# PiTR Serialization Constraint Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Prevent concurrent PiTR restores that share the same log backup from running simultaneously, which would cause duplicate CompactBackup SST files on S3 and exhaust disk space.

**Architecture:** Two-layer defense. Layer 1 is a best-effort check in `makeCreateRestoreEntity` (pitr_restore.go) that rejects creation if an active PiTR restore for the same log backup already exists. Layer 2 is an atomic double-check inside the `execute()` transaction in `run_compact_backup.go` that catches any race (TOCTOU) that slips past Layer 1. Both layers are gated on `override.AllowConcurrentPiTR`; setting this to `true` bypasses all checks (escape hatch for operators).

**Tech Stack:** Go, GORM, testify suites, mockery-generated mocks, `bizerr`/`bizerr_wrapper` error types.

---

## File Map

| File | Change |
|------|--------|
| `internal/config/config.go` | Add `AllowConcurrentPiTR bool` field to `RestoreOverride` struct |
| `app/config/errcode.go` | Add `ErrMsgPiTRConcurrentRestoreNotAllowed` constant |
| `app/models/restore/impl/pitr_restore.go` | Layer 1 check in `makeCreateRestoreEntity()` |
| `app/models/restore/impl/pitr_restore_test.go` | New test file for Layer 1 |
| `app/workflow/action/restore/pitr/create/run_compact_backup.go` | Layer 2 check in `execute()` + pass `override` param |
| `app/workflow/action/restore/pitr/create/run_compact_backup_test.go` | Add Layer 2 `execute()` tests |

---

## Task 1: Add Config Field and Error Constant

**Files:**
- Modify: `internal/config/config.go:801-806`
- Modify: `app/config/errcode.go:247-256`

- [ ] **Step 1: Add `AllowConcurrentPiTR` to `RestoreOverride`**

  In `internal/config/config.go`, find the `RestoreOverride` struct (lines 801–806) and add the new field after `CompactionAcceleratedPiTR`:

  ```go
  // RestoreOverride set Restore override for clusters matched by FeatureGate.
  type RestoreOverride struct {
  	Gate FeatureGate `mapstructure:"gate"`

  	// CompactionAcceleratedPiTR controls whether a backup compaction(compactbackup) is performed before PiTR.
  	CompactionAcceleratedPiTR bool `mapstructure:"compactionAcceleratedPiTR" default:"false"`

  	// AllowConcurrentPiTR bypasses the serialization constraint that prevents multiple PiTR
  	// restores from running concurrently for the same log backup. Set to true only when you
  	// deliberately need to override this safety guard.
  	AllowConcurrentPiTR bool `mapstructure:"allowConcurrentPiTR" default:"false"`
  }
  ```

- [ ] **Step 2: Add error message constant**

  In `app/config/errcode.go`, append to the `restorev2 error message define` const block (after `ErrMsgTikvStorageSizeInvalid`):

  ```go
  ErrMsgPiTRConcurrentRestoreNotAllowed = "another PiTR restore is already running for the same log backup, concurrent PiTR restores are not allowed"
  ```

  The full block becomes:
  ```go
  // restorev2 error message define
  const (
  	ErrMsgBackupNotFound                = "backup(id= %d) is not found"
  	ErrMsgRestoreNotFound               = "restore is not found"
  	ErrMsgBackupStatusIsNotAllowRestore = "backup %d status is not allow restore"
  	ErrMsgComponentStorageSizeInvalid   = "requested storage size of component %s is less than the storage size of backup cluster"
  	ErrMsgComponentNodeQuantityInvalid  = "requested node quantity of component %s is less than the node quantity of backup cluster"
  	ErrMsgRestoreTsNotAllow             = "it can't restore to ts, because it is not in the range of backup ts"
  	ErrMsgTikvStorageSizeInvalid        = "The storage size of each tikv node should be at least '%d' GiB, and it is recommended to reserve more"
  	ErrMsgPiTRConcurrentRestoreNotAllowed = "another PiTR restore is already running for the same log backup, concurrent PiTR restores are not allowed"
  )
  ```

