Reuse a stale non-completed snapshot request when snapshots are repeatable

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

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)
+		})
+	}
+}

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))}))

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)
 }