1package work23import (4 "context"5 "crypto/rand"6 "crypto/sha256"7 "database/sql"8 "encoding/hex"9 "errors"10 "fmt"11 "math"12 "sync"13 "sync/atomic"14 "time"1516 "github.com/FelineStateMachine/tinyrelay/internal/storage"17 "github.com/prometheus/client_golang/prometheus"18)1920const defaultLease = 30 * time.Second2122// ErrClaimLost denotes loss of worker lease ownership.23var ErrClaimLost = errors.New("work: claim lost")2425type stopError struct{ err error }2627func (e *stopError) Error() string {28 if e.err == nil {29 return "work: stopped"30 }31 return e.err.Error()32}3334func (e *stopError) Unwrap() error { return e.err }3536// Stop marks a handler failure as terminal. Workers cancel the claimed intent37// instead of retrying it, while retaining the underlying error for operators.38func Stop(err error) error { return &stopError{err: err} }3940// Intent is a durable unit of work. ClaimToken identifies the current claim;41// ClaimUntil determines when another worker can claim an expired lease.42// State is pending, running, completed or cancelled.43type Intent struct {44 ID string45 Kind string46 EventID string47 Target string48 Payload string49 State string50 Attempts int51 NextAt time.Time52 CreatedAt time.Time53 UpdatedAt time.Time54 LastError string55 ClaimToken string56 ClaimUntil time.Time57}5859// Queue implements durable claiming and fencing over a storage Store.60type Queue struct {61 store *storage.Store62 depth *prometheus.GaugeVec63}6465// New creates a queue. Metrics returns a collector that hosts may register.66func New(store *storage.Store) *Queue {67 return &Queue{store: store, depth: prometheus.NewGaugeVec(prometheus.GaugeOpts{Name: "relay_work_pending", Help: "Pending durable work intents."}, []string{"kind"})}68}6970// Metrics returns the per-kind pending gauge collector.71func (q *Queue) Metrics() prometheus.Collector { return q.depth }7273// Enqueue inserts an intent idempotently.74func (q *Queue) Enqueue(ctx context.Context, intent Intent) (string, error) {75 if intent.Kind == "" || intent.EventID == "" || intent.Target == "" {76 return "", errors.New("work: kind, event id, and target are required")77 }78 if intent.ID == "" {79 intent.ID = storage.IntentID(intent.Kind, intent.EventID, intent.Target)80 }81 now := time.Now().Unix()82 if intent.NextAt.IsZero() {83 intent.NextAt = time.Unix(now, 0)84 }85 err := q.store.WithTx(ctx, func(tx *sql.Tx) error {86 _, err := tx.ExecContext(ctx, `INSERT OR IGNORE INTO work_intents(id,kind,event_id,target,payload,next_at,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?)`, intent.ID, intent.Kind, intent.EventID, intent.Target, intent.Payload, intent.NextAt.Unix(), now, now)87 return err88 })89 if err != nil {90 return "", fmt.Errorf("work: enqueue: %w", err)91 }92 if metricKind(intent.Kind) == "other" {93 q.refresh(ctx, "")94 } else {95 q.refresh(ctx, intent.Kind)96 }97 return intent.ID, nil98}99100// Claim atomically claims the oldest due pending or expired intent.101func (q *Queue) Claim(ctx context.Context, now time.Time, lease time.Duration) (*Intent, error) {102 if lease <= 0 {103 lease = defaultLease104 }105 nowUnix := now.Unix()106 until := now.Add(lease).Unix()107 token, err := token()108 if err != nil {109 return nil, err110 }111 var result *Intent112 err = q.store.WithTx(ctx, func(tx *sql.Tx) error {113 row := tx.QueryRowContext(ctx, `WITH candidate AS (114 SELECT id FROM work_intents115 WHERE (state='pending' AND next_at<=?) OR (state='running' AND claim_until<=?)116 ORDER BY next_at,created_at,id LIMIT 1117)118UPDATE work_intents SET state='running',attempts=attempts+1,updated_at=?,claim_token=?,claim_until=?119WHERE id=(SELECT id FROM candidate)120RETURNING id,kind,event_id,target,payload,state,attempts,next_at,created_at,updated_at,last_error,claim_token,claim_until`, nowUnix, nowUnix, nowUnix, token, until)121 intent := Intent{}122 err := scanIntent(row, &intent)123 if errors.Is(err, sql.ErrNoRows) {124 return nil125 }126 if err != nil {127 return err128 }129 result = &intent130 return nil131 })132 if err != nil {133 return nil, fmt.Errorf("work: claim: %w", err)134 }135 if result != nil {136 if metricKind(result.Kind) == "other" {137 q.refresh(ctx, "")138 } else {139 q.refresh(ctx, result.Kind)140 }141 }142 return result, nil143}144145// Renew extends a claim if the token still fences the caller.146func (q *Queue) Renew(ctx context.Context, id, claimToken string, now time.Time, lease time.Duration) (bool, error) {147 if lease <= 0 {148 lease = defaultLease149 }150 result, err := q.store.DB().ExecContext(ctx, `UPDATE work_intents SET claim_until=?,updated_at=? WHERE id=? AND state='running' AND claim_token=?`, now.Add(lease).Unix(), now.Unix(), id, claimToken)151 if err != nil {152 return false, fmt.Errorf("work: renew: %w", err)153 }154 count, err := result.RowsAffected()155 return count == 1, err156}157158// Complete marks a running intent complete if the token still fences the caller.159func (q *Queue) Complete(ctx context.Context, id, claimToken string, now time.Time) (bool, error) {160 result, err := q.store.DB().ExecContext(ctx, `UPDATE work_intents SET state='completed',updated_at=?,claim_token='',claim_until=0 WHERE id=? AND state='running' AND claim_token=?`, now.Unix(), id, claimToken)161 if err != nil {162 return false, fmt.Errorf("work: complete: %w", err)163 }164 count, err := result.RowsAffected()165 if count == 1 {166 q.refresh(ctx, "")167 }168 return count == 1, err169}170171// CompletePendingTx marks a pending intent complete in the caller's172// transaction. Running intents are left untouched so an active worker keeps173// ownership of its claim.174func CompletePendingTx(ctx context.Context, tx *sql.Tx, id string, now time.Time) (bool, error) {175 result, err := tx.ExecContext(ctx, `UPDATE work_intents SET state='completed',updated_at=?,claim_token='',claim_until=0 WHERE id=? AND state='pending'`, now.Unix(), id)176 if err != nil {177 return false, fmt.Errorf("work: complete pending: %w", err)178 }179 count, err := result.RowsAffected()180 return count == 1, err181}182183// CancelClaim marks a running intent terminal if the claim token still fences184// the caller.185func (q *Queue) CancelClaim(ctx context.Context, id, claimToken string, cause error, now time.Time) (bool, error) {186 message := ""187 if cause != nil {188 message = cause.Error()189 }190 result, err := q.store.DB().ExecContext(ctx, `UPDATE work_intents SET state='cancelled',updated_at=?,last_error=?,claim_token='',claim_until=0 WHERE id=? AND state='running' AND claim_token=?`, now.Unix(), message, id, claimToken)191 if err != nil {192 return false, fmt.Errorf("work: cancel claim: %w", err)193 }194 count, err := result.RowsAffected()195 if count == 1 {196 q.refresh(ctx, "")197 }198 return count == 1, err199}200201// Retry returns a failed claim to pending with unbounded exponential backoff.202func (q *Queue) Retry(ctx context.Context, id, claimToken string, cause error, now time.Time) (bool, error) {203 var attempts int204 err := q.store.DB().QueryRowContext(ctx, `SELECT attempts FROM work_intents WHERE id=? AND state='running' AND claim_token=?`, id, claimToken).Scan(&attempts)205 if errors.Is(err, sql.ErrNoRows) {206 return false, nil207 }208 if err != nil {209 return false, fmt.Errorf("work: retry lookup: %w", err)210 }211 delay := backoff(id, attempts)212 message := ""213 if cause != nil {214 message = cause.Error()215 }216 result, err := q.store.DB().ExecContext(ctx, `UPDATE work_intents SET state='pending',next_at=?,updated_at=?,last_error=?,claim_token='',claim_until=0 WHERE id=? AND state='running' AND claim_token=?`, now.Add(delay).Unix(), now.Unix(), message, id, claimToken)217 if err != nil {218 return false, fmt.Errorf("work: retry: %w", err)219 }220 count, err := result.RowsAffected()221 if count == 1 {222 q.refresh(ctx, "")223 }224 return count == 1, err225}226227// RetryNow makes an intent immediately eligible, regardless of its current pending delay.228func (q *Queue) RetryNow(ctx context.Context, id string, now time.Time) error {229 _, err := q.store.DB().ExecContext(ctx, `UPDATE work_intents SET state='pending',next_at=?,updated_at=?,claim_token='',claim_until=0 WHERE id=? AND state!='completed' AND state!='cancelled'`, now.Unix(), now.Unix(), id)230 if err == nil {231 q.refresh(ctx, "")232 }233 return err234}235236// Cancel prevents future claims. A running claim can be cancelled administratively.237func (q *Queue) Cancel(ctx context.Context, id string, now time.Time) error {238 _, err := q.store.DB().ExecContext(ctx, `UPDATE work_intents SET state='cancelled',updated_at=?,claim_token='',claim_until=0 WHERE id=?`, now.Unix(), id)239 if err == nil {240 q.refresh(ctx, "")241 }242 return err243}244245// Pending returns pending and running intents, optionally filtered by kind.246func (q *Queue) Pending(ctx context.Context, kind string) ([]Intent, error) {247 return q.list(ctx, kind, "pending,running")248}249250// List returns intents in the requested state; an empty state lists all states.251func (q *Queue) List(ctx context.Context, state string) ([]Intent, error) {252 return q.list(ctx, "", state)253}254255func (q *Queue) list(ctx context.Context, kind, state string) ([]Intent, error) {256 query := `SELECT id,kind,event_id,target,payload,state,attempts,next_at,created_at,updated_at,last_error,claim_token,claim_until FROM work_intents WHERE 1=1`257 args := []any{}258 if kind != "" {259 query += " AND kind=?"260 args = append(args, kind)261 }262 if state == "pending,running" {263 query += " AND state IN ('pending','running')"264 } else if state != "" {265 query += " AND state=?"266 args = append(args, state)267 }268 query += " ORDER BY next_at,created_at,id"269 rows, err := q.store.DB().QueryContext(ctx, query, args...)270 if err != nil {271 return nil, err272 }273 defer rows.Close()274 var result []Intent275 for rows.Next() {276 intent := Intent{}277 if err := scanIntent(rows, &intent); err != nil {278 return nil, err279 }280 result = append(result, intent)281 }282 return result, rows.Err()283}284285func (q *Queue) refresh(ctx context.Context, kind string) {286 if kind == "" {287 rows, err := q.store.DB().QueryContext(ctx, `SELECT kind,count(*) FROM work_intents WHERE state='pending' GROUP BY kind`)288 if err != nil {289 return290 }291 counts := make(map[string]int)292 for rows.Next() {293 var value string294 var count int295 if rows.Scan(&value, &count) == nil {296 counts[metricKind(value)] += count297 }298 }299 _ = rows.Close()300 for _, value := range knownMetricKinds {301 if _, ok := counts[value]; !ok {302 counts[value] = 0303 }304 }305 for value, count := range counts {306 q.depth.WithLabelValues(value).Set(float64(count))307 }308 return309 }310 var count int311 if err := q.store.DB().QueryRowContext(ctx, `SELECT count(*) FROM work_intents WHERE kind=? AND state='pending'`, kind).Scan(&count); err == nil {312 q.depth.WithLabelValues(metricKind(kind)).Set(float64(count))313 }314}315316var knownMetricKinds = []string{"delivery", "push", "replication", "webhook", "job", "other"}317318func metricKind(kind string) string {319 for _, known := range knownMetricKinds[:len(knownMetricKinds)-1] {320 if kind == known {321 return known322 }323 }324 return "other"325}326327// Handler processes one claimed intent while the worker renews its lease.328// A nil error completes the claim; an error schedules a retry. Stop marks a329// terminal failure. The handler observes ctx cancellation so the worker can330// finish shutdown or release an attempt whose claim was lost. A retry can331// repeat an external effect after an interrupted acknowledgment; handlers own332// deduplication of those effects.333type Handler func(ctx context.Context, intent Intent) error334335// WorkerOptions controls claim leases and optional handler timing observation.336// Lease and Renew default to 30 seconds and one third of the lease.337type WorkerOptions struct {338 Lease time.Duration339 Renew time.Duration340 Observe func(kind, outcome string, duration time.Duration)341}342343// PoolOptions controls a pool of durable workers. Workers less than or equal344// to zero use the default of four.345type PoolOptions struct {346 Workers int347 Worker WorkerOptions348}349350const defaultWorkers = 4351352// RunPool runs a worker pool until ctx is cancelled or a worker reports a353// queue error. The first queue error cancels the remaining workers and is354// returned after they have stopped.355func RunPool(ctx context.Context, queue *Queue, handlers map[string]Handler, options PoolOptions) error {356 workers := options.Workers357 if workers <= 0 {358 workers = defaultWorkers359 }360 workerCtx, cancel := context.WithCancel(ctx)361 defer cancel()362 errCh := make(chan error, 1)363 var group sync.WaitGroup364 for i := 0; i < workers; i++ {365 group.Add(1)366 go func() {367 defer group.Done()368 if err := NewWorkerWithOptions(queue, handlers, options.Worker).Run(workerCtx); err != nil && workerCtx.Err() == nil {369 select {370 case errCh <- err:371 cancel()372 default:373 }374 }375 }()376 }377 group.Wait()378 select {379 case err := <-errCh:380 return err381 default:382 return ctx.Err()383 }384}385386// Worker owns a queue lifecycle and processes one intent at a time.387type Worker struct {388 queue *Queue389 handlers map[string]Handler390 options WorkerOptions391}392393// NewWorker creates a worker with copied handlers.394func NewWorker(queue *Queue, handlers map[string]Handler) *Worker {395 return NewWorkerWithOptions(queue, handlers, WorkerOptions{})396}397398// NewWorkerWithOptions creates a worker with explicit lease behavior.399func NewWorkerWithOptions(queue *Queue, handlers map[string]Handler, options WorkerOptions) *Worker {400 copyHandlers := make(map[string]Handler, len(handlers))401 for kind, handler := range handlers {402 copyHandlers[kind] = handler403 }404 if options.Lease <= 0 {405 options.Lease = defaultLease406 }407 if options.Renew <= 0 {408 options.Renew = options.Lease / 3409 }410 if options.Renew <= 0 {411 options.Renew = time.Millisecond412 }413 return &Worker{queue: queue, handlers: copyHandlers, options: options}414}415416// Run blocks until the context is cancelled or a durable queue error occurs.417func (w *Worker) Run(ctx context.Context) error {418 for {419 if err := ctx.Err(); err != nil {420 return err421 }422 intent, err := w.queue.Claim(ctx, time.Now(), w.options.Lease)423 if err != nil {424 // SQLite may report a transaction cleanup error when cancellation425 // races a claim. The worker lifecycle is already stopping, so426 // preserve the caller's cancellation outcome.427 if ctx.Err() != nil {428 return ctx.Err()429 }430 return err431 }432 if intent == nil {433 if err := wait(ctx, 100*time.Millisecond); err != nil {434 return err435 }436 continue437 }438 handler := w.handlers[intent.Kind]439 if handler == nil {440 handler = func(context.Context, Intent) error { return errors.New("no handler registered") }441 }442 attemptCtx, cancelAttempt := context.WithCancel(ctx)443 var lost atomic.Bool444 renewDone := make(chan struct{})445 go w.renew(attemptCtx, intent, cancelAttempt, &lost, renewDone)446 started := time.Now()447 err = handler(attemptCtx, *intent)448 cancelAttempt()449 <-renewDone450 if w.options.Observe != nil {451 outcome := "success"452 if err != nil {453 outcome = "error"454 }455 if lost.Load() {456 outcome = "fenced"457 }458 w.options.Observe(intent.Kind, outcome, time.Since(started))459 }460 if ctx.Err() != nil {461 return ctx.Err()462 }463 if lost.Load() {464 continue465 }466 if err != nil {467 var terminal *stopError468 if errors.As(err, &terminal) {469 ok, cancelErr := w.queue.CancelClaim(ctx, intent.ID, intent.ClaimToken, terminal.err, time.Now())470 if cancelErr != nil {471 return cancelErr472 }473 if !ok {474 continue475 }476 continue477 }478 ok, retryErr := w.queue.Retry(ctx, intent.ID, intent.ClaimToken, err, time.Now())479 if retryErr != nil {480 return retryErr481 }482 if !ok {483 // A cancellation or another claimant may have fenced this484 // attempt while the handler was unwinding. Its result is stale;485 // keep the worker alive for other durable work.486 continue487 }488 continue489 }490 ok, err := w.queue.Complete(ctx, intent.ID, intent.ClaimToken, time.Now())491 if err != nil {492 return err493 }494 if !ok {495 // The claim can be fenced after the renewal loop's final check.496 // Discard this stale result and continue processing the queue.497 continue498 }499 }500}501502func (w *Worker) renew(ctx context.Context, intent *Intent, cancel context.CancelFunc, lost *atomic.Bool, done chan<- struct{}) {503 defer close(done)504 ticker := time.NewTicker(w.options.Renew)505 defer ticker.Stop()506 for {507 select {508 case <-ctx.Done():509 return510 case now := <-ticker.C:511 ok, err := w.queue.Renew(ctx, intent.ID, intent.ClaimToken, now, w.options.Lease)512 if err != nil || !ok {513 lost.Store(true)514 cancel()515 return516 }517 }518 }519}520521func wait(ctx context.Context, duration time.Duration) error {522 timer := time.NewTimer(duration)523 defer timer.Stop()524 select {525 case <-ctx.Done():526 return ctx.Err()527 case <-timer.C:528 return nil529 }530}531532type scanner interface{ Scan(...any) error }533534func scanIntent(row scanner, intent *Intent) error {535 var next, created, updated, until int64536 err := row.Scan(&intent.ID, &intent.Kind, &intent.EventID, &intent.Target, &intent.Payload, &intent.State, &intent.Attempts, &next, &created, &updated, &intent.LastError, &intent.ClaimToken, &until)537 if err != nil {538 return err539 }540 intent.NextAt, intent.CreatedAt, intent.UpdatedAt, intent.ClaimUntil = time.Unix(next, 0), time.Unix(created, 0), time.Unix(updated, 0), time.Unix(until, 0)541 return nil542}543544func token() (string, error) {545 var bytes [16]byte546 if _, err := rand.Read(bytes[:]); err != nil {547 return "", fmt.Errorf("work: claim token: %w", err)548 }549 return hex.EncodeToString(bytes[:]), nil550}551552func backoff(id string, attempts int) time.Duration {553 if attempts < 1 {554 attempts = 1555 }556 seconds := math.Pow(2, float64(attempts))557 if seconds > float64(time.Hour/time.Second) {558 seconds = float64(time.Hour / time.Second)559 }560 sum := sha256.Sum256([]byte(id + fmt.Sprint(attempts)))561 jitter := (float64(sum[0])/255 - 0.5) * 0.4562 return time.Duration(seconds*(1+jitter)) * time.Second563}