From e0ab8478a74bdf4806073261aadba9617962efc4 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Mon, 28 Apr 2025 10:27:17 +1000
Subject: [PATCH] Add Reader.SetContext to deprecated ReadContext

---
 client_test.go |  1 +
 reader.go      | 26 +++++++++++++++++++++++++-
 reader_test.go | 16 ++++++++++++++++
 t.go           |  2 ++
 4 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/client_test.go b/client_test.go
index 26714deb..fc83684b 100644
--- 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 a283a2ba..2af99ce2 100644
--- 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 c1017a0c..3f330c76 100644
--- 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 6ec484b6..8a07579f 100644
--- 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)
-- 
2.48.1