- [ ] **Step 3: Verify it compiles**

  ```bash
  make build-dataflow
  ```
  Expected: binary produced, no errors.

- [ ] **Step 4: Commit**

  ```bash
  git add internal/config/config.go app/config/errcode.go
  git commit -m "feat(br): add AllowConcurrentPiTR config field and error constant"
  ```

---

## Task 2: Layer 1 — Best-Effort Check in `makeCreateRestoreEntity`

**Files:**
- Create: `app/models/restore/impl/pitr_restore_test.go`
- Modify: `app/models/restore/impl/pitr_restore.go`

### Step 2a: Write Failing Tests

- [ ] **Step 1: Create `pitr_restore_test.go` with the failing tests**

  ```go
  // app/models/restore/impl/pitr_restore_test.go
  package impl

  import (
  	"context"
  	"testing"

  	"github.com/stretchr/testify/mock"
  	"github.com/stretchr/testify/suite"
  	appconfig "github.com/tidbcloud/dataflow-service/app/config"
  	"github.com/tidbcloud/dataflow-service/app/repo/entity"
  	restore_mock "github.com/tidbcloud/dataflow-service/app/repo/restore/mock"
  	bizerr_wrapper "github.com/tidbcloud/dataflow-service/pkg/bizerr-wrapper"
  	"github.com/tidbcloud/pkgv2/bizerr"
  )

  type pitrRestoreTestSuite struct {
  	suite.Suite
  	ctx          context.Context
  	mockRestoreRepo *restore_mock.RestoreRepo
  	pitrRestore  *PiTRestore
  }

  func TestPiTRestoreTestSuite(t *testing.T) {
  	suite.Run(t, new(pitrRestoreTestSuite))
  }

  func (s *pitrRestoreTestSuite) SetupTest() {
  	s.ctx = context.Background()
  	s.mockRestoreRepo = restore_mock.NewRestoreRepo(s.T())
  	s.pitrRestore = &PiTRestore{
  		baseRestoreModel: baseRestoreModel{
  			restoreRepo: s.mockRestoreRepo,
  		},
  	}
  }

  func (s *pitrRestoreTestSuite) TestMakeCreateRestoreEntity_ConcurrencyCheck() {
  	logBackup := &entity.LogBackup{
  		TenantID:  entity.ID(1),
  		ProjectID: entity.ID(2),
  		ClusterID: entity.ID(3),
  	}
  	logBackup.ID = entity.ID(42)

  	baseCreateParam := &createRestoreParam{
  		logBackup: logBackup,
  		backup:    &entity.Backup{},
  		createRestoreReq: &createRestoreReq{
  			TenantID:  1,
  			ProjectID: 2,
  			ClusterConfig: createClusterConfig{
  				ClusterID: 100,
  				Region:    "us-west-2",
  			},
  		},
  	}

  	tests := map[string]struct {
  		countResult int64
  		countErr    bizerr.BizError
  		wantErr     bool
  		wantErrMsg  string
  		wantErrCode bizerr.ErrCode
  	}{
  		"active PiTR restore exists → returns conflict error": {
  			countResult: 1,
  			wantErr:     true,
  			wantErrMsg:  appconfig.ErrMsgPiTRConcurrentRestoreNotAllowed,
  		},
  		"no active PiTR restore → returns entity": {
  			countResult: 0,
  			wantErr:     false,
  		},
  		"repo Count returns error → returns internal error": {
  			countErr: bizerr_wrapper.NewInternalErr("db error"),
  			wantErr:  true,
  		},
  	}

  	for name, tc := range tests {
  		s.Run(name, func() {
  			s.mockRestoreRepo.On("Count", mock.Anything, mock.Anything).
  				Return(tc.countResult, tc.countErr).Once()

  			entity, err := s.pitrRestore.makeCreateRestoreEntity(s.ctx, baseCreateParam)
  			if tc.wantErr {
  				s.Require().NotNil(err)
  				if tc.wantErrMsg != "" {
  					s.Contains(err.Error(), tc.wantErrMsg)
  				}
  			} else {
  				s.Require().Nil(err)
  				s.Require().NotNil(entity)
  				s.Equal(entity.LogBackupID, logBackup.ID)
  			}
  		})
  	}
  }
  ```

  > **Note on `createRestoreReq` and `createClusterConfig`:** These are internal structs in the `impl` package.
  > Check their actual field names in `restore_base.go` before writing. The test is in the same package
  > (`package impl`) so private fields are accessible. See `createParam.createRestoreReq` usage at
  > `pitr_restore.go:306`.

