fix(reshard): fence maintenance through direct primary

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

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,

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
 	}