diff --git a/spindle/db/mill_state.go b/spindle/db/mill_state.go index c11becda..683ee976 100644 --- a/spindle/db/mill_state.go +++ b/spindle/db/mill_state.go @@ -1,8 +1,10 @@ package db import ( + "context" "database/sql" "fmt" + "time" "tangled.org/core/eventstream" "tangled.org/core/notifier" @@ -293,7 +295,8 @@ func (tx *EventBatchTx) DeleteLease(leaseID string) error { func (tx *EventBatchTx) AdvanceCursor(nodeID, epoch string, seqno uint64) error { _, err := tx.tx.Exec( `insert into mill_executor_cursors (node_id, epoch, acked_seqno) values (?, ?, ?) - on conflict(node_id, epoch) do update set acked_seqno = excluded.acked_seqno`, + on conflict(node_id, epoch) do update set acked_seqno = excluded.acked_seqno + where excluded.acked_seqno > mill_executor_cursors.acked_seqno`, nodeID, epoch, seqno, ) return err @@ -432,8 +435,19 @@ func (d *DB) ListPendingArtifacts() ([]PendingArtifact, error) { // effects must be returned through its result: the result is published only // from the attempt that commits, and a failed attempt's result is discarded. func ApplyEventBatch[T any](d *DB, n *notifier.Notifier, fn func(tx *EventBatchTx) (T, error)) (T, error) { + return ApplyEventBatchCtx(context.Background(), d, n, fn) +} + +func ApplyEventBatchCtx[T any](ctx context.Context, d *DB, n *notifier.Notifier, fn func(tx *EventBatchTx) (T, error)) (T, error) { var result T - err := retrySQLite(func() error { + if err := ctx.Err(); err != nil { + var zero T + return zero, err + } + err := retrySQLiteCtx(ctx, 20*time.Second, func() error { + if err := ctx.Err(); err != nil { + return err + } return eventstream.WithClock(func(insert eventstream.Inserter) error { tx, err := d.Begin() if err != nil { @@ -470,7 +484,11 @@ func ApplyEventBatch[T any](d *DB, n *notifier.Notifier, fn func(tx *EventBatchT // ApplyEventBatch is the package-level ApplyEventBatch for callers without a // result. func (d *DB) ApplyEventBatch(n *notifier.Notifier, fn func(tx *EventBatchTx) error) error { - _, err := ApplyEventBatch(d, n, func(tx *EventBatchTx) (struct{}, error) { + return d.ApplyEventBatchCtx(context.Background(), n, fn) +} + +func (d *DB) ApplyEventBatchCtx(ctx context.Context, n *notifier.Notifier, fn func(tx *EventBatchTx) error) error { + _, err := ApplyEventBatchCtx(ctx, d, n, func(tx *EventBatchTx) (struct{}, error) { return struct{}{}, fn(tx) }) return err diff --git a/spindle/db/mill_state_test.go b/spindle/db/mill_state_test.go index 7ac1b7f6..ccea001b 100644 --- a/spindle/db/mill_state_test.go +++ b/spindle/db/mill_state_test.go @@ -3,6 +3,7 @@ package db import ( "context" "database/sql" + "errors" "fmt" "path/filepath" "sync" @@ -87,6 +88,13 @@ func TestExecutorCursors(t *testing.T) { t.Fatalf("unexpected cursors: %v", cursorMap) } + if err := d.ApplyEventBatch(nil, func(tx *EventBatchTx) error { return tx.AdvanceCursor("node-1", "inc-1", 4) }); err != nil { + t.Fatalf("AdvanceCursor(stale): %v", err) + } + cur, err := d.GetExecutorCursor("node-1", "inc-1") + if err != nil || cur != 9 { + t.Fatalf("cursor regressed to %d (err: %v), want 9", cur, err) + } } func TestApplyEventBatchPublishesResultOnlyOnCommit(t *testing.T) { @@ -138,6 +146,25 @@ func TestApplyEventBatchRetriesBusy(t *testing.T) { } } +func TestApplyEventBatchCtxAbortsOnCanceledContext(t *testing.T) { + d := newTestDB(t) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := ApplyEventBatchCtx(ctx, d, nil, func(tx *EventBatchTx) (string, error) { + return "committed", tx.AdvanceCursor("node-1", "inc-1", 10) + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("ApplyEventBatchCtx err = %v, want context.Canceled", err) + } + cur, err := d.GetExecutorCursor("node-1", "inc-1") + if err != nil { + t.Fatalf("GetExecutorCursor: %v", err) + } + if cur != 0 { + t.Fatalf("cursor advanced to %d on canceled context, want 0", cur) + } +} + func TestCompleteMillLeaseIsAtomic(t *testing.T) { d := newTestDB(t) lease := MillLease{ diff --git a/spindle/db/retry.go b/spindle/db/retry.go index 7146c470..deecf060 100644 --- a/spindle/db/retry.go +++ b/spindle/db/retry.go @@ -1,13 +1,14 @@ package db import ( + "context" "errors" "time" "github.com/mattn/go-sqlite3" ) -const sqliteRetryAttempts = 4 +const defaultSQLiteRetryBudget = 20 * time.Second func isSQLiteBusy(err error) bool { var sqliteErr sqlite3.Error @@ -18,13 +19,34 @@ func isSQLiteBusy(err error) bool { } func retrySQLite(fn func() error) error { - var err error - for attempt := 0; attempt < sqliteRetryAttempts; attempt++ { - err = fn() - if !isSQLiteBusy(err) || attempt == sqliteRetryAttempts-1 { + return retrySQLiteCtx(context.Background(), defaultSQLiteRetryBudget, fn) +} + +func retrySQLiteCtx(ctx context.Context, budget time.Duration, fn func() error) error { + if budget <= 0 { + budget = defaultSQLiteRetryBudget + } + deadline := time.Now().Add(budget) + backoff := 10 * time.Millisecond + const maxBackoff = 100 * time.Millisecond + + for { + if err := ctx.Err(); err != nil { + return err + } + err := fn() + if err == nil || !isSQLiteBusy(err) { + return err + } + if time.Now().After(deadline) { return err } - time.Sleep(time.Duration(10*(1<= dialBackoffMax { + backoff = dialBackoffMin + } e.l.Warn("mill session ended; reconnecting", "err", err, "backoff", backoff) select { case <-ctx.Done(): @@ -664,7 +668,7 @@ func (e *Executor) handleCommit(ctx context.Context, cl *millv1.CommitLease) { runSpan.SetAttributes(attrs...) } - jobCtx, cancel := context.WithCancel(runCtx) + jobCtx, cancel := context.WithCancelCause(runCtx) if cl.GetMillRecordsTerminalMetrics() { jobCtx = observability.WithWorkflowTerminalObserver(jobCtx, func(_, _, failureClass, reason string) { e.mu.Lock() @@ -745,7 +749,7 @@ func (e *Executor) handleCancel(leaseID string) { e.sendCancelAck(leaseID) if cancel != nil { - cancel() + cancel(engine.ErrWorkflowCanceled) } } diff --git a/spindle/mill/mill.go b/spindle/mill/mill.go index a7d7c1c5..3759ec0b 100644 --- a/spindle/mill/mill.go +++ b/spindle/mill/mill.go @@ -276,6 +276,21 @@ func (m *Mill) armSessionRecoveryLocked(sess *millSession, deadline time.Time) { sess.graceTimer = time.AfterFunc(delay, func() { m.failLeasesAfterGrace(sess) }) } +func (m *Mill) extendGraceLocked(sess *millSession) { + if !sess.recovering || sess.recoveryExpired { + return + } + deadline := time.Now().Add(m.cfg.ReconnectGrace) + if deadline.After(sess.graceDeadline) { + sess.graceDeadline = deadline + if sess.graceTimer != nil { + sess.graceTimer.Stop() + } + delay := time.Until(deadline) + sess.graceTimer = time.AfterFunc(delay, func() { m.failLeasesAfterGrace(sess) }) + } +} + func (m *Mill) maybeFinishSessionRecoveryLocked(sess *millSession) { if m.sessions[sess.nodeID] != sess || sess.disconnected || !sess.recovering || sess.recoveryExpired || !sess.recoveryReady || len(sess.pendingCancels) != 0 { @@ -1398,8 +1413,6 @@ func (m *Mill) onEventBatch(sess *millSession, batch *millv1.EventBatch) error { lease *RemoteLease ar *millv1.AttemptResult } - cacheDeletes := make(map[string]db.CacheEntry) - cacheObjectDeletes := make(map[string]struct{}) type startupObservation struct { engine string delay time.Duration @@ -1433,17 +1446,37 @@ func (m *Mill) onEventBatch(sess *millSession, batch *millv1.EventBatch) error { } } + // precompute storage sizes so object-storage stat does not run inside the write tx + precomputedSizes := make(map[string]int64) + for _, entry := range newEntries { + if u := entry.GetCacheUpdate(); u != nil && u.GetAction() == millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_STORED { + var storageKey string + if m.db != nil { + _ = m.db.QueryRow(`select storage_key from mill_cache_capabilities where lease_id = ? and action = 'save' and cache_id = ?`, entry.LeaseId, u.GetId()).Scan(&storageKey) + } + if storageKey != "" { + if sz, err := m.cacheObjectSize(sess.ctx, storageKey); err == nil { + precomputedSizes[storageKey] = sz + } + } + } + } + type batchApply struct { terminalCancelIDs []string pendingTerminals []pendingTerminal artifactLeases []*RemoteLease startupObservations []startupObservation highestSeqno uint64 + cacheDeletes map[string]db.CacheEntry + cacheObjectDeletes map[string]struct{} } applyFunc := func(tx *db.EventBatchTx) (batchApply, error) { var r batchApply r.highestSeqno = current + r.cacheDeletes = make(map[string]db.CacheEntry) + r.cacheObjectDeletes = make(map[string]struct{}) finishedInBatch := make(map[string]struct{}) for _, entry := range newEntries { @@ -1551,9 +1584,13 @@ func (m *Mill) onEventBatch(sess *millSession, batch *millv1.EventBatch) error { return batchApply{}, err } case millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_STORED: - actualSize, err := m.cacheObjectSize(context.Background(), ref) - if err != nil { - return batchApply{}, fmt.Errorf("stat stored cache object %q: %w", ref, err) + actualSize, ok := precomputedSizes[ref] + if !ok { + var err error + actualSize, err = m.cacheObjectSize(context.Background(), ref) + if err != nil { + return batchApply{}, fmt.Errorf("stat stored cache object %q: %w", ref, err) + } } superseded, ready, err := tx.MarkCacheEntryReadyWithChecksum( context.Background(), @@ -1567,13 +1604,13 @@ func (m *Mill) onEventBatch(sess *millSession, batch *millv1.EventBatch) error { return batchApply{}, err } for _, old := range superseded { - cacheDeletes[old.ID] = old + r.cacheDeletes[old.ID] = old } if !ready { if err := tx.QueueCacheObjectDeletion(context.Background(), ref, at); err != nil { return batchApply{}, err } - cacheObjectDeletes[ref] = struct{}{} + r.cacheObjectDeletes[ref] = struct{}{} } case millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_DISCARDED, millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_MISSING: // A missing restore object must not let one executor delete a @@ -1589,7 +1626,7 @@ func (m *Mill) onEventBatch(sess *millSession, batch *millv1.EventBatch) error { return batchApply{}, err } if discarded != nil { - cacheDeletes[discarded.ID] = *discarded + r.cacheDeletes[discarded.ID] = *discarded } } } @@ -1671,7 +1708,7 @@ func (m *Mill) onEventBatch(sess *millSession, batch *millv1.EventBatch) error { var applied batchApply var err error if m.db != nil { - applied, err = db.ApplyEventBatch(m.db, m.n, applyFunc) + applied, err = db.ApplyEventBatchCtx(sess.ctx, m.db, m.n, applyFunc) } else { applied, err = applyFunc(nil) } @@ -1685,7 +1722,7 @@ func (m *Mill) onEventBatch(sess *millSession, batch *millv1.EventBatch) error { m.metrics.RecordWorkflowStartupDelay(context.Background(), observation.engine, "mill", observation.delay) } - m.enqueueCacheDeletions(cacheDeletes, cacheObjectDeletes) + m.enqueueCacheDeletions(applied.cacheDeletes, applied.cacheObjectDeletes) if m.cfg.LogDir != "" { for _, lease := range applied.artifactLeases { @@ -1700,6 +1737,10 @@ func (m *Mill) onEventBatch(sess *millSession, batch *millv1.EventBatch) error { if applied.highestSeqno > m.nodeSeqno[currentKey] { m.nodeSeqno[currentKey] = applied.highestSeqno } + if sess.recovering && !sess.recoveryExpired { + m.extendGraceLocked(sess) + } + m.maybeFinishSessionRecoveryLocked(sess) m.mu.Unlock() for _, leaseID := range applied.terminalCancelIDs { diff --git a/spindle/mill/session.go b/spindle/mill/session.go index 2ce71327..e039ca8b 100644 --- a/spindle/mill/session.go +++ b/spindle/mill/session.go @@ -42,6 +42,9 @@ type millSession struct { mu sync.Mutex pending map[string]chan *millproto.Message // maps lease ID to response waiter + ctx context.Context + cancel context.CancelFunc + closeOnce sync.Once closed chan struct{} } @@ -51,6 +54,7 @@ type messageEncoder interface { } func newSession(nodeID string, epoch string, labels []string, enc messageEncoder, l *slog.Logger) *millSession { + ctx, cancel := context.WithCancel(context.Background()) return &millSession{ nodeID: nodeID, epoch: epoch, @@ -61,6 +65,8 @@ func newSession(nodeID string, epoch string, labels []string, enc messageEncoder pendingCancels: make(map[string]struct{}), closed: make(chan struct{}), lastSeen: time.Now(), + ctx: ctx, + cancel: cancel, } } @@ -75,6 +81,7 @@ func (s *millSession) send(msg *millproto.Message) error { func (s *millSession) close() { s.closeOnce.Do(func() { + s.cancel() close(s.closed) if s.closeTransport != nil { _ = s.closeTransport()