- [ ] **Step 2: Run tests to confirm they fail**

  ```bash
  cd /path/to/repo && go test ./app/models/restore/impl/... -run TestPiTRestoreTestSuite -v 2>&1 | tail -30
  ```
  Expected: FAIL — `makeCreateRestoreEntity` doesn't call `Count` yet.

### Step 2b: Implement Layer 1

- [ ] **Step 3: Add the concurrency check to `makeCreateRestoreEntity`**

  In `app/models/restore/impl/pitr_restore.go`, find `makeCreateRestoreEntity` (line 284).
  Add the check **before** building the `entity.Restore` (before the `req := createParam.createRestoreReq` line, ~line 306).
  Also add `"github.com/tidbcloud/dataflow-service/internal/config"` to imports if not present.

  Insert the following block between the closing `}` of the `if/else` (end of region/createClusterReqV2 block, ~line 304) and the `req :=` line:

  ```go
  // Layer 1: best-effort serialization check for PiTR — prevent concurrent restores for the
  // same log backup. This check has a TOCTOU window; Layer 2 (run_compact_backup.go) provides
  // the hard atomic gate inside a transaction.
  override, _ := config.Global().BackupConfig.RestoreOverrides.GetMatched(
  	createParam.logBackup.TenantID.Uint64(),
  	createParam.logBackup.ProjectID.Uint64(),
  	createParam.logBackup.ClusterID.Uint64(),
  	"", "", // provider/region unknown at this point; omit for override matching
  )
  if !override.AllowConcurrentPiTR {
  	count, countErr := p.restoreRepo.Count(ctx, func(qs entity.RestoreQuerySet) entity.RestoreQuerySet {
  		return qs.LogBackupIDEq(createParam.logBackup.ID).
  			ModeEq(entity.RestoreModePiTR).
  			StatusIn(entity.RestoreStatusPENDING, entity.RestoreStatusRUNNING)
  	})
  	if countErr != nil {
  		log.GetLogger(ctx).Error("failed to count active PiTR restores", zap.Error(countErr))
  		return nil, bizerr_wrapper.NewInternalErr(appconfig.ErrInternalError)
  	}
  	if count > 0 {
  		return nil, bizerr_wrapper.NewInvalidInputError(appconfig.ErrMsgPiTRConcurrentRestoreNotAllowed)
  	}
  }
  ```

  Add `"github.com/tidbcloud/dataflow-service/internal/config"` to the import block in `pitr_restore.go` (it is not currently imported).

- [ ] **Step 4: Run tests to confirm they pass**

  ```bash
  go test ./app/models/restore/impl/... -run TestPiTRestoreTestSuite -v 2>&1 | tail -30
  ```
  Expected: all subtests PASS.

- [ ] **Step 5: Run all restore model tests (regression)**

  ```bash
  make app-unit-test-local 2>&1 | grep -E "FAIL|PASS|ok" | tail -20
  ```
  Expected: no failures.

- [ ] **Step 6: Commit**

  ```bash
  git add app/models/restore/impl/pitr_restore.go app/models/restore/impl/pitr_restore_test.go
  git commit -m "feat(br): Layer 1 — best-effort PiTR serialization check in makeCreateRestoreEntity"
  ```

---

## Task 3: Layer 2 — Atomic Double-Check in `execute()`

