cloudstorage: fix a bug causing context cancelled errors and stuck sink (#12276)

close #12277

Author: zurakutsia

pkg/util/external_storage.go

@@ -197,9 +197,28 @@ func (s *extStorageWithTimeout) DeleteFile(ctx context.Context, name string) err
 func (s *extStorageWithTimeout) Open(
 	ctx context.Context, path string, _ *storage.ReaderOption,
 ) (storage.ExternalFileReader, error) {
+	// Unlike other methods, Open method cannot call cancel() in defer.
+	// This is because the reader's lifetime is bound to the context provided at Open().
+	// Subsequent Read() calls on reader will observe context cancellation.
+	// Instead, we wrap the reader in a struct and cancel it's context in Close().
 	ctx, cancel := context.WithTimeout(ctx, s.timeout)
-	defer cancel()
-	return s.ExternalStorage.Open(ctx, path, nil)
+	r, err := s.ExternalStorage.Open(ctx, path, nil)
+	if err != nil {
+		cancel()
+		return nil, err
+	}
+	return &readerWithCancel{ExternalFileReader: r, cancel: cancel}, nil
+}
+
+type readerWithCancel struct {
+	storage.ExternalFileReader
+	cancel context.CancelFunc
+}
+
+// Close the reader and cancel the context.
+func (r *readerWithCancel) Close() error {
+	defer r.cancel()
+	return r.ExternalFileReader.Close()
 }
 
 // WalkDir traverse all the files in a dir.

pkg/util/external_storage_test.go

@@ -67,6 +67,15 @@ func (m *mockExternalStorage) WriteFile(ctx context.Context, name string, data [
 	return nil
 }
 
+// Implement Open so tests can simulate readers that bind to the Open() context.
+func (m *mockExternalStorage) Open(ctx context.Context, path string, option *storage.ReaderOption) (storage.ExternalFileReader, error) {
+	return &ctxBoundReader{ctx: ctx}, nil
+}
+
+func (m *mockExternalStorage) URI() string { return "mock://" }
+
+func (m *mockExternalStorage) Close() {}
+
 func TestExtStorageWithTimeoutWriteFileTimeout(t *testing.T) {
 	testTimeout := 50 * time.Millisecond
 
@@ -123,3 +132,62 @@ func TestExtStorageWithTimeoutWriteFileSuccess(t *testing.T) {
 	// Assert success
 	require.NoError(t, err, "Expected no error for successful write within timeout")
 }
+
+// ctxBoundReader is a reader that checks the context passed to Open().
+// It simulates backends (e.g., Azure) that bind reader lifetime to the Open() context.
+type ctxBoundReader struct {
+	ctx context.Context
+}
+
+func (r *ctxBoundReader) Read(p []byte) (int, error) {
+	if err := r.ctx.Err(); err != nil {
+		return 0, err
+	}
+	if len(p) == 0 {
+		return 0, nil
+	}
+	p[0] = 'x'
+	return 1, nil
+}
+
+func (r *ctxBoundReader) Seek(offset int64, whence int) (int64, error) {
+	return 0, nil
+}
+
+func (r *ctxBoundReader) Close() error { return nil }
+
+func (r *ctxBoundReader) GetFileSize() (int64, error) { return 1, nil }
+
+func TestExtStorageOpenDoesNotCancelReaderContext(t *testing.T) {
+	timedStore := &extStorageWithTimeout{
+		ExternalStorage: &mockExternalStorage{},
+		timeout:         100 * time.Millisecond,
+	}
+
+	rd, err := timedStore.Open(context.Background(), "file", nil)
+	require.NoError(t, err)
+	defer rd.Close()
+
+	// If Open() had used a derived context with immediate cancel, this would fail with context canceled.
+	_, err = rd.Read(make([]byte, 1))
+	require.NoError(t, err)
+}
+
+func TestExtStorageOpenReaderRespectsCallerCancel(t *testing.T) {
+	timedStore := &extStorageWithTimeout{
+		ExternalStorage: &mockExternalStorage{},
+		timeout:         10 * time.Millisecond,
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	rd, err := timedStore.Open(ctx, "file", nil)
+	require.NoError(t, err)
+	defer rd.Close()
+
+	// This should cause the reader to fail with context canceled.
+	cancel()
+	_, err = rd.Read(make([]byte, 1))
+
+	require.Error(t, err)
+	require.True(t, errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded))
+}