diff --git a/internal/http/webhooks_admin_calls_test.go b/internal/http/webhooks_admin_calls_test.go index dc5b0982..37c21213 100644 --- a/internal/http/webhooks_admin_calls_test.go +++ b/internal/http/webhooks_admin_calls_test.go @@ -90,6 +90,7 @@ func (s *adminCallStore) DeleteOlderThan(context.Context, uuid.UUID, time.Time) return 0, nil } func (s *adminCallStore) ReclaimStale(context.Context, time.Time) (int64, error) { return 0, nil } +func (s *adminCallStore) Heartbeat(context.Context, uuid.UUID, string, time.Time) error { return nil } // ---- stub testers ---- diff --git a/internal/http/webhooks_auth_test.go b/internal/http/webhooks_auth_test.go index 356c3b03..4b58bef3 100644 --- a/internal/http/webhooks_auth_test.go +++ b/internal/http/webhooks_auth_test.go @@ -138,6 +138,9 @@ func (s *stubWebhookCallStore) DeleteOlderThan(_ context.Context, _ uuid.UUID, _ func (s *stubWebhookCallStore) ReclaimStale(_ context.Context, _ time.Time) (int64, error) { return 0, nil } +func (s *stubWebhookCallStore) Heartbeat(_ context.Context, _ uuid.UUID, _ string, _ time.Time) error { + return nil +} // ---- helpers ---- diff --git a/internal/http/webhooks_llm_test.go b/internal/http/webhooks_llm_test.go index 4513623b..2a6700df 100644 --- a/internal/http/webhooks_llm_test.go +++ b/internal/http/webhooks_llm_test.go @@ -82,6 +82,9 @@ func (s *llmCallStore) DeleteOlderThan(_ context.Context, _ uuid.UUID, _ time.Ti func (s *llmCallStore) ReclaimStale(_ context.Context, _ time.Time) (int64, error) { return 0, nil } +func (s *llmCallStore) Heartbeat(_ context.Context, _ uuid.UUID, _ string, _ time.Time) error { + return nil +} // ---- helpers ---- diff --git a/internal/http/webhooks_message_test.go b/internal/http/webhooks_message_test.go index 059362dd..0903e952 100644 --- a/internal/http/webhooks_message_test.go +++ b/internal/http/webhooks_message_test.go @@ -123,6 +123,9 @@ func (s *msgCallStore) DeleteOlderThan(_ context.Context, _ uuid.UUID, _ time.Ti func (s *msgCallStore) ReclaimStale(_ context.Context, _ time.Time) (int64, error) { return 0, nil } +func (s *msgCallStore) Heartbeat(_ context.Context, _ uuid.UUID, _ string, _ time.Time) error { + return nil +} // ---- stub: store.WebhookStore (message handler tests — minimal no-op) ---- diff --git a/internal/store/pg/webhook_calls.go b/internal/store/pg/webhook_calls.go index 65ddcc9c..2c94f3ff 100644 --- a/internal/store/pg/webhook_calls.go +++ b/internal/store/pg/webhook_calls.go @@ -160,7 +160,7 @@ func (s *PGWebhookCallStore) ClaimNext(ctx context.Context, tenantID uuid.UUID, lease := uuid.New().String() row := tx.QueryRowContext(ctx, `UPDATE webhook_calls - SET status = 'running', started_at = $1, lease_token = $2 + SET status = 'running', started_at = $1, lease_token = $2, last_heartbeat_at = $1 WHERE id = $3 RETURNING `+webhookCallColumns, now, lease, callID, @@ -274,15 +274,17 @@ func (s *PGWebhookCallStore) DeleteOlderThan(ctx context.Context, tenantID uuid. } // ReclaimStale resets stale running rows back to queued so the worker can retry them. -// A row is considered stale when started_at < staleThreshold (i.e., the worker that -// claimed it crashed before completing UpdateStatus). +// A row is considered stale when last_heartbeat_at < staleThreshold (the worker stopped +// renewing its lease, indicating it crashed or hung). Rows with NULL last_heartbeat_at +// (legacy/never-heartbeated running rows) are also reclaimed defensively. // Cross-tenant: no tenant_id filter — the retention worker sweeps the whole table. func (s *PGWebhookCallStore) ReclaimStale(ctx context.Context, staleThreshold time.Time) (int64, error) { // Clear lease_token so any in-flight UpdateStatusCAS from the crashed worker returns ErrLeaseExpired. res, err := s.db.ExecContext(ctx, `UPDATE webhook_calls - SET status = 'queued', started_at = NULL, lease_token = NULL - WHERE mode = 'async' AND status = 'running' AND started_at < $1`, + SET status = 'queued', started_at = NULL, lease_token = NULL, last_heartbeat_at = NULL + WHERE mode = 'async' AND status = 'running' + AND (last_heartbeat_at IS NULL OR last_heartbeat_at < $1)`, staleThreshold, ) if err != nil { @@ -291,6 +293,31 @@ func (s *PGWebhookCallStore) ReclaimStale(ctx context.Context, staleThreshold ti return res.RowsAffected() } +// Heartbeat renews the lease for a running call (CAS on lease_token + tenant). +// Returns store.ErrLeaseExpired if 0 rows matched (row reclaimed → caller must stop the run). +// Hand-rolled (not execMapUpdateWhereTenantLease) because we also need AND status='running' +// to reject renewing terminal-state rows — the shared helper has no status guard. +func (s *PGWebhookCallStore) Heartbeat(ctx context.Context, callID uuid.UUID, lease string, now time.Time) error { + tid, err := requireTenantID(ctx) + if err != nil { + return err + } + res, err := s.db.ExecContext(ctx, + `UPDATE webhook_calls + SET last_heartbeat_at = $1 + WHERE id = $2 AND tenant_id = $3 AND lease_token = $4 AND status = 'running'`, + now, callID, tid, lease, + ) + if err != nil { + return err + } + affected, _ := res.RowsAffected() + if affected == 0 { + return store.ErrLeaseExpired + } + return nil +} + // execMapUpdateWhereTenantLease is like execMapUpdateWhereTenantNoUpdatedAt but adds // AND lease_token = $N to the WHERE clause for optimistic concurrency. // Returns store.ErrLeaseExpired when RowsAffected() == 0 (lease mismatch). diff --git a/internal/store/sqlitestore/schema.go b/internal/store/sqlitestore/schema.go index 54d1b906..14487278 100644 --- a/internal/store/sqlitestore/schema.go +++ b/internal/store/sqlitestore/schema.go @@ -16,7 +16,7 @@ var schemaSQL string // SchemaVersion is the current SQLite schema version. // Bump this when adding new migration steps below. -const SchemaVersion = 51 +const SchemaVersion = 52 // migrations maps version → SQL to apply when upgrading FROM that version. // schema.sql always represents the LATEST full schema (for fresh DBs). @@ -881,6 +881,9 @@ CREATE INDEX IF NOT EXISTS idx_skill_user_grants_tenant ON skill_user_grants(ten CREATE UNIQUE INDEX IF NOT EXISTS mcp_oauth_tokens_global_uq ON mcp_oauth_tokens (server_id, tenant_id) WHERE user_id IS NULL; CREATE UNIQUE INDEX IF NOT EXISTS mcp_oauth_tokens_user_uq ON mcp_oauth_tokens (server_id, tenant_id, user_id) WHERE user_id IS NOT NULL; CREATE INDEX IF NOT EXISTS idx_mcp_oauth_tokens_server_tenant ON mcp_oauth_tokens (server_id, tenant_id);`, + // Version 51 → 52: add last_heartbeat_at to webhook_calls for lease heartbeat. + // Mirrors PG migration 000085. Idempotent-guarded via idempotentColumnMigration(51). + 51: `ALTER TABLE webhook_calls ADD COLUMN last_heartbeat_at TEXT;`, } const addUsageEventAnalyticsTables = ` @@ -1479,6 +1482,8 @@ func idempotentColumnMigration(version int) (string, string, bool) { return "secure_cli_user_credentials", "host_scope", true case 41: return "secure_cli_binaries", "adapter_name", true + case 51: + return "webhook_calls", "last_heartbeat_at", true default: return "", "", false } diff --git a/internal/store/sqlitestore/schema.sql b/internal/store/sqlitestore/schema.sql index 93c1d030..c59b6701 100644 --- a/internal/store/sqlitestore/schema.sql +++ b/internal/store/sqlitestore/schema.sql @@ -2072,6 +2072,7 @@ CREATE TABLE IF NOT EXISTS webhook_calls ( delivery_id TEXT NOT NULL, next_attempt_at TEXT, started_at TEXT, + last_heartbeat_at TEXT, lease_token TEXT, request_payload TEXT, response TEXT, diff --git a/internal/store/sqlitestore/webhook_calls.go b/internal/store/sqlitestore/webhook_calls.go index fe566286..59e1a7eb 100644 --- a/internal/store/sqlitestore/webhook_calls.go +++ b/internal/store/sqlitestore/webhook_calls.go @@ -171,8 +171,8 @@ func (s *SQLiteWebhookCallStore) ClaimNext(ctx context.Context, tenantID uuid.UU // Attempts untouched — worker increments post-send. lease := uuid.New().String() _, err = tx.ExecContext(ctx, - `UPDATE webhook_calls SET status = 'running', started_at = ?, lease_token = ? WHERE id = ?`, - now, lease, callID, + `UPDATE webhook_calls SET status = 'running', started_at = ?, lease_token = ?, last_heartbeat_at = ? WHERE id = ?`, + now, lease, now, callID, ) if err != nil { return nil, fmt.Errorf("webhook_calls ClaimNext update: %w", err) @@ -287,13 +287,16 @@ func (s *SQLiteWebhookCallStore) DeleteOlderThan(ctx context.Context, tenantID u } // ReclaimStale resets stale running rows back to queued so the worker can retry them. +// A row is stale when last_heartbeat_at < staleThreshold (worker stopped renewing its lease); +// NULL last_heartbeat_at (legacy/never-heartbeated running rows) is also reclaimed defensively. // Clears lease_token so any in-flight UpdateStatusCAS from the crashed goroutine returns ErrLeaseExpired. // SQLite stores timestamps as ISO-8601 strings; comparison uses standard string ordering. func (s *SQLiteWebhookCallStore) ReclaimStale(ctx context.Context, staleThreshold time.Time) (int64, error) { res, err := s.db.ExecContext(ctx, `UPDATE webhook_calls - SET status = 'queued', started_at = NULL, lease_token = NULL - WHERE mode = 'async' AND status = 'running' AND started_at < ?`, + SET status = 'queued', started_at = NULL, lease_token = NULL, last_heartbeat_at = NULL + WHERE mode = 'async' AND status = 'running' + AND (last_heartbeat_at IS NULL OR last_heartbeat_at < ?)`, staleThreshold, ) if err != nil { @@ -302,6 +305,31 @@ func (s *SQLiteWebhookCallStore) ReclaimStale(ctx context.Context, staleThreshol return res.RowsAffected() } +// Heartbeat renews the lease for a running call (CAS on lease_token + tenant). +// Returns store.ErrLeaseExpired if 0 rows matched (row reclaimed → caller must stop the run). +// Hand-rolled (not execMapUpdateWhereTenantLeaseNoUpdatedAt) because we also need +// AND status='running' to reject renewing terminal-state rows — the shared helper has no status guard. +func (s *SQLiteWebhookCallStore) Heartbeat(ctx context.Context, callID uuid.UUID, lease string, now time.Time) error { + tid, err := requireTenantID(ctx) + if err != nil { + return err + } + res, err := s.db.ExecContext(ctx, + `UPDATE webhook_calls + SET last_heartbeat_at = ? + WHERE id = ? AND tenant_id = ? AND lease_token = ? AND status = 'running'`, + now, callID, tid, lease, + ) + if err != nil { + return err + } + affected, _ := res.RowsAffected() + if affected == 0 { + return store.ErrLeaseExpired + } + return nil +} + // execMapUpdateWhereTenantLeaseNoUpdatedAt is like execMapUpdateWhereTenantNoUpdatedAt but adds // AND lease_token = ? to the WHERE clause for optimistic concurrency. // Returns store.ErrLeaseExpired when RowsAffected() == 0 (lease mismatch). diff --git a/internal/store/sqlitestore/webhooks_test.go b/internal/store/sqlitestore/webhooks_test.go index bff57aa8..ccfade09 100644 --- a/internal/store/sqlitestore/webhooks_test.go +++ b/internal/store/sqlitestore/webhooks_test.go @@ -5,6 +5,7 @@ package sqlitestore import ( "context" "database/sql" + "errors" "testing" "time" @@ -314,3 +315,102 @@ func TestWebhookCallIdempotencyConflict(t *testing.T) { t.Errorf("expected ErrIdempotencyConflict, got: %v", err) } } + +// TestWebhookCallHeartbeat verifies Heartbeat renews last_heartbeat_at when lease matches, +// and returns ErrLeaseExpired when the lease is wrong. +func TestWebhookCallHeartbeat(t *testing.T) { + db := openTestWebhookDB(t) + ws := NewSQLiteWebhookStore(db) + cs := NewSQLiteWebhookCallStore(db) + + tenantID := uuid.New() + ctx := testTenantCtx(tenantID) + wh := &store.WebhookData{ + ID: uuid.New(), TenantID: tenantID, Name: "wh-hb", Kind: "llm", + SecretHash: "h-hb", Scopes: []string{}, IPAllowlist: []string{}, + RateLimitPerMin: 60, CreatedAt: time.Now().UTC(), UpdatedAt: time.Now().UTC(), + } + if err := ws.Create(ctx, wh); err != nil { + t.Fatalf("Create webhook: %v", err) + } + + now := time.Now().UTC() + queued := &store.WebhookCallData{ + ID: uuid.New(), TenantID: tenantID, WebhookID: wh.ID, + DeliveryID: uuid.New(), Mode: "async", Status: "queued", CreatedAt: now, + } + if err := cs.Create(ctx, queued); err != nil { + t.Fatalf("Create queued: %v", err) + } + + claimed, err := cs.ClaimNext(ctx, tenantID, now) + if err != nil { + t.Fatalf("ClaimNext: %v", err) + } + lease := "" + if claimed.LeaseToken != nil { + lease = *claimed.LeaseToken + } + if lease == "" { + t.Fatal("ClaimNext must set lease_token") + } + + // Correct lease → renews, no error. + if err := cs.Heartbeat(ctx, claimed.ID, lease, time.Now().UTC()); err != nil { + t.Fatalf("Heartbeat with valid lease: %v", err) + } + + // Wrong lease → ErrLeaseExpired. + err = cs.Heartbeat(ctx, claimed.ID, "wrong-lease", time.Now().UTC()) + if !errors.Is(err, store.ErrLeaseExpired) { + t.Errorf("Heartbeat wrong lease: got %v, want ErrLeaseExpired", err) + } +} + +// TestWebhookCallReclaimStaleRespectsHeartbeat verifies a row with a fresh heartbeat +// is NOT reclaimed even if started_at is old. +func TestWebhookCallReclaimStaleRespectsHeartbeat(t *testing.T) { + db := openTestWebhookDB(t) + ws := NewSQLiteWebhookStore(db) + cs := NewSQLiteWebhookCallStore(db) + + tenantID := uuid.New() + ctx := testTenantCtx(tenantID) + wh := &store.WebhookData{ + ID: uuid.New(), TenantID: tenantID, Name: "wh-hb2", Kind: "llm", + SecretHash: "h-hb2", Scopes: []string{}, IPAllowlist: []string{}, + RateLimitPerMin: 60, CreatedAt: time.Now().UTC(), UpdatedAt: time.Now().UTC(), + } + if err := ws.Create(ctx, wh); err != nil { + t.Fatalf("Create webhook: %v", err) + } + + oldStart := time.Now().UTC().Add(-time.Hour) + freshHB := time.Now().UTC() // recent heartbeat + id := uuid.New() + _, err := db.ExecContext(ctx, + `INSERT INTO webhook_calls (id,tenant_id,webhook_id,delivery_id,mode,status,attempts,created_at,started_at,last_heartbeat_at,lease_token) + VALUES (?,?,?,?,?,?,?,?,?,?,?)`, + id, tenantID, wh.ID, uuid.New(), "async", "running", 0, oldStart, oldStart, freshHB, "lease-x", + ) + if err != nil { + t.Fatalf("insert running row: %v", err) + } + + // Threshold = now - 90s. Fresh heartbeat (now) is NOT older than threshold → not reclaimed. + n, err := cs.ReclaimStale(ctx, time.Now().UTC().Add(-90*time.Second)) + if err != nil { + t.Fatalf("ReclaimStale: %v", err) + } + if n != 0 { + t.Fatalf("reclaimed %d rows, want 0 (fresh heartbeat must protect row)", n) + } + + var status string + if err := db.QueryRowContext(ctx, `SELECT status FROM webhook_calls WHERE id = ?`, id).Scan(&status); err != nil { + t.Fatalf("select row: %v", err) + } + if status != "running" { + t.Fatalf("status = %q, want running", status) + } +} diff --git a/internal/store/webhook_store.go b/internal/store/webhook_store.go index 691dc633..3e0cb17e 100644 --- a/internal/store/webhook_store.go +++ b/internal/store/webhook_store.go @@ -163,6 +163,11 @@ type WebhookCallStore interface { // Returns sql.ErrNoRows if the queue is empty. ClaimNext(ctx context.Context, tenantID uuid.UUID, now time.Time) (*WebhookCallData, error) + // Heartbeat renews the lease for a running call (CAS on lease_token). + // Sets last_heartbeat_at = now WHERE id=callID AND lease_token=lease AND status='running'. + // Returns store.ErrLeaseExpired if 0 rows match (lease was reclaimed) — the caller MUST stop the current run. + Heartbeat(ctx context.Context, callID uuid.UUID, lease string, now time.Time) error + // List returns calls for the context tenant with optional filters. List(ctx context.Context, f WebhookCallListFilter) ([]WebhookCallData, error) @@ -173,9 +178,9 @@ type WebhookCallStore interface { // If tenantID is uuid.Nil, deletes across all tenants (retention worker). DeleteOlderThan(ctx context.Context, tenantID uuid.UUID, ts time.Time) (int64, error) - // ReclaimStale resets rows stuck in status='running' with started_at older than - // staleThreshold back to status='queued'. Called on worker startup and periodically - // (every 60s) to recover from crashes between ClaimNext and UpdateStatus. - // Returns the number of rows reclaimed. + // ReclaimStale resets rows stuck in status='running' whose last_heartbeat_at is older + // than staleThreshold (or NULL — never-heartbeated/legacy rows) back to status='queued'. + // Called on worker startup and periodically (every 60s) to recover from a worker that + // crashed or hung and stopped renewing its lease. Returns the number of rows reclaimed. ReclaimStale(ctx context.Context, staleThreshold time.Time) (int64, error) } diff --git a/internal/upgrade/version.go b/internal/upgrade/version.go index 08caf60b..6b3efe02 100644 --- a/internal/upgrade/version.go +++ b/internal/upgrade/version.go @@ -2,4 +2,4 @@ package upgrade // RequiredSchemaVersion is the schema migration version this binary requires. // Bump this whenever adding a new SQL migration file. -const RequiredSchemaVersion uint = 84 +const RequiredSchemaVersion uint = 85 diff --git a/internal/webhooks/worker.go b/internal/webhooks/worker.go index 9bb2cd6d..9a32adbc 100644 --- a/internal/webhooks/worker.go +++ b/internal/webhooks/worker.go @@ -36,6 +36,10 @@ const ( // staleRunningWindow is how long a running row must be inactive before being reclaimed. staleRunningWindow = 90 * time.Second + // heartbeatInterval is how often a running agent renews its lease. + // = staleRunningWindow / 3 → safe margin so a live run is never reclaimed. + heartbeatInterval = 30 * time.Second + // reclaimTickInterval is how often the reclaim sweep runs after startup. reclaimTickInterval = 60 * time.Second @@ -340,7 +344,7 @@ func (w *WebhookWorker) execute(ctx context.Context, call *store.WebhookCallData var agentErrMsg string if len(call.Response) == 0 && call.AgentID != nil { - out, usage, invokeErr := w.invokeAgent(tctx, call, req) + out, usage, invokeErr := w.invokeAgentWithHeartbeat(tctx, call, req, lease) if invokeErr != nil { agentErrMsg = invokeErr.Error() slog.Warn("webhook.worker.agent_invoke_failed", @@ -706,9 +710,10 @@ func (w *WebhookWorker) resetToQueued(ctx context.Context, call *store.WebhookCa } tctx := store.WithTenantID(ctx, tenantID) updates := map[string]any{ - "status": "queued", - "started_at": nil, - "lease_token": nil, // clear lease so next claimer can acquire + "status": "queued", + "started_at": nil, + "lease_token": nil, // clear lease so next claimer can acquire + "last_heartbeat_at": nil, // clear heartbeat — row is no longer running // attempts left unchanged — this was not a real send attempt } if err := w.calls.UpdateStatusCAS(tctx, call.ID, lease, updates); err != nil { @@ -724,6 +729,63 @@ func (w *WebhookWorker) resetToQueued(ctx context.Context, call *store.WebhookCa } } +// invokeAgentWithHeartbeat runs the agent alongside a heartbeat goroutine that renews the lease. +// If the lease is lost (row reclaimed), the heartbeat cancels runCtx so the agent stops immediately +// and writes no further MCP side-effects. ctx must already be detached from worker-shutdown +// cancellation (execute passes tctx). +func (w *WebhookWorker) invokeAgentWithHeartbeat( + ctx context.Context, + call *store.WebhookCallData, + req asyncPayload, + lease string, +) (string, *callbackUsage, error) { + runCtx, cancelRun := context.WithCancel(ctx) + defer cancelRun() + + hbStop := make(chan struct{}) + // Deferred so the heartbeat goroutine is stopped even if invokeAgent panics + // (the panic unwinds past this frame; execute's recover() runs afterwards). + defer close(hbStop) + go w.heartbeatLoop(ctx, call.ID, lease, heartbeatInterval, hbStop, cancelRun) + + return w.invokeAgent(runCtx, call, req) +} + +// heartbeatLoop calls Heartbeat every interval until stop is closed. When Heartbeat returns +// ErrLeaseExpired (lease reclaimed), it calls cancelRun() to stop the current run, then returns. +// It uses ctx (detached from shutdown) for the DB call so the heartbeat is not cancelled with runCtx. +func (w *WebhookWorker) heartbeatLoop( + ctx context.Context, + callID uuid.UUID, + lease string, + interval time.Duration, + stop <-chan struct{}, + cancelRun context.CancelFunc, +) { + t := time.NewTicker(interval) + defer t.Stop() + for { + select { + case <-stop: + return + case <-t.C: + // Bound each heartbeat DB call so a stalled connection can't block the goroutine + // past the next tick — it should retry on the next interval, not hang indefinitely. + hbCtx, hbCancel := context.WithTimeout(ctx, 5*time.Second) + err := w.calls.Heartbeat(hbCtx, callID, lease, time.Now()) + hbCancel() + if err != nil { + if errors.Is(err, store.ErrLeaseExpired) { + slog.Warn("webhook.worker.lease_lost_cancel_run", "call_id", callID) + cancelRun() + return + } + slog.Error("webhook.worker.heartbeat_failed", "call_id", callID, "error", err) + } + } + } +} + // invokeAgent runs the agent for an async call and returns (output, usage, error). func (w *WebhookWorker) invokeAgent( ctx context.Context, @@ -770,10 +832,13 @@ func (w *WebhookWorker) invokeAgent( TraceTags: []string{"webhook", "async"}, } - runCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), w.cfg.AsyncAgentTimeout) + // Honor ctx so heartbeatLoop can cancel the run on lease loss. execute already detached + // ctx from worker-shutdown cancellation (tctx = WithoutCancel), so dropping WithoutCancel + // here does NOT let shutdown kill the run — only cancelRun() from the heartbeat does. + agentCtx, cancel := context.WithTimeout(ctx, w.cfg.AsyncAgentTimeout) defer cancel() - result, runErr := ag.Run(runCtx, rr) + result, runErr := ag.Run(agentCtx, rr) if runErr != nil { return "", nil, runErr } diff --git a/internal/webhooks/worker_test.go b/internal/webhooks/worker_test.go index d38572d5..797c07af 100644 --- a/internal/webhooks/worker_test.go +++ b/internal/webhooks/worker_test.go @@ -28,6 +28,8 @@ type stubCallStore struct { claimErr error // if non-nil, returned by ClaimNext reclaimN int64 // count returned by ReclaimStale casLeaseErr error // if non-nil, returned by UpdateStatusCAS + hbCount int32 // số lần Heartbeat được gọi (atomic) + hbErr error // nếu non-nil, Heartbeat trả về lỗi này } func newStubCallStore(initial *store.WebhookCallData) *stubCallStore { @@ -103,6 +105,10 @@ func (s *stubCallStore) DeleteOlderThan(_ context.Context, _ uuid.UUID, _ time.T func (s *stubCallStore) ReclaimStale(_ context.Context, _ time.Time) (int64, error) { return s.reclaimN, nil } +func (s *stubCallStore) Heartbeat(_ context.Context, _ uuid.UUID, _ string, _ time.Time) error { + atomic.AddInt32(&s.hbCount, 1) + return s.hbErr +} // stubWebhookStore returns a fixed webhook on GetByID. type stubWebhookStore struct { @@ -702,6 +708,52 @@ func TestSign(t *testing.T) { } } +// TestHeartbeatLoopRenews verifies the loop calls Heartbeat repeatedly and does NOT +// cancel the run while the lease stays valid. +func TestHeartbeatLoopRenews(t *testing.T) { + callStore := newStubCallStore(nil) + w := newTestWorker(callStore, &stubWebhookStore{}) + + runCtx, cancelRun := context.WithCancel(context.Background()) + defer cancelRun() + stop := make(chan struct{}) + + go w.heartbeatLoop(context.Background(), uuid.New(), "lease-1", 5*time.Millisecond, stop, cancelRun) + + time.Sleep(40 * time.Millisecond) + close(stop) + time.Sleep(10 * time.Millisecond) + + if atomic.LoadInt32(&callStore.hbCount) < 3 { + t.Errorf("expected ≥3 heartbeats, got %d", atomic.LoadInt32(&callStore.hbCount)) + } + if runCtx.Err() != nil { + t.Errorf("runCtx must NOT be cancelled while lease valid; err=%v", runCtx.Err()) + } +} + +// TestHeartbeatLoopCancelsOnLeaseLost verifies that when Heartbeat returns ErrLeaseExpired, +// the loop cancels runCtx and returns. +func TestHeartbeatLoopCancelsOnLeaseLost(t *testing.T) { + callStore := newStubCallStore(nil) + callStore.hbErr = store.ErrLeaseExpired // mất lease ngay tick đầu + w := newTestWorker(callStore, &stubWebhookStore{}) + + runCtx, cancelRun := context.WithCancel(context.Background()) + defer cancelRun() + stop := make(chan struct{}) + defer close(stop) + + go w.heartbeatLoop(context.Background(), uuid.New(), "lease-1", 5*time.Millisecond, stop, cancelRun) + + select { + case <-runCtx.Done(): + // đúng — run bị hủy do mất lease + case <-time.After(200 * time.Millisecond): + t.Fatal("runCtx not cancelled after lease lost within 200ms") + } +} + // ---- test helpers ---- func startsWith(s, pfx string) bool { diff --git a/migrations/000085_webhook_calls_heartbeat.down.sql b/migrations/000085_webhook_calls_heartbeat.down.sql new file mode 100644 index 00000000..bec7d37c --- /dev/null +++ b/migrations/000085_webhook_calls_heartbeat.down.sql @@ -0,0 +1,2 @@ +DROP INDEX IF EXISTS idx_webhook_calls_running_heartbeat; +ALTER TABLE webhook_calls DROP COLUMN IF EXISTS last_heartbeat_at; diff --git a/migrations/000085_webhook_calls_heartbeat.up.sql b/migrations/000085_webhook_calls_heartbeat.up.sql new file mode 100644 index 00000000..9317f631 --- /dev/null +++ b/migrations/000085_webhook_calls_heartbeat.up.sql @@ -0,0 +1,11 @@ +-- Heartbeat (lease renewal) column for async webhook call leases. +-- Lets a long-running agent renew its lease so ReclaimStale won't reclaim a live worker. +ALTER TABLE webhook_calls ADD COLUMN last_heartbeat_at timestamptz; + +-- Backfill running rows so they are not reclaimed immediately after deploy +-- (their lease is still valid until staleRunningWindow elapses from started_at). +UPDATE webhook_calls SET last_heartbeat_at = started_at WHERE status = 'running'; + +-- Partial index: ReclaimStale scans only running rows by heartbeat. +CREATE INDEX IF NOT EXISTS idx_webhook_calls_running_heartbeat + ON webhook_calls (status, last_heartbeat_at) WHERE status = 'running';