github.com/anacrolix/torrent@v1.61.0/internal/ctxrw/ctxrw.go (about)

     1  package ctxrw
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  
     7  	g "github.com/anacrolix/generics"
     8  )
     9  
    10  type contextedReader struct {
    11  	ctx context.Context
    12  	r   io.Reader
    13  }
    14  
    15  func (me contextedReader) Read(p []byte) (n int, err error) {
    16  	return contextedReadOrWrite(me.ctx, me.r.Read, p)
    17  }
    18  
    19  type contextedWriter struct {
    20  	ctx context.Context
    21  	w   io.Writer
    22  }
    23  
    24  // This is problematic. If you return with a context error, a read or write is still pending, and
    25  // could mess up the stream.
    26  func contextedReadOrWrite(ctx context.Context, method func(b []byte) (int, error), b []byte) (_ int, err error) {
    27  	asyncCh := make(chan g.Result[int], 1)
    28  	go func() {
    29  		asyncCh <- g.ResultFromTuple(method(b))
    30  	}()
    31  	select {
    32  	case <-ctx.Done():
    33  		err = context.Cause(ctx)
    34  		return
    35  	case res := <-asyncCh:
    36  		return res.AsTuple()
    37  	}
    38  
    39  }
    40  
    41  func (me contextedWriter) Write(p []byte) (n int, err error) {
    42  	return contextedReadOrWrite(me.ctx, me.w.Write, p)
    43  }
    44  
    45  func WrapReadWriter(ctx context.Context, rw io.ReadWriter) io.ReadWriter {
    46  	return struct {
    47  		io.Reader
    48  		io.Writer
    49  	}{
    50  		contextedReader{
    51  			ctx: ctx,
    52  			r:   rw,
    53  		},
    54  		contextedWriter{
    55  			ctx: ctx,
    56  			w:   rw,
    57  		},
    58  	}
    59  }