Add a Read timeout for readbuffer

readbuffer.go

@@ -5,23 +5,28 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"os"
 	"sync"
 	"time"
 
 	"github.com/sirupsen/logrus"
 )
 
 const (
-	MaxBuffer = 1 << 21
+	MaxBuffer          = 1 << 21
+	DefaultReadTimeout = 1 * time.Minute
 )
 
+var ErrReadTimeoutExceeded = errors.New("read timeout exceeded")
+
 type readBuffer struct {
 	id, readCount, offerCount int64
 	cond                      sync.Cond
 	deadline                  time.Time
 	buf                       bytes.Buffer
 	err                       error
 	backPressure              *backPressure
+	readTimeout               time.Duration
 }
 
 func newReadBuffer(id int64, backPressure *backPressure) *readBuffer {
@@ -31,6 +36,7 @@ func newReadBuffer(id int64, backPressure *backPressure) *readBuffer {
 		cond: sync.Cond{
 			L: &sync.Mutex{},
 		},
+		readTimeout: DefaultReadTimeout,
 	}
 }
 
@@ -100,13 +106,15 @@ func (r *readBuffer) Read(b []byte) (int, error) {
 		now := time.Now()
 		if !r.deadline.IsZero() {
 			if now.After(r.deadline) {
-				return 0, errors.New("deadline exceeded")
+				return 0, fmt.Errorf("remotedialer readbuffer Read deadline exceeded: %w", os.ErrDeadlineExceeded)
 			}
 		}
 
 		var t *time.Timer
 		if !r.deadline.IsZero() {
 			t = time.AfterFunc(r.deadline.Sub(now), func() { r.cond.Broadcast() })
+		} else {
+			t = time.AfterFunc(r.readTimeout, func() { r.Close(ErrReadTimeoutExceeded) })
 		}
 		r.cond.Wait()
 		if t != nil {

readbuffer_test.go

@@ -0,0 +1,102 @@
+package remotedialer
+
+import (
+	"bytes"
+	"fmt"
+	"os"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_readBuffer_Read(t *testing.T) {
+	type args struct {
+		b []byte
+	}
+	tests := []struct {
+		name     string
+		args     args
+		deadline time.Duration
+		timeout  time.Duration
+		offer    []byte
+		want     int
+		wantErr  assert.ErrorAssertionFunc
+	}{
+		{
+			name: "Read",
+			args: args{
+				b: make([]byte, 10),
+			},
+			offer:   []byte("test"),
+			want:    4,
+			wantErr: assert.NoError,
+		},
+		{
+			name: "Read with timeout",
+			args: args{
+				b: make([]byte, 10),
+			},
+			timeout: 100 * time.Millisecond,
+			want:    0,
+			wantErr: func(t assert.TestingT, err error, msgAndArgs ...interface{}) bool {
+				return assert.ErrorIs(t, err, ErrReadTimeoutExceeded)
+			},
+		},
+		{
+			name: "Read with deadline",
+			args: args{
+				b: make([]byte, 10),
+			},
+			deadline: 100 * time.Millisecond,
+			offer:    nil,
+			want:     0,
+			wantErr: func(t assert.TestingT, err error, msgAndArgs ...interface{}) bool {
+				return assert.ErrorIs(t, err, os.ErrDeadlineExceeded)
+			},
+		},
+		{
+			name: "Read with timeout and deadline",
+			args: args{
+				b: make([]byte, 10),
+			},
+			deadline: 50 * time.Millisecond,
+			offer:    nil,
+			want:     0,
+			wantErr: func(t assert.TestingT, err error, msgAndArgs ...interface{}) bool {
+				return assert.ErrorIs(t, err, os.ErrDeadlineExceeded)
+			},
+		},
+	}
+
+	t.Parallel()
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			r := &readBuffer{
+				id:           1,
+				backPressure: newBackPressure(nil),
+				cond: sync.Cond{
+					L: &sync.Mutex{},
+				},
+			}
+			if tt.deadline > 0 {
+				r.deadline = time.Now().Add(tt.deadline)
+			}
+			if tt.timeout > 0 {
+				r.readTimeout = tt.timeout
+			}
+			if tt.offer != nil {
+				err := r.Offer(bytes.NewReader(tt.offer))
+				if err != nil {
+					t.Errorf("Offer() returned error: %v", err)
+				}
+			}
+			got, err := r.Read(tt.args.b)
+			if !tt.wantErr(t, err, fmt.Sprintf("Read(%v)", tt.args.b)) {
+				return
+			}
+			assert.Equalf(t, tt.want, got, "Read(%v)", tt.args.b)
+		})
+	}
+}