add ReadFileStream deadline for disk call (#17745)

timeout the reader side if hung via disk max timeout

internal/ioutil/ioutil.go

@@ -73,7 +73,54 @@ type ioret struct {
 	err error
 }
 
-// DeadlineWriter deadline writer with context
+// DeadlineReader deadline reader with timeout
+type DeadlineReader struct {
+	io.ReadCloser
+	timeout time.Duration
+	err     error
+}
+
+// NewDeadlineReader wraps a writer to make it respect given deadline
+// value per Write(). If there is a blocking write, the returned Reader
+// will return whenever the timer hits (the return values are n=0
+// and err=context.DeadlineExceeded.)
+func NewDeadlineReader(r io.ReadCloser, timeout time.Duration) io.ReadCloser {
+	return &DeadlineReader{ReadCloser: r, timeout: timeout}
+}
+
+func (r *DeadlineReader) Read(buf []byte) (int, error) {
+	if r.err != nil {
+		return 0, r.err
+	}
+
+	c := make(chan ioret, 1)
+	t := time.NewTimer(r.timeout)
+	go func() {
+		n, err := r.ReadCloser.Read(buf)
+		c <- ioret{n, err}
+		close(c)
+	}()
+
+	select {
+	case res := <-c:
+		if !t.Stop() {
+			<-t.C
+		}
+		r.err = res.err
+		return res.n, res.err
+	case <-t.C:
+		r.ReadCloser.Close()
+		r.err = context.DeadlineExceeded
+		return 0, context.DeadlineExceeded
+	}
+}
+
+// Close closer interface to close the underlying closer
+func (r *DeadlineReader) Close() error {
+	return r.ReadCloser.Close()
+}
+
+// DeadlineWriter deadline writer with timeout
 type DeadlineWriter struct {
 	io.WriteCloser
 	timeout time.Duration
@@ -119,15 +166,15 @@ func (d *DeadlineWorker) Run(work func() error) error {
 		d.err = r.err
 		return r.err
 	case <-t.C:
-		d.err = context.Canceled
-		return context.Canceled
+		d.err = context.DeadlineExceeded
+		return context.DeadlineExceeded
 	}
 }
 
 // NewDeadlineWriter wraps a writer to make it respect given deadline
 // value per Write(). If there is a blocking write, the returned Writer
 // will return whenever the timer hits (the return values are n=0
-// and err=context.Canceled.)
+// and err=context.DeadlineExceeded.)
 func NewDeadlineWriter(w io.WriteCloser, timeout time.Duration) io.WriteCloser {
 	return &DeadlineWriter{WriteCloser: w, timeout: timeout}
 }
@@ -154,8 +201,8 @@ func (w *DeadlineWriter) Write(buf []byte) (int, error) {
 		return r.n, r.err
 	case <-t.C:
 		w.WriteCloser.Close()
-		w.err = context.Canceled
-		return 0, context.Canceled
+		w.err = context.DeadlineExceeded
+		return 0, context.DeadlineExceeded
 	}
 }
 

internal/ioutil/ioutil_test.go

@@ -26,6 +26,19 @@ import (
 	"time"
 )
 
+type sleepReader struct {
+	timeout time.Duration
+}
+
+func (r *sleepReader) Read(p []byte) (n int, err error) {
+	time.Sleep(r.timeout)
+	return len(p), nil
+}
+
+func (r *sleepReader) Close() error {
+	return nil
+}
+
 type sleepWriter struct {
 	timeout time.Duration
 }
@@ -39,16 +52,39 @@ func (w *sleepWriter) Close() error {
 	return nil
 }
 
+func TestDeadlineReader(t *testing.T) {
+	r := NewDeadlineReader(&sleepReader{timeout: 500 * time.Millisecond}, 450*time.Millisecond)
+	buf := make([]byte, 4)
+	_, err := r.Read(buf)
+	r.Close()
+	if err != context.DeadlineExceeded {
+		t.Errorf("DeadlineReader shouldn't be successful %v - should return context.DeadlineExceeded", err)
+	}
+	_, err = r.Read(buf)
+	if err != context.DeadlineExceeded {
+		t.Errorf("DeadlineReader shouldn't be successful %v - should return context.DeadlineExceeded", err)
+	}
+	r = NewDeadlineReader(&sleepReader{timeout: 100 * time.Millisecond}, 600*time.Millisecond)
+	n, err := r.Read(buf)
+	r.Close()
+	if err != nil {
+		t.Errorf("DeadlineReader should succeed but failed with %s", err)
+	}
+	if n != 4 {
+		t.Errorf("DeadlineReader should succeed but should have only read 4 bytes, but returned %d instead", n)
+	}
+}
+
 func TestDeadlineWriter(t *testing.T) {
 	w := NewDeadlineWriter(&sleepWriter{timeout: 500 * time.Millisecond}, 450*time.Millisecond)
 	_, err := w.Write([]byte("1"))
 	w.Close()
-	if err != context.Canceled {
-		t.Error("DeadlineWriter shouldn't be successful - should return context.Canceled")
+	if err != context.DeadlineExceeded {
+		t.Error("DeadlineWriter shouldn't be successful - should return context.DeadlineExceeded")
 	}
 	_, err = w.Write([]byte("1"))
-	if err != context.Canceled {
-		t.Error("DeadlineWriter shouldn't be successful - should return context.Canceled")
+	if err != context.DeadlineExceeded {
+		t.Error("DeadlineWriter shouldn't be successful - should return context.DeadlineExceeded")
 	}
 	w = NewDeadlineWriter(&sleepWriter{timeout: 100 * time.Millisecond}, 600*time.Millisecond)
 	n, err := w.Write([]byte("abcd"))