**Files:**
- Modify: `app/workflow/action/restore/pitr/create/run_compact_backup.go`
- Modify: `app/workflow/action/restore/pitr/create/run_compact_backup_test.go`

### Step 3a: Write Failing Tests

- [ ] **Step 1: Add `execute()` test cases to `run_compact_backup_test.go`**

  Append the following test method to `runCompactBackupForRestoreTestSuite` (after the existing `TestPrepareTimeRange` method):

  ```go
  func (s *runCompactBackupForRestoreTestSuite) TestExecute_ConcurrentPiTRCheck() {
  	restoreID := entity.ID(99)
  	logBackupID := entity.ID(42)

  	activeRestore := &entity.Restore{
  		Mode:        entity.RestoreModePiTR,
  		LogBackupID: logBackupID,
  	}
  	activeRestore.ID = restoreID

  	tests := map[string]struct {
  		override         config.RestoreOverride
  		concurrentCount  int64
  		countErr         error
  		wantErr          bool
  		wantErrContains  string
  	}{
  		"concurrent PiTR restore detected → returns error": {
  			override:        config.RestoreOverride{AllowConcurrentPiTR: false},
  			concurrentCount: 1,
  			wantErr:         true,
  			wantErrContains: "concurrent PiTR restore",
  		},
  		"no concurrent restore → proceeds (count=0)": {
  			override:        config.RestoreOverride{AllowConcurrentPiTR: false},
  			concurrentCount: 0,
  			wantErr:         false, // may fail further due to missing deps, but not on count check
  		},
  		"AllowConcurrentPiTR=true → check skipped": {
  			override:        config.RestoreOverride{AllowConcurrentPiTR: true},
  			concurrentCount: 5, // count is non-zero but should be ignored
  			wantErr:         false,
  		},
  		"count returns error → propagates error": {
  			override:        config.RestoreOverride{AllowConcurrentPiTR: false},
  			countErr:        fmt.Errorf("db error"),
  			wantErr:         true,
  			wantErrContains: "db error",
  		},
  	}

  	for name, tc := range tests {
  		s.Run(name, func() {
  			mockRestoreRepo := restore_mock.NewRestoreRepo(s.T())
  			mockBackupRepo := backup_mock.NewBackupRepo(s.T())
  			mockCompactBackupModel := compact_mock.NewCompactBackupModel(s.T())
  			mockTransProvider := db_mock.NewTransProvider(s.T())

  			action := &RunCompactBackupForRestore{
  				restoreRepo:        mockRestoreRepo,
  				backupRepo:         mockBackupRepo,
  				compactBackupModel: mockCompactBackupModel,
  				transProvider:      mockTransProvider,
  			}

  			// execute() reloads the restore first
  			mockRestoreRepo.On("Load", mock.Anything, restoreID.Uint64()).
  				Return(activeRestore, nil).Once()

  			if !tc.override.AllowConcurrentPiTR {
  				if tc.countErr != nil {
  					mockRestoreRepo.On("Count", mock.Anything, mock.Anything).
  						Return(int64(0), tc.countErr).Once()
  				} else {
  					mockRestoreRepo.On("Count", mock.Anything, mock.Anything).
  						Return(tc.concurrentCount, nil).Once()
  				}
  			}

  			if !tc.wantErr && tc.concurrentCount == 0 && !tc.override.AllowConcurrentPiTR {
  				// execute() will continue to Load backup — mock it to cause a predictable early exit
  				mockBackupRepo.On("Load", mock.Anything, mock.Anything).
  					Return(nil, bizerr_wrapper.NewNotFoundErr("backup not found")).Once()
  			}
  			if !tc.wantErr && tc.override.AllowConcurrentPiTR {
  				mockBackupRepo.On("Load", mock.Anything, mock.Anything).
  					Return(nil, bizerr_wrapper.NewNotFoundErr("backup not found")).Once()
  			}

  			logger := zap.NewNop()
  			err := action.execute(s.ctx, logger, activeRestore, nil, tc.override)

  			if tc.wantErr {
  				s.Require().Error(err)
  				if tc.wantErrContains != "" {
  					s.Contains(err.Error(), tc.wantErrContains)
  				}
  			} else {
  				// error is expected from backup.Load mock, not from the concurrency check
  				s.Error(err) // backup load will fail since mock returns not-found
  				s.NotContains(err.Error(), "concurrent PiTR restore")
  			}
  		})
  	}
  }
  ```

  Add required imports to the test file:
  ```go
  import (
  	// existing imports …
  	backup_mock   "github.com/tidbcloud/dataflow-service/app/repo/backup/mock"
  	restore_mock  "github.com/tidbcloud/dataflow-service/app/repo/restore/mock"
  	compact_mock  "github.com/tidbcloud/dataflow-service/app/models/backup/mock"
  	db_mock       "github.com/tidbcloud/pkgv2/db/mock"
  	bizerr_wrapper "github.com/tidbcloud/dataflow-service/pkg/bizerr-wrapper"
  )
  ```

  > **Finding mock package paths:** Run `find app/repo/backup/mock app/models/backup/mock -name "*.go" | head -5`
  > to verify the import paths. The CompactBackupModel mock lives in `app/models/backup/mock`.
  > The TransProvider mock is from `pkgv2/db/mock` — verify with `grep -r "TransProvider" --include="*_mock*" .`.

