diff --git a/controlplane/provisioner/reshard_runner.go b/controlplane/provisioner/reshard_runner.go index b41bfa4e..a7d59973 100644 --- a/controlplane/provisioner/reshard_runner.go +++ b/controlplane/provisioner/reshard_runner.go @@ -1026,6 +1026,15 @@ func (o *opRun) disableMaintenance(ctx context.Context, beforeDisable func() err if m.OperationID != o.op.ID || m.User == "" || m.Password == "" { return fmt.Errorf("maintenance identity for operation %d is not available before cleanup", o.op.ID) } + if !isValidCnpgShardName(o.op.FromShard) { + return fmt.Errorf("invalid source shard %q before disabling maintenance identity", o.op.FromShard) + } + resolveCtx, resolveCancel := context.WithTimeout(ctx, o.sourceMaintenanceTimeout()) + direct, err := o.r.duckling.CnpgProvisionerEndpoint(resolveCtx, o.op.FromShard) + resolveCancel() + if err != nil { + return fmt.Errorf("resolve direct source endpoint before disabling maintenance identity: %w", err) + } if m.MaintenanceNoLogin { // A prior runner can crash after PostgreSQL applied NOLOGIN and drained // sessions but before it persisted maintenance_disabled_at. Re-enable @@ -1040,13 +1049,9 @@ func (o *opRun) disableMaintenance(ctx context.Context, beforeDisable func() err } m = st.ReshardMaintenance } - host := o.op.SourceEndpoint - if host == "" { - host = st.MetadataStore.Endpoint - } admin := CatalogEndpoint{ - Host: host, Port: 5432, User: m.User, Password: m.Password, - Database: "postgres", SSLMode: sslModeFor(o.op.SourceKind), + Host: direct.Host, Port: direct.Port, User: m.User, Password: m.Password, + Database: "postgres", SSLMode: direct.SSLMode, } disableCtx, cancel := context.WithTimeout(ctx, o.sourceMaintenanceTimeout()) defer cancel() diff --git a/controlplane/provisioner/reshard_runner_test.go b/controlplane/provisioner/reshard_runner_test.go index 636356a1..65033866 100644 --- a/controlplane/provisioner/reshard_runner_test.go +++ b/controlplane/provisioner/reshard_runner_test.go @@ -6,6 +6,7 @@ import ( "context" "errors" "fmt" + "reflect" "strings" "sync" "sync/atomic" @@ -308,9 +309,17 @@ type fakeDuckling struct { cleanupChecks int provisionerEndpoint CatalogEndpoint provisionerErr error + provisionerShards []string + provisionerDeadline time.Time } -func (f *fakeDuckling) CnpgProvisionerEndpoint(context.Context, string) (CatalogEndpoint, error) { +func (f *fakeDuckling) CnpgProvisionerEndpoint(ctx context.Context, shard string) (CatalogEndpoint, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.provisionerShards = append(f.provisionerShards, shard) + if deadline, ok := ctx.Deadline(); ok { + f.provisionerDeadline = deadline + } if f.provisionerErr != nil { return CatalogEndpoint{}, f.provisionerErr } @@ -640,15 +649,21 @@ type fakeBackuper struct { type fakeSourceFencer struct { mu sync.Mutex calls int + disableEndpoint CatalogEndpoint before func() error err error disableCalls int disableErrsAfter int + disableDeadline time.Time } -func (f *fakeSourceFencer) DisableMaintenanceAndTerminate(_ context.Context, _ CatalogEndpoint, disableAndWait func() error) error { +func (f *fakeSourceFencer) DisableMaintenanceAndTerminate(ctx context.Context, endpoint CatalogEndpoint, disableAndWait func() error) error { f.mu.Lock() f.disableCalls++ + f.disableEndpoint = endpoint + if deadline, ok := ctx.Deadline(); ok { + f.disableDeadline = deadline + } f.mu.Unlock() if err := disableAndWait(); err != nil { return err @@ -1492,6 +1507,106 @@ func TestSourceMaintenanceTimeoutIgnoresShortCutoverOverride(t *testing.T) { } } +func TestDisableMaintenanceUsesDirectSourceEndpoint(t *testing.T) { + op := cnpgOp() + store := newFakeReshardStore(op) + duckling := &fakeDuckling{ + status: cnpgSourceStatus(), + provisionerEndpoint: CatalogEndpoint{ + Host: "shard-001-rw.cnpg-shards.svc.cluster.local", Port: 6432, + User: "provisioner", Password: "provisioner-password", Database: "postgres", SSLMode: "require", + }, + } + if err := duckling.SetReshardMaintenance(context.Background(), op.DucklingName, &ReshardMaintenanceSpec{ + OperationID: op.ID, + SourceShard: op.FromShard, + Phase: ReshardMaintenancePhasePrepared, + }); err != nil { + t.Fatalf("prepare maintenance: %v", err) + } + fencer := &fakeSourceFencer{} + runner := testRunner(store, duckling, &fakeCopier{}) + runner.fencer = fencer + + run := &opRun{r: runner, op: op} + if err := run.disableMaintenance(context.Background(), nil); err != nil { + t.Fatalf("disable maintenance: %v", err) + } + + if want := []string{"shard-001"}; !reflect.DeepEqual(duckling.provisionerShards, want) { + t.Fatalf("provisioner endpoint shards = %v, want %v", duckling.provisionerShards, want) + } + got := fencer.disableEndpoint + if got.Host != "shard-001-rw.cnpg-shards.svc.cluster.local" || got.Port != 6432 || got.SSLMode != "require" { + t.Fatalf("maintenance endpoint = %s, want direct source primary with TLS", got.Redacted()) + } + if got.User != "reshard_1" || got.Password != "maintenance-password" || got.Database != "postgres" { + t.Fatalf("maintenance credentials = user %q password %q database %q, want operation-scoped credentials", got.User, got.Password, got.Database) + } + if !fencer.disableDeadline.After(duckling.provisionerDeadline) { + t.Fatalf("disable deadline = %s, want a fresh budget after endpoint lookup deadline %s", fencer.disableDeadline, duckling.provisionerDeadline) + } +} + +func TestDisableMaintenanceFailsClosedWhenDirectEndpointUnavailable(t *testing.T) { + op := cnpgOp() + store := newFakeReshardStore(op) + duckling := &fakeDuckling{status: cnpgSourceStatus(), provisionerErr: errors.New("source primary unavailable")} + if err := duckling.SetReshardMaintenance(context.Background(), op.DucklingName, &ReshardMaintenanceSpec{ + OperationID: op.ID, + SourceShard: op.FromShard, + Phase: ReshardMaintenancePhasePrepared, + }); err != nil { + t.Fatalf("prepare maintenance: %v", err) + } + fencer := &fakeSourceFencer{} + runner := testRunner(store, duckling, &fakeCopier{}) + runner.fencer = fencer + + run := &opRun{r: runner, op: op} + err := run.disableMaintenance(context.Background(), nil) + if err == nil || !strings.Contains(err.Error(), "source primary unavailable") { + t.Fatalf("disable maintenance error = %v, want direct endpoint failure", err) + } + if fencer.disableCalls != 0 { + t.Fatalf("disable calls = %d, want 0 without direct endpoint", fencer.disableCalls) + } + if got := duckling.status.ReshardMaintenance.MaintenanceNoLogin; got { + t.Fatal("maintenance role set NOLOGIN without a direct cleanup endpoint") + } +} + +func TestDisableMaintenanceKeepsNoLoginWhenDirectEndpointUnavailable(t *testing.T) { + op := cnpgOp() + store := newFakeReshardStore(op) + duckling := &fakeDuckling{status: cnpgSourceStatus(), provisionerErr: errors.New("source primary unavailable")} + if err := duckling.SetReshardMaintenance(context.Background(), op.DucklingName, &ReshardMaintenanceSpec{ + OperationID: op.ID, + SourceShard: op.FromShard, + Phase: ReshardMaintenancePhaseDisabled, + }); err != nil { + t.Fatalf("disable maintenance for takeover: %v", err) + } + fencer := &fakeSourceFencer{} + runner := testRunner(store, duckling, &fakeCopier{}) + runner.fencer = fencer + + run := &opRun{r: runner, op: op} + err := run.disableMaintenance(context.Background(), nil) + if err == nil || !strings.Contains(err.Error(), "source primary unavailable") { + t.Fatalf("disable maintenance error = %v, want direct endpoint failure", err) + } + if got := duckling.maintenancePhases(); !reflect.DeepEqual(got, []string{ReshardMaintenancePhaseDisabled}) { + t.Fatalf("maintenance phases = %v, want role left NOLOGIN", got) + } + if !duckling.status.ReshardMaintenance.MaintenanceNoLogin { + t.Fatal("maintenance role re-enabled before direct endpoint was available") + } + if fencer.disableCalls != 0 { + t.Fatalf("disable calls = %d, want 0 without direct endpoint", fencer.disableCalls) + } +} + func TestMaintenanceWaitUsesFullSourceBudgetInsteadOfCutoverOverride(t *testing.T) { op := &configstore.ReshardOperation{ ID: 1, diff --git a/controlplane/provisioner/reshard_source_fence.go b/controlplane/provisioner/reshard_source_fence.go index 8b3204e9..9b1a0e00 100644 --- a/controlplane/provisioner/reshard_source_fence.go +++ b/controlplane/provisioner/reshard_source_fence.go @@ -32,15 +32,16 @@ func (PGReshardSourceFencer) TerminateAndWait(ctx context.Context, maintenance C } func (PGReshardSourceFencer) DisableMaintenanceAndTerminate(ctx context.Context, maintenance CatalogEndpoint, disableAndWait func() error) error { - // Establish the final administrative session before requesting NOLOGIN. - // Once PostgreSQL confirms NOLOGIN, this already-authenticated session can - // terminate every other maintenance session and then close itself. There - // is no gap in which a new privileged session can race cleanup. + // Callers provide a direct PostgreSQL endpoint so the final administrative + // session is never returned to a pool after NOLOGIN. Once PostgreSQL + // confirms NOLOGIN, this already-authenticated session can terminate every + // other maintenance session and then close its own backend. conn, err := pgx.Connect(ctx, maintenance.DSN()) if err != nil { return fmt.Errorf("connect before disabling reshard maintenance identity %s: %w", maintenance.Redacted(), err) } defer conn.Close(context.WithoutCancel(ctx)) + if err := disableAndWait(); err != nil { return err }