MM-46610 Ignore the reader timeout for S3 files when importing (#21578)

shared/filestore/s3store.go

@@ -208,26 +208,41 @@ func (b *S3FileBackend) MakeBucket() error {
 // s3WithCancel is a wrapper struct which cancels the context
 // when the object is closed.
 type s3WithCancel struct {
-	*s3.Object
+	io.ReadSeekCloser
+	timer  *time.Timer
 	cancel context.CancelFunc
 }
 
 func (sc *s3WithCancel) Close() error {
+	sc.timer.Stop()
 	sc.cancel()
-	return sc.Object.Close()
+	return sc.ReadSeekCloser.Close()
+}
+
+// CancelTimeout attempts to cancel the timeout for this reader. It allows calling
+// code to ignore the timeout in case of longer running operations. The methods returns
+// false if the timeout has already fired.
+func (sc *s3WithCancel) CancelTimeout() bool {
+	return sc.timer.Stop()
 }
 
 // Caller must close the first return value
 func (b *S3FileBackend) Reader(path string) (ReadCloseSeeker, error) {
 	path = filepath.Join(b.pathPrefix, path)
-	ctx, cancel := context.WithTimeout(context.Background(), b.timeout)
+	ctx, cancel := context.WithCancel(context.Background())
 	minioObject, err := b.client.GetObject(ctx, b.bucket, path, s3.GetObjectOptions{})
 	if err != nil {
 		cancel()
 		return nil, errors.Wrapf(err, "unable to open file %s", path)
 	}
 
-	return &s3WithCancel{Object: minioObject, cancel: cancel}, nil
+	sc := &s3WithCancel{
+		ReadSeekCloser: minioObject,
+		timer:          time.AfterFunc(b.timeout, cancel),
+		cancel:         cancel,
+	}
+
+	return sc, nil
 }
 
 func (b *S3FileBackend) ReadFile(path string) ([]byte, error) {

shared/filestore/s3store_test.go

@@ -10,12 +10,14 @@ import (
 	"encoding/base64"
 	"errors"
 	"fmt"
+	"io"
 	"net/http/httptest"
 	"net/http/httputil"
 	"net/url"
 	"os"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 )
@@ -196,3 +198,105 @@ func TestInsecureMakeBucket(t *testing.T) {
 func newTLSProxyServer(backend *url.URL) *httptest.Server {
 	return httptest.NewTLSServer(httputil.NewSingleHostReverseProxy(backend))
 }
+
+func TestS3WithCancel(t *testing.T) {
+	// Some of these tests use time.Sleep to wait for the timeout to expire.
+	// They are run in parallel to reduce wait times.
+
+	t.Run("zero timeout", func(t *testing.T) {
+		t.Parallel()
+		r, ctx := newMockS3WithCancel(0, nil)
+
+		time.Sleep(10 * time.Millisecond) // give the context time to cancel
+
+		require.False(t, r.CancelTimeout())
+		require.Error(t, ctx.Err())
+	})
+
+	t.Run("timeout", func(t *testing.T) {
+		t.Parallel()
+		r, ctx := newMockS3WithCancel(50*time.Millisecond, nil)
+
+		time.Sleep(100 * time.Millisecond) // give the context time to cancel
+
+		require.False(t, r.CancelTimeout())
+		require.Error(t, ctx.Err())
+	})
+
+	t.Run("timeout cancel", func(t *testing.T) {
+		t.Parallel()
+		r, ctx := newMockS3WithCancel(50*time.Millisecond, nil)
+
+		time.Sleep(10 * time.Millisecond) // give the context time to cancel
+
+		require.True(t, r.CancelTimeout())
+		require.NoError(t, ctx.Err())
+
+		time.Sleep(100 * time.Millisecond) // wait for the original (canceled) timeout to expire
+
+		require.False(t, r.CancelTimeout())
+		require.NoError(t, ctx.Err())
+		require.NoError(t, r.Close())
+	})
+
+	t.Run("timeout closed", func(t *testing.T) {
+		t.Parallel()
+		r, ctx := newMockS3WithCancel(50*time.Millisecond, nil)
+
+		time.Sleep(10 * time.Millisecond) // give the context time to cancel
+
+		require.True(t, r.CancelTimeout())
+		require.NoError(t, ctx.Err())
+		require.NoError(t, r.Close())
+
+		time.Sleep(100 * time.Millisecond) // wait for the original (canceled) timeout to expire
+
+		require.False(t, r.CancelTimeout())
+		require.Error(t, ctx.Err())
+		require.NoError(t, r.Close())
+	})
+
+	t.Run("close cancel close", func(t *testing.T) {
+		t.Parallel()
+		r, ctx := newMockS3WithCancel(50*time.Millisecond, nil)
+
+		time.Sleep(10 * time.Millisecond) // give the context time to cancel
+
+		require.True(t, r.CancelTimeout())
+		require.NoError(t, r.Close())
+		require.Error(t, ctx.Err())
+		require.False(t, r.CancelTimeout())
+		require.Error(t, ctx.Err())
+		require.NoError(t, r.Close())
+	})
+
+	t.Run("close error", func(t *testing.T) {
+		t.Parallel()
+		r, ctx := newMockS3WithCancel(50*time.Millisecond, errors.New("test error"))
+
+		time.Sleep(10 * time.Millisecond) // give the context time to cancel
+
+		require.NoError(t, ctx.Err())
+		require.Error(t, r.Close())
+		require.False(t, r.CancelTimeout())
+		require.Error(t, ctx.Err())
+	})
+}
+
+func newMockS3WithCancel(timeout time.Duration, closeErr error) (*s3WithCancel, context.Context) {
+	ctx, cancel := context.WithCancel(context.Background())
+	return &s3WithCancel{
+		ReadSeekCloser: fauxCloser{strings.NewReader("testdata"), closeErr},
+		timer:          time.AfterFunc(timeout, cancel),
+		cancel:         cancel,
+	}, ctx
+}
+
+type fauxCloser struct {
+	io.ReadSeeker
+	closeErr error
+}
+
+func (fc fauxCloser) Close() error {
+	return fc.closeErr
+}