fix: checkpoint sync progress on SIGTERM before exit

pkg/connectorrunner/runner.go

@@ -9,6 +9,7 @@ import (
 	"path/filepath"
 	"strings"
 	"sync"
+	"syscall"
 	"time"
 
 	"github.com/conductorone/baton-sdk/pkg/bid"
@@ -48,6 +49,42 @@ type connectorRunner struct {
 
 var ErrSigTerm = errors.New("context cancelled by process shutdown")
 
+// shutdownDrainTimeout bounds how long run() waits, on cancellation, for
+// in-flight task goroutines to finish before returning. Each task goroutine
+// is untracked (dispatched via a bare `go func`) and Close() — called
+// immediately after Run() returns by every caller of this runner — tears
+// down the connector client via c.cw.Close() with no synchronization against
+// those goroutines. Without a drain, a SIGTERM can still race the very
+// checkpoint this package exists to protect: Run() returns as soon as ctx is
+// Done, Close() runs, and the in-flight sync's forced checkpoint write (see
+// sync/parallel_syncer.go's handleOperationError, bounded to 15s) can be torn
+// out from under it.
+//
+// This budget intentionally only targets that checkpoint write, not a full
+// c1z finalize: syncer.Close() (called by the task handler on every path,
+// including sync failure) bounds its own detached finalize by
+// dotc1z.FinalizeTimeout, which defaults to 1 hour — no realistic
+// termination grace period can wait that out, and this drain does not try
+// to. It only aims to let the smaller, higher-priority checkpoint write
+// complete before the connector client is torn down.
+const shutdownDrainTimeout = 25 * time.Second
+
+// drainInFlightTasks waits for every currently-held semaphore slot to be
+// released — i.e. every dispatched task goroutine (see the `go func(t
+// *v1.Task)` dispatch below) to return — bounded by shutdownDrainTimeout.
+// Called right before run() returns on cancellation, so the caller's
+// deferred Close() (which tears down the connector client) doesn't race an
+// in-flight sync that's still writing its forced checkpoint.
+func (c *connectorRunner) drainInFlightTasks(ctx context.Context, sem *semaphore.Weighted, l *zap.Logger) {
+	drainCtx, cancelDrain := context.WithTimeout(context.WithoutCancel(ctx), shutdownDrainTimeout)
+	defer cancelDrain()
+	if err := sem.Acquire(drainCtx, int64(c.taskConcurrency)); err != nil {
+		l.Warn("runner: shutdown drain timed out; in-flight tasks may not have finished checkpointing", zap.Error(err))
+		return
+	}
+	sem.Release(int64(c.taskConcurrency))
+}
+
 // setupPersistentLog ensures that a log file on disk is created,
 // when required by either the stored Manager or by a Task.
 // A log file created by a stored Manager persists for our entire run,
@@ -138,7 +175,7 @@ func (c *connectorRunner) Run(ctx context.Context) error {
 	}
 
 	sigChan := make(chan os.Signal, 1)
-	signal.Notify(sigChan, os.Interrupt)
+	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
 	go func() {
 		for range sigChan {
 			cancel(ErrSigTerm)
@@ -218,13 +255,15 @@ func (c *connectorRunner) run(ctx context.Context) error {
 	for !stopForLoop {
 		select {
 		case <-ctx.Done():
+			c.drainInFlightTasks(ctx, sem, l)
 			return c.handleContextCancel(ctx)
 		case <-time.After(nextCheckAfter):
 			l.Debug("runner: claiming worker")
 			// Acquire a worker slot before we call Next() so we don't claim a task before we can actually process it.
 			err = sem.Acquire(ctx, 1)
 			if err != nil {
 				l.Error("runner: error acquiring semaphore to claim worker", zap.Error(err))
+				c.drainInFlightTasks(ctx, sem, l)
 				return c.handleContextCancel(ctx)
 			}
 			l.Debug("runner: worker claimed, checking for next task")

pkg/connectorrunner/runner_test.go

@@ -101,6 +101,58 @@ func TestRun_TaskConcurrencyCapsProcessing(t *testing.T) {
 	}
 }
 
+// TestRun_DrainsInFlightTaskBeforeReturningOnCancellation covers the gap
+// between catching SIGTERM and actually protecting an in-flight sync's
+// forced checkpoint: every task is dispatched to an untracked goroutine
+// (`go func(t *v1.Task)` in run()), and every caller of this runner closes
+// the connector client (Close -> c.cw.Close()) immediately after Run()
+// returns. Without draining, run() returns as soon as ctx is Done
+// regardless of what that goroutine is still doing -- including a
+// checkpoint write deliberately running on a context.WithoutCancel(ctx)
+// scope specifically so cancellation doesn't cut it off.
+//
+// detachedTaskManager's Process ignores ctx.Done() (mirroring that
+// detached-context write) and only returns when explicitly released, so
+// this test can assert run() actually waits for it instead of just
+// hoping the timing works out.
+func TestRun_DrainsInFlightTaskBeforeReturningOnCancellation(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithCancelCause(context.Background())
+	defer cancel(ErrSigTerm)
+
+	tm := newDetachedTaskManager()
+	runner := &connectorRunner{
+		cw:              noopClientWrapper{},
+		tasks:           tm,
+		taskConcurrency: 1,
+	}
+
+	errCh := make(chan error, 1)
+	go func() {
+		errCh <- runner.run(ctx)
+	}()
+
+	waitForStartedTasks(t, tm.started, 1)
+
+	cancel(ErrSigTerm)
+
+	select {
+	case <-errCh:
+		require.Fail(t, "run() returned before the in-flight task finished; the shutdown drain did not wait for it")
+	case <-time.After(100 * time.Millisecond):
+	}
+
+	close(tm.release)
+
+	select {
+	case err := <-errCh:
+		require.NoError(t, err)
+	case <-time.After(time.Second):
+		require.Fail(t, "timed out waiting for runner to stop after releasing the in-flight task")
+	}
+}
+
 func waitForStartedTasks(t *testing.T, started <-chan struct{}, count int) {
 	t.Helper()
 	for range count {
@@ -233,6 +285,51 @@ func (m *blockingTaskManager) GetTempDir() string {
 	return ""
 }
 
+// detachedTaskManager's Process deliberately does not select on ctx.Done()
+// -- it mirrors a sync's forced checkpoint write, which runs on a
+// context.WithoutCancel(ctx) scope specifically so the caller's cancellation
+// doesn't cut it off. It only returns once explicitly released.
+type detachedTaskManager struct {
+	nextCalls atomic.Int64
+	started   chan struct{}
+	release   chan struct{}
+}
+
+func newDetachedTaskManager() *detachedTaskManager {
+	return &detachedTaskManager{
+		started: make(chan struct{}, 1),
+		release: make(chan struct{}),
+	}
+}
+
+func (m *detachedTaskManager) Next(ctx context.Context) (*v1.Task, time.Duration, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, 0, err
+	}
+	if m.nextCalls.Add(1) > 1 {
+		return nil, time.Hour, nil
+	}
+	return v1.Task_builder{
+		Id:     "task-1",
+		Status: v1.Task_STATUS_PENDING,
+		Hello:  &v1.Task_HelloTask{},
+	}.Build(), 0, nil
+}
+
+func (m *detachedTaskManager) Process(ctx context.Context, task *v1.Task, cc types.ConnectorClient) error {
+	m.started <- struct{}{}
+	<-m.release
+	return nil
+}
+
+func (m *detachedTaskManager) ShouldDebug() bool {
+	return false
+}
+
+func (m *detachedTaskManager) GetTempDir() string {
+	return ""
+}
+
 type noopClientWrapper struct{}
 
 func (noopClientWrapper) C(ctx context.Context) (types.ConnectorClient, error) {

pkg/sync/parallel_syncer.go

@@ -161,9 +161,22 @@ func (s *syncer) parallelSync(
 			break
 		}
 
+		// Checked before the throttled checkpoint below: that checkpoint runs
+		// on the raw ctx, which may already be Done (a SIGTERM-driven
+		// cancellation cancels ctx itself, not just runCtx's derived
+		// deadline) -- checking cancellation first means a cancelled ctx
+		// never gets a chance to fail that write and return a raw,
+		// non-resumable error before this iteration reaches
+		// handleOperationError.
+		select {
+		case <-runCtx.Done():
+			return s.handleOperationError(ctx, runCtx, warnings, nil)
+		default:
+		}
+
 		err := s.Checkpoint(ctx, false)
 		if err != nil {
-			return warnings, err
+			return s.handleOperationError(ctx, runCtx, warnings, err)
 		}
 
 		// If we have more than 10 warnings and more than 10% of actions ended in a warning, exit the sync.
@@ -173,28 +186,6 @@ func (s *syncer) parallelSync(
 				return warnings, fmt.Errorf("%w: warnings: %v completed actions: %d", ErrTooManyWarnings, warnings, completedActionsCount)
 			}
 		}
-		select {
-		case <-runCtx.Done():
-			err = context.Cause(runCtx)
-			switch {
-			case errors.Is(err, context.DeadlineExceeded):
-				if s.recordStats {
-					l.Info("sync run duration has expired, exiting sync early", s.syncSummaryFields(trace.SpanFromContext(ctx))...)
-				} else {
-					l.Info("sync run duration has expired, exiting sync early", zap.String("sync_id", s.syncID))
-				}
-				// It would be nice to remove this once we're more confident in the checkpointing logic.
-				checkpointErr := s.Checkpoint(ctx, true)
-				if checkpointErr != nil {
-					l.Error("error checkpointing before exiting sync", zap.Error(checkpointErr))
-				}
-				return warnings, errors.Join(checkpointErr, ErrSyncNotComplete)
-			default:
-				l.Error("sync context cancelled", zap.String("sync_id", s.syncID), zap.Error(err))
-				return warnings, err
-			}
-		default:
-		}
 
 		switch stateAction.Op {
 		case InitOp:
@@ -220,7 +211,7 @@ func (s *syncer) parallelSync(
 				s.state.PushAction(ctx, Action{Op: SyncResourceTypesOp})
 				err = s.Checkpoint(ctx, true)
 				if err != nil {
-					return warnings, err
+					return s.handleOperationError(ctx, runCtx, warnings, err)
 				}
 				// Don't do grant expansion or external resources in partial syncs, as we likely lack related resources/entitlements/grants
 				continue
@@ -238,7 +229,7 @@ func (s *syncer) parallelSync(
 				s.state.SetNeedsExpansion()
 				err = s.Checkpoint(ctx, true)
 				if err != nil {
-					return warnings, err
+					return s.handleOperationError(ctx, runCtx, warnings, err)
 				}
 				continue
 			}
@@ -256,7 +247,7 @@ func (s *syncer) parallelSync(
 
 			err = s.Checkpoint(ctx, true)
 			if err != nil {
-				return warnings, err
+				return s.handleOperationError(ctx, runCtx, warnings, err)
 			}
 			continue
 
@@ -400,7 +391,7 @@ func (s *syncer) parallelSync(
 				}
 				if err := s.store.SyncMeta().MarkSyncSupportsDiff(ctx, s.syncID); err != nil {
 					l.Error("failed to set supports_diff marker", zap.Error(err))
-					return warnings, err
+					return s.handleOperationError(ctx, runCtx, warnings, err)
 				}
 			}
 
@@ -422,29 +413,90 @@ func (s *syncer) parallelSync(
 	return warnings, nil
 }
 
+// forcedCheckpointTimeout bounds the detached (context.WithoutCancel) context
+// used for the forced checkpoint write below. This runs precisely during an
+// orchestrator's termination grace period, so it must stay well short of a
+// typical grace period (Kubernetes defaults to 30s) — unlike
+// dotc1z.FinalizeTimeout, which bounds the much heavier WAL
+// checkpoint+save+upload tail, this is a single small state write.
+const forcedCheckpointTimeout = 15 * time.Second
+
+// handleOperationError checks whether runCtx has been cancelled — by a
+// deadline timeout, a SIGTERM/SIGINT shutdown (connectorrunner.ErrSigTerm),
+// a server-directed task cancellation, a heartbeat failure, or any other
+// cause — and if so treats it uniformly as resumable: force a checkpoint so
+// the next resume picks up from here instead of restarting, and return an
+// error that both IsSyncPreservable recognizes (via ErrSyncNotComplete) and
+// still lets a caller distinguish *why* via errors.Is(err, cause).
+//
+// batchErr is the error an in-flight operation returned (nil when called
+// from the top-of-loop check, where cancellation was observed between
+// actions rather than surfaced by a connector call). If runCtx was not
+// actually cancelled, batchErr is a genuine, unrelated operation error and
+// is returned unchanged — this function only special-cases cancellation.
 func (s *syncer) handleOperationError(
 	ctx context.Context,
 	runCtx context.Context,
 	warnings []error,
 	batchErr error,
 ) ([]error, error) {
-	if !errors.Is(context.Cause(runCtx), context.DeadlineExceeded) {
+	cause := context.Cause(runCtx)
+	if cause == nil {
 		return warnings, batchErr
 	}
-	// The run duration expired. The operation error is usually just the
-	// resulting cancellation, but a genuine connector/store failure can
-	// land in the same window — log it so expiry never masks a real bug.
-	// The sync resumes from the checkpoint, so the work is retried either
-	// way; only ErrSyncNotComplete is surfaced to keep the resumable
-	// contract for callers.
+	// The operation error is usually just the resulting cancellation, but a
+	// genuine connector/store failure can land in the same window — log it
+	// so cancellation never masks a real bug. The sync resumes from the
+	// checkpoint, so the work is retried either way; only ErrSyncNotComplete
+	// (plus cause) is surfaced to keep the resumable contract for callers.
 	if batchErr != nil && !errors.Is(batchErr, context.Canceled) && !errors.Is(batchErr, context.DeadlineExceeded) {
 		ctxzap.Extract(ctx).Error(
-			"sync operation failed while run duration expired; exiting early for resume",
+			"sync operation failed while the sync was being cancelled; exiting early for resume",
 			zap.Error(batchErr),
+			zap.NamedError("cancel_cause", cause),
 		)
 	}
-	checkpointErr := s.Checkpoint(ctx, true)
-	return warnings, errors.Join(checkpointErr, ErrSyncNotComplete)
+	l := ctxzap.Extract(ctx)
+	switch {
+	case errors.Is(cause, context.DeadlineExceeded):
+		if s.recordStats {
+			l.Info("sync run duration has expired, exiting sync early", s.syncSummaryFields(trace.SpanFromContext(ctx))...)
+		} else {
+			l.Info("sync run duration has expired, exiting sync early", zap.String("sync_id", s.syncID))
+		}
+	case s.recordStats:
+		// Same per-step/per-resource-type/retry-wait fields as the deadline
+		// branch above — these are exactly the counters used to reconstruct
+		// a sync's timeline after the fact (see CXH-2121), so a
+		// SIGTERM-interrupted sync must not produce less diagnostic output
+		// than a deadline-expired one.
+		l.Info("sync context cancelled, exiting sync early",
+			append(s.syncSummaryFields(trace.SpanFromContext(ctx)), zap.NamedError("cancel_cause", cause))...)
+	default:
+		l.Info("sync context cancelled, exiting sync early", zap.String("sync_id", s.syncID), zap.Error(cause))
+	}
+	// ctx itself may already be Done here (e.g. a SIGTERM-driven cancellation
+	// cancels ctx directly, not just runCtx's derived deadline), so this
+	// checkpoint must run on a context that survives that cancellation —
+	// same pattern as the other finalize/cleanup writes in this codebase
+	// (c1file.go, clone_sync.go, syncer.go) — but still bounded, so a wedged
+	// store write during the grace period doesn't just burn the whole window
+	// before getting SIGKILLed anyway.
+	// It would be nice to remove this once we're more confident in the checkpointing logic.
+	cpCtx, cpCancel := context.WithTimeout(context.WithoutCancel(ctx), forcedCheckpointTimeout)
+	defer cpCancel()
+	checkpointErr := s.Checkpoint(cpCtx, true)
+	if checkpointErr != nil {
+		l.Error("error checkpointing before exiting sync", zap.Error(checkpointErr))
+	}
+	// cause and batchErr are joined in (not just logged) so a caller with
+	// more context — e.g. pkg/tasks/c1api distinguishing ErrTaskCancelled
+	// from a plain shutdown, or a retry loop that wants to know a genuine
+	// store/connector failure happened to land in the same window — can
+	// still inspect errors.Is(err, ...) instead of it being lost behind
+	// ErrSyncNotComplete. errors.Join drops nils, so this is a no-op when
+	// called from the top-of-loop check (batchErr is nil there).
+	return warnings, errors.Join(checkpointErr, ErrSyncNotComplete, cause, batchErr)
 }
 
 func tooManyWarnings(warningCount int, completedActionsCount uint64) bool {