From 68ac58fdb70553e905162b04683e68f25ed404fd Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 30 Jul 2026 06:05:58 -0400 Subject: [PATCH 1/4] ci: reproduce reshard cancel failure From 01535340a00cf49ec83fc702876fbce0a72a0ed4 Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 30 Jul 2026 06:11:50 -0400 Subject: [PATCH 2/4] fix(reshard): pin maintenance backend before NOLOGIN --- .../provisioner/reshard_source_fence.go | 34 +++++++-- .../provisioner/reshard_source_fence_test.go | 74 +++++++++++++++++++ 2 files changed, 102 insertions(+), 6 deletions(-) create mode 100644 controlplane/provisioner/reshard_source_fence_test.go diff --git a/controlplane/provisioner/reshard_source_fence.go b/controlplane/provisioner/reshard_source_fence.go index 8b3204e9..4619d4cc 100644 --- a/controlplane/provisioner/reshard_source_fence.go +++ b/controlplane/provisioner/reshard_source_fence.go @@ -8,6 +8,7 @@ import ( "time" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" ) // ReshardSourceFencer removes sessions which bypass Duckgres after the @@ -32,22 +33,43 @@ 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. 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)) + return disableMaintenanceAndTerminateOnSession(ctx, conn, maintenance.User, disableAndWait) +} + +type reshardSQLSession interface { + Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) + QueryRow(ctx context.Context, sql string, args ...any) pgx.Row +} + +func disableMaintenanceAndTerminateOnSession(ctx context.Context, session reshardSQLSession, username string, disableAndWait func() error) error { + // PgBouncer accepts a client connection before assigning a PostgreSQL + // backend. Pin one with an explicit transaction while the maintenance role + // still has LOGIN; otherwise the first SQL can be delayed until after the + // role becomes NOLOGIN and fail with PgBouncer's cached + // server_login_retry error. The transaction also keeps transaction-pooling + // PgBouncer from returning that authenticated backend before cleanup. + if _, err := session.Exec(ctx, "BEGIN"); err != nil { + return fmt.Errorf("pin backend before disabling reshard maintenance identity: %w", err) + } + defer func() { + _, _ = session.Exec(context.WithoutCancel(ctx), "ROLLBACK") + }() + + // Once PostgreSQL confirms NOLOGIN, this already-authenticated backend can + // terminate every other maintenance session and then close itself. There + // is no gap in which a new privileged session can race cleanup. if err := disableAndWait(); err != nil { return err } - return terminateRoleSessionsAndWait(ctx, conn, maintenance.User, "") + return terminateRoleSessionsAndWait(ctx, session, username, "") } -func terminateRoleSessionsAndWait(ctx context.Context, conn *pgx.Conn, username, database string) error { +func terminateRoleSessionsAndWait(ctx context.Context, conn reshardSQLSession, username, database string) error { const terminate = ` SELECT pg_terminate_backend(pid) FROM pg_stat_activity diff --git a/controlplane/provisioner/reshard_source_fence_test.go b/controlplane/provisioner/reshard_source_fence_test.go new file mode 100644 index 00000000..869e7024 --- /dev/null +++ b/controlplane/provisioner/reshard_source_fence_test.go @@ -0,0 +1,74 @@ +//go:build kubernetes + +package provisioner + +import ( + "context" + "fmt" + "reflect" + "strings" + "testing" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +type recordingReshardSession struct { + events []string +} + +func (s *recordingReshardSession) Exec(_ context.Context, query string, _ ...any) (pgconn.CommandTag, error) { + switch { + case strings.TrimSpace(query) == "BEGIN": + s.events = append(s.events, "begin") + case strings.TrimSpace(query) == "ROLLBACK": + s.events = append(s.events, "rollback") + case strings.Contains(query, "pg_terminate_backend"): + s.events = append(s.events, "terminate") + default: + return pgconn.CommandTag{}, fmt.Errorf("unexpected query: %s", query) + } + return pgconn.CommandTag{}, nil +} + +func (s *recordingReshardSession) QueryRow(_ context.Context, query string, _ ...any) pgx.Row { + if !strings.Contains(query, "pg_stat_activity") { + return reshardCountRow{err: fmt.Errorf("unexpected query: %s", query)} + } + s.events = append(s.events, "count") + return reshardCountRow{} +} + +type reshardCountRow struct { + err error +} + +func (r reshardCountRow) Scan(dest ...any) error { + if r.err != nil { + return r.err + } + if len(dest) != 1 { + return fmt.Errorf("scan destinations = %d, want 1", len(dest)) + } + remaining, ok := dest[0].(*int64) + if !ok { + return fmt.Errorf("scan destination type = %T, want *int64", dest[0]) + } + *remaining = 0 + return nil +} + +func TestDisableMaintenancePinsBackendBeforeNoLogin(t *testing.T) { + session := &recordingReshardSession{} + err := disableMaintenanceAndTerminateOnSession(context.Background(), session, "reshard_op", func() error { + session.events = append(session.events, "disable") + return nil + }) + if err != nil { + t.Fatalf("disable maintenance: %v", err) + } + want := []string{"begin", "disable", "terminate", "count", "rollback"} + if !reflect.DeepEqual(session.events, want) { + t.Fatalf("events = %v, want %v", session.events, want) + } +} From 4fc21ab84e83d42fc9df946d062bcc9b348616f1 Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 30 Jul 2026 06:50:24 -0400 Subject: [PATCH 3/4] fix(reshard): avoid transaction during maintenance fence --- .../provisioner/reshard_source_fence.go | 17 ++++----- .../provisioner/reshard_source_fence_test.go | 36 +++++++++++++++---- 2 files changed, 35 insertions(+), 18 deletions(-) diff --git a/controlplane/provisioner/reshard_source_fence.go b/controlplane/provisioner/reshard_source_fence.go index 4619d4cc..b24243d4 100644 --- a/controlplane/provisioner/reshard_source_fence.go +++ b/controlplane/provisioner/reshard_source_fence.go @@ -47,18 +47,13 @@ type reshardSQLSession interface { } func disableMaintenanceAndTerminateOnSession(ctx context.Context, session reshardSQLSession, username string, disableAndWait func() error) error { - // PgBouncer accepts a client connection before assigning a PostgreSQL - // backend. Pin one with an explicit transaction while the maintenance role - // still has LOGIN; otherwise the first SQL can be delayed until after the - // role becomes NOLOGIN and fail with PgBouncer's cached - // server_login_retry error. The transaction also keeps transaction-pooling - // PgBouncer from returning that authenticated backend before cleanup. - if _, err := session.Exec(ctx, "BEGIN"); err != nil { - return fmt.Errorf("pin backend before disabling reshard maintenance identity: %w", err) + // PgBouncer can accept the client before it has authenticated a PostgreSQL + // backend. Force that authentication while the maintenance role still has + // LOGIN. CNPG's shard Pooler uses session mode, so this backend remains + // attached without holding a transaction across the reconciliation wait. + if _, err := session.Exec(ctx, "SELECT 1"); err != nil { + return fmt.Errorf("authenticate backend before disabling reshard maintenance identity: %w", err) } - defer func() { - _, _ = session.Exec(context.WithoutCancel(ctx), "ROLLBACK") - }() // Once PostgreSQL confirms NOLOGIN, this already-authenticated backend can // terminate every other maintenance session and then close itself. There diff --git a/controlplane/provisioner/reshard_source_fence_test.go b/controlplane/provisioner/reshard_source_fence_test.go index 869e7024..8b1f15fd 100644 --- a/controlplane/provisioner/reshard_source_fence_test.go +++ b/controlplane/provisioner/reshard_source_fence_test.go @@ -4,6 +4,7 @@ package provisioner import ( "context" + "errors" "fmt" "reflect" "strings" @@ -14,15 +15,17 @@ import ( ) type recordingReshardSession struct { - events []string + events []string + probeErr error } func (s *recordingReshardSession) Exec(_ context.Context, query string, _ ...any) (pgconn.CommandTag, error) { switch { - case strings.TrimSpace(query) == "BEGIN": - s.events = append(s.events, "begin") - case strings.TrimSpace(query) == "ROLLBACK": - s.events = append(s.events, "rollback") + case strings.TrimSpace(query) == "SELECT 1": + s.events = append(s.events, "probe") + if s.probeErr != nil { + return pgconn.CommandTag{}, s.probeErr + } case strings.Contains(query, "pg_terminate_backend"): s.events = append(s.events, "terminate") default: @@ -58,7 +61,7 @@ func (r reshardCountRow) Scan(dest ...any) error { return nil } -func TestDisableMaintenancePinsBackendBeforeNoLogin(t *testing.T) { +func TestDisableMaintenanceAuthenticatesBackendBeforeNoLogin(t *testing.T) { session := &recordingReshardSession{} err := disableMaintenanceAndTerminateOnSession(context.Background(), session, "reshard_op", func() error { session.events = append(session.events, "disable") @@ -67,8 +70,27 @@ func TestDisableMaintenancePinsBackendBeforeNoLogin(t *testing.T) { if err != nil { t.Fatalf("disable maintenance: %v", err) } - want := []string{"begin", "disable", "terminate", "count", "rollback"} + want := []string{"probe", "disable", "terminate", "count"} if !reflect.DeepEqual(session.events, want) { t.Fatalf("events = %v, want %v", session.events, want) } } + +func TestDisableMaintenanceStopsBeforeNoLoginWhenBackendProbeFails(t *testing.T) { + probeErr := errors.New("backend unavailable") + session := &recordingReshardSession{probeErr: probeErr} + disableCalled := false + err := disableMaintenanceAndTerminateOnSession(context.Background(), session, "reshard_op", func() error { + disableCalled = true + return nil + }) + if !errors.Is(err, probeErr) { + t.Fatalf("disable maintenance error = %v, want %v", err, probeErr) + } + if disableCalled { + t.Fatal("disable callback called without an authenticated backend") + } + if want := []string{"probe"}; !reflect.DeepEqual(session.events, want) { + t.Fatalf("events = %v, want %v", session.events, want) + } +} From 8b9afa183876a38d32fe4767c536691c9193942e Mon Sep 17 00:00:00 2001 From: Bill Guowei Yang Date: Thu, 30 Jul 2026 09:53:38 -0400 Subject: [PATCH 4/4] fix(reshard): fence maintenance through primary --- controlplane/provisioner/reshard_runner.go | 17 ++- .../provisioner/reshard_runner_test.go | 119 +++++++++++++++++- .../provisioner/reshard_source_fence.go | 28 +---- .../provisioner/reshard_source_fence_test.go | 96 -------------- 4 files changed, 134 insertions(+), 126 deletions(-) delete mode 100644 controlplane/provisioner/reshard_source_fence_test.go 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 b24243d4..9b1a0e00 100644 --- a/controlplane/provisioner/reshard_source_fence.go +++ b/controlplane/provisioner/reshard_source_fence.go @@ -8,7 +8,6 @@ import ( "time" "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgconn" ) // ReshardSourceFencer removes sessions which bypass Duckgres after the @@ -33,38 +32,23 @@ func (PGReshardSourceFencer) TerminateAndWait(ctx context.Context, maintenance C } func (PGReshardSourceFencer) DisableMaintenanceAndTerminate(ctx context.Context, maintenance CatalogEndpoint, disableAndWait func() error) error { + // 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)) - return disableMaintenanceAndTerminateOnSession(ctx, conn, maintenance.User, disableAndWait) -} - -type reshardSQLSession interface { - Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) - QueryRow(ctx context.Context, sql string, args ...any) pgx.Row -} - -func disableMaintenanceAndTerminateOnSession(ctx context.Context, session reshardSQLSession, username string, disableAndWait func() error) error { - // PgBouncer can accept the client before it has authenticated a PostgreSQL - // backend. Force that authentication while the maintenance role still has - // LOGIN. CNPG's shard Pooler uses session mode, so this backend remains - // attached without holding a transaction across the reconciliation wait. - if _, err := session.Exec(ctx, "SELECT 1"); err != nil { - return fmt.Errorf("authenticate backend before disabling reshard maintenance identity: %w", err) - } - // Once PostgreSQL confirms NOLOGIN, this already-authenticated backend can - // terminate every other maintenance session and then close itself. There - // is no gap in which a new privileged session can race cleanup. if err := disableAndWait(); err != nil { return err } - return terminateRoleSessionsAndWait(ctx, session, username, "") + return terminateRoleSessionsAndWait(ctx, conn, maintenance.User, "") } -func terminateRoleSessionsAndWait(ctx context.Context, conn reshardSQLSession, username, database string) error { +func terminateRoleSessionsAndWait(ctx context.Context, conn *pgx.Conn, username, database string) error { const terminate = ` SELECT pg_terminate_backend(pid) FROM pg_stat_activity diff --git a/controlplane/provisioner/reshard_source_fence_test.go b/controlplane/provisioner/reshard_source_fence_test.go deleted file mode 100644 index 8b1f15fd..00000000 --- a/controlplane/provisioner/reshard_source_fence_test.go +++ /dev/null @@ -1,96 +0,0 @@ -//go:build kubernetes - -package provisioner - -import ( - "context" - "errors" - "fmt" - "reflect" - "strings" - "testing" - - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/pgconn" -) - -type recordingReshardSession struct { - events []string - probeErr error -} - -func (s *recordingReshardSession) Exec(_ context.Context, query string, _ ...any) (pgconn.CommandTag, error) { - switch { - case strings.TrimSpace(query) == "SELECT 1": - s.events = append(s.events, "probe") - if s.probeErr != nil { - return pgconn.CommandTag{}, s.probeErr - } - case strings.Contains(query, "pg_terminate_backend"): - s.events = append(s.events, "terminate") - default: - return pgconn.CommandTag{}, fmt.Errorf("unexpected query: %s", query) - } - return pgconn.CommandTag{}, nil -} - -func (s *recordingReshardSession) QueryRow(_ context.Context, query string, _ ...any) pgx.Row { - if !strings.Contains(query, "pg_stat_activity") { - return reshardCountRow{err: fmt.Errorf("unexpected query: %s", query)} - } - s.events = append(s.events, "count") - return reshardCountRow{} -} - -type reshardCountRow struct { - err error -} - -func (r reshardCountRow) Scan(dest ...any) error { - if r.err != nil { - return r.err - } - if len(dest) != 1 { - return fmt.Errorf("scan destinations = %d, want 1", len(dest)) - } - remaining, ok := dest[0].(*int64) - if !ok { - return fmt.Errorf("scan destination type = %T, want *int64", dest[0]) - } - *remaining = 0 - return nil -} - -func TestDisableMaintenanceAuthenticatesBackendBeforeNoLogin(t *testing.T) { - session := &recordingReshardSession{} - err := disableMaintenanceAndTerminateOnSession(context.Background(), session, "reshard_op", func() error { - session.events = append(session.events, "disable") - return nil - }) - if err != nil { - t.Fatalf("disable maintenance: %v", err) - } - want := []string{"probe", "disable", "terminate", "count"} - if !reflect.DeepEqual(session.events, want) { - t.Fatalf("events = %v, want %v", session.events, want) - } -} - -func TestDisableMaintenanceStopsBeforeNoLoginWhenBackendProbeFails(t *testing.T) { - probeErr := errors.New("backend unavailable") - session := &recordingReshardSession{probeErr: probeErr} - disableCalled := false - err := disableMaintenanceAndTerminateOnSession(context.Background(), session, "reshard_op", func() error { - disableCalled = true - return nil - }) - if !errors.Is(err, probeErr) { - t.Fatalf("disable maintenance error = %v, want %v", err, probeErr) - } - if disableCalled { - t.Fatal("disable callback called without an authenticated backend") - } - if want := []string{"probe"}; !reflect.DeepEqual(session.events, want) { - t.Fatalf("events = %v, want %v", session.events, want) - } -}