- [ ] **Step 2: Run tests to confirm they fail (compile error expected)**

  ```bash
  go test ./app/workflow/action/restore/pitr/create/... -run TestRunCompactBackupForRestoreTestSuite/TestExecute -v 2>&1 | tail -20
  ```
  Expected: compile error — `execute` has wrong number of arguments.

### Step 3b: Implement Layer 2

- [ ] **Step 3: Change `execute()` signature to accept `override`**

  In `app/workflow/action/restore/pitr/create/run_compact_backup.go`, change line 185:

  Before:
  ```go
  func (a *RunCompactBackupForRestore) execute(ctx context.Context, logger *zap.Logger, restore *entity.Restore, cluster *core.Cluster) error {
  ```

  After:
  ```go
  func (a *RunCompactBackupForRestore) execute(ctx context.Context, logger *zap.Logger, restore *entity.Restore, cluster *core.Cluster, override config.RestoreOverride) error {
  ```

- [ ] **Step 4: Update the call site in `Run()` to pass `override`**

  In the same file, around line 113, change:

  Before:
  ```go
  	if err := action.InvokeRunActionWithRetry(ctx.Context(), logger, func() error {
  		return a.execute(ctx.Context(), logger, restore, cluster)
  	}); err != nil {
  ```

  After:
  ```go
  	if err := action.InvokeRunActionWithRetry(ctx.Context(), logger, func() error {
  		return a.execute(ctx.Context(), logger, restore, cluster, override)
  	}); err != nil {
  ```

