diff --git a/pkg/snapshot/generator/snapshot_generator_recorder.go b/pkg/snapshot/generator/snapshot_generator_recorder.go index 0b6d8578d..0b42251e4 100644 --- a/pkg/snapshot/generator/snapshot_generator_recorder.go +++ b/pkg/snapshot/generator/snapshot_generator_recorder.go @@ -82,11 +82,24 @@ func (s *SnapshotRecorder) CreateSnapshot(ctx context.Context, ss *snapshot.Snap requests := s.createRequests(ss) + // only one run at a time may work on a given schema and table list. A + // leftover non-completed request is reused below only while holding this + // lock: a live run would still hold it, since the lock is released only + // when that run ends or its process dies. + release, err := s.store.LockSnapshotRequests(ctx, requests) + if err != nil { + if errors.Is(err, snapshotstore.ErrSnapshotRequestLocked) { + return fmt.Errorf("another snapshot of the same schema and tables is running, wait for it to finish: %w", err) + } + return err + } + defer release() + if err := s.markSnapshotInProgress(ctx, requests); err != nil { return err } - err := s.wrapped.CreateSnapshot(ctx, ss) + err = s.wrapped.CreateSnapshot(ctx, ss) return s.markSnapshotCompleted(ctx, requests, err) } @@ -159,18 +172,53 @@ func (s *SnapshotRecorder) markSnapshotInProgress(ctx context.Context, requests // create one request per schema for _, req := range requests { eg.Go(func() error { - if err := s.store.CreateSnapshotRequest(ctx, req); err != nil { - return err - } - // the snapshot will start immediately - req.MarkInProgress() - return s.store.UpdateSnapshotRequest(ctx, req) + return s.startRequest(ctx, req) }) } return eg.Wait() } +// startRequest records the request as in progress. A non-completed request +// for the same schema and table list can only be left behind by a run that +// has ended, since the caller holds the request lock. Repeatable snapshots +// take it over, otherwise the conflict is reported. +func (s *SnapshotRecorder) startRequest(ctx context.Context, req *snapshot.Request) error { + err := s.store.CreateSnapshotRequest(ctx, req) + if err == nil { + // the snapshot will start immediately + req.MarkInProgress() + return s.store.UpdateSnapshotRequest(ctx, req) + } + if !errors.Is(err, snapshotstore.ErrNonCompletedRequestExists) { + return err + } + if !s.repeatableSnapshots { + return fmt.Errorf("a non-completed snapshot request for schema %q already exists in %s.%s; if no snapshot is running, delete that row or mark it completed: %w", req.Schema, snapshotstore.SchemaName, snapshotstore.TableName, err) + } + + resumed, resumeErr := s.store.ResumeSnapshotRequest(ctx, req) + if resumeErr != nil { + return resumeErr + } + if resumed { + s.logger.Warn(err, "reusing non-completed snapshot request left by a previous run", loglib.Fields{ + "schema": req.Schema, + "mode": req.GetMode(), + }) + req.MarkInProgress() + return nil + } + + // the leftover request was completed between the insert and the take + // over, so the insert can go through now + if err := s.store.CreateSnapshotRequest(ctx, req); err != nil { + return err + } + req.MarkInProgress() + return s.store.UpdateSnapshotRequest(ctx, req) +} + func (s *SnapshotRecorder) markSnapshotCompleted(ctx context.Context, requests []*snapshot.Request, err error) error { // make sure we can update the request status in the store regardless of // context cancelations diff --git a/pkg/snapshot/generator/snapshot_generator_recorder_test.go b/pkg/snapshot/generator/snapshot_generator_recorder_test.go index 2cf5339ba..c79413d29 100644 --- a/pkg/snapshot/generator/snapshot_generator_recorder_test.go +++ b/pkg/snapshot/generator/snapshot_generator_recorder_test.go @@ -1007,3 +1007,164 @@ func TestSnapshotRecorder_filterOutExistingSnapshots(t *testing.T) { }) } } + +// errStaleRequest is what the store returns when a non-completed request for +// the same schema and tables is already recorded +var errStaleRequest = fmt.Errorf("error creating snapshot request: %w: %w", snapshotstore.ErrNonCompletedRequestExists, + errors.New(`constraint violation: duplicate key value violates unique constraint "schema_table_status_hash_unique_index"`)) + +func TestSnapshotRecorder_CreateSnapshot_staleRequest(t *testing.T) { + t.Parallel() + + errTest := errors.New("oh noes") + + tests := []struct { + name string + repeatable bool + createErrs []error + lockErr error + resumed bool + resumeErr error + generateErr error + + wantErr error + wantErrContains string + wantCreates int + wantResumes int + wantGenerated bool + wantStatuses []snapshot.Status + }{ + { + name: "repeatable takes over the stale request", + repeatable: true, + createErrs: []error{errStaleRequest}, + resumed: true, + wantCreates: 1, + wantResumes: 1, + wantGenerated: true, + wantStatuses: []snapshot.Status{snapshot.StatusCompleted}, + }, + { + name: "repeatable inserts again when the stale request completed before the take over", + repeatable: true, + createErrs: []error{errStaleRequest, nil}, + resumed: false, + wantCreates: 2, + wantResumes: 1, + wantGenerated: true, + wantStatuses: []snapshot.Status{snapshot.StatusInProgress, snapshot.StatusCompleted}, + }, + { + name: "repeatable fails when the take over fails", + repeatable: true, + createErrs: []error{errStaleRequest}, + resumeErr: errTest, + wantErr: errTest, + wantCreates: 1, + wantResumes: 1, + }, + { + name: "non repeatable conflict is actionable", + repeatable: false, + createErrs: []error{errStaleRequest}, + wantErr: snapshotstore.ErrNonCompletedRequestExists, + wantErrContains: "pgstream.snapshot_requests", + wantCreates: 1, + }, + { + name: "other constraint errors are not taken for a stale request", + repeatable: true, + createErrs: []error{errTest}, + wantErr: errTest, + wantCreates: 1, + }, + { + name: "another run holds the lock", + repeatable: true, + lockErr: fmt.Errorf("schema \"public\": %w", snapshotstore.ErrSnapshotRequestLocked), + wantErr: snapshotstore.ErrSnapshotRequestLocked, + wantErrContains: "another snapshot of the same schema and tables is running", + }, + { + name: "lock is released when the snapshot fails", + repeatable: true, + createErrs: []error{nil}, + generateErr: errTest, + wantErr: errTest, + wantCreates: 1, + wantGenerated: true, + wantStatuses: []snapshot.Status{snapshot.StatusInProgress, snapshot.StatusCompleted}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + creates, resumes, locks, releases := 0, 0, 0, 0 + var statuses []snapshot.Status + store := &snapshotstoremocks.Store{ + LockSnapshotRequestsFn: func(ctx context.Context, reqs []*snapshot.Request) (func(), error) { + locks++ + require.Len(t, reqs, 1) + if tc.lockErr != nil { + return nil, tc.lockErr + } + return func() { releases++ }, nil + }, + CreateSnapshotRequestFn: func(ctx context.Context, r *snapshot.Request) error { + creates++ + require.LessOrEqual(t, creates, len(tc.createErrs)) + return tc.createErrs[creates-1] + }, + ResumeSnapshotRequestFn: func(ctx context.Context, r *snapshot.Request) (bool, error) { + resumes++ + require.Equal(t, "public", r.Schema) + require.Equal(t, snapshot.RequestModeData, r.GetMode()) + return tc.resumed, tc.resumeErr + }, + UpdateSnapshotRequestFn: func(ctx context.Context, _ uint, r *snapshot.Request) error { + statuses = append(statuses, r.Status) + return nil + }, + GetSnapshotRequestsBySchemaFn: func(ctx context.Context, s string) ([]*snapshot.Request, error) { + return []*snapshot.Request{}, nil + }, + } + + generated := false + generator := &mockGenerator{ + createSnapshotFn: func(ctx context.Context, ss *snapshot.Snapshot) error { + generated = true + return tc.generateErr + }, + } + + sr := NewSnapshotRecorder(&Config{ + RepeatableSnapshots: tc.repeatable, + SnapshotWorkers: 1, + }, store, generator) + defer sr.Close() + + err := sr.CreateSnapshot(context.Background(), &snapshot.Snapshot{ + SchemaTables: map[string][]string{"public": {"t1"}}, + }) + if tc.wantErr != nil { + require.ErrorIs(t, err, tc.wantErr) + } else { + require.NoError(t, err) + } + if tc.wantErrContains != "" { + require.ErrorContains(t, err, tc.wantErrContains) + } + require.Equal(t, 1, locks) + if tc.lockErr == nil { + require.Equal(t, 1, releases) + } + require.Equal(t, tc.wantCreates, creates) + require.Equal(t, tc.wantResumes, resumes) + require.Equal(t, tc.wantGenerated, generated) + require.Equal(t, tc.wantStatuses, statuses) + }) + } +} diff --git a/pkg/snapshot/store/instrumentation/instrumented_snapshot_store.go b/pkg/snapshot/store/instrumentation/instrumented_snapshot_store.go index 90eac3e2d..afaac64ab 100644 --- a/pkg/snapshot/store/instrumentation/instrumented_snapshot_store.go +++ b/pkg/snapshot/store/instrumentation/instrumented_snapshot_store.go @@ -43,6 +43,21 @@ func (i *Store) UpdateSnapshotRequest(ctx context.Context, req *snapshot.Request return i.inner.UpdateSnapshotRequest(ctx, req) } +func (i *Store) LockSnapshotRequests(ctx context.Context, reqs []*snapshot.Request) (release func(), err error) { + ctx, span := otel.StartSpan(ctx, i.tracer, "snapshotStore.LockSnapshotRequests", + trace.WithAttributes(attribute.KeyValue{Key: "requests", Value: attribute.IntValue(len(reqs))})) + defer otel.CloseSpan(span, err) + + return i.inner.LockSnapshotRequests(ctx, reqs) +} + +func (i *Store) ResumeSnapshotRequest(ctx context.Context, req *snapshot.Request) (resumed bool, err error) { + ctx, span := otel.StartSpan(ctx, i.tracer, "snapshotStore.ResumeSnapshotRequest", trace.WithAttributes(i.requestAttributes(req)...)) + defer otel.CloseSpan(span, err) + + return i.inner.ResumeSnapshotRequest(ctx, req) +} + func (i *Store) GetSnapshotRequestsByStatus(ctx context.Context, status snapshot.Status) (req []*snapshot.Request, err error) { ctx, span := otel.StartSpan(ctx, i.tracer, "snapshotStore.GetSnapshotRequestsByStatus", trace.WithAttributes(attribute.KeyValue{Key: "status", Value: attribute.StringValue(string(status))})) diff --git a/pkg/snapshot/store/mocks/mock_snapshot_store.go b/pkg/snapshot/store/mocks/mock_snapshot_store.go index d2b8a8e46..b1ca0ba2a 100644 --- a/pkg/snapshot/store/mocks/mock_snapshot_store.go +++ b/pkg/snapshot/store/mocks/mock_snapshot_store.go @@ -13,7 +13,10 @@ type Store struct { UpdateSnapshotRequestFn func(context.Context, uint, *snapshot.Request) error GetSnapshotRequestsByStatusFn func(ctx context.Context, status snapshot.Status) ([]*snapshot.Request, error) GetSnapshotRequestsBySchemaFn func(ctx context.Context, s string) ([]*snapshot.Request, error) - updateSnapshotRequestCalls uint + // LockSnapshotRequestsFn defaults to granting every lock when nil + LockSnapshotRequestsFn func(context.Context, []*snapshot.Request) (func(), error) + ResumeSnapshotRequestFn func(context.Context, *snapshot.Request) (bool, error) + updateSnapshotRequestCalls uint } func (m *Store) CreateSnapshotRequest(ctx context.Context, s *snapshot.Request) error { @@ -25,6 +28,17 @@ func (m *Store) UpdateSnapshotRequest(ctx context.Context, s *snapshot.Request) return m.UpdateSnapshotRequestFn(ctx, m.updateSnapshotRequestCalls, s) } +func (m *Store) LockSnapshotRequests(ctx context.Context, reqs []*snapshot.Request) (func(), error) { + if m.LockSnapshotRequestsFn == nil { + return func() {}, nil + } + return m.LockSnapshotRequestsFn(ctx, reqs) +} + +func (m *Store) ResumeSnapshotRequest(ctx context.Context, req *snapshot.Request) (bool, error) { + return m.ResumeSnapshotRequestFn(ctx, req) +} + func (m *Store) GetSnapshotRequestsByStatus(ctx context.Context, status snapshot.Status) ([]*snapshot.Request, error) { return m.GetSnapshotRequestsByStatusFn(ctx, status) } diff --git a/pkg/snapshot/store/postgres/integration/pg_snapshot_store_recovery_integration_test.go b/pkg/snapshot/store/postgres/integration/pg_snapshot_store_recovery_integration_test.go new file mode 100644 index 000000000..282886a4e --- /dev/null +++ b/pkg/snapshot/store/postgres/integration/pg_snapshot_store_recovery_integration_test.go @@ -0,0 +1,255 @@ +// SPDX-License-Identifier: Apache-2.0 + +package integration + +import ( + "context" + "fmt" + "os" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + pglib "github.com/xataio/pgstream/internal/postgres" + "github.com/xataio/pgstream/pkg/snapshot" + "github.com/xataio/pgstream/pkg/snapshot/generator" + "github.com/xataio/pgstream/pkg/snapshot/store" + pgstore "github.com/xataio/pgstream/pkg/snapshot/store/postgres" +) + +// Test_SnapshotStore_LockSnapshotRequests checks that only one holder can lock +// a schema and table list at a time, whatever the request mode, and that the +// lock is free again once released. +func Test_SnapshotStore_LockSnapshotRequests(t *testing.T) { + if os.Getenv("PGSTREAM_INTEGRATION_TESTS") == "" { + t.Skip("skipping integration test...") + } + + ctx := context.Background() + s, err := pgstore.New(ctx, sharedPGURL) + require.NoError(t, err) + defer s.Close() + + dataReq := &snapshot.Request{Schema: "lock_schema", Tables: []string{"t1"}, Mode: snapshot.RequestModeData} + schemaOnlyReq := &snapshot.Request{Schema: "lock_schema", Tables: []string{"t1"}, Mode: snapshot.RequestModeSchemaOnly} + otherReq := &snapshot.Request{Schema: "lock_schema", Tables: []string{"t2"}} + + release, err := s.LockSnapshotRequests(ctx, []*snapshot.Request{dataReq}) + require.NoError(t, err) + + // the unique index ignores mode, so the lock does too + _, err = s.LockSnapshotRequests(ctx, []*snapshot.Request{schemaOnlyReq}) + require.ErrorIs(t, err, store.ErrSnapshotRequestLocked) + + // a different table list is a different scope + releaseOther, err := s.LockSnapshotRequests(ctx, []*snapshot.Request{otherReq}) + require.NoError(t, err) + releaseOther() + + release() + release, err = s.LockSnapshotRequests(ctx, []*snapshot.Request{schemaOnlyReq}) + require.NoError(t, err) + release() +} + +// Test_SnapshotStore_LockReleasedWhenHolderDies checks that the server frees +// the lock of a run whose connection is gone, as after a kill -9, so a later +// run can recover its request. +func Test_SnapshotStore_LockReleasedWhenHolderDies(t *testing.T) { + if os.Getenv("PGSTREAM_INTEGRATION_TESTS") == "" { + t.Skip("skipping integration test...") + } + + ctx := context.Background() + s, err := pgstore.New(ctx, sharedPGURL) + require.NoError(t, err) + defer s.Close() + + req := &snapshot.Request{Schema: "lock_dies_schema", Tables: []string{"t1"}} + release, err := s.LockSnapshotRequests(ctx, []*snapshot.Request{req}) + require.NoError(t, err) + defer release() + + _, err = s.LockSnapshotRequests(ctx, []*snapshot.Request{req}) + require.ErrorIs(t, err, store.ErrSnapshotRequestLocked) + + // end the holder's session from outside, like a dead process would + conn, err := pglib.NewConn(ctx, sharedPGURL) + require.NoError(t, err) + defer conn.Close(ctx) + var terminated int + err = conn.QueryRow(ctx, []any{&terminated}, fmt.Sprintf(`SELECT count(pg_terminate_backend(l.pid))::int FROM pg_locks l + WHERE l.locktype = 'advisory' AND l.granted + AND ((l.classid::bigint << 32) | l.objid::bigint) = hashtextextended($1::text || '/' || %s.snapshot_table_names_hash($2::text[]), 0)`, store.SchemaName), + req.Schema, req.Tables) + require.NoError(t, err) + require.Equal(t, 1, terminated) + + require.Eventually(t, func() bool { + release, err := s.LockSnapshotRequests(ctx, []*snapshot.Request{req}) + if err != nil { + return false + } + release() + return true + }, 10*time.Second, 100*time.Millisecond) +} + +// Test_SnapshotStore_ResumeSnapshotRequest checks the take over of a leftover +// request: it matches the unique index key whatever the mode, records the new +// mode so the completion update finds the row, and reports when there is +// nothing left to take over. +func Test_SnapshotStore_ResumeSnapshotRequest(t *testing.T) { + if os.Getenv("PGSTREAM_INTEGRATION_TESTS") == "" { + t.Skip("skipping integration test...") + } + + ctx := context.Background() + clearRequests(t, "resume_schema") + + s, err := pgstore.New(ctx, sharedPGURL) + require.NoError(t, err) + defer s.Close() + + // a schema-only run left its request behind + leftover := &snapshot.Request{Schema: "resume_schema", Tables: []string{"t1"}, Mode: snapshot.RequestModeSchemaOnly} + require.NoError(t, s.CreateSnapshotRequest(ctx, leftover)) + + // a data run for the same tables conflicts on the index, which ignores mode + req := &snapshot.Request{Schema: "resume_schema", Tables: []string{"t1"}, Mode: snapshot.RequestModeData} + require.ErrorIs(t, s.CreateSnapshotRequest(ctx, req), store.ErrNonCompletedRequestExists) + + resumed, err := s.ResumeSnapshotRequest(ctx, req) + require.NoError(t, err) + require.True(t, resumed) + requireRequests(t, s, "resume_schema", []snapshot.Status{snapshot.StatusInProgress}, snapshot.RequestModeData) + + // the completion update is keyed on mode and must find the row + req.MarkCompleted(req.Schema, nil) + require.NoError(t, s.UpdateSnapshotRequest(ctx, req)) + requireRequests(t, s, "resume_schema", []snapshot.Status{snapshot.StatusCompleted}, snapshot.RequestModeData) + + // nothing non-completed is left, so there is nothing to take over + resumed, err = s.ResumeSnapshotRequest(ctx, req) + require.NoError(t, err) + require.False(t, resumed) +} + +// Test_SnapshotRecorder_ConcurrentRuns runs the recorder against real postgres +// with repeatable snapshots. While a slow run is copying, two retries started +// together are refused and leave its request alone, so neither can complete +// it early and the slow run's own completion updates its own row. A run that +// died (its request is left in progress and its lock is gone) is recovered. +func Test_SnapshotRecorder_ConcurrentRuns(t *testing.T) { + if os.Getenv("PGSTREAM_INTEGRATION_TESTS") == "" { + t.Skip("skipping integration test...") + } + + ctx := context.Background() + const schema = "concurrent_schema" + clearRequests(t, schema) + + newRecorder := func(gen generator.SnapshotGenerator) *generator.SnapshotRecorder { + s, err := pgstore.New(ctx, sharedPGURL) + require.NoError(t, err) + return generator.NewSnapshotRecorder(&generator.Config{RepeatableSnapshots: true}, s, gen) + } + newSnapshot := func() *snapshot.Snapshot { + return &snapshot.Snapshot{SchemaTables: map[string][]string{schema: {"t1"}}} + } + + s, err := pgstore.New(ctx, sharedPGURL) + require.NoError(t, err) + defer s.Close() + + // slow writer: blocks until told to finish + started := make(chan struct{}) + finish := make(chan struct{}) + slow := newRecorder(&blockingGenerator{started: started, finish: finish}) + defer slow.Close() + slowErr := make(chan error, 1) + go func() { slowErr <- slow.CreateSnapshot(ctx, newSnapshot()) }() + <-started + requireRequests(t, s, schema, []snapshot.Status{snapshot.StatusInProgress}, snapshot.RequestModeData) + + // two retries at the same time while the slow run is still copying. The + // recorders are built first, since building a store runs its table setup. + retryGens := []*recordingGenerator{{}, {}} + retries := make([]*generator.SnapshotRecorder, 0, len(retryGens)) + for _, gen := range retryGens { + r := newRecorder(gen) + defer r.Close() + retries = append(retries, r) + } + var wg sync.WaitGroup + retryErrs := make([]error, len(retries)) + for i, r := range retries { + wg.Add(1) + go func() { + defer wg.Done() + retryErrs[i] = r.CreateSnapshot(ctx, newSnapshot()) + }() + } + wg.Wait() + for i := range retries { + require.ErrorIs(t, retryErrs[i], store.ErrSnapshotRequestLocked) + require.False(t, retryGens[i].called, "a refused retry must not copy") + } + // the slow run's request is untouched by the refused retries + requireRequests(t, s, schema, []snapshot.Status{snapshot.StatusInProgress}, snapshot.RequestModeData) + + close(finish) + require.NoError(t, <-slowErr) + requireRequests(t, s, schema, []snapshot.Status{snapshot.StatusCompleted}, snapshot.RequestModeData) + + // a run that died: its request stays in progress and nobody holds its lock + dead := &snapshot.Request{Schema: schema, Tables: []string{"t1"}, Mode: snapshot.RequestModeData} + require.NoError(t, s.CreateSnapshotRequest(ctx, dead)) + dead.MarkInProgress() + require.NoError(t, s.UpdateSnapshotRequest(ctx, dead)) + + gen := &recordingGenerator{} + recovery := newRecorder(gen) + defer recovery.Close() + require.NoError(t, recovery.CreateSnapshot(ctx, newSnapshot())) + require.True(t, gen.called) + requireRequests(t, s, schema, []snapshot.Status{snapshot.StatusCompleted, snapshot.StatusCompleted}, snapshot.RequestModeData) +} + +type blockingGenerator struct { + started chan struct{} + finish chan struct{} +} + +func (g *blockingGenerator) CreateSnapshot(ctx context.Context, _ *snapshot.Snapshot) error { + close(g.started) + <-g.finish + return nil +} + +func (g *blockingGenerator) Close() error { return nil } + +type recordingGenerator struct { + called bool +} + +func (g *recordingGenerator) CreateSnapshot(context.Context, *snapshot.Snapshot) error { + g.called = true + return nil +} + +func (g *recordingGenerator) Close() error { return nil } + +func requireRequests(t *testing.T, s *pgstore.Store, schema string, wantStatuses []snapshot.Status, wantMode snapshot.RequestMode) { + t.Helper() + + reqs, err := s.GetSnapshotRequestsBySchema(context.Background(), schema) + require.NoError(t, err) + statuses := make([]snapshot.Status, 0, len(reqs)) + for _, req := range reqs { + statuses = append(statuses, req.Status) + require.Equal(t, wantMode, req.GetMode()) + } + require.Equal(t, wantStatuses, statuses) +} diff --git a/pkg/snapshot/store/postgres/pg_snapshot_store.go b/pkg/snapshot/store/postgres/pg_snapshot_store.go index 7a0fb94b8..ca0981f71 100644 --- a/pkg/snapshot/store/postgres/pg_snapshot_store.go +++ b/pkg/snapshot/store/postgres/pg_snapshot_store.go @@ -4,7 +4,10 @@ package postgres import ( "context" + "errors" "fmt" + "slices" + "strings" "github.com/lib/pq" "github.com/xataio/pgstream/internal/postgres" @@ -14,9 +17,17 @@ import ( type Store struct { conn postgres.Querier + // lockConnBuilder opens the dedicated connection that holds the request + // locks. Session level advisory locks belong to one connection, so they + // can't be taken through the pool. + lockConnBuilder func(context.Context) (postgres.Querier, error) } -const queryLimit = 1000 +const ( + queryLimit = 1000 + + uniqueIndexName = "schema_table_status_hash_unique_index" +) func New(ctx context.Context, url string) (*Store, error) { conn, err := postgres.NewConnPool(ctx, url) @@ -26,6 +37,9 @@ func New(ctx context.Context, url string) (*Store, error) { s := &Store{ conn: conn, + lockConnBuilder: func(ctx context.Context) (postgres.Querier, error) { + return postgres.NewConn(ctx, url) + }, } // create snapshots table if it doesn't exist @@ -45,12 +59,79 @@ func (s *Store) CreateSnapshotRequest(ctx context.Context, req *snapshot.Request VALUES($1, $2, now(), now(),'requested', $3)`, snapshotsTable()) _, err := s.conn.Exec(ctx, query, req.Schema, pq.StringArray(req.Tables), req.GetMode()) if err != nil { + if isNonCompletedRequestConflict(err) { + return fmt.Errorf("error creating snapshot request: %w: %w", store.ErrNonCompletedRequestExists, err) + } return fmt.Errorf("error creating snapshot request: %w", err) } return nil } +// isNonCompletedRequestConflict reports whether err is the unique violation +// on the index that allows one non-completed request per schema and table +// list. Other integrity errors (check or not null constraints) are not. +func isNonCompletedRequestConflict(err error) bool { + constraintErr := &postgres.ErrConstraintViolation{} + return errors.As(err, &constraintErr) && strings.Contains(constraintErr.Details, uniqueIndexName) +} + +// LockSnapshotRequests takes a session level advisory lock per request, keyed +// like the unique index on the schema name and the table names hash (mode is +// not part of the key). The locks live on a dedicated connection, so they are +// released when release closes it, or by the server when the process dies and +// its connection drops. +func (s *Store) LockSnapshotRequests(ctx context.Context, reqs []*snapshot.Request) (func(), error) { + conn, err := s.lockConnBuilder(ctx) + if err != nil { + return nil, fmt.Errorf("opening snapshot request lock connection: %w", err) + } + release := func() { + conn.Close(context.Background()) + } + + // take the locks in a fixed order so two runs over the same scopes + // contend on the same first lock + sorted := slices.Clone(reqs) + slices.SortFunc(sorted, func(a, b *snapshot.Request) int { + return strings.Compare(lockKey(a), lockKey(b)) + }) + + query := fmt.Sprintf(`SELECT pg_try_advisory_lock(hashtextextended($1::text || '/' || %s.snapshot_table_names_hash($2::text[]), 0))`, store.SchemaName) + for _, req := range sorted { + var locked bool + if err := conn.QueryRow(ctx, []any{&locked}, query, req.Schema, pq.StringArray(req.Tables)); err != nil { + release() + return nil, fmt.Errorf("locking snapshot request for schema %q: %w", req.Schema, err) + } + if !locked { + release() + return nil, fmt.Errorf("schema %q tables %v: %w", req.Schema, req.Tables, store.ErrSnapshotRequestLocked) + } + } + + return release, nil +} + +func lockKey(req *snapshot.Request) string { + return req.Schema + "/" + strings.Join(req.Tables, ",") +} + +// ResumeSnapshotRequest takes over the non-completed request matching the +// unique index key (schema name and table names, any mode), marks it in +// progress and records the mode on input, so later updates keyed on mode +// match it. +func (s *Store) ResumeSnapshotRequest(ctx context.Context, req *snapshot.Request) (bool, error) { + query := fmt.Sprintf(`UPDATE %s SET status = 'in progress', mode = $1, errors = NULL, updated_at = now() + WHERE schema_name = $2 and table_names = $3 and status != 'completed'`, snapshotsTable()) + tag, err := s.conn.Exec(ctx, query, req.GetMode(), req.Schema, pq.StringArray(req.Tables)) + if err != nil { + return false, fmt.Errorf("error resuming snapshot request: %w", err) + } + + return tag.RowsAffected() > 0, nil +} + func (s *Store) UpdateSnapshotRequest(ctx context.Context, req *snapshot.Request) error { query := fmt.Sprintf(`UPDATE %s SET status = $1, errors = $2, updated_at = now() WHERE schema_name = $3 and table_names = $4 and mode = $5 and status != 'completed'`, snapshotsTable()) diff --git a/pkg/snapshot/store/postgres/pg_snapshot_store_test.go b/pkg/snapshot/store/postgres/pg_snapshot_store_test.go index 5ac6039d1..1c38d1602 100644 --- a/pkg/snapshot/store/postgres/pg_snapshot_store_test.go +++ b/pkg/snapshot/store/postgres/pg_snapshot_store_test.go @@ -8,6 +8,7 @@ import ( "fmt" "testing" + "github.com/jackc/pgx/v5/pgconn" "github.com/lib/pq" "github.com/stretchr/testify/require" pglib "github.com/xataio/pgstream/internal/postgres" @@ -585,3 +586,160 @@ func TestStore_createTable(t *testing.T) { }) } } + +func TestStore_CreateSnapshotRequest_conflict(t *testing.T) { + t.Parallel() + + req := &snapshot.Request{Schema: "public", Tables: []string{"t1"}} + + tests := []struct { + name string + execErr error + wantStale bool + }{ + { + name: "unique violation on the non-completed request index", + execErr: &pglib.ErrConstraintViolation{Details: `duplicate key value violates unique constraint "schema_table_status_hash_unique_index"`}, + wantStale: true, + }, + { + name: "check constraint violation", + execErr: &pglib.ErrConstraintViolation{Details: `new row for relation "snapshot_requests" violates check constraint "snapshot_requests_mode_check"`}, + wantStale: false, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + s := &Store{ + conn: &postgresmocks.Querier{ + ExecFn: func(ctx context.Context, _ uint, _ string, _ ...any) (pglib.CommandTag, error) { + return pglib.CommandTag{}, tc.execErr + }, + }, + } + err := s.CreateSnapshotRequest(context.Background(), req) + require.ErrorIs(t, err, tc.execErr) + require.Equal(t, tc.wantStale, errors.Is(err, store.ErrNonCompletedRequestExists)) + }) + } +} + +func TestStore_ResumeSnapshotRequest(t *testing.T) { + t.Parallel() + + req := &snapshot.Request{Schema: "public", Tables: []string{"t1"}, Mode: snapshot.RequestModeSchemaOnly} + errTest := errors.New("oh noes") + + tests := []struct { + name string + tag string + execErr error + + wantResumed bool + wantErr error + }{ + {name: "takes over the leftover request", tag: "UPDATE 1", wantResumed: true}, + {name: "no leftover request", tag: "UPDATE 0", wantResumed: false}, + {name: "error", execErr: errTest, wantErr: errTest}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + s := &Store{ + conn: &postgresmocks.Querier{ + ExecFn: func(ctx context.Context, _ uint, q string, a ...any) (pglib.CommandTag, error) { + wantQuery := fmt.Sprintf(`UPDATE %s SET status = 'in progress', mode = $1, errors = NULL, updated_at = now() + WHERE schema_name = $2 and table_names = $3 and status != 'completed'`, snapshotsTable()) + require.Equal(t, wantQuery, q) + require.Equal(t, []any{snapshot.RequestModeSchemaOnly, "public", pq.StringArray{"t1"}}, a) + return pglib.CommandTag{CommandTag: pgconn.NewCommandTag(tc.tag)}, tc.execErr + }, + }, + } + resumed, err := s.ResumeSnapshotRequest(context.Background(), req) + require.ErrorIs(t, err, tc.wantErr) + require.Equal(t, tc.wantResumed, resumed) + }) + } +} + +func TestStore_LockSnapshotRequests(t *testing.T) { + t.Parallel() + + reqs := []*snapshot.Request{ + {Schema: "b", Tables: []string{"t1"}}, + {Schema: "a", Tables: []string{"t1"}}, + } + errTest := errors.New("oh noes") + + tests := []struct { + name string + locked []bool + queryErr error + connErr error + wantErr error + wantCalls int + }{ + {name: "all locks taken", locked: []bool{true, true}, wantCalls: 2}, + {name: "second lock held by another run", locked: []bool{true, false}, wantErr: store.ErrSnapshotRequestLocked, wantCalls: 2}, + {name: "lock query fails", queryErr: errTest, wantErr: errTest, wantCalls: 1}, + {name: "lock connection fails", connErr: errTest, wantErr: errTest}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + calls, closes := 0, 0 + var schemas []string + conn := &postgresmocks.Querier{ + QueryRowFn: func(ctx context.Context, dest []any, q string, a ...any) error { + calls++ + schema, ok := a[0].(string) + require.True(t, ok) + schemas = append(schemas, schema) + if tc.queryErr != nil { + return tc.queryErr + } + locked, ok := dest[0].(*bool) + require.True(t, ok) + *locked = tc.locked[calls-1] + return nil + }, + CloseFn: func(context.Context) error { + closes++ + return nil + }, + } + s := &Store{ + lockConnBuilder: func(context.Context) (pglib.Querier, error) { + if tc.connErr != nil { + return nil, tc.connErr + } + return conn, nil + }, + } + + release, err := s.LockSnapshotRequests(context.Background(), reqs) + require.ErrorIs(t, err, tc.wantErr) + require.Equal(t, tc.wantCalls, calls) + if tc.wantCalls > 0 { + // locks are taken in a fixed order + require.Equal(t, "a", schemas[0]) + } + if err != nil { + require.Nil(t, release) + require.Equal(t, min(tc.wantCalls, 1), closes) + return + } + require.Equal(t, 0, closes) + release() + require.Equal(t, 1, closes) + }) + } +} diff --git a/pkg/snapshot/store/snapshot_store.go b/pkg/snapshot/store/snapshot_store.go index 4f31bd113..f03d63f0e 100644 --- a/pkg/snapshot/store/snapshot_store.go +++ b/pkg/snapshot/store/snapshot_store.go @@ -4,18 +4,37 @@ package store import ( "context" + "errors" "github.com/xataio/pgstream/pkg/snapshot" ) type Store interface { + // CreateSnapshotRequest records a new request. It returns an error + // wrapping ErrNonCompletedRequestExists when a non-completed request for + // the same schema and table list is already recorded. CreateSnapshotRequest(context.Context, *snapshot.Request) error UpdateSnapshotRequest(context.Context, *snapshot.Request) error + // LockSnapshotRequests takes a lock per request on its schema and table + // list, held until release is called or the process exits, so only one + // run can work on a given scope at a time. It returns an error wrapping + // ErrSnapshotRequestLocked when another run holds one of the locks. + LockSnapshotRequests(context.Context, []*snapshot.Request) (release func(), err error) + // ResumeSnapshotRequest takes over the non-completed request recorded for + // the same schema and table list as the one on input, whatever its mode, + // and marks it in progress with the mode on input. It reports whether a + // request was taken over. The caller must hold the request lock. + ResumeSnapshotRequest(context.Context, *snapshot.Request) (bool, error) GetSnapshotRequestsByStatus(ctx context.Context, status snapshot.Status) ([]*snapshot.Request, error) GetSnapshotRequestsBySchema(ctx context.Context, schema string) ([]*snapshot.Request, error) Close() error } +var ( + ErrNonCompletedRequestExists = errors.New("a non-completed snapshot request already exists") + ErrSnapshotRequestLocked = errors.New("snapshot request locked by another run") +) + const ( SchemaName = "pgstream" TableName = "snapshot_requests"