From 91f46eb25d7fcfceeac8b0e7eddc804309ebea51 Mon Sep 17 00:00:00 2001 From: Eli Mallon Date: Wed, 10 Sep 2025 13:41:48 -0700 Subject: [PATCH] statedb: much better postgres distributed locking --- pkg/atproto/chat_message_test.go | 2 +- pkg/atproto/lexicon_repo_test.go | 2 +- pkg/cmd/streamplace.go | 5 ++-- pkg/statedb/locks.go | 40 +++++++++++++++++++++++-- pkg/statedb/locks_test.go | 51 ++++++++++++++------------------ pkg/statedb/migrate.go | 2 +- pkg/statedb/statedb.go | 21 +++++++++---- 7 files changed, 81 insertions(+), 42 deletions(-) diff --git a/pkg/atproto/chat_message_test.go b/pkg/atproto/chat_message_test.go index 0bd184f9..d6ea5d87 100644 --- a/pkg/atproto/chat_message_test.go +++ b/pkg/atproto/chat_message_test.go @@ -35,7 +35,7 @@ func TestChatMessage(t *testing.T) { cli.DataDir = t.TempDir() mod, err := model.MakeDB(":memory:") require.NoError(t, err) - state, err := statedb.MakeDB(&cli, nil, mod) + state, err := statedb.MakeDB(context.Background(), &cli, nil, mod) require.NoError(t, err) atsync := &ATProtoSynchronizer{ CLI: &cli, diff --git a/pkg/atproto/lexicon_repo_test.go b/pkg/atproto/lexicon_repo_test.go index 307327d9..364b60d6 100644 --- a/pkg/atproto/lexicon_repo_test.go +++ b/pkg/atproto/lexicon_repo_test.go @@ -22,7 +22,7 @@ func TestLexiconRepo(t *testing.T) { cli.DataDir = t.TempDir() mod, err := model.MakeDB(":memory:") require.NoError(t, err) - state, err := statedb.MakeDB(&cli, nil, mod) + state, err := statedb.MakeDB(context.Background(), &cli, nil, mod) require.NoError(t, err) // creating a new repo diff --git a/pkg/cmd/streamplace.go b/pkg/cmd/streamplace.go index 0e2dc55f..4bc022a3 100644 --- a/pkg/cmd/streamplace.go +++ b/pkg/cmd/streamplace.go @@ -318,12 +318,14 @@ func start(build *config.BuildFlags, platformJobs []jobFunc) error { } } + group, ctx := TimeoutGroupWithContext(ctx) + out := carstore.SQLiteStore{} err = out.Open(":memory:") if err != nil { return err } - state, err := statedb.MakeDB(&cli, noter, mod) + state, err := statedb.MakeDB(ctx, &cli, noter, mod) if err != nil { return err } @@ -394,7 +396,6 @@ func start(build *config.BuildFlags, platformJobs []jobFunc) error { return err } - group, ctx := TimeoutGroupWithContext(ctx) ctx = log.WithLogValues(ctx, "version", build.Version) group.Go(func() error { diff --git a/pkg/statedb/locks.go b/pkg/statedb/locks.go index 6635ca83..ddf59ead 100644 --- a/pkg/statedb/locks.go +++ b/pkg/statedb/locks.go @@ -1,10 +1,14 @@ package statedb import ( + "context" "crypto/sha256" "encoding/binary" "fmt" "sync" + + "gorm.io/gorm" + "stream.place/streamplace/pkg/log" ) func (state *StatefulDB) GetNamedLock(name string) (func(), error) { @@ -21,26 +25,56 @@ func (state *StatefulDB) getNamedLockPostgres(name string) (func(), error) { // we also use a local lock here - whoever is locking wants exclusive access even within the node lock := state.locks.GetLock(name) lock.Lock() + state.pgLockConnMu.Lock() + defer state.pgLockConnMu.Unlock() // Convert string to sha256 hash and use decimal value for advisory lock h := sha256.Sum256([]byte(name)) nameInt := int64(binary.BigEndian.Uint64(h[:8])) - err := state.DB.Exec("SELECT pg_advisory_lock($1)", nameInt).Error + log.Debug(context.Background(), fmt.Sprintf("starting SELECT pg_advisory_lock(%d)", nameInt)) + err := state.pgLockConn.Exec("SELECT pg_advisory_lock($1)", nameInt).Error if err != nil { lock.Unlock() return nil, err } return func() { - err := state.DB.Exec("SELECT pg_advisory_unlock($1)", nameInt).Error - lock.Unlock() + state.pgLockConnMu.Lock() + defer state.pgLockConnMu.Unlock() + log.Debug(context.Background(), fmt.Sprintf("starting SELECT pg_advisory_unlock(%d)", nameInt)) + var unlocked bool + err := state.pgLockConn.Raw("SELECT pg_advisory_unlock($1)", nameInt).Scan(&unlocked).Error + if err == nil && !unlocked { + err = fmt.Errorf("pg_advisory_unlock returned false") + } if err != nil { // unfortunate, but the risk is that we're holding on to the lock forever, // so it's responsible to crash in this case panic(fmt.Errorf("error unlocking named lock: %w", err)) } + lock.Unlock() }, nil } +// startLockerConn starts a dedicated connection to the database for locking +func (state *StatefulDB) startPostgresLockerConn(ctx context.Context) error { + done := make(chan struct{}) + var err error + go func() { + err = state.DB.Connection(func(tx *gorm.DB) error { + state.pgLockConn = tx + close(done) + // hold this open until the context is done + <-ctx.Done() + return nil + }) + if err != nil { + close(done) + } + }() + <-done + return err +} + func (state *StatefulDB) getNamedLockSQLite(name string) (func(), error) { lock := state.locks.GetLock(name) lock.Lock() diff --git a/pkg/statedb/locks_test.go b/pkg/statedb/locks_test.go index b112e491..0aff3cf9 100644 --- a/pkg/statedb/locks_test.go +++ b/pkg/statedb/locks_test.go @@ -1,16 +1,19 @@ package statedb import ( + "context" "fmt" "net/url" "os" "os/exec" "strings" + "sync/atomic" "testing" "time" "github.com/google/uuid" "github.com/stretchr/testify/require" + "golang.org/x/sync/errgroup" "gorm.io/driver/postgres" "stream.place/streamplace/pkg/config" "stream.place/streamplace/pkg/model" @@ -86,6 +89,8 @@ func makePostgresURL(t *testing.T) string { return u.String() } +var lockRuns = 50000 + func TestPostgresLocks(t *testing.T) { if postgresURL == "" { t.Skip("no postgres url, skipping postgres tests") @@ -97,40 +102,28 @@ func TestPostgresLocks(t *testing.T) { } mod, err := model.MakeDB(":memory:") require.NoError(t, err) - state, err := MakeDB(&cli, nil, mod) - require.NoError(t, err) - - unlock, err := state.GetNamedLock("test") - t.Log("got lock") + state, err := MakeDB(context.Background(), &cli, nil, mod) require.NoError(t, err) - require.NotNil(t, unlock) - shouldBeLocked := true + var g errgroup.Group + var count atomic.Uint64 - done := make(chan struct{}) - - go func() { - unlock2, err := state.GetNamedLock("test") - t.Log("got lock 2") - require.Equal(t, shouldBeLocked, false) + doLock := func() error { + unlock, err := state.GetNamedLock("test") require.NoError(t, err) - require.NotNil(t, unlock2) - unlock2() - close(done) - }() - - time.Sleep(1 * time.Second) - - t.Log("unlocking") - shouldBeLocked = false - unlock() - t.Log("unlocked") - - select { - case <-done: - case <-time.After(1 * time.Second): - require.Fail(t, "lock not released") + defer unlock() + count.Add(1) + return nil } + + for i := 0; i < lockRuns; i++ { + g.Go(doLock) + } + + err = g.Wait() + require.NoError(t, err) + require.Equal(t, int(count.Load()), int(uint64(lockRuns))) + sqlDB, err := state.DB.DB() require.NoError(t, err) diff --git a/pkg/statedb/migrate.go b/pkg/statedb/migrate.go index f1cd94a5..43243efb 100644 --- a/pkg/statedb/migrate.go +++ b/pkg/statedb/migrate.go @@ -23,7 +23,7 @@ func Migrate(cli *config.CLI) error { // slogGorm.WithTraceAll(), ) - newDB, err := MakeDB(cli, nil, nil) + newDB, err := MakeDB(context.Background(), cli, nil, nil) if err != nil { return err } diff --git a/pkg/statedb/statedb.go b/pkg/statedb/statedb.go index d7784a5f..28d089ca 100644 --- a/pkg/statedb/statedb.go +++ b/pkg/statedb/statedb.go @@ -6,6 +6,7 @@ import ( "net/url" "os" "strings" + "sync" "time" "github.com/lmittmann/tint" @@ -36,6 +37,9 @@ type StatefulDB struct { model model.Model // pokeQueue is used to wake up the queue processor when a new task is enqueued pokeQueue chan struct{} + // pgLockConn is used to hold a connection to the database for locking + pgLockConn *gorm.DB + pgLockConnMu sync.Mutex } // list tables here so we can migrate them @@ -51,9 +55,9 @@ var StatefulDBModels = []any{ var NoPostgresDatabaseCode = "3D000" // Stateful database for storing private streamplace state -func MakeDB(cli *config.CLI, noter notificationpkg.FirebaseNotifier, model model.Model) (*StatefulDB, error) { +func MakeDB(ctx context.Context, cli *config.CLI, noter notificationpkg.FirebaseNotifier, model model.Model) (*StatefulDB, error) { dbURL := cli.DBURL - log.Log(context.Background(), "starting stateful database", "dbURL", redactDBURL(dbURL)) + log.Log(ctx, "starting stateful database", "dbURL", redactDBURL(dbURL)) var dial gorm.Dialector var dbType DBType if dbURL == ":memory:" { @@ -98,14 +102,21 @@ func MakeDB(cli *config.CLI, noter notificationpkg.FirebaseNotifier, model model return nil, err } } - return &StatefulDB{ + state := &StatefulDB{ DB: db, CLI: cli, Type: dbType, locks: NewNamedLocks(), model: model, pokeQueue: make(chan struct{}, 1), - }, nil + } + if state.Type == DBTypePostgres { + err = state.startPostgresLockerConn(ctx) + if err != nil { + return nil, fmt.Errorf("error starting postgres locker connection: %w", err) + } + } + return state, nil } func openDB(dial gorm.Dialector) (*gorm.DB, error) { @@ -113,7 +124,7 @@ func openDB(dial gorm.Dialector) (*gorm.DB, error) { slogGorm.WithHandler(tint.NewHandler(os.Stderr, &tint.Options{ TimeFormat: time.RFC3339, })), - // slogGorm.WithTraceAll(), + slogGorm.WithTraceAll(), ) return gorm.Open(dial, &gorm.Config{ -- 2.51.2