Merge pull request #2529 from keboola/devin/PSGO-184-1772005379-fix-file-close-lock-timeout

fix: use dedicated context for lock acquisition in closeFile (PSGO-184)

Author: Matovidlo

internal/pkg/service/stream/storage/node/coordinator/filerotation/operator.go

@@ -476,8 +476,12 @@ func (o *operator) closeFile(ctx context.Context, file *fileData) {
 	dbCtx, dbCancel := context.WithTimeoutCause(context.WithoutCancel(ctx), dbOperationTimeout, errors.New("switch to importing timeout"))
 	defer dbCancel()
 
-	// Lock all file operations
-	lock, unlock, lockErr := clusterlock.LockFile(ctx, o.locks, o.logger, file.FileKey)
+	// Lock all file operations.
+	// Use a fresh context - waitForFileClosing may have consumed the FileCloseTimeout budget.
+	lockCtx, lockCancel := context.WithTimeoutCause(context.WithoutCancel(ctx), dbOperationTimeout, errors.New("file close lock acquisition timeout"))
+	defer lockCancel()
+
+	lock, unlock, lockErr := clusterlock.LockFile(lockCtx, o.locks, o.logger, file.FileKey)
 	if lockErr != nil {
 		o.logger.Errorf(ctx, `file close error: %s`, lockErr)
 		return
@@ -507,8 +511,12 @@ func (o *operator) closeFile(ctx context.Context, file *fileData) {
 
 	// If there is an error, increment retry delay
 	if err != nil {
-		o.logger.Error(dbCtx, err.Error())
-		fileEntity, rErr := o.storage.File().IncrementRetryAttempt(file.FileKey, o.clock.Now(), err.Error()).RequireLock(lock).Do(dbCtx).ResultOrErr()
+		retryCtx, retryCancel := context.WithTimeoutCause(context.WithoutCancel(ctx), dbOperationTimeout, errors.New("retry increment timeout"))
+		defer retryCancel()
+
+		o.logger.Error(ctx, err.Error())
+
+		fileEntity, rErr := o.storage.File().IncrementRetryAttempt(file.FileKey, o.clock.Now(), err.Error()).RequireLock(lock).Do(retryCtx).ResultOrErr()
 		if rErr != nil {
 			o.logger.Errorf(ctx, "cannot increment file close retry: %s", rErr)
 			return

internal/pkg/service/stream/storage/node/coordinator/filerotation/operator_test.go

@@ -33,6 +33,8 @@ import (
 // fileExpirationDiff defines the time between file opening and file import trigger.
 const fileExpirationDiff = time.Minute
 
+const shortFileCloseTimeout = 10 * time.Second
+
 func TestFileRotation(t *testing.T) {
 	t.Parallel()
 
@@ -182,6 +184,71 @@ func TestFileRotation(t *testing.T) {
 	ts.logger.AssertNoErrorMessage(t)
 }
 
+// TestCloseFile_LockAcquisitionAfterWaitTimeout is a regression test for PSGO-184: when
+// waitForFileClosing exhausts the FileCloseTimeout, the subsequent lock acquisition must use
+// its own fresh context, not the expired one, so IncrementRetryAttempt can still be called.
+func TestCloseFile_LockAcquisitionAfterWaitTimeout(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(t.Context(), 60*time.Second)
+	defer cancel()
+
+	ts := setupWithCloseTimeout(t, ctx, shortFileCloseTimeout)
+	defer ts.teardown(t)
+	ts.prepareFixtures(t, ctx)
+
+	// Trigger file rotation to put the old file into FileClosing state.
+	ts.clk.Advance(fileExpirationDiff)
+	ts.triggerCheck(t, true, `
+{"level":"info","message":"rotating file, import conditions met: expiration threshold met, expiration: 2000-01-01T00:31:00.000Z, remains: %s, threshold: 30m0s","file.id":"%s","component":"storage.node.operator.file.rotation"}
+`)
+
+	files, err := ts.dependencies.StorageRepository().File().ListIn(ts.sink.SinkKey).Do(ctx).All()
+	require.NoError(t, err)
+	require.Len(t, files, 2)
+	require.Equal(t, model.FileClosing, files[0].State)
+	require.Equal(t, model.FileWriting, files[1].State)
+
+	slices, err := ts.dependencies.StorageRepository().Slice().ListIn(ts.sink.SinkKey).Do(ctx).All()
+	require.NoError(t, err)
+	require.Len(t, slices, 4)
+	// ListIn returns slices in lexicographic order by etcd key (oldest file first).
+	require.Equal(t, files[0].FileKey, slices[0].FileKey)
+	require.Equal(t, model.SliceClosing, slices[0].State)
+	require.Equal(t, files[0].FileKey, slices[1].FileKey)
+	require.Equal(t, model.SliceClosing, slices[1].State)
+	require.Equal(t, files[1].FileKey, slices[2].FileKey)
+	require.Equal(t, model.SliceWriting, slices[2].State)
+	require.Equal(t, files[1].FileKey, slices[3].FileKey)
+	require.Equal(t, model.SliceWriting, slices[3].State)
+
+	closingFileKey := files[0].FileKey
+	ts.logger.Truncate()
+
+	// Trigger a close attempt; waitForFileClosing will block until FileCloseTimeout (10s) expires.
+	// Afterwards, lock acquisition and IncrementRetryAttempt must still succeed.
+	// Uses real wall-clock time - the 10s timeout cannot be faked with clockwork.FakeClock.
+	ts.clk.Advance(ts.interval)
+
+	assert.EventuallyWithT(t, func(c *assert.CollectT) {
+		ts.logger.AssertJSONMessages(c, `
+{"level":"error","message":"error when waiting for file slices upload:\n- context deadline exceeded","component":"storage.node.operator.file.rotation"}
+{"level":"info","message":"file closing will be retried after %s","component":"storage.node.operator.file.rotation"}
+`)
+	}, 15*time.Second, 100*time.Millisecond)
+
+	file, err := ts.dependencies.StorageRepository().File().Get(closingFileKey).Do(ctx).ResultOrErr()
+	require.NoError(t, err)
+	require.Equal(t, model.FileClosing, file.State)
+	require.Positive(t, file.RetryAttempt, "retry attempt should be incremented, proving lock was acquired after waitForFileClosing timeout")
+
+	warnAndErrorLogs := ts.logger.WarnAndErrorMessages()
+	assert.NotContains(t, warnAndErrorLogs, "cannot acquire lock", "lock acquisition should not fail when using a dedicated context")
+
+	ts.dependencies.Process().Shutdown(ctx, errors.New("bye bye"))
+	ts.dependencies.Process().WaitForShutdown()
+}
+
 type testState struct {
 	interval      time.Duration
 	importTrigger targetConfig.ImportTrigger
@@ -196,6 +263,11 @@ type testState struct {
 
 func setup(t *testing.T, ctx context.Context) *testState {
 	t.Helper()
+	return setupWithCloseTimeout(t, ctx, 0)
+}
+
+func setupWithCloseTimeout(t *testing.T, ctx context.Context, closeTimeout time.Duration) *testState {
+	t.Helper()
 
 	importTrigger := targetConfig.ImportTrigger{
 		Count:       50000,
@@ -219,6 +291,9 @@ func setup(t *testing.T, ctx context.Context) *testState {
 			Trigger:     importTrigger,
 		}
 		cfg.Storage.Level.Target.Operator.FileRotationCheckInterval = duration.From(conditionsCheckInterval)
+		if closeTimeout > 0 {
+			cfg.Storage.Level.Target.Operator.FileCloseTimeout = duration.From(closeTimeout)
+		}
 	}, commonDeps.WithClock(clk))
 	client := mock.TestEtcdClient()
 
@@ -289,4 +364,5 @@ func (ts *testState) prepareFixtures(t *testing.T, ctx context.Context) {
 	require.NoError(t, ts.dependencies.DefinitionRepository().Source().Create(&source, ts.clk.Now(), test.ByUser(), "create").Do(ctx).Err())
 	require.NoError(t, ts.dependencies.DefinitionRepository().Sink().Create(&ts.sink, ts.clk.Now(), test.ByUser(), "create").Do(ctx).Err())
 	ts.waitForFilesSync(t)
+	ts.logger.Truncate()
 }