[-] fix active_chain row leak when a chain fails before its tasks run

internal/log/log.go

@@ -3,6 +3,7 @@ package log
 import (
 	"context"
 	"os"
+	"time"
 
 	"github.com/cybertec-postgresql/pg_timetable/internal/config"
 	"github.com/jackc/pgx/v5/tracelog"
@@ -119,6 +120,15 @@ func WithLogger(ctx context.Context, logger LoggerIface) context.Context {
 	return context.WithValue(ctx, loggerKey{}, logger)
 }
 
+// DetachedTimeout bounds cleanup work that must run even when the caller's context is cancelled.
+var DetachedTimeout = 30 * time.Second // var so tests can shorten it
+
+// DetachedContext returns a context that survives cancellation of ctx but keeps its logger,
+// bounded by DetachedTimeout so a dead database cannot block the caller forever.
+func DetachedContext(ctx context.Context) (context.Context, context.CancelFunc) {
+	return context.WithTimeout(context.WithoutCancel(ctx), DetachedTimeout)
+}
+
 // FallbackLogger is an alias for the standard logger
 var FallbackLogger = logrus.StandardLogger()
 

internal/pgengine/access.go

@@ -66,8 +66,13 @@ SELECT $1, $2 WHERE
 	return res.RowsAffected() == 1
 }
 
