package database import ( "context" "database/sql" "database/sql/driver" "fmt" "hash/fnv" "time" ) // Several instances can share one Postgres database. Advisory locks keep // them from creating the schema, migrating it or running a background job at // the same time. SQLite is single-instance, so there every lock is granted // at once. // lockKey maps a lock name to the 64-bit key Postgres advisory locks take. func lockKey(name string) int64 { h := fnv.New64a() _, _ = h.Write([]byte("git-pkgs/proxy:" + name)) return int64(h.Sum64()) //nolint:gosec // a key, wrapping is fine } // WithLock runs fn while holding the named lock, waiting for another // instance that holds it. func (db *DB) WithLock(ctx context.Context, name string, fn func() error) error { if db.dialect != DialectPostgres { return fn() } conn, err := db.Conn(ctx) if err != nil { return fmt.Errorf("taking connection for lock %s: %w", name, err) } defer func() { _ = conn.Close() }() if _, err := conn.ExecContext(ctx, "SELECT pg_advisory_lock($1)", lockKey(name)); err != nil { return fmt.Errorf("taking lock %s: %w", name, err) } defer unlock(conn, name) return fn() } // TryLock takes the named lock if no other instance holds it. When ok is // true the caller must call release once done. The lock lives in a // connection held for that time, and Postgres frees it if the process dies. func (db *DB) TryLock(ctx context.Context, name string) (release func(), ok bool, err error) { if db.dialect != DialectPostgres { return func() {}, true, nil } conn, err := db.Conn(ctx) if err != nil { return nil, false, fmt.Errorf("taking connection for lock %s: %w", name, err) } if err := conn.QueryRowContext(ctx, "SELECT pg_try_advisory_lock($1)", lockKey(name)).Scan(&ok); err != nil { _ = conn.Close() return nil, false, fmt.Errorf("taking lock %s: %w", name, err) } if !ok { _ = conn.Close() return nil, false, nil } return func() { unlock(conn, name) _ = conn.Close() }, true, nil } // unlockTimeout bounds the unlock. If it fails, closing the connection still // ends the session that holds the lock. const unlockTimeout = 10 * time.Second func unlock(conn *sql.Conn, name string) { ctx, cancel := context.WithTimeout(context.Background(), unlockTimeout) defer cancel() if _, err := conn.ExecContext(ctx, "SELECT pg_advisory_unlock($1)", lockKey(name)); err != nil { // Drop the connection so the session, and with it the lock, ends. _ = conn.Raw(func(any) error { return driver.ErrBadConn }) } } // ClaimJobRun records that the named job starts now unless it started less // than minGap ago, and reports whether the caller should run it. Instances // sharing the database each keep their own timer for a scheduled job, and // this lets only the first of them run it per interval. func (db *DB) ClaimJobRun(name string, minGap time.Duration) (bool, error) { now := time.Now().UTC() res, err := db.Exec(db.Rebind(` INSERT INTO job_runs (name, started_at) VALUES (?, ?) ON CONFLICT(name) DO UPDATE SET started_at = excluded.started_at WHERE job_runs.started_at <= ?`), name, now, now.Add(-minGap)) if err != nil { return false, fmt.Errorf("claiming job run %s: %w", name, err) } n, err := res.RowsAffected() if err != nil { return false, fmt.Errorf("claiming job run %s: %w", name, err) } return n > 0, nil } // ReleaseJobRun drops the named job's last claim, so the next attempt on any // instance runs it. It is for a run that was interrupted before it finished. func (db *DB) ReleaseJobRun(name string) error { if _, err := db.Exec(db.Rebind(`DELETE FROM job_runs WHERE name = ?`), name); err != nil { return fmt.Errorf("releasing job run %s: %w", name, err) } return nil }