- [ ] **Step 5: Add the Layer 2 concurrency check inside `execute()`**

  In `execute()`, the body starts by reloading the restore and checking the idempotency guard (lines 188–198). The count check must run **after** the reload + idempotency guard (so we use the fresh `restore` and skip already-done work) but **before** entering `transProvider.Transaction()`.

  Insert after the idempotency guard block (after line 198, before line 200's blank line before "Load snapshot backup"):

  ```go
  	// Layer 2: atomic gate — double-check for concurrent PiTR restores inside the action.
  	// This catches TOCTOU races that slipped past the best-effort check in makeCreateRestoreEntity.
  	// Only applies when compaction is in play (this action only runs when needCompaction() is true).
  	if !override.AllowConcurrentPiTR {
  		count, countErr := a.restoreRepo.Count(ctx, func(qs entity.RestoreQuerySet) entity.RestoreQuerySet {
  			return qs.LogBackupIDEq(entity.ID(restore.LogBackupID.Uint64())).
  				ModeEq(entity.RestoreModePiTR).
  				StatusIn(entity.RestoreStatusRUNNING).
  				IDNotIn(restore.ID)
  		})
  		if countErr != nil {
  			logger.Error("failed to count concurrent PiTR restores", zap.Error(countErr))
  			return countErr
  		}
  		if count > 0 {
  			logger.Warn("concurrent PiTR restore detected, aborting compaction",
  				zap.Uint64("restore_id", restore.ID.Uint64()),
  				zap.Uint64("log_backup_id", restore.LogBackupID.Uint64()))
  			return fmt.Errorf("concurrent PiTR restore detected for log backup %d: "+
  				"another PiTR restore is already running, aborting to prevent duplicate compaction files",
  				restore.LogBackupID.Uint64())
  		}
  	}
  ```

  The `entity` package is already imported. No new imports needed in this file.

- [ ] **Step 6: Run the new tests**

  ```bash
  go test ./app/workflow/action/restore/pitr/create/... -run TestRunCompactBackupForRestoreTestSuite -v 2>&1 | tail -30
  ```
  Expected: all subtests PASS (including new `TestExecute_ConcurrentPiTRCheck` subtests).

- [ ] **Step 7: Full test suite (regression)**

  ```bash
  make app-unit-test-local 2>&1 | grep -E "FAIL|ok" | tail -20
  ```
  Expected: no failures.

- [ ] **Step 8: Commit**

  ```bash
  git add app/workflow/action/restore/pitr/create/run_compact_backup.go \
          app/workflow/action/restore/pitr/create/run_compact_backup_test.go
  git commit -m "feat(br): Layer 2 — atomic PiTR serialization gate in RunCompactBackupForRestore.execute()"
  ```

---

## Task 4: Final Verification

- [ ] **Step 1: Full build**

  ```bash
  make build-dataflow
  ```
  Expected: binary produced.

- [ ] **Step 2: Full unit test run**

  ```bash
  make app-unit-test-local
  ```
  Expected: all tests pass, no races detected.

- [ ] **Step 3: Verify spec coverage**

  Check that `docs/superpowers/specs/2026-05-18-pitr-serialization-design.md` requirements are all covered:
  - Config field `AllowConcurrentPiTR` added → Task 1 ✓
  - Layer 1 best-effort check in `makeCreateRestoreEntity` → Task 2 ✓
  - Layer 2 atomic gate in `execute()` with `override` param → Task 3 ✓
  - Snapshot restore path untouched (no changes to `restore_base.go`) → verified by code review ✓
  - Error message returns HTTP 422 via `bizerr_wrapper.NewInvalidInputError` → Layer 1 ✓
  - Layer 2 fails the DAG with a plain `error` → restore goes to FAILED status ✓

---

## Notes

### Why no `restore_base.go` changes?
The user explicitly said not to touch `restore_base.go`. The check is PiTR-specific so it belongs in `pitr_restore.go`.

### TOCTOU race acceptance
Layer 1 has a time-of-check-time-of-use window: two simultaneous PiTR creates can both see count=0 and both proceed. This is accepted (best-effort semantics). Layer 2 is not inside a single serialized DB transaction either, but it runs inside `execute()` which is called from the workflow action — at that point only one compaction attempt is being made per restore (the idempotency guard at the top of `execute()` ensures re-runs are no-ops). The race window at Layer 2 is much narrower.

### Why empty provider/region in Layer 1?
At `makeCreateRestoreEntity` time we don't have the target cluster object (it may not exist yet for restore-to-new-cluster flows). We use the log backup's `TenantID/ProjectID/ClusterID` with empty provider/region. `GetMatched` will still match overrides that are gated only on tenant/project/cluster IDs. Provider/region-specific overrides won't match — acceptable for this safety gate.

### Test helpers
- `restore_mock.NewRestoreRepo(s.T())` — from `app/repo/restore/mock`
- `backup_mock.NewBackupRepo(s.T())` — from `app/repo/backup/mock`
- CompactBackupModel mock — from `app/models/backup/mock`
- TransProvider mock — verify exact package with `grep -r "NewTransProvider" --include="*.go" vendor/`