Add Reader.SetContext to deprecated ReadContext

diff --git a/client_test.go b/client_test.go
index 26714deba7ea9d40abae391866d7f4c3b59d3133..fc83684b6ef6d7e9422aabf5e1bb141dbd58e219 100644 (file)
--- a/client_test.go
+++ b/client_test.go
@@ -194,6 +194,7 @@ func TestCompletedPieceWrongSize(t *testing.T) {
        assert.True(t, new)
        r := tt.NewReader()
        defer r.Close()
+       r.SetContext(t.Context())
        qt.Check(t, qt.IsNil(iotest.TestReader(r, []byte(testutil.GreetingFileContents))))
 }
 

diff --git a/reader.go b/reader.go
index a283a2ba91b323c8da9acc64dd9db85c138e046a..2af99ce2664227af0f04a1ab4bc226547a81403b 100644 (file)
--- a/reader.go
+++ b/reader.go
@@ -15,6 +15,9 @@ import (
 // also drive Client behaviour. Not safe for concurrent use. There are Torrent, File and Piece
 // constructors for this.
 type Reader interface {
+       // Set the context for reads. When done, reads should get cancelled so they don't get stuck
+       // waiting for data.
+       SetContext(context.Context)
        // Read/Seek and not ReadAt because we want to return data as soon as it's available, and
        // because we want a single read head.
        io.ReadSeekCloser
@@ -55,6 +58,10 @@ type reader struct {
        // Function to dynamically calculate readahead. If nil, readahead is static.
        readaheadFunc ReadaheadFunc
 
+       // This is not protected by a lock because you should be coordinating setting this. If you want
+       // different contexts, you should have different Readers.
+       ctx context.Context
+
        // Required when modifying pos and readahead.
        mu sync.Locker
 
@@ -72,6 +79,10 @@ type reader struct {
        responsive bool
 }
 
+func (r *reader) SetContext(ctx context.Context) {
+       r.ctx = ctx
+}
+
 var _ io.ReadSeekCloser = (*reader)(nil)
 
 func (r *reader) SetResponsive() {
@@ -151,10 +162,23 @@ func (r *reader) piecesUncached() (ret pieceRange) {
 }
 
 func (r *reader) Read(b []byte) (n int, err error) {
-       return r.ReadContext(context.Background(), b)
+       return r.read(b)
 }
 
+func (r *reader) read(b []byte) (n int, err error) {
+       return r.readContext(r.ctx, b)
+}
+
+// Deprecated: Use SetContext and Read. TODO: I've realised this breaks the ability to pass through
+// optional interfaces like io.WriterTo and io.ReaderFrom. Go sux. Context should be provided
+// somewhere else.
 func (r *reader) ReadContext(ctx context.Context, b []byte) (n int, err error) {
+       r.ctx = ctx
+       return r.Read(b)
+}
+
+// We still pass ctx here, although it's a reader field now.
+func (r *reader) readContext(ctx context.Context, b []byte) (n int, err error) {
        if len(b) > 0 {
                r.reading = true
                // TODO: Rework reader piece priorities so we don't have to push updates in to the Client

diff --git a/reader_test.go b/reader_test.go
index c1017a0cb8af6488d8a037312f31b3be84125fb7..3f330c76bfdb719d950c36671bd566f7356af6bf 100644 (file)
--- a/reader_test.go
+++ b/reader_test.go
@@ -24,3 +24,19 @@ func TestReaderReadContext(t *testing.T) {
        _, err = r.ReadContext(ctx, make([]byte, 1))
        require.EqualValues(t, context.DeadlineExceeded, err)
 }
+
+func TestReaderSetContextAndRead(t *testing.T) {
+       cl, err := NewClient(TestingConfig(t))
+       require.NoError(t, err)
+       defer cl.Close()
+       tt, err := cl.AddTorrent(testutil.GreetingMetaInfo())
+       require.NoError(t, err)
+       defer tt.Drop()
+       ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(time.Millisecond))
+       defer cancel()
+       r := tt.Files()[0].NewReader()
+       defer r.Close()
+       r.SetContext(ctx)
+       _, err = r.Read(make([]byte, 1))
+       require.EqualValues(t, context.DeadlineExceeded, err)
+}

diff --git a/t.go b/t.go
index 6ec484b6fb783d587ccaa02f4c6ee5c93da84c56..8a07579f403823de3c59eb99f5f0ed1d70575af0 100644 (file)
--- a/t.go
+++ b/t.go
@@ -1,6 +1,7 @@
 package torrent
 
 import (
+       "context"
        "strconv"
        "strings"
 
@@ -43,6 +44,7 @@ func (t *Torrent) newReader(offset, length int64) Reader {
                t:      t,
                offset: offset,
                length: length,
+               ctx:    context.Background(),
        }
        r.readaheadFunc = defaultReadaheadFunc
        t.addReader(&r)