From af21716be80267715186b7a700412311f58b362f Mon Sep 17 00:00:00 2001 From: Anuj Parihar Date: Tue, 7 Jul 2026 15:31:27 +0530 Subject: [PATCH 1/2] feat: add sidekiq job runner --- framework/configstore/migrations.go | 67 ++++++ framework/configstore/sidekiq.go | 172 +++++++++++++++ framework/configstore/store.go | 10 + framework/configstore/tables/sidekiq.go | 37 ++++ framework/sidekiq/sidekiq.go | 240 +++++++++++++++++++++ framework/sidekiq/sidekiq_test.go | 234 ++++++++++++++++++++ transports/bifrost-http/lib/config_test.go | 33 +++ transports/bifrost-http/server/server.go | 29 ++- 8 files changed, 821 insertions(+), 1 deletion(-) create mode 100644 framework/configstore/sidekiq.go create mode 100644 framework/configstore/tables/sidekiq.go create mode 100644 framework/sidekiq/sidekiq.go create mode 100644 framework/sidekiq/sidekiq_test.go diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index 74a6c785db0..aba1db485ee 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -437,6 +437,7 @@ var configstoreMigrationSteps = []migrationStep{ {IDs: []string{"add_mcp_client_tool_execution_timeout_column"}, run: migrationAddMCPClientToolExecutionTimeoutColumn}, {IDs: []string{"add_virtual_key_expires_at_column"}, run: migrationAddVirtualKeyExpiresAtColumn}, {IDs: []string{"add_vertex_force_single_region_column"}, run: migrationAddVertexForceSingleRegionColumn}, + {IDs: []string{"add_sidekiq_table"}, run: migrationAddSidekiqTable}, } // quoteSQLiteIdentifier quotes a SQLite identifier, escaping any double quotes. @@ -10340,3 +10341,69 @@ func migrationAddVertexForceSingleRegionColumn(ctx context.Context, db *gorm.DB, } return nil } + +// migrationAddSidekiqTable creates the generic `sidekiq` background-job table. Uses raw SQL +// (not GORM auto-DDL) so the schema is explicit and stable across GORM versions. +// Idempotent via CREATE TABLE IF NOT EXISTS; covers postgres and sqlite dialects. +func migrationAddSidekiqTable(ctx context.Context, db *gorm.DB, logger schemas.Logger) error { + migrationName := "add_sidekiq_table" + logger.Info("[configstore] starting migration %s", migrationName) + defer logger.Info("[configstore] finished migration %s", migrationName) + m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ + ID: migrationName, + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + + var createTable string + switch tx.Dialector.Name() { + case "postgres": + createTable = ` + CREATE TABLE IF NOT EXISTS sidekiq ( + id TEXT PRIMARY KEY, + kind TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + metadata TEXT DEFAULT '{}', + attempts INTEGER NOT NULL DEFAULT 0, + last_error TEXT, + created_at TIMESTAMPTZ NOT NULL, + updated_at TIMESTAMPTZ NOT NULL, + started_at TIMESTAMPTZ, + completed_at TIMESTAMPTZ + )` + case "sqlite": + createTable = ` + CREATE TABLE IF NOT EXISTS sidekiq ( + id TEXT PRIMARY KEY, + kind TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + metadata TEXT DEFAULT '{}', + attempts INTEGER NOT NULL DEFAULT 0, + last_error TEXT, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL, + started_at DATETIME, + completed_at DATETIME + )` + default: + // Fall back to GORM for any other dialect so the migration does not + // hard-fail on an unsupported backend. + return tx.Migrator().CreateTable(&tables.TableSidekiqJob{}) + } + + if err := tx.Exec(createTable).Error; err != nil { + return err + } + // Index supports the reaper / recovery scan that filters by status and + // orders/filters by updated_at. + return tx.Exec(`CREATE INDEX IF NOT EXISTS idx_sidekiq_status_updated ON sidekiq (status, updated_at)`).Error + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + return tx.Exec(`DROP TABLE IF EXISTS sidekiq`).Error + }, + }}) + if err := m.Migrate(); err != nil { + return fmt.Errorf("error running %s migration: %w", migrationName, err) + } + return nil +} diff --git a/framework/configstore/sidekiq.go b/framework/configstore/sidekiq.go new file mode 100644 index 00000000000..b3aedfcf1fa --- /dev/null +++ b/framework/configstore/sidekiq.go @@ -0,0 +1,172 @@ +package configstore + +import ( + "context" + "errors" + "strings" + "time" + + "github.com/maximhq/bifrost/framework/configstore/tables" + "gorm.io/gorm" +) + +// CreateSidekiqJob inserts a new background job. The caller supplies the id, kind +// and metadata; status defaults to pending and timestamps are stamped here. +func (s *RDBConfigStore) CreateSidekiqJob(ctx context.Context, job *tables.TableSidekiqJob) error { + if job == nil { + return errors.New("sidekiq job is required") + } + if strings.TrimSpace(job.ID) == "" { + return errors.New("sidekiq job id is required") + } + if strings.TrimSpace(job.Kind) == "" { + return errors.New("sidekiq job kind is required") + } + now := time.Now() + if job.Status == "" { + job.Status = tables.SidekiqStatusPending + } + if job.Metadata == "" { + job.Metadata = "{}" + } + job.CreatedAt = now + job.UpdatedAt = now + return s.DB().WithContext(ctx).Create(job).Error +} + +// GetSidekiqJob returns a single job by id, or nil when it does not exist. +func (s *RDBConfigStore) GetSidekiqJob(ctx context.Context, id string) (*tables.TableSidekiqJob, error) { + var job tables.TableSidekiqJob + err := s.DB().WithContext(ctx).Where("id = ?", id).First(&job).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + if err != nil { + return nil, err + } + return &job, nil +} + +// MarkSidekiqJobRunning transitions a job to running, stamps started_at, bumps the +// heartbeat (updated_at), and increments the attempt counter. Safe to call on +// resume: each resumed run counts as a fresh attempt. +func (s *RDBConfigStore) MarkSidekiqJobRunning(ctx context.Context, id string) error { + now := time.Now() + res := s.DB().WithContext(ctx). + Model(&tables.TableSidekiqJob{}). + Where("id = ?", id). + Updates(map[string]any{ + "status": tables.SidekiqStatusRunning, + "started_at": now, + "updated_at": now, + "attempts": gorm.Expr("attempts + 1"), + }) + if res.Error != nil { + return res.Error + } + if res.RowsAffected == 0 { + return errors.New("sidekiq job not found or already in terminal state") + } + return nil +} + +// UpdateSidekiqJobProgress persists a progress checkpoint: it replaces the metadata +// blob and bumps the heartbeat (updated_at) so the reaper does not treat the job as +// stale. Called after each processed page. +func (s *RDBConfigStore) UpdateSidekiqJobProgress(ctx context.Context, id, metadata string) error { + res := s.DB().WithContext(ctx). + Model(&tables.TableSidekiqJob{}). + Where("id = ?", id). + Updates(map[string]any{ + "metadata": metadata, + "updated_at": time.Now(), + }) + if res.Error != nil { + return res.Error + } + if res.RowsAffected == 0 { + return errors.New("sidekiq job not found") + } + return nil +} + +// CompleteSidekiqJob marks a job completed, stamps completed_at, and stores the +// final metadata (counts, summary). +func (s *RDBConfigStore) CompleteSidekiqJob(ctx context.Context, id, metadata string) error { + now := time.Now() + res := s.DB().WithContext(ctx). + Model(&tables.TableSidekiqJob{}). + Where("id = ?", id). + Updates(map[string]any{ + "status": tables.SidekiqStatusCompleted, + "metadata": metadata, + "updated_at": now, + "completed_at": now, + }) + if res.Error != nil { + return res.Error + } + if res.RowsAffected == 0 { + return errors.New("sidekiq job not found") + } + return nil +} + +// FailSidekiqJob marks a job failed, records the error, stamps completed_at, and +// preserves the latest metadata so a later resume can read the checkpoint cursor. +func (s *RDBConfigStore) FailSidekiqJob(ctx context.Context, id, metadata, lastErr string) error { + now := time.Now() + updates := map[string]any{ + "status": tables.SidekiqStatusFailed, + "last_error": lastErr, + "updated_at": now, + "completed_at": now, + } + if metadata != "" { + updates["metadata"] = metadata + } + res := s.DB().WithContext(ctx). + Model(&tables.TableSidekiqJob{}). + Where("id = ?", id). + Updates(updates) + if res.Error != nil { + return res.Error + } + if res.RowsAffected == 0 { + return errors.New("sidekiq job not found") + } + return nil +} + +// ListIncompleteSidekiqJobs returns jobs that are not in a terminal state +// (pending or running). Used by startup recovery to resume work that was +// interrupted by a restart or crash. +func (s *RDBConfigStore) ListIncompleteSidekiqJobs(ctx context.Context) ([]tables.TableSidekiqJob, error) { + var jobs []tables.TableSidekiqJob + err := s.DB().WithContext(ctx). + Where("status IN ?", []string{tables.SidekiqStatusPending, tables.SidekiqStatusRunning}). + Order("created_at ASC"). + Find(&jobs).Error + if err != nil { + return nil, err + } + return jobs, nil +} + +// MarkStaleSidekiqJobsFailed flips any running job whose heartbeat (updated_at) is +// older than staleBefore to failed. This is the safety net for a goroutine or node +// that died without marking its job: the job stops looking "running" and becomes +// eligible for inspection or a manual resume. Returns the number of jobs reaped. +func (s *RDBConfigStore) MarkStaleSidekiqJobsFailed(ctx context.Context, staleBefore time.Time) (int64, error) { + now := time.Now() + res := s.DB().WithContext(ctx). + Model(&tables.TableSidekiqJob{}). + Where("status = ? AND updated_at < ?", tables.SidekiqStatusRunning, staleBefore). + Updates(map[string]any{ + "status": tables.SidekiqStatusFailed, + "last_error": "job timed out: no heartbeat before stale threshold", + "updated_at": now, + "completed_at": now, + }) + return res.RowsAffected, res.Error +} diff --git a/framework/configstore/store.go b/framework/configstore/store.go index d6cb6f69c6f..7069b656a14 100644 --- a/framework/configstore/store.go +++ b/framework/configstore/store.go @@ -663,6 +663,16 @@ type ConfigStore interface { RenamePromptSession(ctx context.Context, id uint, name string) error DeletePromptSession(ctx context.Context, id uint) error + // Sidekiq - generic durable background jobs + CreateSidekiqJob(ctx context.Context, job *tables.TableSidekiqJob) error + GetSidekiqJob(ctx context.Context, id string) (*tables.TableSidekiqJob, error) + MarkSidekiqJobRunning(ctx context.Context, id string) error + UpdateSidekiqJobProgress(ctx context.Context, id, metadata string) error + CompleteSidekiqJob(ctx context.Context, id, metadata string) error + FailSidekiqJob(ctx context.Context, id, metadata, lastErr string) error + ListIncompleteSidekiqJobs(ctx context.Context) ([]tables.TableSidekiqJob, error) + MarkStaleSidekiqJobsFailed(ctx context.Context, staleBefore time.Time) (int64, error) + // DB returns the underlying database connection. DB() *gorm.DB diff --git a/framework/configstore/tables/sidekiq.go b/framework/configstore/tables/sidekiq.go new file mode 100644 index 00000000000..bf114a77620 --- /dev/null +++ b/framework/configstore/tables/sidekiq.go @@ -0,0 +1,37 @@ +package tables + +import "time" + +// Sidekiq job status values. +const ( + // SidekiqStatusPending marks a job that has been enqueued but not yet started. + SidekiqStatusPending = "pending" + // SidekiqStatusRunning marks a job whose goroutine is actively processing it. + SidekiqStatusRunning = "running" + // SidekiqStatusCompleted marks a job that finished successfully. + SidekiqStatusCompleted = "completed" + // SidekiqStatusFailed marks a job that errored or was reaped as stale. + SidekiqStatusFailed = "failed" +) + +// TableSidekiqJob is a generic, durable background-job record. It is intentionally +// not tied to any feature: callers store all job-specific data (provider, filters, +// resume cursor, running counts, errors) in the Metadata JSON blob. The runner only +// reads Status and UpdatedAt; everything else is opaque to it. +type TableSidekiqJob struct { + ID string `gorm:"column:id;primaryKey;type:text" json:"id"` + Kind string `gorm:"column:kind;not null;type:text;index:idx_sidekiq_status_updated,priority:3" json:"kind"` + Status string `gorm:"column:status;not null;default:pending;type:text;index:idx_sidekiq_status_updated,priority:1" json:"status"` + Metadata string `gorm:"column:metadata;type:text;default:'{}'" json:"metadata"` + Attempts int `gorm:"column:attempts;not null;default:0" json:"attempts"` + LastError string `gorm:"column:last_error;type:text" json:"last_error,omitempty"` + CreatedAt time.Time `gorm:"column:created_at;not null" json:"created_at"` + UpdatedAt time.Time `gorm:"column:updated_at;not null;index:idx_sidekiq_status_updated,priority:2" json:"updated_at"` + StartedAt *time.Time `gorm:"column:started_at" json:"started_at,omitempty"` + CompletedAt *time.Time `gorm:"column:completed_at" json:"completed_at,omitempty"` +} + +// TableName returns the backing table name for sidekiq jobs. +func (TableSidekiqJob) TableName() string { + return "sidekiq" +} diff --git a/framework/sidekiq/sidekiq.go b/framework/sidekiq/sidekiq.go new file mode 100644 index 00000000000..d94fd44c73f --- /dev/null +++ b/framework/sidekiq/sidekiq.go @@ -0,0 +1,240 @@ +package sidekiq + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/framework/configstore/tables" +) + +// Store is the subset of the configstore the runner needs. Keeping it narrow lets +// the runner be tested with a fake and avoids a hard dependency on the full store. +type Store interface { + CreateSidekiqJob(ctx context.Context, job *tables.TableSidekiqJob) error + GetSidekiqJob(ctx context.Context, id string) (*tables.TableSidekiqJob, error) + MarkSidekiqJobRunning(ctx context.Context, id string) error + UpdateSidekiqJobProgress(ctx context.Context, id, metadata string) error + CompleteSidekiqJob(ctx context.Context, id, metadata string) error + FailSidekiqJob(ctx context.Context, id, metadata, lastErr string) error + ListIncompleteSidekiqJobs(ctx context.Context) ([]tables.TableSidekiqJob, error) + MarkStaleSidekiqJobsFailed(ctx context.Context, staleBefore time.Time) (int64, error) +} + +// ProgressFunc persists a checkpoint: it replaces the job's metadata blob and +// bumps the heartbeat. Handlers call it after each unit of work (e.g. each page). +type ProgressFunc func(metadata string) error + +// HandlerFunc processes one job. It is given the job (read its Metadata for the +// resume cursor) and a progress callback to checkpoint after each unit of work. +// It returns the final metadata to persist and an error. A nil error completes the +// job; a non-nil error fails it (the returned metadata is still stored so a later +// resume can read the last cursor). +type HandlerFunc func(ctx context.Context, job tables.TableSidekiqJob, progress ProgressFunc) (finalMetadata string, err error) + +const ( + ReaperInterval = 1 * time.Minute + StaleAfter = 15 * time.Minute +) + +// Runner owns the handler registry and the lifecycle of job goroutines. +type Runner struct { + store Store + logger schemas.Logger + handlers map[string]HandlerFunc + mu sync.RWMutex + + baseCtx context.Context + cancel context.CancelFunc + sem chan struct{} + wg sync.WaitGroup +} + +// New creates a Runner. maxConcurrent bounds how many job goroutines run at once +// (<=0 defaults to 4). Jobs run on a background context derived here, never on a +// request context, so they outlive the HTTP request that enqueued them. +func New(store Store, logger schemas.Logger, maxConcurrent int) *Runner { + if maxConcurrent <= 0 { + maxConcurrent = 4 + } + ctx, cancel := context.WithCancel(context.Background()) + return &Runner{ + store: store, + logger: logger, + handlers: make(map[string]HandlerFunc), + baseCtx: ctx, + cancel: cancel, + sem: make(chan struct{}, maxConcurrent), + } +} + +// Register binds a handler to a job kind. Call during startup, before enqueuing. +func (r *Runner) Register(kind string, fn HandlerFunc) { + r.mu.Lock() + defer r.mu.Unlock() + r.handlers[kind] = fn +} + +// handlerFor returns the registered handler for a kind, if any. +func (r *Runner) handlerFor(kind string) (HandlerFunc, bool) { + r.mu.RLock() + defer r.mu.RUnlock() + fn, ok := r.handlers[kind] + return fn, ok +} + +// Enqueue persists a new pending job and starts its goroutine. The caller supplies +// the id (also usable as a UI operation id), the kind, and the initial metadata +// JSON. It returns once the row is committed, so the HTTP handler can respond +// immediately while processing continues in the background. +func (r *Runner) Enqueue(ctx context.Context, id, kind, metadata string) error { + if _, ok := r.handlerFor(kind); !ok { + return fmt.Errorf("sidekiq: no handler registered for kind %q", kind) + } + job := &tables.TableSidekiqJob{ + ID: id, + Kind: kind, + Status: tables.SidekiqStatusPending, + Metadata: metadata, + } + if err := r.store.CreateSidekiqJob(ctx, job); err != nil { + return err + } + r.spawn(*job) + return nil +} + +// spawn runs a job in its own goroutine, bounded by the concurrency semaphore. +func (r *Runner) spawn(job tables.TableSidekiqJob) { + r.wg.Add(1) + go func() { + defer r.wg.Done() + select { + case r.sem <- struct{}{}: + case <-r.baseCtx.Done(): + return + } + defer func() { <-r.sem }() + r.execute(job) + }() +} + +// execute marks the job running, invokes its handler, and records the terminal +// state. A panic in the handler is recovered and recorded as a failure so one bad +// job cannot crash the process. +func (r *Runner) execute(job tables.TableSidekiqJob) { + defer func() { + if rec := recover(); rec != nil { + r.logger.Error("sidekiq: job %s (%s) panicked: %v", job.ID, job.Kind, rec) + if err := r.store.FailSidekiqJob(r.baseCtx, job.ID, "", fmt.Sprintf("panic: %v", rec)); err != nil { + r.logger.Error("sidekiq: failed to mark panicked job %s failed: %v", job.ID, err) + } + } + }() + + fn, ok := r.handlerFor(job.Kind) + if !ok { + if err := r.store.FailSidekiqJob(r.baseCtx, job.ID, "", "no handler registered for kind "+job.Kind); err != nil { + r.logger.Error("sidekiq: failed to fail unhandled job %s: %v", job.ID, err) + } + return + } + + if err := r.store.MarkSidekiqJobRunning(r.baseCtx, job.ID); err != nil { + r.logger.Error("sidekiq: failed to mark job %s running: %v", job.ID, err) + if ferr := r.store.FailSidekiqJob(r.baseCtx, job.ID, job.Metadata, err.Error()); ferr != nil { + r.logger.Error("sidekiq: failed to fail job %s after running-mark failure: %v", job.ID, ferr) + } + return + } + + progress := func(metadata string) error { + return r.store.UpdateSidekiqJobProgress(r.baseCtx, job.ID, metadata) + } + + finalMetadata, err := fn(r.baseCtx, job, progress) + if err != nil { + r.logger.Error("sidekiq: job %s (%s) failed: %v", job.ID, job.Kind, err) + if ferr := r.store.FailSidekiqJob(r.baseCtx, job.ID, finalMetadata, err.Error()); ferr != nil { + r.logger.Error("sidekiq: failed to mark job %s failed: %v", job.ID, ferr) + } + return + } + if cerr := r.store.CompleteSidekiqJob(r.baseCtx, job.ID, finalMetadata); cerr != nil { + r.logger.Error("sidekiq: failed to mark job %s completed: %v", job.ID, cerr) + } +} + +// RecoverIncomplete re-runs jobs left pending or running by a previous process +// (a restart or crash). Each handler resumes from the cursor stored in its +// metadata; because per-item work is idempotent, reprocessing the in-flight unit +// is safe. In a multi-node cluster this may double-run a job across nodes, which +// idempotency tolerates; it does not de-duplicate work across nodes by design, +// matching the choice to keep the runner simple (no leader election). +func (r *Runner) RecoverIncomplete(ctx context.Context) error { + jobs, err := r.store.ListIncompleteSidekiqJobs(ctx) + if err != nil { + return err + } + for _, job := range jobs { + if _, ok := r.handlerFor(job.Kind); !ok { + r.logger.Warn("sidekiq: skipping recovery of job %s, no handler for kind %s", job.ID, job.Kind) + continue + } + r.logger.Info("sidekiq: recovering incomplete job %s (%s)", job.ID, job.Kind) + r.spawn(job) + } + return nil +} + +// StartReaper periodically marks running jobs whose heartbeat is older than +// staleAfter as failed, catching goroutines or nodes that died without recording a +// terminal state. It returns a stop function. Both interval and staleAfter must be +// positive; staleAfter should comfortably exceed the handler's per-checkpoint time. +func (r *Runner) StartReaper(interval, staleAfter time.Duration) (stop func()) { + // Guard against invalid durations: time.NewTicker panics on interval <= 0, and + // a non-positive staleAfter would make every running job look stale. Fall back + // to the package defaults rather than crashing or reaping live jobs. + if interval <= 0 { + interval = ReaperInterval + } + if staleAfter <= 0 { + staleAfter = StaleAfter + } + ticker := time.NewTicker(interval) + done := make(chan struct{}) + r.wg.Add(1) + go func() { + defer r.wg.Done() + defer ticker.Stop() + for { + select { + case <-done: + return + case <-r.baseCtx.Done(): + return + case <-ticker.C: + n, err := r.store.MarkStaleSidekiqJobsFailed(r.baseCtx, time.Now().Add(-staleAfter)) + if err != nil { + r.logger.Error("sidekiq: reaper failed: %v", err) + continue + } + if n > 0 { + r.logger.Warn("sidekiq: reaper marked %d stale job(s) as failed", n) + } + } + } + }() + var once sync.Once + return func() { once.Do(func() { close(done) }) } +} + +// Shutdown cancels the background context and waits for in-flight goroutines to +// return. In-flight jobs observe baseCtx cancellation and stop at their next +// checkpoint, leaving a resumable cursor in metadata. +func (r *Runner) Shutdown() { + r.cancel() + r.wg.Wait() +} diff --git a/framework/sidekiq/sidekiq_test.go b/framework/sidekiq/sidekiq_test.go new file mode 100644 index 00000000000..4c062311319 --- /dev/null +++ b/framework/sidekiq/sidekiq_test.go @@ -0,0 +1,234 @@ +package sidekiq + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + bifrost "github.com/maximhq/bifrost/core" + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/framework/configstore/tables" +) + +// fakeStore is an in-memory Store for exercising the runner without a database. +type fakeStore struct { + mu sync.Mutex + created []tables.TableSidekiqJob + running map[string]int + progress map[string]string + completed map[string]string + failedMeta map[string]string + failedErr map[string]string + incomplete []tables.TableSidekiqJob + staleCalls int + terminal chan string +} + +func newFakeStore() *fakeStore { + return &fakeStore{ + running: map[string]int{}, + progress: map[string]string{}, + completed: map[string]string{}, + failedMeta: map[string]string{}, + failedErr: map[string]string{}, + terminal: make(chan string, 16), + } +} + +func (f *fakeStore) CreateSidekiqJob(_ context.Context, job *tables.TableSidekiqJob) error { + f.mu.Lock() + defer f.mu.Unlock() + f.created = append(f.created, *job) + return nil +} + +func (f *fakeStore) GetSidekiqJob(_ context.Context, _ string) (*tables.TableSidekiqJob, error) { + return nil, nil +} + +func (f *fakeStore) MarkSidekiqJobRunning(_ context.Context, id string) error { + f.mu.Lock() + defer f.mu.Unlock() + f.running[id]++ + return nil +} + +func (f *fakeStore) UpdateSidekiqJobProgress(_ context.Context, id, metadata string) error { + f.mu.Lock() + defer f.mu.Unlock() + f.progress[id] = metadata + return nil +} + +func (f *fakeStore) CompleteSidekiqJob(_ context.Context, id, metadata string) error { + f.mu.Lock() + f.completed[id] = metadata + f.mu.Unlock() + f.terminal <- id + return nil +} + +func (f *fakeStore) FailSidekiqJob(_ context.Context, id, metadata, lastErr string) error { + f.mu.Lock() + f.failedMeta[id] = metadata + f.failedErr[id] = lastErr + f.mu.Unlock() + f.terminal <- id + return nil +} + +func (f *fakeStore) ListIncompleteSidekiqJobs(_ context.Context) ([]tables.TableSidekiqJob, error) { + return f.incomplete, nil +} + +func (f *fakeStore) MarkStaleSidekiqJobsFailed(_ context.Context, _ time.Time) (int64, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.staleCalls++ + return 0, nil +} + +// waitTerminal blocks until a job reaches a terminal state or the test times out. +func waitTerminal(t *testing.T, f *fakeStore) string { + t.Helper() + select { + case id := <-f.terminal: + return id + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for job to reach a terminal state") + return "" + } +} + +func testRunner(store Store) *Runner { + return New(store, bifrost.NewDefaultLogger(schemas.LogLevelError), 4) +} + +func TestEnqueueRunsHandlerAndCompletes(t *testing.T) { + store := newFakeStore() + r := testRunner(store) + r.Register("k", func(_ context.Context, job tables.TableSidekiqJob, progress ProgressFunc) (string, error) { + _ = progress("checkpoint") + return "final", nil + }) + + if err := r.Enqueue(context.Background(), "job1", "k", "{}"); err != nil { + t.Fatalf("Enqueue: %v", err) + } + id := waitTerminal(t, store) + if id != "job1" { + t.Fatalf("terminal id = %q, want job1", id) + } + + store.mu.Lock() + defer store.mu.Unlock() + if store.running["job1"] != 1 { + t.Errorf("MarkSidekiqJobRunning called %d times, want 1", store.running["job1"]) + } + if store.progress["job1"] != "checkpoint" { + t.Errorf("progress = %q, want checkpoint", store.progress["job1"]) + } + if store.completed["job1"] != "final" { + t.Errorf("completed metadata = %q, want final", store.completed["job1"]) + } +} + +func TestHandlerErrorMarksFailed(t *testing.T) { + store := newFakeStore() + r := testRunner(store) + r.Register("k", func(_ context.Context, job tables.TableSidekiqJob, _ ProgressFunc) (string, error) { + return "partial", errors.New("boom") + }) + + if err := r.Enqueue(context.Background(), "job2", "k", "{}"); err != nil { + t.Fatalf("Enqueue: %v", err) + } + waitTerminal(t, store) + + store.mu.Lock() + defer store.mu.Unlock() + if store.failedErr["job2"] != "boom" { + t.Errorf("failed err = %q, want boom", store.failedErr["job2"]) + } + if store.failedMeta["job2"] != "partial" { + t.Errorf("failed metadata = %q, want partial", store.failedMeta["job2"]) + } +} + +func TestHandlerPanicRecovered(t *testing.T) { + store := newFakeStore() + r := testRunner(store) + r.Register("k", func(_ context.Context, job tables.TableSidekiqJob, _ ProgressFunc) (string, error) { + panic("kaboom") + }) + + if err := r.Enqueue(context.Background(), "job3", "k", "{}"); err != nil { + t.Fatalf("Enqueue: %v", err) + } + waitTerminal(t, store) + + store.mu.Lock() + defer store.mu.Unlock() + if store.failedErr["job3"] == "" { + t.Errorf("expected a failure recorded for a panicking handler") + } +} + +func TestEnqueueUnknownKindErrors(t *testing.T) { + store := newFakeStore() + r := testRunner(store) + if err := r.Enqueue(context.Background(), "job4", "missing", "{}"); err == nil { + t.Fatal("expected error enqueuing an unregistered kind") + } + store.mu.Lock() + defer store.mu.Unlock() + if len(store.created) != 0 { + t.Errorf("no job should be created for an unknown kind, got %d", len(store.created)) + } +} + +func TestRecoverIncompleteResumesJobs(t *testing.T) { + store := newFakeStore() + store.incomplete = []tables.TableSidekiqJob{ + {ID: "r1", Kind: "k", Status: tables.SidekiqStatusRunning, Metadata: "{}"}, + {ID: "r2", Kind: "k", Status: tables.SidekiqStatusPending, Metadata: "{}"}, + } + r := testRunner(store) + r.Register("k", func(_ context.Context, job tables.TableSidekiqJob, _ ProgressFunc) (string, error) { + return "done", nil + }) + + if err := r.RecoverIncomplete(context.Background()); err != nil { + t.Fatalf("RecoverIncomplete: %v", err) + } + got := map[string]bool{} + got[waitTerminal(t, store)] = true + got[waitTerminal(t, store)] = true + if !got["r1"] || !got["r2"] { + t.Errorf("expected both r1 and r2 to be recovered, got %v", got) + } +} + +func TestReaperInvokesStaleSweep(t *testing.T) { + store := newFakeStore() + r := testRunner(store) + stop := r.StartReaper(10*time.Millisecond, time.Millisecond) + defer stop() + + deadline := time.After(2 * time.Second) + for { + store.mu.Lock() + n := store.staleCalls + store.mu.Unlock() + if n > 0 { + return + } + select { + case <-deadline: + t.Fatal("reaper never invoked the stale sweep") + case <-time.After(5 * time.Millisecond): + } + } +} diff --git a/transports/bifrost-http/lib/config_test.go b/transports/bifrost-http/lib/config_test.go index 1e23526436f..03e475d37b8 100644 --- a/transports/bifrost-http/lib/config_test.go +++ b/transports/bifrost-http/lib/config_test.go @@ -1553,6 +1553,39 @@ func (m *MockConfigStore) DeleteRoutingRule(ctx context.Context, id string, tx . return nil } +// Sidekiq +func (m *MockConfigStore) CreateSidekiqJob(ctx context.Context, job *tables.TableSidekiqJob) error { + return nil +} + +func (m *MockConfigStore) GetSidekiqJob(ctx context.Context, id string) (*tables.TableSidekiqJob, error) { + return nil, nil +} + +func (m *MockConfigStore) MarkSidekiqJobRunning(ctx context.Context, id string) error { + return nil +} + +func (m *MockConfigStore) CompleteSidekiqJob(ctx context.Context, id, metadata string) error { + return nil +} + +func (m *MockConfigStore) UpdateSidekiqJobProgress(ctx context.Context, id, metadata string) error { + return nil +} + +func (m *MockConfigStore) FailSidekiqJob(ctx context.Context, id, metadata, lastErr string) error { + return nil +} + +func (m *MockConfigStore) ListIncompleteSidekiqJobs(ctx context.Context) ([]tables.TableSidekiqJob, error) { + return nil, nil +} + +func (m *MockConfigStore) MarkStaleSidekiqJobsFailed(ctx context.Context, staleBefore time.Time) (int64, error) { + return 0, nil +} + func TestMergeGovernanceConfig_SyncsComplexityAnalyzerConfig(t *testing.T) { initTestLogger() diff --git a/transports/bifrost-http/server/server.go b/transports/bifrost-http/server/server.go index 4aa631bd42a..7c0026ba3f6 100644 --- a/transports/bifrost-http/server/server.go +++ b/transports/bifrost-http/server/server.go @@ -24,6 +24,7 @@ import ( "github.com/maximhq/bifrost/framework/encrypt" "github.com/maximhq/bifrost/framework/logstore" dynamicPlugins "github.com/maximhq/bifrost/framework/plugins" + "github.com/maximhq/bifrost/framework/sidekiq" "github.com/maximhq/bifrost/framework/temptoken" "github.com/maximhq/bifrost/framework/tracing" "github.com/maximhq/bifrost/plugins/governance" @@ -179,6 +180,9 @@ type BifrostHTTPServer struct { // otherwise left nil so the quota endpoint reads the VK's own budget rows. ExternalQuotaBudgetResolver handlers.ExternalQuotaBudgetResolver + SidekiqRunner *sidekiq.Runner + SidekiqReaperStop func() + wsPool *bfws.Pool } @@ -1698,7 +1702,6 @@ func (s *BifrostHTTPServer) Bootstrap(ctx context.Context) error { s.Ctx, s.cancel = schemas.NewBifrostContextWithCancel(ctx) handlers.SetVersion(s.Version) configDir := GetDefaultConfigDir(s.AppDir) - // Ensure app directory exists if err := os.MkdirAll(configDir, 0o755); err != nil { return fmt.Errorf("failed to create app directory %s: %v", configDir, err) @@ -1892,6 +1895,11 @@ func (s *BifrostHTTPServer) Bootstrap(ctx context.Context) error { semanticCachePlugin.SetEmbeddingRequestExecutor(s.Client.EmbeddingRequest) } + // Initialize Sidekiq runner for background jobs + if s.Config != nil && s.Config.ConfigStore != nil { + s.SidekiqRunner = sidekiq.New(s.Config.ConfigStore, logger, 4) + } + // Register routes err = s.RegisterAPIRoutes(s.Ctx, s, apiMiddlewares...) if err != nil { @@ -1962,6 +1970,17 @@ func (s *BifrostHTTPServer) Bootstrap(ctx context.Context) error { }) // Register UI handler s.RegisterUIRoutes() + + // Start Sidekiq reaper to clean up stale jobs + if s.SidekiqRunner != nil { + s.SidekiqReaperStop = s.SidekiqRunner.StartReaper(sidekiq.ReaperInterval, sidekiq.StaleAfter) + go func() { + if err := s.SidekiqRunner.RecoverIncomplete(ctx); err != nil { + logger.Error("sidekiq: failed to recover incomplete provisioning jobs: %v", err) + } + }() + } + // Checking if config has server config and use it to set read buffer size logger.Debug("server read buffer size: %d", s.Config.ServerConfig.ReadBufferSize) // Create fasthttp server instance @@ -2049,6 +2068,14 @@ func (s *BifrostHTTPServer) Start() error { s.OAuth2SweepWorker.stop() s.OAuth2SweepWorker = nil } + if s.SidekiqReaperStop != nil { + logger.Info("stopping sidekiq reaper...") + s.SidekiqReaperStop() + } + if s.SidekiqRunner != nil { + logger.Info("stopping sidekiq runner...") + s.SidekiqRunner.Shutdown() + } if s.devPprofHandler != nil { logger.Info("stopping dev pprof handler...") s.devPprofHandler.Cleanup() From 220539ff0870dad6104d4d4dbd8c489aab23ad37 Mon Sep 17 00:00:00 2001 From: Anuj Parihar Date: Tue, 7 Jul 2026 18:05:27 +0530 Subject: [PATCH 2/2] feat: respect sidekiq over clustering --- framework/configstore/migrations.go | 12 +- framework/configstore/sidekiq.go | 98 ++-- framework/configstore/sidekiq_test.go | 364 +++++++++++++++ framework/configstore/store.go | 11 +- framework/configstore/tables/sidekiq.go | 5 + framework/sidekiq/sidekiq.go | 289 ++++++++---- framework/sidekiq/sidekiq_test.go | 512 +++++++++++++++++++-- transports/bifrost-http/lib/config_test.go | 16 +- transports/bifrost-http/server/server.go | 23 +- 9 files changed, 1143 insertions(+), 187 deletions(-) create mode 100644 framework/configstore/sidekiq_test.go diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index aba1db485ee..2724f10ddfa 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -10362,6 +10362,7 @@ func migrationAddSidekiqTable(ctx context.Context, db *gorm.DB, logger schemas.L id TEXT PRIMARY KEY, kind TEXT NOT NULL, status TEXT NOT NULL DEFAULT 'pending', + runner_id TEXT, metadata TEXT DEFAULT '{}', attempts INTEGER NOT NULL DEFAULT 0, last_error TEXT, @@ -10376,6 +10377,7 @@ func migrationAddSidekiqTable(ctx context.Context, db *gorm.DB, logger schemas.L id TEXT PRIMARY KEY, kind TEXT NOT NULL, status TEXT NOT NULL DEFAULT 'pending', + runner_id TEXT, metadata TEXT DEFAULT '{}', attempts INTEGER NOT NULL DEFAULT 0, last_error TEXT, @@ -10393,9 +10395,12 @@ func migrationAddSidekiqTable(ctx context.Context, db *gorm.DB, logger schemas.L if err := tx.Exec(createTable).Error; err != nil { return err } - // Index supports the reaper / recovery scan that filters by status and - // orders/filters by updated_at. - return tx.Exec(`CREATE INDEX IF NOT EXISTS idx_sidekiq_status_updated ON sidekiq (status, updated_at)`).Error + // idx_sidekiq_status_updated supports the reaper/recovery scan. + if err := tx.Exec(`CREATE INDEX IF NOT EXISTS idx_sidekiq_status_updated ON sidekiq (status, updated_at)`).Error; err != nil { + return err + } + // idx_sidekiq_runner supports fencing lookups by runner_id. + return tx.Exec(`CREATE INDEX IF NOT EXISTS idx_sidekiq_runner ON sidekiq (runner_id)`).Error }, Rollback: func(tx *gorm.DB) error { tx = tx.WithContext(ctx) @@ -10407,3 +10412,4 @@ func migrationAddSidekiqTable(ctx context.Context, db *gorm.DB, logger schemas.L } return nil } + diff --git a/framework/configstore/sidekiq.go b/framework/configstore/sidekiq.go index b3aedfcf1fa..b1cecc4e28d 100644 --- a/framework/configstore/sidekiq.go +++ b/framework/configstore/sidekiq.go @@ -47,36 +47,71 @@ func (s *RDBConfigStore) GetSidekiqJob(ctx context.Context, id string) (*tables. return &job, nil } -// MarkSidekiqJobRunning transitions a job to running, stamps started_at, bumps the -// heartbeat (updated_at), and increments the attempt counter. Safe to call on -// resume: each resumed run counts as a fresh attempt. -func (s *RDBConfigStore) MarkSidekiqJobRunning(ctx context.Context, id string) error { +// ClaimSidekiqJob atomically claims a job for runnerID and transitions it to +// running. It is the cluster-wide mutual-exclusion primitive: the conditional +// UPDATE means at most one claim affects a row, so exactly one node — and exactly +// one goroutine — runs each job. A claim succeeds when the job is: +// - pending (never started), or +// - running but stale (updated_at < staleBefore, i.e. the owner's heartbeat +// lapsed, so it is presumed dead and the job is orphaned/resumable). +// +// A job running under a live owner (fresh heartbeat) yields RowsAffected == 0, so +// it is not claimed. Note there is deliberately no "runner_id = runnerID" escape: +// resume after a crash is covered by the stale condition (a restarted process has +// a new runnerID anyway), and omitting it means a second concurrent claim on the +// same node (e.g. Enqueue racing a dispatcher tick) loses instead of double-running. +// The claim stamps runner_id, bumps the heartbeat, and increments the attempt +// counter (each claimed run is a fresh attempt). started_at is only set on first +// start; a resume keeps the original. Returns true when this claim won. +// +// In OSS (single-node) mode runnerID is empty and staleBefore is time.Now(), so +// any running job (e.g. from a crashed previous process) is immediately claimable. +func (s *RDBConfigStore) ClaimSidekiqJob(ctx context.Context, id, runnerID string, staleBefore time.Time) (bool, error) { now := time.Now() res := s.DB().WithContext(ctx). Model(&tables.TableSidekiqJob{}). - Where("id = ?", id). + Where("id = ? AND (status = ? OR (status = ? AND updated_at < ?))", + id, + tables.SidekiqStatusPending, + tables.SidekiqStatusRunning, staleBefore). Updates(map[string]any{ "status": tables.SidekiqStatusRunning, - "started_at": now, + "runner_id": runnerID, + "started_at": gorm.Expr("COALESCE(started_at, ?)", now), "updated_at": now, "attempts": gorm.Expr("attempts + 1"), }) if res.Error != nil { - return res.Error + return false, res.Error } - if res.RowsAffected == 0 { - return errors.New("sidekiq job not found or already in terminal state") + return res.RowsAffected == 1, nil +} + +// HeartbeatSidekiqJob bumps the heartbeat (updated_at) for a job the caller still +// owns and is still running. Called on a fixed interval by the owning runner so a +// slow-but-alive job (one whose handler has not checkpointed recently) is not +// judged stale and re-claimed elsewhere. Fenced on runner_id: returns false when +// the caller no longer owns the job (it was reaped and re-claimed), which the +// runner treats as a signal to cancel its in-flight work. +func (s *RDBConfigStore) HeartbeatSidekiqJob(ctx context.Context, id, runnerID string) (bool, error) { + res := s.DB().WithContext(ctx). + Model(&tables.TableSidekiqJob{}). + Where("id = ? AND runner_id = ? AND status = ?", id, runnerID, tables.SidekiqStatusRunning). + Update("updated_at", time.Now()) + if res.Error != nil { + return false, res.Error } - return nil + return res.RowsAffected == 1, nil } // UpdateSidekiqJobProgress persists a progress checkpoint: it replaces the metadata // blob and bumps the heartbeat (updated_at) so the reaper does not treat the job as -// stale. Called after each processed page. -func (s *RDBConfigStore) UpdateSidekiqJobProgress(ctx context.Context, id, metadata string) error { +// stale. Called after each processed page. Fenced on runner_id so only the current +// owner can advance the job; a stale runner that revives affects 0 rows. +func (s *RDBConfigStore) UpdateSidekiqJobProgress(ctx context.Context, id, runnerID, metadata string) error { res := s.DB().WithContext(ctx). Model(&tables.TableSidekiqJob{}). - Where("id = ?", id). + Where("id = ? AND runner_id = ?", id, runnerID). Updates(map[string]any{ "metadata": metadata, "updated_at": time.Now(), @@ -85,18 +120,22 @@ func (s *RDBConfigStore) UpdateSidekiqJobProgress(ctx context.Context, id, metad return res.Error } if res.RowsAffected == 0 { - return errors.New("sidekiq job not found") + return errors.New("sidekiq job not found or no longer owned by caller") } return nil } // CompleteSidekiqJob marks a job completed, stamps completed_at, and stores the -// final metadata (counts, summary). -func (s *RDBConfigStore) CompleteSidekiqJob(ctx context.Context, id, metadata string) error { +// final metadata (counts, summary). Fenced on runner_id AND status = running so a +// job that was reaped and re-claimed elsewhere is not marked complete by its former +// runner, and — critically — so a job the reaper already flipped to failed (because +// this runner ran past the stale threshold) is not silently resurrected to completed, +// which would mask the staleness signal. +func (s *RDBConfigStore) CompleteSidekiqJob(ctx context.Context, id, runnerID, metadata string) error { now := time.Now() res := s.DB().WithContext(ctx). Model(&tables.TableSidekiqJob{}). - Where("id = ?", id). + Where("id = ? AND runner_id = ? AND status = ?", id, runnerID, tables.SidekiqStatusRunning). Updates(map[string]any{ "status": tables.SidekiqStatusCompleted, "metadata": metadata, @@ -107,14 +146,17 @@ func (s *RDBConfigStore) CompleteSidekiqJob(ctx context.Context, id, metadata st return res.Error } if res.RowsAffected == 0 { - return errors.New("sidekiq job not found") + return errors.New("sidekiq job not found, no longer owned by caller, or no longer running") } return nil } // FailSidekiqJob marks a job failed, records the error, stamps completed_at, and // preserves the latest metadata so a later resume can read the checkpoint cursor. -func (s *RDBConfigStore) FailSidekiqJob(ctx context.Context, id, metadata, lastErr string) error { +// Fenced on runner_id AND status = running so a former runner cannot overwrite a +// re-claimed job's state, and so the execute/panic paths cannot overwrite the +// last_error the reaper already wrote when it failed this job for going stale. +func (s *RDBConfigStore) FailSidekiqJob(ctx context.Context, id, runnerID, metadata, lastErr string) error { now := time.Now() updates := map[string]any{ "status": tables.SidekiqStatusFailed, @@ -127,24 +169,28 @@ func (s *RDBConfigStore) FailSidekiqJob(ctx context.Context, id, metadata, lastE } res := s.DB().WithContext(ctx). Model(&tables.TableSidekiqJob{}). - Where("id = ?", id). + Where("id = ? AND runner_id = ? AND status = ?", id, runnerID, tables.SidekiqStatusRunning). Updates(updates) if res.Error != nil { return res.Error } if res.RowsAffected == 0 { - return errors.New("sidekiq job not found") + return errors.New("sidekiq job not found, no longer owned by caller, or no longer running") } return nil } -// ListIncompleteSidekiqJobs returns jobs that are not in a terminal state -// (pending or running). Used by startup recovery to resume work that was -// interrupted by a restart or crash. -func (s *RDBConfigStore) ListIncompleteSidekiqJobs(ctx context.Context) ([]tables.TableSidekiqJob, error) { +// ListClaimableSidekiqJobs returns jobs eligible to be picked up: those that are +// pending (never started), or running but stale (heartbeat older than staleBefore, +// i.e. their owner is presumed dead). Ordered oldest-first. The dispatcher scans +// this list and attempts to claim each; the atomic ClaimSidekiqJob decides the one +// winner, so listing on every node is safe and needs no cross-node coordination. +func (s *RDBConfigStore) ListClaimableSidekiqJobs(ctx context.Context, staleBefore time.Time) ([]tables.TableSidekiqJob, error) { var jobs []tables.TableSidekiqJob err := s.DB().WithContext(ctx). - Where("status IN ?", []string{tables.SidekiqStatusPending, tables.SidekiqStatusRunning}). + Where("status = ? OR (status = ? AND updated_at < ?)", + tables.SidekiqStatusPending, + tables.SidekiqStatusRunning, staleBefore). Order("created_at ASC"). Find(&jobs).Error if err != nil { diff --git a/framework/configstore/sidekiq_test.go b/framework/configstore/sidekiq_test.go new file mode 100644 index 00000000000..ad58eac08c2 --- /dev/null +++ b/framework/configstore/sidekiq_test.go @@ -0,0 +1,364 @@ +package configstore + +import ( + "context" + "testing" + "time" + + "github.com/maximhq/bifrost/framework/configstore/tables" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// setupSidekiqTestStore extends the base test store with the sidekiq table. +func setupSidekiqTestStore(t *testing.T) *RDBConfigStore { + store := setupRDBTestStore(t) + require.NoError(t, store.DB().AutoMigrate(&tables.TableSidekiqJob{}), "migrate sidekiq table") + return store +} + +// setUpdatedAt forces a job's updated_at to a fixed time so staleness can be +// exercised deterministically without sleeping. +func setUpdatedAt(t *testing.T, store *RDBConfigStore, id string, ts time.Time) { + t.Helper() + require.NoError(t, store.DB().Model(&tables.TableSidekiqJob{}). + Where("id = ?", id).Update("updated_at", ts).Error) +} + +func getJob(t *testing.T, store *RDBConfigStore, id string) *tables.TableSidekiqJob { + t.Helper() + job, err := store.GetSidekiqJob(context.Background(), id) + require.NoError(t, err) + require.NotNil(t, job, "job %s should exist", id) + return job +} + +func TestCreateSidekiqJobValidation(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + + assert.Error(t, store.CreateSidekiqJob(ctx, nil), "nil job") + assert.Error(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{Kind: "k"}), "empty id") + assert.Error(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "x"}), "empty kind") +} + +func TestCreateSidekiqJobDefaults(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + + job := &tables.TableSidekiqJob{ID: "j1", Kind: "sync"} + require.NoError(t, store.CreateSidekiqJob(ctx, job)) + + got := getJob(t, store, "j1") + assert.Equal(t, tables.SidekiqStatusPending, got.Status, "status defaults to pending") + assert.Equal(t, "{}", got.Metadata, "metadata defaults to {}") + assert.Equal(t, 0, got.Attempts) + assert.False(t, got.CreatedAt.IsZero(), "created_at stamped") + assert.False(t, got.UpdatedAt.IsZero(), "updated_at stamped") + assert.Nil(t, got.StartedAt, "started_at nil until claimed") +} + +func TestCreateSidekiqJobHonoursExplicitFields(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + + job := &tables.TableSidekiqJob{ID: "j1", Kind: "sync", Status: tables.SidekiqStatusRunning, Metadata: `{"cursor":5}`} + require.NoError(t, store.CreateSidekiqJob(ctx, job)) + + got := getJob(t, store, "j1") + assert.Equal(t, tables.SidekiqStatusRunning, got.Status) + assert.Equal(t, `{"cursor":5}`, got.Metadata) +} + +func TestGetSidekiqJobMissingReturnsNil(t *testing.T) { + store := setupSidekiqTestStore(t) + job, err := store.GetSidekiqJob(context.Background(), "nope") + require.NoError(t, err) + assert.Nil(t, job) +} + +func TestClaimSidekiqJobPending(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "j1", Kind: "k"})) + + ok, err := store.ClaimSidekiqJob(ctx, "j1", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + assert.True(t, ok, "pending job is claimable") + + got := getJob(t, store, "j1") + assert.Equal(t, tables.SidekiqStatusRunning, got.Status) + assert.Equal(t, "owner-A", got.RunnerID) + assert.Equal(t, 1, got.Attempts, "claim increments attempts") + require.NotNil(t, got.StartedAt, "started_at set on first claim") +} + +func TestClaimSidekiqJobFreshRunningRejected(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "j1", Kind: "k"})) + + ok, err := store.ClaimSidekiqJob(ctx, "j1", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + require.True(t, ok) + + // A second owner cannot claim while the heartbeat is fresh. + ok2, err := store.ClaimSidekiqJob(ctx, "j1", "owner-B", time.Now().Add(-time.Minute)) + require.NoError(t, err) + assert.False(t, ok2, "fresh running job is not re-claimable") + assert.Equal(t, "owner-A", getJob(t, store, "j1").RunnerID, "owner unchanged") +} + +func TestClaimSidekiqJobStaleReclaimPreservesStartedAt(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "j1", Kind: "k"})) + + ok, err := store.ClaimSidekiqJob(ctx, "j1", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + require.True(t, ok) + firstStart := *getJob(t, store, "j1").StartedAt + + // Age the heartbeat past the stale window. + setUpdatedAt(t, store, "j1", time.Now().Add(-30*time.Minute)) + + ok2, err := store.ClaimSidekiqJob(ctx, "j1", "owner-B", time.Now().Add(-15*time.Minute)) + require.NoError(t, err) + assert.True(t, ok2, "stale running job is re-claimable") + + got := getJob(t, store, "j1") + assert.Equal(t, "owner-B", got.RunnerID, "ownership transferred") + assert.Equal(t, 2, got.Attempts, "re-claim increments attempts") + require.NotNil(t, got.StartedAt) + assert.WithinDuration(t, firstStart, *got.StartedAt, time.Millisecond, "started_at preserved across resume") +} + +func TestClaimSidekiqJobMissing(t *testing.T) { + store := setupSidekiqTestStore(t) + ok, err := store.ClaimSidekiqJob(context.Background(), "ghost", "owner-A", time.Now()) + require.NoError(t, err) + assert.False(t, ok) +} + +func TestHeartbeatSidekiqJob(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "j1", Kind: "k"})) + _, err := store.ClaimSidekiqJob(ctx, "j1", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + + setUpdatedAt(t, store, "j1", time.Now().Add(-5*time.Minute)) + before := getJob(t, store, "j1").UpdatedAt + + ok, err := store.HeartbeatSidekiqJob(ctx, "j1", "owner-A") + require.NoError(t, err) + assert.True(t, ok, "owner heartbeat succeeds") + assert.True(t, getJob(t, store, "j1").UpdatedAt.After(before), "heartbeat bumps updated_at") + + // Wrong owner cannot heartbeat. + ok, err = store.HeartbeatSidekiqJob(ctx, "j1", "owner-B") + require.NoError(t, err) + assert.False(t, ok, "non-owner heartbeat rejected") +} + +func TestHeartbeatSidekiqJobRejectedWhenNotRunning(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "j1", Kind: "k"})) + _, err := store.ClaimSidekiqJob(ctx, "j1", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + require.NoError(t, store.CompleteSidekiqJob(ctx, "j1", "owner-A", "{}")) + + ok, err := store.HeartbeatSidekiqJob(ctx, "j1", "owner-A") + require.NoError(t, err) + assert.False(t, ok, "heartbeat on a completed job is rejected") +} + +func TestUpdateSidekiqJobProgress(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "j1", Kind: "k"})) + _, err := store.ClaimSidekiqJob(ctx, "j1", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + + require.NoError(t, store.UpdateSidekiqJobProgress(ctx, "j1", "owner-A", `{"cursor":42}`)) + assert.Equal(t, `{"cursor":42}`, getJob(t, store, "j1").Metadata) + + // A stale/non-owner cannot advance progress. + assert.Error(t, store.UpdateSidekiqJobProgress(ctx, "j1", "owner-B", `{"cursor":99}`)) + assert.Equal(t, `{"cursor":42}`, getJob(t, store, "j1").Metadata, "metadata unchanged by non-owner") +} + +func TestCompleteSidekiqJob(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "j1", Kind: "k"})) + _, err := store.ClaimSidekiqJob(ctx, "j1", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + + require.NoError(t, store.CompleteSidekiqJob(ctx, "j1", "owner-A", `{"done":true}`)) + got := getJob(t, store, "j1") + assert.Equal(t, tables.SidekiqStatusCompleted, got.Status) + assert.Equal(t, `{"done":true}`, got.Metadata) + require.NotNil(t, got.CompletedAt) +} + +func TestCompleteSidekiqJobRejectsNonOwner(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "j1", Kind: "k"})) + _, err := store.ClaimSidekiqJob(ctx, "j1", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + + assert.Error(t, store.CompleteSidekiqJob(ctx, "j1", "owner-B", "{}"), "non-owner cannot complete") + assert.Equal(t, tables.SidekiqStatusRunning, getJob(t, store, "j1").Status) +} + +// TestCompleteSidekiqJobRejectsReapedJob covers the status guard: once the reaper +// has flipped a running job to failed, its former owner must not resurrect it. +func TestCompleteSidekiqJobRejectsReapedJob(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "j1", Kind: "k"})) + _, err := store.ClaimSidekiqJob(ctx, "j1", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + + // Reaper fails the job (owner_id left intact) while owner-A is still running. + setUpdatedAt(t, store, "j1", time.Now().Add(-30*time.Minute)) + n, err := store.MarkStaleSidekiqJobsFailed(ctx, time.Now().Add(-15*time.Minute)) + require.NoError(t, err) + require.Equal(t, int64(1), n) + + assert.Error(t, store.CompleteSidekiqJob(ctx, "j1", "owner-A", "{}"), + "complete must fail once the job is no longer running") + got := getJob(t, store, "j1") + assert.Equal(t, tables.SidekiqStatusFailed, got.Status, "reaped failure must not be resurrected") +} + +func TestFailSidekiqJob(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "j1", Kind: "k"})) + _, err := store.ClaimSidekiqJob(ctx, "j1", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + + require.NoError(t, store.FailSidekiqJob(ctx, "j1", "owner-A", `{"cursor":7}`, "boom")) + got := getJob(t, store, "j1") + assert.Equal(t, tables.SidekiqStatusFailed, got.Status) + assert.Equal(t, "boom", got.LastError) + assert.Equal(t, `{"cursor":7}`, got.Metadata, "checkpoint metadata preserved for resume") + require.NotNil(t, got.CompletedAt) +} + +func TestFailSidekiqJobEmptyMetadataPreservesExisting(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "j1", Kind: "k", Metadata: `{"cursor":3}`})) + _, err := store.ClaimSidekiqJob(ctx, "j1", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + + require.NoError(t, store.FailSidekiqJob(ctx, "j1", "owner-A", "", "panic")) + got := getJob(t, store, "j1") + assert.Equal(t, "panic", got.LastError) + assert.Equal(t, `{"cursor":3}`, got.Metadata, "empty metadata does not clobber last checkpoint") +} + +// TestFailSidekiqJobRejectsReapedJob covers the status guard on the fail path: the +// panic/execute path must not overwrite a last_error the reaper already wrote. +func TestFailSidekiqJobRejectsReapedJob(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "j1", Kind: "k"})) + _, err := store.ClaimSidekiqJob(ctx, "j1", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + + setUpdatedAt(t, store, "j1", time.Now().Add(-30*time.Minute)) + _, err = store.MarkStaleSidekiqJobsFailed(ctx, time.Now().Add(-15*time.Minute)) + require.NoError(t, err) + reapedErr := getJob(t, store, "j1").LastError + + assert.Error(t, store.FailSidekiqJob(ctx, "j1", "owner-A", "", "late handler error")) + assert.Equal(t, reapedErr, getJob(t, store, "j1").LastError, "reaper's last_error preserved") +} + +func TestListClaimableSidekiqJobs(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + + // pending → claimable + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "pending", Kind: "k"})) + + // running + fresh → not claimable + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "fresh", Kind: "k"})) + _, err := store.ClaimSidekiqJob(ctx, "fresh", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + + // running + stale → claimable + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "stale", Kind: "k"})) + _, err = store.ClaimSidekiqJob(ctx, "stale", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + setUpdatedAt(t, store, "stale", time.Now().Add(-30*time.Minute)) + + // completed → not claimable + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "done", Kind: "k"})) + _, err = store.ClaimSidekiqJob(ctx, "done", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + require.NoError(t, store.CompleteSidekiqJob(ctx, "done", "owner-A", "{}")) + + jobs, err := store.ListClaimableSidekiqJobs(ctx, time.Now().Add(-15*time.Minute)) + require.NoError(t, err) + + ids := map[string]bool{} + for _, j := range jobs { + ids[j.ID] = true + } + assert.True(t, ids["pending"], "pending is claimable") + assert.True(t, ids["stale"], "stale running is claimable") + assert.False(t, ids["fresh"], "fresh running is not claimable") + assert.False(t, ids["done"], "completed is not claimable") +} + +func TestListClaimableSidekiqJobsOrderedOldestFirst(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "a", Kind: "k"})) + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "b", Kind: "k"})) + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "c", Kind: "k"})) + // Force a deterministic created_at ordering: c, a, b. + require.NoError(t, store.DB().Model(&tables.TableSidekiqJob{}).Where("id = ?", "c").Update("created_at", time.Now().Add(-3*time.Hour)).Error) + require.NoError(t, store.DB().Model(&tables.TableSidekiqJob{}).Where("id = ?", "a").Update("created_at", time.Now().Add(-2*time.Hour)).Error) + require.NoError(t, store.DB().Model(&tables.TableSidekiqJob{}).Where("id = ?", "b").Update("created_at", time.Now().Add(-1*time.Hour)).Error) + + jobs, err := store.ListClaimableSidekiqJobs(ctx, time.Now().Add(-15*time.Minute)) + require.NoError(t, err) + require.Len(t, jobs, 3) + assert.Equal(t, []string{"c", "a", "b"}, []string{jobs[0].ID, jobs[1].ID, jobs[2].ID}) +} + +func TestMarkStaleSidekiqJobsFailed(t *testing.T) { + store := setupSidekiqTestStore(t) + ctx := context.Background() + + // stale running → reaped + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "stale", Kind: "k"})) + _, err := store.ClaimSidekiqJob(ctx, "stale", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + setUpdatedAt(t, store, "stale", time.Now().Add(-30*time.Minute)) + + // fresh running → left alone + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "fresh", Kind: "k"})) + _, err = store.ClaimSidekiqJob(ctx, "fresh", "owner-A", time.Now().Add(-time.Minute)) + require.NoError(t, err) + + // pending → left alone (not running) + require.NoError(t, store.CreateSidekiqJob(ctx, &tables.TableSidekiqJob{ID: "pending", Kind: "k"})) + + n, err := store.MarkStaleSidekiqJobsFailed(ctx, time.Now().Add(-15*time.Minute)) + require.NoError(t, err) + assert.Equal(t, int64(1), n, "only the stale running job is reaped") + + assert.Equal(t, tables.SidekiqStatusFailed, getJob(t, store, "stale").Status) + assert.NotEmpty(t, getJob(t, store, "stale").LastError) + assert.Equal(t, tables.SidekiqStatusRunning, getJob(t, store, "fresh").Status) + assert.Equal(t, tables.SidekiqStatusPending, getJob(t, store, "pending").Status) +} diff --git a/framework/configstore/store.go b/framework/configstore/store.go index 7069b656a14..81401aeb689 100644 --- a/framework/configstore/store.go +++ b/framework/configstore/store.go @@ -666,11 +666,12 @@ type ConfigStore interface { // Sidekiq - generic durable background jobs CreateSidekiqJob(ctx context.Context, job *tables.TableSidekiqJob) error GetSidekiqJob(ctx context.Context, id string) (*tables.TableSidekiqJob, error) - MarkSidekiqJobRunning(ctx context.Context, id string) error - UpdateSidekiqJobProgress(ctx context.Context, id, metadata string) error - CompleteSidekiqJob(ctx context.Context, id, metadata string) error - FailSidekiqJob(ctx context.Context, id, metadata, lastErr string) error - ListIncompleteSidekiqJobs(ctx context.Context) ([]tables.TableSidekiqJob, error) + ClaimSidekiqJob(ctx context.Context, id, runnerID string, staleBefore time.Time) (bool, error) + HeartbeatSidekiqJob(ctx context.Context, id, runnerID string) (bool, error) + UpdateSidekiqJobProgress(ctx context.Context, id, runnerID, metadata string) error + CompleteSidekiqJob(ctx context.Context, id, runnerID, metadata string) error + FailSidekiqJob(ctx context.Context, id, runnerID, metadata, lastErr string) error + ListClaimableSidekiqJobs(ctx context.Context, staleBefore time.Time) ([]tables.TableSidekiqJob, error) MarkStaleSidekiqJobsFailed(ctx context.Context, staleBefore time.Time) (int64, error) // DB returns the underlying database connection. diff --git a/framework/configstore/tables/sidekiq.go b/framework/configstore/tables/sidekiq.go index bf114a77620..cc785d2f52d 100644 --- a/framework/configstore/tables/sidekiq.go +++ b/framework/configstore/tables/sidekiq.go @@ -22,6 +22,11 @@ type TableSidekiqJob struct { ID string `gorm:"column:id;primaryKey;type:text" json:"id"` Kind string `gorm:"column:kind;not null;type:text;index:idx_sidekiq_status_updated,priority:3" json:"kind"` Status string `gorm:"column:status;not null;default:pending;type:text;index:idx_sidekiq_status_updated,priority:1" json:"status"` + // RunnerID identifies the runner process that currently owns (claimed) this job. + // Empty in OSS (single-node) mode; set to the node ID in enterprise cluster mode. + // Progress/complete/fail/heartbeat writes are fenced on this value so a revived + // stale node cannot stomp a job another node has re-claimed. + RunnerID string `gorm:"column:runner_id;type:text;index" json:"runner_id,omitempty"` Metadata string `gorm:"column:metadata;type:text;default:'{}'" json:"metadata"` Attempts int `gorm:"column:attempts;not null;default:0" json:"attempts"` LastError string `gorm:"column:last_error;type:text" json:"last_error,omitempty"` diff --git a/framework/sidekiq/sidekiq.go b/framework/sidekiq/sidekiq.go index d94fd44c73f..96492878708 100644 --- a/framework/sidekiq/sidekiq.go +++ b/framework/sidekiq/sidekiq.go @@ -4,80 +4,102 @@ import ( "context" "fmt" "sync" + "sync/atomic" "time" "github.com/maximhq/bifrost/core/schemas" "github.com/maximhq/bifrost/framework/configstore/tables" ) -// Store is the subset of the configstore the runner needs. Keeping it narrow lets -// the runner be tested with a fake and avoids a hard dependency on the full store. +// Store is the narrow subset of configstore the runner needs. type Store interface { CreateSidekiqJob(ctx context.Context, job *tables.TableSidekiqJob) error GetSidekiqJob(ctx context.Context, id string) (*tables.TableSidekiqJob, error) - MarkSidekiqJobRunning(ctx context.Context, id string) error - UpdateSidekiqJobProgress(ctx context.Context, id, metadata string) error - CompleteSidekiqJob(ctx context.Context, id, metadata string) error - FailSidekiqJob(ctx context.Context, id, metadata, lastErr string) error - ListIncompleteSidekiqJobs(ctx context.Context) ([]tables.TableSidekiqJob, error) - MarkStaleSidekiqJobsFailed(ctx context.Context, staleBefore time.Time) (int64, error) + // ClaimSidekiqJob atomically claims a job for runnerID; returns true only for the winner. + ClaimSidekiqJob(ctx context.Context, id, runnerID string, staleBefore time.Time) (bool, error) + // HeartbeatSidekiqJob bumps updated_at for a job still owned by runnerID; returns false on lost ownership. + HeartbeatSidekiqJob(ctx context.Context, id, runnerID string) (bool, error) + UpdateSidekiqJobProgress(ctx context.Context, id, runnerID, metadata string) error + CompleteSidekiqJob(ctx context.Context, id, runnerID, metadata string) error + FailSidekiqJob(ctx context.Context, id, runnerID, metadata, lastErr string) error + // ListClaimableSidekiqJobs returns pending jobs and running jobs whose heartbeat is older than staleBefore. + ListClaimableSidekiqJobs(ctx context.Context, staleBefore time.Time) ([]tables.TableSidekiqJob, error) } -// ProgressFunc persists a checkpoint: it replaces the job's metadata blob and -// bumps the heartbeat. Handlers call it after each unit of work (e.g. each page). +// ProgressFunc persists a checkpoint and bumps the heartbeat. Handlers call it after each unit of work. type ProgressFunc func(metadata string) error -// HandlerFunc processes one job. It is given the job (read its Metadata for the -// resume cursor) and a progress callback to checkpoint after each unit of work. -// It returns the final metadata to persist and an error. A nil error completes the -// job; a non-nil error fails it (the returned metadata is still stored so a later -// resume can read the last cursor). +// HandlerFunc processes one job. Receives the job row (read Metadata for the resume cursor) and a +// progress callback. Returns final metadata and an error. Nil error completes the job; non-nil fails it. +// The context is cancelled if this node loses ownership or on shutdown. type HandlerFunc func(ctx context.Context, job tables.TableSidekiqJob, progress ProgressFunc) (finalMetadata string, err error) const ( - ReaperInterval = 1 * time.Minute - StaleAfter = 15 * time.Minute + DispatchInterval = 30 * time.Second + HeartbeatInterval = 1 * time.Minute + StaleAfter = 15 * time.Minute + // MaxAttempts caps re-claims before a job is permanently failed, preventing poison-job loops. + MaxAttempts = 5 ) -// Runner owns the handler registry and the lifecycle of job goroutines. +// Runner owns the handler registry and job goroutine lifecycle. type Runner struct { store Store logger schemas.Logger handlers map[string]HandlerFunc mu sync.RWMutex + // runnerID fences job mutations to the node that claimed the job. + // Empty string disables the stale window so crashed jobs are immediately re-claimable. + runnerID string + staleAfter atomic.Int64 // nanoseconds + + heartbeatInterval time.Duration + + // inflight tracks job IDs being processed on this node to prevent duplicate goroutines. + inflightMu sync.Mutex + inflight map[string]struct{} + baseCtx context.Context cancel context.CancelFunc sem chan struct{} wg sync.WaitGroup } -// New creates a Runner. maxConcurrent bounds how many job goroutines run at once -// (<=0 defaults to 4). Jobs run on a background context derived here, never on a -// request context, so they outlive the HTTP request that enqueued them. -func New(store Store, logger schemas.Logger, maxConcurrent int) *Runner { +// New creates a Runner. maxConcurrent bounds simultaneous job goroutines (<=0 defaults to 4). +// Pass the node ID as runnerID in cluster mode. Pass "" to make any running job immediately +// re-claimable on restart (no stale window). +func New(store Store, logger schemas.Logger, maxConcurrent int, runnerID string) *Runner { if maxConcurrent <= 0 { maxConcurrent = 4 } ctx, cancel := context.WithCancel(context.Background()) - return &Runner{ - store: store, - logger: logger, - handlers: make(map[string]HandlerFunc), - baseCtx: ctx, - cancel: cancel, - sem: make(chan struct{}, maxConcurrent), + r := &Runner{ + store: store, + logger: logger, + handlers: make(map[string]HandlerFunc), + runnerID: runnerID, + heartbeatInterval: HeartbeatInterval, + inflight: make(map[string]struct{}), + baseCtx: ctx, + cancel: cancel, + sem: make(chan struct{}, maxConcurrent), } + if runnerID == "" { + r.staleAfter.Store(0) + } else { + r.staleAfter.Store(int64(StaleAfter)) + } + return r } -// Register binds a handler to a job kind. Call during startup, before enqueuing. +// Register binds a handler to a job kind. Call before enqueuing. func (r *Runner) Register(kind string, fn HandlerFunc) { r.mu.Lock() defer r.mu.Unlock() r.handlers[kind] = fn } -// handlerFor returns the registered handler for a kind, if any. func (r *Runner) handlerFor(kind string) (HandlerFunc, bool) { r.mu.RLock() defer r.mu.RUnlock() @@ -85,10 +107,8 @@ func (r *Runner) handlerFor(kind string) (HandlerFunc, bool) { return fn, ok } -// Enqueue persists a new pending job and starts its goroutine. The caller supplies -// the id (also usable as a UI operation id), the kind, and the initial metadata -// JSON. It returns once the row is committed, so the HTTP handler can respond -// immediately while processing continues in the background. +// Enqueue persists a new pending job and starts it as soon as a concurrency slot is free. +// Returns once the DB row is committed so the caller can respond immediately. func (r *Runner) Enqueue(ctx context.Context, id, kind, metadata string) error { if _, ok := r.handlerFor(kind); !ok { return fmt.Errorf("sidekiq: no handler registered for kind %q", kind) @@ -106,11 +126,36 @@ func (r *Runner) Enqueue(ctx context.Context, id, kind, metadata string) error { return nil } -// spawn runs a job in its own goroutine, bounded by the concurrency semaphore. +func (r *Runner) staleBefore() time.Time { + return time.Now().Add(-time.Duration(r.staleAfter.Load())) +} + +func (r *Runner) tryMarkInflight(id string) bool { + r.inflightMu.Lock() + defer r.inflightMu.Unlock() + if _, ok := r.inflight[id]; ok { + return false + } + r.inflight[id] = struct{}{} + return true +} + +func (r *Runner) clearInflight(id string) { + r.inflightMu.Lock() + delete(r.inflight, id) + r.inflightMu.Unlock() +} + +// spawn runs a job in its own goroutine, blocking until a concurrency slot is free. +// Used by Enqueue so an explicitly triggered job starts as soon as possible. func (r *Runner) spawn(job tables.TableSidekiqJob) { + if !r.tryMarkInflight(job.ID) { + return + } r.wg.Add(1) go func() { defer r.wg.Done() + defer r.clearInflight(job.ID) select { case r.sem <- struct{}{}: case <-r.baseCtx.Done(): @@ -121,94 +166,128 @@ func (r *Runner) spawn(job tables.TableSidekiqJob) { }() } -// execute marks the job running, invokes its handler, and records the terminal -// state. A panic in the handler is recovered and recorded as a failure so one bad -// job cannot crash the process. +// execute claims the job and runs its handler. Uses a non-blocking claim so multiple nodes +// racing for the same job are safe: only the winner (RowsAffected == 1) proceeds. +// Panics are recovered and recorded as failures. func (r *Runner) execute(job tables.TableSidekiqJob) { + fn, ok := r.handlerFor(job.Kind) + if !ok { + r.logger.Warn("sidekiq: no handler for kind %s, skipping job %s", job.Kind, job.ID) + return + } + + jobCtx, cancel := context.WithCancel(r.baseCtx) + defer cancel() + defer func() { if rec := recover(); rec != nil { r.logger.Error("sidekiq: job %s (%s) panicked: %v", job.ID, job.Kind, rec) - if err := r.store.FailSidekiqJob(r.baseCtx, job.ID, "", fmt.Sprintf("panic: %v", rec)); err != nil { + if err := r.store.FailSidekiqJob(r.baseCtx, job.ID, r.runnerID, "", fmt.Sprintf("panic: %v", rec)); err != nil { r.logger.Error("sidekiq: failed to mark panicked job %s failed: %v", job.ID, err) } } }() - fn, ok := r.handlerFor(job.Kind) - if !ok { - if err := r.store.FailSidekiqJob(r.baseCtx, job.ID, "", "no handler registered for kind "+job.Kind); err != nil { - r.logger.Error("sidekiq: failed to fail unhandled job %s: %v", job.ID, err) - } + claimed, err := r.store.ClaimSidekiqJob(r.baseCtx, job.ID, r.runnerID, r.staleBefore()) + if err != nil { + r.logger.Error("sidekiq: failed to claim job %s: %v", job.ID, err) + return + } + if !claimed { + return + } + + // Re-fetch so the handler sees the latest metadata/cursor, not the snapshot from dispatch time. + fresh, err := r.store.GetSidekiqJob(r.baseCtx, job.ID) + if err != nil { + r.logger.Error("sidekiq: failed to fetch job %s after claim: %v", job.ID, err) + return + } + if fresh == nil { + r.logger.Error("sidekiq: job %s vanished after claim", job.ID) return } + job = *fresh - if err := r.store.MarkSidekiqJobRunning(r.baseCtx, job.ID); err != nil { - r.logger.Error("sidekiq: failed to mark job %s running: %v", job.ID, err) - if ferr := r.store.FailSidekiqJob(r.baseCtx, job.ID, job.Metadata, err.Error()); ferr != nil { - r.logger.Error("sidekiq: failed to fail job %s after running-mark failure: %v", job.ID, ferr) + if job.Attempts >= MaxAttempts { + r.logger.Warn("sidekiq: job %s (%s) exceeded max attempts (%d)", job.ID, job.Kind, MaxAttempts) + if ferr := r.store.FailSidekiqJob(r.baseCtx, job.ID, r.runnerID, job.Metadata, fmt.Sprintf("exceeded max attempts (%d)", MaxAttempts)); ferr != nil { + r.logger.Error("sidekiq: failed to fail exhausted job %s: %v", job.ID, ferr) } return } + stopHeartbeat := r.startHeartbeat(jobCtx, cancel, job.ID) + defer stopHeartbeat() + progress := func(metadata string) error { - return r.store.UpdateSidekiqJobProgress(r.baseCtx, job.ID, metadata) + return r.store.UpdateSidekiqJobProgress(r.baseCtx, job.ID, r.runnerID, metadata) } - finalMetadata, err := fn(r.baseCtx, job, progress) + finalMetadata, err := fn(jobCtx, job, progress) if err != nil { r.logger.Error("sidekiq: job %s (%s) failed: %v", job.ID, job.Kind, err) - if ferr := r.store.FailSidekiqJob(r.baseCtx, job.ID, finalMetadata, err.Error()); ferr != nil { + if ferr := r.store.FailSidekiqJob(r.baseCtx, job.ID, r.runnerID, finalMetadata, err.Error()); ferr != nil { r.logger.Error("sidekiq: failed to mark job %s failed: %v", job.ID, ferr) } return } - if cerr := r.store.CompleteSidekiqJob(r.baseCtx, job.ID, finalMetadata); cerr != nil { + if cerr := r.store.CompleteSidekiqJob(r.baseCtx, job.ID, r.runnerID, finalMetadata); cerr != nil { r.logger.Error("sidekiq: failed to mark job %s completed: %v", job.ID, cerr) } } -// RecoverIncomplete re-runs jobs left pending or running by a previous process -// (a restart or crash). Each handler resumes from the cursor stored in its -// metadata; because per-item work is idempotent, reprocessing the in-flight unit -// is safe. In a multi-node cluster this may double-run a job across nodes, which -// idempotency tolerates; it does not de-duplicate work across nodes by design, -// matching the choice to keep the runner simple (no leader election). -func (r *Runner) RecoverIncomplete(ctx context.Context) error { - jobs, err := r.store.ListIncompleteSidekiqJobs(ctx) - if err != nil { - return err - } - for _, job := range jobs { - if _, ok := r.handlerFor(job.Kind); !ok { - r.logger.Warn("sidekiq: skipping recovery of job %s, no handler for kind %s", job.ID, job.Kind) - continue +// startHeartbeat periodically bumps updated_at so the job isn't judged stale. +// Cancels jobCtx if ownership is lost (job reaped and re-claimed elsewhere). +func (r *Runner) startHeartbeat(jobCtx context.Context, cancel context.CancelFunc, id string) (stop func()) { + ticker := time.NewTicker(r.heartbeatInterval) + done := make(chan struct{}) + r.wg.Add(1) + go func() { + defer r.wg.Done() + defer ticker.Stop() + for { + select { + case <-done: + return + case <-jobCtx.Done(): + return + case <-ticker.C: + ok, err := r.store.HeartbeatSidekiqJob(r.baseCtx, id, r.runnerID) + if err != nil { + r.logger.Error("sidekiq: heartbeat for job %s failed: %v", id, err) + continue + } + if !ok { + r.logger.Warn("sidekiq: lost ownership of job %s, cancelling", id) + cancel() + return + } + } } - r.logger.Info("sidekiq: recovering incomplete job %s (%s)", job.ID, job.Kind) - r.spawn(job) - } - return nil + }() + var once sync.Once + return func() { once.Do(func() { close(done) }) } } -// StartReaper periodically marks running jobs whose heartbeat is older than -// staleAfter as failed, catching goroutines or nodes that died without recording a -// terminal state. It returns a stop function. Both interval and staleAfter must be -// positive; staleAfter should comfortably exceed the handler's per-checkpoint time. -func (r *Runner) StartReaper(interval, staleAfter time.Duration) (stop func()) { - // Guard against invalid durations: time.NewTicker panics on interval <= 0, and - // a non-positive staleAfter would make every running job look stale. Fall back - // to the package defaults rather than crashing or reaping live jobs. +// StartDispatcher scans for claimable jobs on an interval. Uses non-blocking semaphore +// acquisition so it never spawns more goroutines than available concurrency slots — +// remaining jobs are left for the next tick. Runs one scan immediately on start. +func (r *Runner) StartDispatcher(interval, staleAfter time.Duration) (stop func()) { if interval <= 0 { - interval = ReaperInterval + interval = DispatchInterval } if staleAfter <= 0 { staleAfter = StaleAfter } + r.staleAfter.Store(int64(staleAfter)) ticker := time.NewTicker(interval) done := make(chan struct{}) r.wg.Add(1) go func() { defer r.wg.Done() defer ticker.Stop() + r.dispatchOnce() for { select { case <-done: @@ -216,14 +295,7 @@ func (r *Runner) StartReaper(interval, staleAfter time.Duration) (stop func()) { case <-r.baseCtx.Done(): return case <-ticker.C: - n, err := r.store.MarkStaleSidekiqJobsFailed(r.baseCtx, time.Now().Add(-staleAfter)) - if err != nil { - r.logger.Error("sidekiq: reaper failed: %v", err) - continue - } - if n > 0 { - r.logger.Warn("sidekiq: reaper marked %d stale job(s) as failed", n) - } + r.dispatchOnce() } } }() @@ -231,9 +303,42 @@ func (r *Runner) StartReaper(interval, staleAfter time.Duration) (stop func()) { return func() { once.Do(func() { close(done) }) } } -// Shutdown cancels the background context and waits for in-flight goroutines to -// return. In-flight jobs observe baseCtx cancellation and stop at their next -// checkpoint, leaving a resumable cursor in metadata. +// dispatchOnce lists claimable jobs and spawns goroutines only for slots that are +// immediately available. Stops as soon as the semaphore is full so 1000 pending +// jobs never produce 1000 parked goroutines — the remainder are picked up next tick. +func (r *Runner) dispatchOnce() { + jobs, err := r.store.ListClaimableSidekiqJobs(r.baseCtx, r.staleBefore()) + if err != nil { + r.logger.Error("sidekiq: dispatcher failed to list claimable jobs: %v", err) + return + } + for _, job := range jobs { + if _, ok := r.handlerFor(job.Kind); !ok { + r.logger.Warn("sidekiq: skipping job %s, no handler for kind %s", job.ID, job.Kind) + continue + } + if !r.tryMarkInflight(job.ID) { + continue + } + select { + case r.sem <- struct{}{}: + default: + // Semaphore full; release inflight and stop — next tick picks up the rest. + r.clearInflight(job.ID) + return + } + job := job + r.wg.Add(1) + go func() { + defer r.wg.Done() + defer r.clearInflight(job.ID) + defer func() { <-r.sem }() + r.execute(job) + }() + } +} + +// Shutdown cancels the background context and waits for in-flight goroutines to return. func (r *Runner) Shutdown() { r.cancel() r.wg.Wait() diff --git a/framework/sidekiq/sidekiq_test.go b/framework/sidekiq/sidekiq_test.go index 4c062311319..cb34afa6ee5 100644 --- a/framework/sidekiq/sidekiq_test.go +++ b/framework/sidekiq/sidekiq_test.go @@ -3,6 +3,7 @@ package sidekiq import ( "context" "errors" + "fmt" "sync" "testing" "time" @@ -10,24 +11,37 @@ import ( bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" "github.com/maximhq/bifrost/framework/configstore/tables" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) -// fakeStore is an in-memory Store for exercising the runner without a database. +// jobState is the fake's in-memory model of one row. +type jobState struct { + kind string + status string + owner string + metadata string + attempts int + updatedAt time.Time +} + +// fakeStore is an in-memory Store for exercising the runner without a database. It +// models the atomic claim and owner fencing so multi-owner behaviour can be tested. type fakeStore struct { mu sync.Mutex + jobs map[string]*jobState created []tables.TableSidekiqJob - running map[string]int + running map[string]int // successful claim count per id progress map[string]string completed map[string]string failedMeta map[string]string failedErr map[string]string - incomplete []tables.TableSidekiqJob - staleCalls int terminal chan string } func newFakeStore() *fakeStore { return &fakeStore{ + jobs: map[string]*jobState{}, running: map[string]int{}, progress: map[string]string{}, completed: map[string]string{}, @@ -37,41 +51,102 @@ func newFakeStore() *fakeStore { } } +// seed inserts a job directly, bypassing Create (used to simulate rows left by a +// previous process). staleAge sets how long ago the last heartbeat was. +func (f *fakeStore) seed(id, kind, status string, staleAge time.Duration) { + f.mu.Lock() + defer f.mu.Unlock() + f.jobs[id] = &jobState{kind: kind, status: status, metadata: "{}", updatedAt: time.Now().Add(-staleAge)} +} + func (f *fakeStore) CreateSidekiqJob(_ context.Context, job *tables.TableSidekiqJob) error { f.mu.Lock() defer f.mu.Unlock() f.created = append(f.created, *job) + f.jobs[job.ID] = &jobState{kind: job.Kind, status: tables.SidekiqStatusPending, metadata: job.Metadata, updatedAt: time.Now()} return nil } -func (f *fakeStore) GetSidekiqJob(_ context.Context, _ string) (*tables.TableSidekiqJob, error) { - return nil, nil +func (f *fakeStore) GetSidekiqJob(_ context.Context, id string) (*tables.TableSidekiqJob, error) { + f.mu.Lock() + defer f.mu.Unlock() + j, ok := f.jobs[id] + if !ok { + return nil, nil + } + return &tables.TableSidekiqJob{ID: id, Kind: j.kind, Status: j.status, RunnerID: j.owner, Metadata: j.metadata, Attempts: j.attempts}, nil } -func (f *fakeStore) MarkSidekiqJobRunning(_ context.Context, id string) error { +func (f *fakeStore) ClaimSidekiqJob(_ context.Context, id, runnerID string, staleBefore time.Time) (bool, error) { f.mu.Lock() defer f.mu.Unlock() + j, ok := f.jobs[id] + if !ok { + return false, nil + } + claimable := j.status == tables.SidekiqStatusPending || + (j.status == tables.SidekiqStatusRunning && j.updatedAt.Before(staleBefore)) + if !claimable { + return false, nil + } + j.status = tables.SidekiqStatusRunning + j.owner = runnerID + j.attempts++ + j.updatedAt = time.Now() f.running[id]++ - return nil + return true, nil +} + +func (f *fakeStore) HeartbeatSidekiqJob(_ context.Context, id, runnerID string) (bool, error) { + f.mu.Lock() + defer f.mu.Unlock() + j, ok := f.jobs[id] + if !ok || j.owner != runnerID || j.status != tables.SidekiqStatusRunning { + return false, nil + } + j.updatedAt = time.Now() + return true, nil } -func (f *fakeStore) UpdateSidekiqJobProgress(_ context.Context, id, metadata string) error { +func (f *fakeStore) UpdateSidekiqJobProgress(_ context.Context, id, runnerID, metadata string) error { f.mu.Lock() defer f.mu.Unlock() + j, ok := f.jobs[id] + if !ok || j.owner != runnerID { + return errors.New("not owned by caller") + } + j.metadata = metadata + j.updatedAt = time.Now() f.progress[id] = metadata return nil } -func (f *fakeStore) CompleteSidekiqJob(_ context.Context, id, metadata string) error { +func (f *fakeStore) CompleteSidekiqJob(_ context.Context, id, runnerID, metadata string) error { f.mu.Lock() + j, ok := f.jobs[id] + if !ok || j.owner != runnerID || j.status != tables.SidekiqStatusRunning { + f.mu.Unlock() + return errors.New("not owned by caller or no longer running") + } + j.status = tables.SidekiqStatusCompleted + j.metadata = metadata f.completed[id] = metadata f.mu.Unlock() f.terminal <- id return nil } -func (f *fakeStore) FailSidekiqJob(_ context.Context, id, metadata, lastErr string) error { +func (f *fakeStore) FailSidekiqJob(_ context.Context, id, runnerID, metadata, lastErr string) error { f.mu.Lock() + j, ok := f.jobs[id] + if !ok || j.owner != runnerID || j.status != tables.SidekiqStatusRunning { + f.mu.Unlock() + return errors.New("not owned by caller or no longer running") + } + j.status = tables.SidekiqStatusFailed + if metadata != "" { + j.metadata = metadata + } f.failedMeta[id] = metadata f.failedErr[id] = lastErr f.mu.Unlock() @@ -79,15 +154,17 @@ func (f *fakeStore) FailSidekiqJob(_ context.Context, id, metadata, lastErr stri return nil } -func (f *fakeStore) ListIncompleteSidekiqJobs(_ context.Context) ([]tables.TableSidekiqJob, error) { - return f.incomplete, nil -} - -func (f *fakeStore) MarkStaleSidekiqJobsFailed(_ context.Context, _ time.Time) (int64, error) { +func (f *fakeStore) ListClaimableSidekiqJobs(_ context.Context, staleBefore time.Time) ([]tables.TableSidekiqJob, error) { f.mu.Lock() defer f.mu.Unlock() - f.staleCalls++ - return 0, nil + var out []tables.TableSidekiqJob + for id, j := range f.jobs { + if j.status == tables.SidekiqStatusPending || + (j.status == tables.SidekiqStatusRunning && j.updatedAt.Before(staleBefore)) { + out = append(out, tables.TableSidekiqJob{ID: id, Kind: j.kind, Status: j.status, Metadata: j.metadata, Attempts: j.attempts}) + } + } + return out, nil } // waitTerminal blocks until a job reaches a terminal state or the test times out. @@ -103,7 +180,7 @@ func waitTerminal(t *testing.T, f *fakeStore) string { } func testRunner(store Store) *Runner { - return New(store, bifrost.NewDefaultLogger(schemas.LogLevelError), 4) + return New(store, bifrost.NewDefaultLogger(schemas.LogLevelError), 4, "") } func TestEnqueueRunsHandlerAndCompletes(t *testing.T) { @@ -125,7 +202,7 @@ func TestEnqueueRunsHandlerAndCompletes(t *testing.T) { store.mu.Lock() defer store.mu.Unlock() if store.running["job1"] != 1 { - t.Errorf("MarkSidekiqJobRunning called %d times, want 1", store.running["job1"]) + t.Errorf("claimed %d times, want 1", store.running["job1"]) } if store.progress["job1"] != "checkpoint" { t.Errorf("progress = %q, want checkpoint", store.progress["job1"]) @@ -189,46 +266,397 @@ func TestEnqueueUnknownKindErrors(t *testing.T) { } } -func TestRecoverIncompleteResumesJobs(t *testing.T) { +// TestDispatcherRunsClaimableJobs verifies both a pending job and a stale running +// job (orphaned by a dead owner) are picked up and run. +func TestDispatcherRunsClaimableJobs(t *testing.T) { store := newFakeStore() - store.incomplete = []tables.TableSidekiqJob{ - {ID: "r1", Kind: "k", Status: tables.SidekiqStatusRunning, Metadata: "{}"}, - {ID: "r2", Kind: "k", Status: tables.SidekiqStatusPending, Metadata: "{}"}, - } + store.seed("r1", "k", tables.SidekiqStatusRunning, 30*time.Minute) // stale, orphaned + store.seed("r2", "k", tables.SidekiqStatusPending, 0) // fresh pending r := testRunner(store) r.Register("k", func(_ context.Context, job tables.TableSidekiqJob, _ ProgressFunc) (string, error) { return "done", nil }) - if err := r.RecoverIncomplete(context.Background()); err != nil { - t.Fatalf("RecoverIncomplete: %v", err) - } + r.dispatchOnce() + got := map[string]bool{} got[waitTerminal(t, store)] = true got[waitTerminal(t, store)] = true if !got["r1"] || !got["r2"] { - t.Errorf("expected both r1 and r2 to be recovered, got %v", got) + t.Errorf("expected both r1 and r2 to be dispatched, got %v", got) + } +} + +// TestClaimIsExclusiveAcrossOwners verifies only one owner wins a claim while the +// job is running with a fresh heartbeat. +func TestClaimIsExclusiveAcrossOwners(t *testing.T) { + store := newFakeStore() + store.seed("j", "k", tables.SidekiqStatusPending, 0) + staleBefore := time.Now().Add(-StaleAfter) + + ok1, err := store.ClaimSidekiqJob(context.Background(), "j", "runner-A", staleBefore) + if err != nil || !ok1 { + t.Fatalf("first claim should win: ok=%v err=%v", ok1, err) + } + ok2, err := store.ClaimSidekiqJob(context.Background(), "j", "runner-B", staleBefore) + if err != nil || ok2 { + t.Fatalf("second claim should lose while job is running fresh: ok=%v err=%v", ok2, err) } } -func TestReaperInvokesStaleSweep(t *testing.T) { +// TestRunnerRaceSingleWinner runs two runners (distinct owners) against one shared +// store and one pending job; exactly one must run it. +func TestRunnerRaceSingleWinner(t *testing.T) { store := newFakeStore() + store.seed("race", "k", tables.SidekiqStatusPending, 0) + handler := func(_ context.Context, job tables.TableSidekiqJob, _ ProgressFunc) (string, error) { + return "done", nil + } + r1, r2 := testRunner(store), testRunner(store) + r1.Register("k", handler) + r2.Register("k", handler) + + go r1.dispatchOnce() + go r2.dispatchOnce() + + waitTerminal(t, store) + + // No second terminal should arrive. + select { + case id := <-store.terminal: + t.Fatalf("job %s reached a terminal state twice; not a single winner", id) + case <-time.After(200 * time.Millisecond): + } + store.mu.Lock() + defer store.mu.Unlock() + if store.running["race"] != 1 { + t.Errorf("job claimed %d times, want exactly 1 winner", store.running["race"]) + } +} + +// TestMaxAttemptsExhaustedFailsWithoutRunningHandler verifies that a job whose +// pre-claim attempt count has already reached MaxAttempts is failed permanently and +// its handler is never invoked. This guards the poison-job boundary in execute(). +func TestMaxAttemptsExhaustedFailsWithoutRunningHandler(t *testing.T) { + store := newFakeStore() + store.seed("poison", "k", tables.SidekiqStatusPending, 0) + // Simulate a job that has already been claimed MaxAttempts times. + store.mu.Lock() + store.jobs["poison"].attempts = MaxAttempts + store.mu.Unlock() + r := testRunner(store) - stop := r.StartReaper(10*time.Millisecond, time.Millisecond) - defer stop() + handlerCalled := make(chan struct{}, 1) + r.Register("k", func(_ context.Context, job tables.TableSidekiqJob, _ ProgressFunc) (string, error) { + handlerCalled <- struct{}{} + return "done", nil + }) + + r.dispatchOnce() + + if id := waitTerminal(t, store); id != "poison" { + t.Fatalf("terminal id = %q, want poison", id) + } + select { + case <-handlerCalled: + t.Fatal("handler must not run for a job that has exhausted its attempts") + default: + } + + store.mu.Lock() + defer store.mu.Unlock() + if store.jobs["poison"].status != tables.SidekiqStatusFailed { + t.Errorf("status = %q, want failed", store.jobs["poison"].status) + } + if _, ok := store.failedErr["poison"]; !ok { + t.Error("expected a failure to be recorded for the exhausted job") + } + if _, ok := store.completed["poison"]; ok { + t.Error("exhausted job must not be completed") + } +} + +// TestCompleteDoesNotOverwriteReapedFailure verifies the status guard on the +// terminal writes: once the reaper has flipped a running job to failed, a late +// CompleteSidekiqJob from its former owner must not resurrect it to completed. +func TestCompleteDoesNotOverwriteReapedFailure(t *testing.T) { + store := newFakeStore() + store.seed("stale", "k", tables.SidekiqStatusPending, 0) + staleBefore := time.Now().Add(-StaleAfter) + + // Owner claims the job (status -> running). + if ok, err := store.ClaimSidekiqJob(context.Background(), "stale", "runner-A", staleBefore); err != nil || !ok { + t.Fatalf("claim should win: ok=%v err=%v", ok, err) + } + // Reaper fails the job out from under the still-running owner. + store.mu.Lock() + store.jobs["stale"].status = tables.SidekiqStatusFailed + store.mu.Unlock() + + // The former owner finishing late must not overwrite the failed state. + if err := store.CompleteSidekiqJob(context.Background(), "stale", "runner-A", "final"); err == nil { + t.Fatal("CompleteSidekiqJob should fail once the job is no longer running") + } + store.mu.Lock() + defer store.mu.Unlock() + if store.jobs["stale"].status != tables.SidekiqStatusFailed { + t.Errorf("status = %q, want failed (must not be resurrected)", store.jobs["stale"].status) + } +} + +// TestHeartbeatCancelsOnLostOwnership verifies the owning runner cancels its +// in-flight handler when it discovers the job was re-claimed elsewhere. +func TestHeartbeatCancelsOnLostOwnership(t *testing.T) { + store := newFakeStore() + r := testRunner(store) + r.heartbeatInterval = 5 * time.Millisecond // beat frequently for the test + + started := make(chan struct{}) + cancelled := make(chan struct{}) + r.Register("k", func(ctx context.Context, job tables.TableSidekiqJob, _ ProgressFunc) (string, error) { + close(started) + <-ctx.Done() // block until the heartbeat cancels us + close(cancelled) + return "", ctx.Err() + }) + + if err := r.Enqueue(context.Background(), "hb", "k", "{}"); err != nil { + t.Fatalf("Enqueue: %v", err) + } + <-started + + // Steal ownership out from under the runner; the next heartbeat sees the mismatch. + store.mu.Lock() + store.jobs["hb"].owner = "someone-else" + store.mu.Unlock() - deadline := time.After(2 * time.Second) - for { - store.mu.Lock() - n := store.staleCalls - store.mu.Unlock() - if n > 0 { - return + select { + case <-cancelled: + case <-time.After(2 * time.Second): + t.Fatal("handler was not cancelled after losing ownership") + } +} + +// TestConcurrencySemaphoreBound verifies that at most maxConcurrent handlers run at +// the same time. Extra jobs wait behind the semaphore and run once a slot frees. +func TestConcurrencySemaphoreBound(t *testing.T) { + const maxConcurrent = 2 + store := newFakeStore() + r := New(store, testRunner(store).logger, maxConcurrent, "") + + // Gate that keeps the first batch of handlers running until we release them. + gate := make(chan struct{}) + var inFlight sync.WaitGroup + + r.Register("k", func(_ context.Context, _ tables.TableSidekiqJob, _ ProgressFunc) (string, error) { + inFlight.Done() // signal that we entered the handler + <-gate // block until the test releases us + return "", nil + }) + + // Enqueue maxConcurrent jobs and wait until all are inside the handler. + inFlight.Add(maxConcurrent) + for i := range maxConcurrent { + require.NoError(t, r.Enqueue(context.Background(), fmt.Sprintf("job-%d", i), "k", "{}")) + } + inFlight.Wait() // both slots are now busy + + // Enqueue a third job — it should not start while the semaphore is full. + require.NoError(t, r.Enqueue(context.Background(), "job-extra", "k", "{}")) + // Give it a moment to (not) run. + time.Sleep(50 * time.Millisecond) + + store.mu.Lock() + extraStatus := store.jobs["job-extra"].status + store.mu.Unlock() + assert.Equal(t, tables.SidekiqStatusPending, extraStatus, "extra job must still be pending while semaphore is full") + + // Release the first batch; the extra job should now run. + close(gate) + waitTerminal(t, store) + waitTerminal(t, store) + waitTerminal(t, store) +} + +// TestInflightDedupPreventsDoubleSpawn verifies that repeated dispatchOnce calls while +// the semaphore is full do not double-claim a job once a slot eventually frees. +func TestInflightDedupPreventsDoubleSpawn(t *testing.T) { + const maxConcurrent = 1 + store := newFakeStore() + r := New(store, testRunner(store).logger, maxConcurrent, "") + + gate := make(chan struct{}) + + store.seed("blocker", "k", tables.SidekiqStatusPending, 0) + r.Register("k", func(_ context.Context, job tables.TableSidekiqJob, _ ProgressFunc) (string, error) { + if job.ID == "blocker" { + <-gate } + return "", nil + }) + r.dispatchOnce() // claims the sole slot with "blocker" + time.Sleep(20 * time.Millisecond) + + // Seed waiter while semaphore is full; repeated ticks must not cause a double-claim later. + store.seed("waiter", "k", tables.SidekiqStatusPending, 0) + r.dispatchOnce() // semaphore full — skips waiter + r.dispatchOnce() // still full + r.dispatchOnce() // still full + + // Release blocker and wait for its slot to free before dispatching again. + close(gate) + waitTerminal(t, store) // blocker done + time.Sleep(10 * time.Millisecond) // let deferred semaphore release run + + r.dispatchOnce() // slot is now free — picks up waiter + waitTerminal(t, store) // waiter done + + store.mu.Lock() + defer store.mu.Unlock() + if store.running["waiter"] != 1 { + t.Errorf("waiter claimed %d time(s), want exactly 1", store.running["waiter"]) + } +} + +// TestDispatcherSkipsUnknownKind verifies that a job whose kind has no registered +// handler is skipped by the dispatcher without claiming or failing it. +func TestDispatcherSkipsUnknownKind(t *testing.T) { + store := newFakeStore() + store.seed("j", "unregistered", tables.SidekiqStatusPending, 0) + r := testRunner(store) + // No handler registered for "unregistered". + r.dispatchOnce() + time.Sleep(50 * time.Millisecond) + + store.mu.Lock() + defer store.mu.Unlock() + assert.Equal(t, tables.SidekiqStatusPending, store.jobs["j"].status, "job with no handler stays pending") + if _, ok := store.completed["j"]; ok { + t.Error("job with no handler must not be completed") + } + if _, ok := store.failedErr["j"]; ok { + t.Error("job with no handler must not be failed") + } +} + +// TestResumePicksUpCheckpointCursor verifies the dispatcher-driven recovery flow: +// a job left running-but-stale (simulating a dead node) is re-claimed and its +// handler receives the metadata cursor from the last progress checkpoint. +func TestResumePicksUpCheckpointCursor(t *testing.T) { + store := newFakeStore() + // Seed a stale running job with a progress cursor in its metadata. + store.seed("resume-job", "k", tables.SidekiqStatusRunning, 30*time.Minute) + store.mu.Lock() + store.jobs["resume-job"].metadata = `{"cursor":42}` + store.mu.Unlock() + + r := testRunner(store) + var gotCursor string + r.Register("k", func(_ context.Context, job tables.TableSidekiqJob, _ ProgressFunc) (string, error) { + gotCursor = job.Metadata + return job.Metadata, nil + }) + r.dispatchOnce() + + id := waitTerminal(t, store) + assert.Equal(t, "resume-job", id) + assert.Equal(t, `{"cursor":42}`, gotCursor, "handler must receive the checkpoint cursor") +} + +// TestShutdownDrainsInFlight verifies that Shutdown waits for in-flight handlers to +// return rather than orphaning goroutines. +func TestShutdownDrainsInFlight(t *testing.T) { + store := newFakeStore() + r := testRunner(store) + + started := make(chan struct{}) + finished := make(chan struct{}) + r.Register("k", func(_ context.Context, _ tables.TableSidekiqJob, _ ProgressFunc) (string, error) { + close(started) + time.Sleep(50 * time.Millisecond) + close(finished) + return "", nil + }) + + require.NoError(t, r.Enqueue(context.Background(), "j", "k", "{}")) + <-started + + done := make(chan struct{}) + go func() { + r.Shutdown() + close(done) + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("Shutdown did not return after in-flight handler finished") + } + select { + case <-finished: + default: + t.Error("Shutdown returned before the handler finished") + } +} + +// TestDispatcherQueuesOverflowOnSubsequentTicks verifies that when there are more +// pending jobs than concurrency slots, the dispatcher picks up only what fits now +// and leaves the rest for the next tick — without spawning a goroutine per job. +func TestDispatcherQueuesOverflowOnSubsequentTicks(t *testing.T) { + const ( + maxConcurrent = 2 + totalJobs = 6 + ) + store := newFakeStore() + r := New(store, testRunner(store).logger, maxConcurrent, "node-1") + + for i := range totalJobs { + store.seed(fmt.Sprintf("job-%d", i), "k", tables.SidekiqStatusPending, 0) + } + + gate := make(chan struct{}) + started := make(chan struct{}, totalJobs) + + r.Register("k", func(_ context.Context, _ tables.TableSidekiqJob, _ ProgressFunc) (string, error) { + started <- struct{}{} + <-gate + return "", nil + }) + + stop := r.StartDispatcher(30*time.Millisecond, StaleAfter) + defer stop() + + // Wait for the first batch (maxConcurrent) to enter their handlers. + for range maxConcurrent { select { - case <-deadline: - t.Fatal("reaper never invoked the stale sweep") - case <-time.After(5 * time.Millisecond): + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for first batch to start") } } + + // Semaphore is full — remaining jobs must still be pending. + store.mu.Lock() + pendingCount := 0 + for i := range totalJobs { + if store.jobs[fmt.Sprintf("job-%d", i)].status == tables.SidekiqStatusPending { + pendingCount++ + } + } + store.mu.Unlock() + require.Equal(t, totalJobs-maxConcurrent, pendingCount, "overflow jobs must remain pending while semaphore is full") + + // Release all handlers; the dispatcher picks up the remaining jobs in subsequent ticks. + close(gate) + for range totalJobs { + waitTerminal(t, store) + } + + // Every job must have been claimed exactly once. + store.mu.Lock() + defer store.mu.Unlock() + for i := range totalJobs { + id := fmt.Sprintf("job-%d", i) + assert.Equal(t, 1, store.running[id], "job %s must be claimed exactly once", id) + } } diff --git a/transports/bifrost-http/lib/config_test.go b/transports/bifrost-http/lib/config_test.go index 03e475d37b8..260b1cd343d 100644 --- a/transports/bifrost-http/lib/config_test.go +++ b/transports/bifrost-http/lib/config_test.go @@ -1562,23 +1562,27 @@ func (m *MockConfigStore) GetSidekiqJob(ctx context.Context, id string) (*tables return nil, nil } -func (m *MockConfigStore) MarkSidekiqJobRunning(ctx context.Context, id string) error { - return nil +func (m *MockConfigStore) ClaimSidekiqJob(ctx context.Context, id, runnerID string, staleBefore time.Time) (bool, error) { + return false, nil +} + +func (m *MockConfigStore) HeartbeatSidekiqJob(ctx context.Context, id, runnerID string) (bool, error) { + return false, nil } -func (m *MockConfigStore) CompleteSidekiqJob(ctx context.Context, id, metadata string) error { +func (m *MockConfigStore) CompleteSidekiqJob(ctx context.Context, id, runnerID, metadata string) error { return nil } -func (m *MockConfigStore) UpdateSidekiqJobProgress(ctx context.Context, id, metadata string) error { +func (m *MockConfigStore) UpdateSidekiqJobProgress(ctx context.Context, id, runnerID, metadata string) error { return nil } -func (m *MockConfigStore) FailSidekiqJob(ctx context.Context, id, metadata, lastErr string) error { +func (m *MockConfigStore) FailSidekiqJob(ctx context.Context, id, runnerID, metadata, lastErr string) error { return nil } -func (m *MockConfigStore) ListIncompleteSidekiqJobs(ctx context.Context) ([]tables.TableSidekiqJob, error) { +func (m *MockConfigStore) ListClaimableSidekiqJobs(ctx context.Context, staleBefore time.Time) ([]tables.TableSidekiqJob, error) { return nil, nil } diff --git a/transports/bifrost-http/server/server.go b/transports/bifrost-http/server/server.go index 7c0026ba3f6..c81e560a068 100644 --- a/transports/bifrost-http/server/server.go +++ b/transports/bifrost-http/server/server.go @@ -180,8 +180,8 @@ type BifrostHTTPServer struct { // otherwise left nil so the quota endpoint reads the VK's own budget rows. ExternalQuotaBudgetResolver handlers.ExternalQuotaBudgetResolver - SidekiqRunner *sidekiq.Runner - SidekiqReaperStop func() + SidekiqRunner *sidekiq.Runner + SidekiqDispatcherStop func() wsPool *bfws.Pool } @@ -1897,7 +1897,7 @@ func (s *BifrostHTTPServer) Bootstrap(ctx context.Context) error { // Initialize Sidekiq runner for background jobs if s.Config != nil && s.Config.ConfigStore != nil { - s.SidekiqRunner = sidekiq.New(s.Config.ConfigStore, logger, 4) + s.SidekiqRunner = sidekiq.New(s.Config.ConfigStore, logger, 4, "") } // Register routes @@ -1971,14 +1971,11 @@ func (s *BifrostHTTPServer) Bootstrap(ctx context.Context) error { // Register UI handler s.RegisterUIRoutes() - // Start Sidekiq reaper to clean up stale jobs + // Start the Sidekiq dispatcher: on every node it periodically claims pending and + // stale (orphaned) jobs, with an atomic claim guaranteeing exactly one node runs + // each job. Subsumes startup recovery of jobs left behind by a crash or restart. if s.SidekiqRunner != nil { - s.SidekiqReaperStop = s.SidekiqRunner.StartReaper(sidekiq.ReaperInterval, sidekiq.StaleAfter) - go func() { - if err := s.SidekiqRunner.RecoverIncomplete(ctx); err != nil { - logger.Error("sidekiq: failed to recover incomplete provisioning jobs: %v", err) - } - }() + s.SidekiqDispatcherStop = s.SidekiqRunner.StartDispatcher(sidekiq.DispatchInterval, sidekiq.StaleAfter) } // Checking if config has server config and use it to set read buffer size @@ -2068,9 +2065,9 @@ func (s *BifrostHTTPServer) Start() error { s.OAuth2SweepWorker.stop() s.OAuth2SweepWorker = nil } - if s.SidekiqReaperStop != nil { - logger.Info("stopping sidekiq reaper...") - s.SidekiqReaperStop() + if s.SidekiqDispatcherStop != nil { + logger.Info("stopping sidekiq dispatcher...") + s.SidekiqDispatcherStop() } if s.SidekiqRunner != nil { logger.Info("stopping sidekiq runner...")