+// RemoveChainRunStatus frees one run's slot, so other running instances keep theirs
 func (pge *PgEngine) RemoveChainRunStatus(ctx context.Context, chainID int) {
-	const sqlRemoveRunStatus = `DELETE FROM timetable.active_chain WHERE chain_id = $1 and client_name = $2`
+	const sqlRemoveRunStatus = `DELETE FROM timetable.active_chain WHERE ctid = (
+	SELECT ctid FROM timetable.active_chain
+	WHERE chain_id = $1 AND client_name = $2
+	LIMIT 1 FOR UPDATE SKIP LOCKED
+)`
 	_, err := pge.ConfigDb.Exec(ctx, sqlRemoveRunStatus, chainID, pge.ClientName)
 	if err != nil {
 		pge.l.WithError(err).Error("Cannot save information about the chain run status")

internal/pgengine/pgengine_test.go

@@ -109,7 +109,7 @@ func TestSchedulerFunctions(t *testing.T) {
 		assert.Greater(t, txid, int64(0), "Should return transaction id")
 		assert.NoError(t, pge.GetChainElements(ctx, &chains, 0), "Should no error in clean database")
 		assert.Empty(t, chains, "Should be empty in clean database")
-		pge.CommitTransaction(ctx, tx)
+		assert.NoError(t, tx.Commit(ctx))
 	})
 
 	t.Run("Check GetChainElements returns only live tasks", func(t *testing.T) {
@@ -137,7 +137,7 @@ func TestSchedulerFunctions(t *testing.T) {
 			TaskID:  0,
 			ChainID: 0}), "Should no error in clean database")
 		assert.Empty(t, paramVals, "Should be empty in clean database")
-		pge.CommitTransaction(ctx, tx)
+		assert.NoError(t, tx.Commit(ctx))
 	})
 
 	t.Run("Check InsertChainRunStatus function", func(t *testing.T) {
@@ -147,6 +147,18 @@ func TestSchedulerFunctions(t *testing.T) {
 		assert.True(t, res, "Active chain should be inserted")
 	})
 
+	t.Run("Check RemoveChainRunStatus removes one instance only", func(t *testing.T) {
+		const chainID = 4242
+		require.True(t, pge.InsertChainRunStatus(ctx, chainID, 16))
+		require.True(t, pge.InsertChainRunStatus(ctx, chainID, 16))
+		pge.RemoveChainRunStatus(ctx, chainID)
+		var count int
+		require.NoError(t, pge.ConfigDb.QueryRow(ctx,
+			"SELECT count(*) FROM timetable.active_chain WHERE chain_id = $1", chainID).Scan(&count))
+		assert.Equal(t, 1, count, "Other running instance must keep its slot")
+		pge.RemoveChainRunStatus(ctx, chainID)
+	})
+
 	t.Run("Check ExecuteSQLCommand function", func(t *testing.T) {
 		tx, txid, err := pge.StartTransaction(ctx)
 		assert.NoError(t, err, "Should start transaction")
@@ -163,7 +175,7 @@ func TestSchedulerFunctions(t *testing.T) {
 		assert.NoError(t, f("SELECT $1::int4", []string{"[42]", `[14]`}), "Simple query with doubled parameters")
 		assert.NoError(t, f("SELECT $1::int4, $2::text", []string{`[42, "hey"]`}), "Simple query with two parameters")
 
-		pge.CommitTransaction(ctx, tx)
+		assert.NoError(t, tx.Commit(ctx))
 	})
 
 }

internal/pgengine/transaction.go

@@ -16,21 +16,19 @@ import (
 // StartTransaction returns transaction object, virtual transaction id and error
 func (pge *PgEngine) StartTransaction(ctx context.Context) (tx pgx.Tx, vxid int64, err error) {
 	if tx, err = pge.ConfigDb.Begin(ctx); err != nil {
-		return
+		return nil, 0, err
 	}
 	err = tx.QueryRow(ctx, `SELECT 
 (split_part(virtualxid, '/', 1)::int8 << 32) | split_part(virtualxid, '/', 2)::int8
 FROM pg_locks 
 WHERE pid = pg_backend_pid() AND virtualxid IS NOT NULL`).Scan(&vxid)
-	return
-}
-
-// CommitTransaction commits transaction and log error in the case of error
-func (pge *PgEngine) CommitTransaction(ctx context.Context, tx pgx.Tx) {
-	err := tx.Commit(ctx)
 	if err != nil {
-		log.GetLogger(ctx).WithError(err).Error("Application cannot commit after job finished")
+		bctx, cancel := log.DetachedContext(ctx)
+		defer cancel()
+		pge.RollbackTransaction(bctx, tx)
+		return nil, 0, err
 	}
+	return tx, vxid, nil
 }
 
 // RollbackTransaction rollbacks transaction and log error in the case of error

internal/pgengine/transaction_test.go

@@ -44,8 +44,17 @@ func TestStartTransaction(t *testing.T) {
 	assert.Error(t, err)
 
 	mockPool.ExpectBegin()
-	mockPool.ExpectQuery("SELECT").WillReturnRows(pgxmock.NewRows([]string{"txid"}).AddRow(int64(42)))
+	mockPool.ExpectQuery("SELECT").WillReturnError(errors.New("foo"))
+	mockPool.ExpectRollback()
 	tx, txid, err := pge.StartTransaction(ctx)
+	assert.Nil(t, tx, "half-started transaction must be rolled back, not returned")
+	assert.Zero(t, txid)
+	assert.Error(t, err)
+	assert.NoError(t, mockPool.ExpectationsWereMet())
+
+	mockPool.ExpectBegin()
+	mockPool.ExpectQuery("SELECT").WillReturnRows(pgxmock.NewRows([]string{"txid"}).AddRow(int64(42)))
+	tx, txid, err = pge.StartTransaction(ctx)
 	assert.NotNil(t, tx)
 	assert.EqualValues(t, 42, txid)
 	assert.NoError(t, err)
@@ -56,15 +65,9 @@ func TestMustTransaction(t *testing.T) {
 
 	pge := pgengine.NewDB(mockPool, "pgengine_unit_test")
 
-	mockPool.ExpectBegin()
-	mockPool.ExpectCommit().WillReturnError(errors.New("error"))
-	tx, err := mockPool.Begin(context.Background())
-	assert.NoError(t, err)
-	pge.CommitTransaction(ctx, tx)
-
 	mockPool.ExpectBegin()
 	mockPool.ExpectRollback().WillReturnError(errors.New("error"))
-	tx, err = mockPool.Begin(context.Background())
+	tx, err := mockPool.Begin(context.Background())
 	assert.NoError(t, err)
 	pge.RollbackTransaction(ctx, tx)
 

internal/scheduler/chain.go

@@ -159,6 +159,7 @@ func (sch *Scheduler) chainWorker(ctx context.Context, chains <-chan Chain) {
 				chainContext, cancel := context.WithCancel(chainContext)
 				sch.addActiveChain(chain.ChainID, cancel)
 				sch.executeChain(chainContext, chain)
+				sch.removeChainRunStatus(chainContext, chain.ChainID)
 				sch.deleteActiveChain(chain.ChainID)
 				cancel()
 				sch.Unlock(chain.ExclusiveExecution)
@@ -178,6 +179,14 @@ func getTimeoutContext(ctx context.Context, globalTimeout int, customTimeout int
 	return ctx, nil
 }
 
+// removeChainRunStatus frees the chain's active_chain slot. Runs detached so it
+// survives a cancelled scheduler context and cannot hang on a dead database.
+func (sch *Scheduler) removeChainRunStatus(ctx context.Context, chainID int) {
+	bctx, cancel := log.DetachedContext(ctx)
+	defer cancel()
+	sch.pgengine.RemoveChainRunStatus(bctx, chainID)
+}
+
 func (sch *Scheduler) executeOnErrorHandler(ctx context.Context, chain Chain) {
 	if ctx.Err() != nil || chain.OnError == "" {
 		return
@@ -191,9 +200,8 @@ func (sch *Scheduler) executeOnErrorHandler(ctx context.Context, chain Chain) {
 	}
 	sch.pgengine.SetCurrentTaskContext(ctx, tx, chain.ChainID, 0)
 	_, err = tx.Exec(ctx, chain.OnError)
-	sch.pgengine.CommitTransaction(ctx, tx)
-	if err != nil {
-		l.Info("Error handler failed")
+	if err = cmp.Or(err, tx.Commit(ctx)); err != nil {
+		l.WithError(err).Info("Error handler failed")
 		return
 	}
 	l.Info("Error handler executed successfully")
@@ -202,8 +210,6 @@ func (sch *Scheduler) executeOnErrorHandler(ctx context.Context, chain Chain) {
 /* execute a chain of tasks */
 func (sch *Scheduler) executeChain(ctx context.Context, chain Chain) {
 	var ChainTasks []pgengine.ChainTask
-	var bctx context.Context
-	var cancel context.CancelFunc
 	var vxid int64
 
 	// OTel tracing: create root span for this chain execution
@@ -224,17 +230,29 @@ func (sch *Scheduler) executeChain(ctx context.Context, chain Chain) {
 	}
 
 	chainL := sch.l.WithField("chain", chain)
+	chainFailed := func(l log.LoggerIface, tx pgx.Tx, err error, msg string) {
+		l.WithError(err).Error(msg)
+		chainSpan.SetStatus(codes.Error, msg)
+		lctx := log.WithLogger(ctx, l)
+		if tx != nil {
+			bctx, cancel := log.DetachedContext(lctx)
+			sch.pgengine.RollbackTransaction(bctx, tx)
+			cancel()
+		}
+		sch.provider.RecordChainFailed(context.WithoutCancel(lctx), sch.Config().ClientName)
+		sch.executeOnErrorHandler(lctx, chain) // skipped if the chain was stopped
+	}
+
 	tx, vxid, err := sch.pgengine.StartTransaction(chainCtx)
 	if err != nil {
-		chainL.WithError(err).Error("Cannot start transaction")
+		chainFailed(chainL, nil, err, "Chain failed to start")
 		return
 	}
 	chainL = chainL.WithField("vxid", vxid)
 
 	err = sch.pgengine.GetChainElements(chainCtx, &ChainTasks, chain.ChainID)
 	if err != nil {
-		chainL.WithError(err).Error("Failed to retrieve chain elements")
-		sch.pgengine.RollbackTransaction(chainCtx, tx)
+		chainFailed(chainL, tx, err, "Chain failed to load tasks")
 		return
 	}
 
@@ -252,30 +270,25 @@ func (sch *Scheduler) executeChain(ctx context.Context, chain Chain) {
 			l.Info("Task executed successfully")
 		}
 
-		// we detach the context from cancellation here because the current one
-		// (chainCtx and its parent ctx) might be cancelled, e.g. by notify_chain_stop().
-		// Cleanup operations below must still run to keep timetable.active_chain consistent.
-		bctx = log.WithLogger(context.WithoutCancel(ctx), l)
 		if err != nil {
 			if !task.IgnoreError {
-				chainL.Error("Chain failed")
-				sch.pgengine.RemoveChainRunStatus(bctx, chain.ChainID)
-				sch.pgengine.RollbackTransaction(bctx, tx)
-				chainSpan.SetStatus(codes.Error, "chain failed")
-				sch.provider.RecordChainFailed(bctx, sch.Config().ClientName)
-				sch.executeOnErrorHandler(bctx, chain)
+				chainFailed(l, tx, err, "Chain failed")
 				return
 			}
 			l.Info("Ignoring task failure")
 		}
 	}
-	bctx = log.WithLogger(context.WithoutCancel(chainCtx), chainL)
-	sch.pgengine.CommitTransaction(bctx, tx)
+	// commit even if the chain was stopped meanwhile: all its tasks have succeeded
+	if err = tx.Commit(log.WithLogger(context.WithoutCancel(ctx), chainL)); err != nil {
+		chainFailed(chainL, nil, err, "Chain failed to commit")
+		return
+	}
 	sch.provider.RecordChainDuration(ctx, time.Since(chainStart).Seconds(), sch.Config().ClientName)
 	chainL.Info("Chain executed successfully")
 	sch.provider.RecordChainCompleted(ctx, sch.Config().ClientName)
-	sch.pgengine.RemoveChainRunStatus(bctx, chain.ChainID)
 	if chain.SelfDestruct {
+		bctx, cancel := log.DetachedContext(log.WithLogger(ctx, chainL))
+		defer cancel()
 		sch.pgengine.DeleteChain(bctx, chain.ChainID)
 	}
 }