From 7b0389dd192599f71843a6b0074b625eef79d0ef Mon Sep 17 00:00:00 2001 From: scanash00 Date: Wed, 24 Jun 2026 13:36:17 -0800 Subject: [PATCH] db leader fixes --- backend/cmd/server/main.go | 190 +++++++++++------- backend/internal/db/migrate.go | 7 +- .../db/migrations/00012_leader_election.sql | 10 + backend/internal/leader/leader.go | 133 ++++++++++++ backend/internal/leader/leader_test.go | 105 ++++++++++ 5 files changed, 364 insertions(+), 81 deletions(-) create mode 100644 backend/internal/db/migrations/00012_leader_election.sql create mode 100644 backend/internal/leader/leader.go create mode 100644 backend/internal/leader/leader_test.go diff --git a/backend/cmd/server/main.go b/backend/cmd/server/main.go index a559fae..2533dae 100644 --- a/backend/cmd/server/main.go +++ b/backend/cmd/server/main.go @@ -8,6 +8,7 @@ import ( "os" "os/signal" "strings" + stdsync "sync" "syscall" "time" @@ -21,6 +22,7 @@ import ( "margin.at/internal/db" "margin.at/internal/embeddings" "margin.at/internal/firehose" + "margin.at/internal/leader" "margin.at/internal/logger" internalMiddleware "margin.at/internal/middleware" "margin.at/internal/oauth" @@ -70,28 +72,6 @@ func main() { } database.MigrateUnifiedNotes(context.Background()) - go func() { - ticker := time.NewTicker(1 * time.Hour) - defer ticker.Stop() - - runCleanup := func() { - _, err := database.TryAdvisoryLock(context.Background(), db.LockSessionCleanup, func() error { - if err := database.DeleteExpiredOAuthSessions(context.Background()); err != nil { - return err - } - return database.DeleteExpiredPendingAuthsOAuth(context.Background()) - }) - if err != nil { - logger.Error("Failed to run cleanup of expired sessions: %v", err) - } - } - - runCleanup() - for range ticker.C { - runCleanup() - } - }() - embeddingClient := embeddings.NewClient() recService := recommendations.NewService(database, embeddingClient) logger.Info("Recommendation engine initialized (embeddings enabled: %v)", embeddingClient.IsEnabled()) @@ -109,63 +89,54 @@ func main() { ingester := firehose.NewIngester(database, syncSvc) firehose.RelayURL = getEnv("BLOCK_RELAY_URL", "wss://jetstream2.us-east.bsky.network/subscribe") logger.Info("Firehose URL: %s", firehose.RelayURL) - - backfillCtx, backfillCancel := context.WithCancel(context.Background()) - defer backfillCancel() - if recService.IsEnabled() { ingester.SetOnAnnotation(recService.OnAnnotation) ingester.SetOnDocument(recService.OnDocument) + } + + workerCtx, workerCancel := context.WithCancel(context.Background()) + defer workerCancel() - if getEnv("DISABLE_BACKFILL", "") == "" { + var workersWG stdsync.WaitGroup + workersWG.Add(1) + go func() { + defer workersWG.Done() + for workerCtx.Err() == nil { + logger.Info("leader: campaigning for leadership...") + leaderCtx, releaseLeader, err := leader.Acquire(workerCtx, database.Pool()) + if err != nil { + return + } + logger.Info("leader: elected — starting background workers") + + var leaderWG stdsync.WaitGroup + + ingester.Start(leaderCtx) //nolint:errcheck + logger.Info("firehose: ingester started") + + if recService.IsEnabled() && getEnv("DISABLE_BACKFILL", "") == "" { + leaderWG.Add(1) + go func() { + defer leaderWG.Done() + runBackfill(leaderCtx, recService) + }() + } else if !recService.IsEnabled() { + logger.Info("Recommendation backfill skipped (embeddings disabled)") + } else { + logger.Info("Recommendation backfill disabled (DISABLE_BACKFILL is set)") + } + + leaderWG.Add(1) go func() { - time.Sleep(5 * time.Second) - select { - case <-backfillCtx.Done(): - return - default: - } - database.TryAdvisoryLock(backfillCtx, db.LockBackfill, func() error { //nolint:errcheck - logger.Info("Starting recommendation backfill...") - if err := recService.BackfillDocumentEmbeddings(200); err != nil { - logger.Error("Document embedding backfill error: %v", err) - } - if backfillCtx.Err() != nil { - return nil - } - annCount, err := recService.BackfillAnnotationEmbeddings(200) - if err != nil { - logger.Error("Annotation embedding backfill error: %v", err) - } - if backfillCtx.Err() != nil { - return nil - } - hlCount, err := recService.BackfillHighlightEmbeddings(200) - if err != nil { - logger.Error("Highlight embedding backfill error: %v", err) - } - if backfillCtx.Err() != nil { - return nil - } - profileCount, err := recService.RebuildAllProfiles() - if err != nil { - logger.Error("Profile rebuild error: %v", err) - } - logger.Info("Recommendation backfill complete (annotations: %d, highlights: %d, profiles: %d)", annCount, hlCount, profileCount) - return nil - }) + defer leaderWG.Done() + runSessionCleanup(leaderCtx, database) }() - } else { - logger.Info("Recommendation backfill disabled (DISABLE_BACKFILL is set)") - } - } - go func() { - _, err := database.TryAdvisoryLock(context.Background(), db.LockFirehose, func() error { - return ingester.Start(context.Background()) - }) - if err != nil { - logger.Error("Firehose ingester error: %v", err) + <-leaderCtx.Done() + logger.Info("leader: stepping down — draining workers") + ingester.Stop() + releaseLeader() + leaderWG.Wait() } }() @@ -233,19 +204,88 @@ func main() { <-quit logger.Infoln("Shutting down server...") - backfillCancel() - ingester.Stop() ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() if err := server.Shutdown(ctx); err != nil { - logger.Fatal("Server forced to shutdown: %v", err) + logger.Error("Server forced to shutdown: %v", err) + } + + workerCancel() + drained := make(chan struct{}) + go func() { + workersWG.Wait() + close(drained) + }() + select { + case <-drained: + logger.Infoln("Background workers drained cleanly") + case <-time.After(30 * time.Second): + logger.Error("Background workers did not drain within 30s — exiting anyway") } logger.Infoln("Server exited") } +func runBackfill(ctx context.Context, recService *recommendations.Service) { + select { + case <-ctx.Done(): + return + case <-time.After(5 * time.Second): + } + + logger.Info("Starting recommendation backfill...") + if err := recService.BackfillDocumentEmbeddings(200); err != nil { + logger.Error("Document embedding backfill error: %v", err) + } + if ctx.Err() != nil { + return + } + annCount, err := recService.BackfillAnnotationEmbeddings(200) + if err != nil { + logger.Error("Annotation embedding backfill error: %v", err) + } + if ctx.Err() != nil { + return + } + hlCount, err := recService.BackfillHighlightEmbeddings(200) + if err != nil { + logger.Error("Highlight embedding backfill error: %v", err) + } + if ctx.Err() != nil { + return + } + profileCount, err := recService.RebuildAllProfiles() + if err != nil { + logger.Error("Profile rebuild error: %v", err) + } + logger.Info("Recommendation backfill complete (annotations: %d, highlights: %d, profiles: %d)", annCount, hlCount, profileCount) +} + +func runSessionCleanup(ctx context.Context, database *db.DB) { + cleanup := func() { + if err := database.DeleteExpiredOAuthSessions(ctx); err != nil { + logger.Error("Failed to delete expired OAuth sessions: %v", err) + } + if err := database.DeleteExpiredPendingAuthsOAuth(ctx); err != nil { + logger.Error("Failed to delete expired pending auths: %v", err) + } + } + + cleanup() + ticker := time.NewTicker(1 * time.Hour) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + cleanup() + } + } +} + func getEnv(key, fallback string) string { if value, ok := os.LookupEnv(key); ok { return value diff --git a/backend/internal/db/migrate.go b/backend/internal/db/migrate.go index 2111ef9..81ed84a 100644 --- a/backend/internal/db/migrate.go +++ b/backend/internal/db/migrate.go @@ -14,12 +14,7 @@ import ( //go:embed migrations var migrationsFS embed.FS -const ( - LockMigrate int64 = 0x6D617267696E - LockFirehose int64 = 0x6669726568736F65 - LockSessionCleanup int64 = 0x73657373696F6E - LockBackfill int64 = 0x6261636B66696C6C -) +const LockMigrate int64 = 0x6D617267696E func (db *DB) Migrate() error { ctx := context.Background() diff --git a/backend/internal/db/migrations/00012_leader_election.sql b/backend/internal/db/migrations/00012_leader_election.sql new file mode 100644 index 0000000..3e0d5f5 --- /dev/null +++ b/backend/internal/db/migrations/00012_leader_election.sql @@ -0,0 +1,10 @@ +-- +goose Up +CREATE TABLE IF NOT EXISTS leader_election ( + name TEXT PRIMARY KEY, + holder TEXT NOT NULL, + acquired_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + expires_at TIMESTAMPTZ NOT NULL +); + +-- +goose Down +DROP TABLE IF EXISTS leader_election; diff --git a/backend/internal/leader/leader.go b/backend/internal/leader/leader.go new file mode 100644 index 0000000..805157d --- /dev/null +++ b/backend/internal/leader/leader.go @@ -0,0 +1,133 @@ +package leader + +import ( + "context" + "crypto/rand" + "encoding/hex" + "errors" + "fmt" + "log" + "os" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" +) + +const ( + lockName = "background_workers" + leaseTTL = 30 * time.Second + heartbeat = 10 * time.Second + acquireInterval = 5 * time.Second +) + +func instanceID() string { + host, _ := os.Hostname() + b := make([]byte, 4) + _, _ = rand.Read(b) + return host + "-" + hex.EncodeToString(b) +} + +func Acquire(ctx context.Context, pool *pgxpool.Pool) (leaderCtx context.Context, release func(), err error) { + id := instanceID() + + waiting := false + for { + ok, aerr := tryAcquire(ctx, pool, id) + if aerr != nil { + log.Printf("leader: acquire query failed: %v", aerr) + } else if ok { + break + } else if !waiting { + log.Printf("leader: another instance is leader — standing by") + waiting = true + } + + select { + case <-ctx.Done(): + return nil, func() {}, ctx.Err() + case <-time.After(acquireInterval): + } + } + + leaderCtx, cancel := context.WithCancel(ctx) + go renewLoop(leaderCtx, cancel, pool, id) + + release = func() { + cancel() + rctx, rcancel := context.WithTimeout(context.Background(), 5*time.Second) + defer rcancel() + _, _ = pool.Exec(rctx, + `DELETE FROM leader_election WHERE name = $1 AND holder = $2`, lockName, id) + } + return leaderCtx, release, nil +} + +func tryAcquire(ctx context.Context, pool *pgxpool.Pool, id string) (bool, error) { + var holder string + err := pool.QueryRow(ctx, ` + INSERT INTO leader_election (name, holder, acquired_at, expires_at) + VALUES ($1, $2, now(), now() + $3::interval) + ON CONFLICT (name) DO UPDATE + SET holder = EXCLUDED.holder, + acquired_at = now(), + expires_at = EXCLUDED.expires_at + WHERE leader_election.expires_at < now() OR leader_election.holder = EXCLUDED.holder + RETURNING holder`, + lockName, id, intervalString(leaseTTL), + ).Scan(&holder) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return false, nil + } + return false, err + } + return holder == id, nil +} + +func renew(ctx context.Context, pool *pgxpool.Pool, id string) (bool, error) { + tag, err := pool.Exec(ctx, ` + UPDATE leader_election + SET expires_at = now() + $3::interval + WHERE name = $1 AND holder = $2`, + lockName, id, intervalString(leaseTTL), + ) + if err != nil { + return false, err + } + return tag.RowsAffected() == 1, nil +} + +func renewLoop(ctx context.Context, cancel context.CancelFunc, pool *pgxpool.Pool, id string) { + ticker := time.NewTicker(heartbeat) + defer ticker.Stop() + misses := 0 + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + held, err := renew(ctx, pool, id) + if err != nil { + misses++ + log.Printf("leader: renew failed (%d/3): %v", misses, err) + if misses >= 3 { + log.Printf("leader: lost lease after repeated renew failures — stopping workers") + cancel() + return + } + continue + } + misses = 0 + if !held { + log.Printf("leader: lease taken over by another instance — stopping workers") + cancel() + return + } + } + } +} + +func intervalString(d time.Duration) string { + return fmt.Sprintf("%d seconds", int(d.Seconds())) +} diff --git a/backend/internal/leader/leader_test.go b/backend/internal/leader/leader_test.go new file mode 100644 index 0000000..809b01a --- /dev/null +++ b/backend/internal/leader/leader_test.go @@ -0,0 +1,105 @@ +package leader + +import ( + "context" + "os" + "testing" + + "github.com/jackc/pgx/v5/pgxpool" +) + +func testPool(t *testing.T) *pgxpool.Pool { + t.Helper() + dsn := os.Getenv("TEST_DATABASE_URL") + if dsn == "" { + dsn = os.Getenv("DATABASE_URL") + } + if dsn == "" { + t.Skip("set TEST_DATABASE_URL or DATABASE_URL to run leader election tests") + } + pool, err := pgxpool.New(context.Background(), dsn) + if err != nil { + t.Fatalf("pool: %v", err) + } + t.Cleanup(pool.Close) + _, err = pool.Exec(context.Background(), ` + CREATE TABLE IF NOT EXISTS leader_election ( + name TEXT PRIMARY KEY, + holder TEXT NOT NULL, + acquired_at TIMESTAMPTZ NOT NULL DEFAULT now(), + expires_at TIMESTAMPTZ NOT NULL + )`) + if err != nil { + t.Fatalf("create table: %v", err) + } + _, _ = pool.Exec(context.Background(), `DELETE FROM leader_election WHERE name = $1`, lockName) + return pool +} + +func TestTryAcquire_SingleWinner(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + + a, err := tryAcquire(ctx, pool, "pod-a") + if err != nil || !a { + t.Fatalf("pod-a should acquire: ok=%v err=%v", a, err) + } + + b, err := tryAcquire(ctx, pool, "pod-b") + if err != nil { + t.Fatalf("pod-b acquire err: %v", err) + } + if b { + t.Fatal("pod-b acquired a lease already held by pod-a") + } + + again, err := tryAcquire(ctx, pool, "pod-a") + if err != nil || !again { + t.Fatalf("pod-a should re-acquire its own lease: ok=%v err=%v", again, err) + } +} + +func TestTryAcquire_TakeoverAfterExpiry(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + + if ok, err := tryAcquire(ctx, pool, "pod-a"); err != nil || !ok { + t.Fatalf("pod-a should acquire: ok=%v err=%v", ok, err) + } + + if _, err := pool.Exec(ctx, + `UPDATE leader_election SET expires_at = now() - interval '1 second' WHERE name = $1`, lockName); err != nil { + t.Fatalf("expire: %v", err) + } + + ok, err := tryAcquire(ctx, pool, "pod-b") + if err != nil || !ok { + t.Fatalf("pod-b should take over expired lease: ok=%v err=%v", ok, err) + } + + if ok, err := tryAcquire(ctx, pool, "pod-a"); err != nil || ok { + t.Fatalf("pod-a should not reclaim now that pod-b holds it: ok=%v err=%v", ok, err) + } +} + +func TestRenew(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + + if ok, err := tryAcquire(ctx, pool, "pod-a"); err != nil || !ok { + t.Fatalf("acquire: ok=%v err=%v", ok, err) + } + + held, err := renew(ctx, pool, "pod-a") + if err != nil || !held { + t.Fatalf("pod-a should renew its lease: held=%v err=%v", held, err) + } + + held, err = renew(ctx, pool, "pod-b") + if err != nil { + t.Fatalf("renew err: %v", err) + } + if held { + t.Fatal("pod-b renewed a lease it does not hold") + } +} -- knotmirror 0.0.1