From 1ca33611b6a99b6a4bf9334757082a14669d9259 Mon Sep 17 00:00:00 2001 From: Vladimir Iliakov Date: Fri, 28 Aug 2026 12:42:08 +0200 Subject: [PATCH] STAC-25639: elasticsearch restore wait until all shards and indices are healthy --- cmd/elasticsearch/check_and_finalize.go | 287 +++++++++++++--- cmd/elasticsearch/check_and_finalize_test.go | 312 ++++++++++++++++++ cmd/elasticsearch/restore.go | 14 +- cmd/elasticsearch/restore_test.go | 4 +- internal/clients/elasticsearch/client.go | 59 ++-- internal/clients/elasticsearch/client_test.go | 90 +++++ internal/clients/elasticsearch/interface.go | 2 +- internal/orchestration/restore/apirestore.go | 5 +- 8 files changed, 686 insertions(+), 87 deletions(-) create mode 100644 cmd/elasticsearch/check_and_finalize_test.go diff --git a/cmd/elasticsearch/check_and_finalize.go b/cmd/elasticsearch/check_and_finalize.go index f01ec44..cd5cf5c 100644 --- a/cmd/elasticsearch/check_and_finalize.go +++ b/cmd/elasticsearch/check_and_finalize.go @@ -2,12 +2,14 @@ package elasticsearch import ( "fmt" + "time" "github.com/spf13/cobra" "github.com/stackvista/stackstate-backup-cli/cmd/cmdutils" "github.com/stackvista/stackstate-backup-cli/internal/app" es "github.com/stackvista/stackstate-backup-cli/internal/clients/elasticsearch" "github.com/stackvista/stackstate-backup-cli/internal/foundation/config" + "github.com/stackvista/stackstate-backup-cli/internal/foundation/logger" "github.com/stackvista/stackstate-backup-cli/internal/orchestration/portforward" "github.com/stackvista/stackstate-backup-cli/internal/orchestration/restore" "github.com/stackvista/stackstate-backup-cli/internal/orchestration/scale" @@ -15,8 +17,10 @@ import ( // Check-and-finalize command flags var ( - checkOperationID string - checkWait bool + checkOperationID string + checkWait bool + checkNoProgressIn time.Duration + checkFinalizeOnly bool ) func checkAndFinalizeCmd(globalFlags *config.CLIGlobalFlags) *cobra.Command { @@ -25,19 +29,38 @@ func checkAndFinalizeCmd(globalFlags *config.CLIGlobalFlags) *cobra.Command { Short: "Check restore status and finalize if complete", Long: `Check the status of a restore operation and perform finalization (scale up deployments) if complete. If the restore is still running and --wait is specified, wait for completion before finalizing.`, + PreRunE: func(_ *cobra.Command, _ []string) error { + if !checkFinalizeOnly && checkOperationID == "" { + return fmt.Errorf("--operation-id is required unless --finalize-only is set") + } + + return nil + }, Run: func(_ *cobra.Command, _ []string) { cmdutils.Run(globalFlags, runCheckAndFinalize, cmdutils.StorageIsRequired) }, } - cmd.Flags().StringVar(&checkOperationID, "operation-id", "", "Operation ID of the restore operation (required)") + cmd.Flags().StringVar(&checkOperationID, "operation-id", "", + "Snapshot name of the restore operation (required unless --finalize-only)") cmd.Flags().BoolVar(&checkWait, "wait", false, "Wait for restore to complete if still running") - _ = cmd.MarkFlagRequired("operation-id") + cmd.Flags().DurationVar(&checkNoProgressIn, "no-progress-timeout", defaultNoProgressTimeout, noProgressTimeoutUsage) + cmd.Flags().BoolVar(&checkFinalizeOnly, "finalize-only", false, + "Scale the deployments back up and release the restore lock without checking restore status. "+ + "Use when the snapshot can no longer be read and the restore is known to be finished") + cmd.MarkFlagsMutuallyExclusive("finalize-only", "wait") return cmd } func runCheckAndFinalize(appCtx *app.Context) error { + // Finalizing needs Kubernetes only, so it stays available when Elasticsearch or the snapshot + // cannot be reached at all - otherwise nothing in the CLI can release the restore lock. + if checkFinalizeOnly { + appCtx.Logger.Warningf("Skipping the restore status check on request; finalizing") + return finalizeRestore(appCtx) + } + // Setup port-forward to Elasticsearch serviceName := appCtx.Config.Elasticsearch.Service.Name remotePort := appCtx.Config.Elasticsearch.Service.Port @@ -56,13 +79,223 @@ func runCheckAndFinalize(appCtx *app.Context) error { repository := appCtx.Config.Elasticsearch.Restore.Repository - return checkAndFinalize(esClient, appCtx, repository, checkOperationID, checkWait) + return checkAndFinalize(esClient, appCtx, repository, checkOperationID, checkWait, checkNoProgressIn) +} + +const ( + // defaultNoProgressTimeout bounds the wait so a stuck restore fails instead of polling until + // the caller's own timeout with the workloads scaled down and the restore lock held. Generous + // on purpose: a spurious failure aborts a working restore, while the only cost of waiting too + // long is a later failure. Progress is counted per primary shard, so tripping this means no + // shard at all completed in the window. + defaultNoProgressTimeout = 2 * time.Hour + + noProgressTimeoutUsage = "Fail if no primary shard finishes restoring within this duration. " + + "Bounds inactivity, not total restore time; 0 waits indefinitely" + + // restoreStatusMaxErrors tolerates a port-forward dropping mid-restore. Polling now spans the + // whole restore, so a pod restart must not end a restore that is still running server-side. + restoreStatusMaxErrors = 5 +) + +// reconnectingHealthClient rebuilds the port-forward and Elasticsearch client after a failed +// health check. It starts from the caller's client so the happy path opens no extra port-forward. +type reconnectingHealthClient struct { + appCtx *app.Context + client indicesHealthGetter + pf *portforward.Conn +} + +func (r *reconnectingHealthClient) GetIndicesHealth() (map[string]es.IndexHealth, error) { + if r.client == nil { + if err := r.connect(); err != nil { + return nil, err + } + } + + health, err := r.client.GetIndicesHealth() + if err != nil { + // The tunnel is bound to a fixed local port, so a broken one never recovers. Drop it and + // let the next poll dial a fresh pod. + r.disconnect() + return nil, err + } + + return health, nil +} + +func (r *reconnectingHealthClient) connect() error { + pf, err := portforward.SetupPortForward( + r.appCtx.K8sClient, + r.appCtx.Namespace, + r.appCtx.Config.Elasticsearch.Service.Name, + r.appCtx.Config.Elasticsearch.Service.Port, + r.appCtx.Logger, + ) + if err != nil { + return err + } + + client, err := r.appCtx.NewESClient(pf.LocalPort) + if err != nil { + close(pf.StopChan) + return fmt.Errorf("failed to create Elasticsearch client: %w", err) + } + + r.pf, r.client = pf, client + + return nil +} + +// disconnect drops the client and closes only a port-forward this type opened itself; the one the +// caller passed in is closed by the caller. +func (r *reconnectingHealthClient) disconnect() { + if r.pf != nil { + close(r.pf.StopChan) + r.pf = nil + } + + r.client = nil +} + +// expectedRestoredIndices returns the snapshot's indices that this restore recreates. The snapshot +// is re-read rather than passed in so that check-and-finalize works from just a snapshot name; +// snapshots are immutable, so the list cannot drift. +func expectedRestoredIndices(esClient es.Interface, appCtx *app.Context, repository, snapshotName string) ([]string, error) { + snapshot, err := esClient.GetSnapshot(repository, snapshotName) + if err != nil { + return nil, fmt.Errorf("failed to get snapshot details: %w", err) + } + + expected := filterSTSIndices( + snapshot.Indices, + appCtx.Config.Elasticsearch.Restore.IndexPrefix, + appCtx.Config.Elasticsearch.Restore.DatastreamIndexPrefix, + ) + if len(expected) == 0 { + return nil, fmt.Errorf("snapshot %s contains no indices matching the configured STS prefixes", snapshotName) + } + + return expected, nil +} + +// restoreProgress summarises how far a restore has got. Replicas are deliberately ignored: a +// cluster with fewer nodes than replicas keeps them unassigned forever, so requiring green would +// never complete there. +type restoreProgress struct { + // indicesRestored counts expected indices with every primary shard active, and decides completion. + indicesRestored int + // primariesActive counts active primary shards, and drives stall detection. Whole indices are + // too coarse for that: a large multi-shard index can restore for a long time without finishing. + primariesActive int } -func checkAndFinalize(esClient es.Interface, appCtx *app.Context, repository, snapshotName string, waitForComplete bool) error { +func measureRestore(expected []string, health map[string]es.IndexHealth) restoreProgress { + var progress restoreProgress + + for _, index := range expected { + indexHealth, exists := health[index] + if !exists { + continue + } + + progress.primariesActive += indexHealth.ActivePrimaryShards + if indexHealth.NumberOfShards > 0 && indexHealth.ActivePrimaryShards == indexHealth.NumberOfShards { + progress.indicesRestored++ + } + } + + return progress +} + +type indicesHealthGetter interface { + GetIndicesHealth() (map[string]es.IndexHealth, error) +} + +// newRestoreStatusFn builds the status callback used for both the single check and the wait loop. +// Completion is derived from the restored indices themselves: a restore that has been accepted but +// not yet applied is indistinguishable from a finished one when judged by recovery activity alone. +func newRestoreStatusFn( + esClient indicesHealthGetter, + log *logger.Logger, + expected []string, + noProgressTimeout time.Duration, + maxErrors int, +) func() (string, bool, error) { + lastPrimariesActive := -1 + lastProgressAt := time.Now() + errCount := 0 + + return func() (string, bool, error) { + health, err := esClient.GetIndicesHealth() + if err != nil { + errCount++ + if errCount >= maxErrors { + return "", false, err + } + + log.Warningf("Restore status check failed (%d/%d), retrying: %v", errCount, maxErrors, err) + + return es.StatusInProgress, false, nil + } + errCount = 0 + + progress := measureRestore(expected, health) + if progress.indicesRestored == len(expected) { + return es.StatusSuccess, true, nil + } + + if progress.primariesActive != lastPrimariesActive { + lastPrimariesActive = progress.primariesActive + lastProgressAt = time.Now() + } else if noProgressTimeout > 0 && time.Since(lastProgressAt) > noProgressTimeout { + // Deliberately not reported as a failed restore: all this establishes is that no progress + // was observed from here. Restarting a restore deletes every STS index first, so a caller + // that reads this as "failed" and retries would destroy a restore that is merely slow. + return "", false, fmt.Errorf( + "elasticsearch restore stalled: no primary shard finished restoring in %s "+ + "(%d of %d indices complete, %d primaries active); "+ + "it may still be running server-side, so check before restarting it", + noProgressTimeout, progress.indicesRestored, len(expected), progress.primariesActive, + ) + } + + log.Debugf( + "Restored %d of %d indices (%d primaries active)", + progress.indicesRestored, len(expected), progress.primariesActive, + ) + + return es.StatusInProgress, false, nil + } +} + +func checkAndFinalize( + esClient es.Interface, + appCtx *app.Context, + repository, snapshotName string, + waitForComplete bool, + noProgressTimeout time.Duration, +) error { + expected, err := expectedRestoredIndices(esClient, appCtx, repository, snapshotName) + if err != nil { + return err + } + + healthClient := &reconnectingHealthClient{appCtx: appCtx, client: esClient} + defer healthClient.disconnect() + + // Retrying only pays off while polling. A one-shot check that swallowed the error would report + // a dead tunnel as a running restore and exit 0. + maxErrors := restoreStatusMaxErrors + if !waitForComplete { + maxErrors = 1 + } + + statusFn := newRestoreStatusFn(healthClient, appCtx.Logger, expected, noProgressTimeout, maxErrors) + // Get restore status - appCtx.Logger.Infof("Checking restore status for snapshot: %s", snapshotName) - status, isComplete, err := esClient.GetRestoreStatus(repository, snapshotName) + appCtx.Logger.Infof("Checking restore status for snapshot: %s (%d indices)", snapshotName, len(expected)) + status, isComplete, err := statusFn() if err != nil { return fmt.Errorf("failed to get restore status: %w", err) } @@ -72,17 +305,9 @@ func checkAndFinalize(esClient es.Interface, appCtx *app.Context, repository, sn // Handle different scenarios if isComplete { switch status { - case "SUCCESS": + case es.StatusSuccess: appCtx.Logger.Successf("Restore completed successfully") return finalizeRestore(appCtx) - case "NOT_FOUND": - appCtx.Logger.Infof("No restore operation found for snapshot: %s", snapshotName) - appCtx.Logger.Infof("The restore may have already been finalized") - appCtx.Logger.Println() - appCtx.Logger.Infof("Checking if deployments need to be scaled up...") - return attemptScaleUp(appCtx) - case "FAILED": - return fmt.Errorf("restore failed with status: %s", status) default: return fmt.Errorf("restore completed with unexpected status: %s", status) } @@ -93,7 +318,7 @@ func checkAndFinalize(esClient es.Interface, appCtx *app.Context, repository, sn if waitForComplete { appCtx.Logger.Println() - return waitAndFinalize(esClient, appCtx, repository, snapshotName) + return waitAndFinalize(statusFn, appCtx, snapshotName) } // Not waiting - print status and exit @@ -103,15 +328,10 @@ func checkAndFinalize(esClient es.Interface, appCtx *app.Context, repository, sn } // waitAndFinalize waits for restore to complete and finalizes (scale up) -func waitAndFinalize(esClient es.Interface, appCtx *app.Context, repository, snapshotName string) error { +func waitAndFinalize(statusFn func() (string, bool, error), appCtx *app.Context, snapshotName string) error { restore.PrintAPIWaitingMessage("elasticsearch", snapshotName, appCtx.Namespace, appCtx.Logger) - // Wait for restore to complete - checkStatusFn := func() (string, bool, error) { - return esClient.GetRestoreStatus(repository, snapshotName) - } - - if err := restore.WaitForAPIRestore(checkStatusFn, 0, appCtx.Logger); err != nil { + if err := restore.WaitForAPIRestore(statusFn, 0, appCtx.Logger); err != nil { return err } @@ -129,20 +349,3 @@ func finalizeRestore(appCtx *app.Context) error { return restore.FinalizeRestore(scaleUpFn, appCtx.Logger) } - -// attemptScaleUp tries to scale up deployments and release lock (used when restore is not found/already complete) -func attemptScaleUp(appCtx *app.Context) error { - labelSelector := appCtx.Config.Elasticsearch.Restore.ScaleDownLabelSelector - scaleUpFn := func() error { - return scale.ScaleUpAndReleaseLock(appCtx.K8sClient, appCtx.Namespace, labelSelector, appCtx.Logger) - } - - if err := scaleUpFn(); err != nil { - // Don't fail if no deployments found to scale up - appCtx.Logger.Infof("No deployments found to scale up (this is normal if already finalized)") - return nil - } - - appCtx.Logger.Successf("Finalization completed successfully") - return nil -} diff --git a/cmd/elasticsearch/check_and_finalize_test.go b/cmd/elasticsearch/check_and_finalize_test.go new file mode 100644 index 0000000..b55ae03 --- /dev/null +++ b/cmd/elasticsearch/check_and_finalize_test.go @@ -0,0 +1,312 @@ +package elasticsearch + +import ( + "errors" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + es "github.com/stackvista/stackstate-backup-cli/internal/clients/elasticsearch" + "github.com/stackvista/stackstate-backup-cli/internal/foundation/logger" +) + +type healthResult struct { + health map[string]es.IndexHealth + err error +} + +// fakeHealthClient replays results in order, repeating the last one once exhausted. +type fakeHealthClient struct { + results []healthResult + calls int +} + +func (f *fakeHealthClient) GetIndicesHealth() (map[string]es.IndexHealth, error) { + result := f.results[min(f.calls, len(f.results)-1)] + f.calls++ + + return result.health, result.err +} + +func ok(health map[string]es.IndexHealth) healthResult { + return healthResult{health: health} +} + +func restored(shards int) es.IndexHealth { + return es.IndexHealth{Status: "green", NumberOfShards: shards, ActivePrimaryShards: shards} +} + +func TestMeasureRestore(t *testing.T) { + tests := []struct { + name string + expected []string + health map[string]es.IndexHealth + wantIndices int + wantPrimaries int + }{ + { + name: "index not recreated yet", + expected: []string{"sts_topology", "sts_events"}, + health: map[string]es.IndexHealth{"sts_topology": restored(1)}, + wantIndices: 1, + wantPrimaries: 1, + }, + { + name: "primaries still recovering are not counted", + expected: []string{"sts_topology"}, + health: map[string]es.IndexHealth{ + "sts_topology": {Status: "red", NumberOfShards: 3, ActivePrimaryShards: 2, InitializingShards: 1}, + }, + wantIndices: 0, + wantPrimaries: 2, + }, + { + name: "index present with no active primaries is not counted", + expected: []string{"sts_topology"}, + health: map[string]es.IndexHealth{ + "sts_topology": {Status: "red", NumberOfShards: 3, ActivePrimaryShards: 0, UnassignedShards: 3}, + }, + wantIndices: 0, + wantPrimaries: 0, + }, + { + name: "unassigned replicas do not block completion", + expected: []string{"sts_topology"}, + health: map[string]es.IndexHealth{ + "sts_topology": { + Status: "yellow", NumberOfShards: 1, NumberOfReplicas: 1, + ActivePrimaryShards: 1, ActiveShards: 1, UnassignedShards: 1, + }, + }, + wantIndices: 1, + wantPrimaries: 1, + }, + { + name: "all primaries active across several indices", + expected: []string{"sts_topology", ".ds-sts_k8s_logs-000001"}, + health: map[string]es.IndexHealth{ + "sts_topology": restored(3), + ".ds-sts_k8s_logs-000001": restored(2), + "unrelated": restored(1), + }, + wantIndices: 2, + wantPrimaries: 5, + }, + { + name: "index reporting zero shards is not counted", + expected: []string{"sts_topology"}, + health: map[string]es.IndexHealth{"sts_topology": {Status: "green"}}, + wantIndices: 0, + wantPrimaries: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + progress := measureRestore(tt.expected, tt.health) + assert.Equal(t, tt.wantIndices, progress.indicesRestored, "indicesRestored") + assert.Equal(t, tt.wantPrimaries, progress.primariesActive, "primariesActive") + }) + } +} + +var statusFnExpected = []string{"sts_topology", "sts_events"} + +func TestNewRestoreStatusFn_Progress(t *testing.T) { + expected := statusFnExpected + log := logger.New(true, false) + + t.Run("accepted but not yet applied reports in progress", func(t *testing.T) { + client := &fakeHealthClient{results: []healthResult{ok(map[string]es.IndexHealth{})}} + + status, isComplete, err := newRestoreStatusFn(client, log, expected, time.Minute, 5)() + + require.NoError(t, err) + assert.Equal(t, es.StatusInProgress, status) + assert.False(t, isComplete) + }) + + t.Run("complete once every index has all primaries active", func(t *testing.T) { + client := &fakeHealthClient{results: []healthResult{ + ok(map[string]es.IndexHealth{"sts_topology": restored(1)}), + ok(map[string]es.IndexHealth{"sts_topology": restored(1), "sts_events": restored(1)}), + }} + statusFn := newRestoreStatusFn(client, log, expected, time.Minute, 5) + + status, isComplete, err := statusFn() + require.NoError(t, err) + assert.Equal(t, es.StatusInProgress, status) + assert.False(t, isComplete) + + status, isComplete, err = statusFn() + require.NoError(t, err) + assert.Equal(t, es.StatusSuccess, status) + assert.True(t, isComplete) + }) +} + +func TestNewRestoreStatusFn_Deadline(t *testing.T) { + expected := statusFnExpected + log := logger.New(true, false) + + t.Run("stalled restore errors instead of polling forever", func(t *testing.T) { + client := &fakeHealthClient{results: []healthResult{ + ok(map[string]es.IndexHealth{"sts_topology": restored(1)}), + }} + statusFn := newRestoreStatusFn(client, log, expected, time.Nanosecond, 5) + + // First call records the progress it observed; the second sees no change past the deadline. + _, isComplete, err := statusFn() + require.NoError(t, err) + assert.False(t, isComplete) + + _, isComplete, err = statusFn() + require.Error(t, err) + assert.False(t, isComplete) + // Must not read as a failed restore: the caller's retry deletes every STS index first. + assert.ErrorContains(t, err, "stalled") + assert.ErrorContains(t, err, "may still be running") + assert.NotContains(t, err.Error(), "restore failed") + }) + + t.Run("shard progress inside one index is not a stall", func(t *testing.T) { + // A large multi-shard index can restore for a long time without completing. Counting whole + // indices would read this as stalled; counting primaries does not. + client := &fakeHealthClient{results: []healthResult{ + ok(map[string]es.IndexHealth{"sts_topology": {NumberOfShards: 5, ActivePrimaryShards: 1}}), + ok(map[string]es.IndexHealth{"sts_topology": {NumberOfShards: 5, ActivePrimaryShards: 2}}), + ok(map[string]es.IndexHealth{"sts_topology": {NumberOfShards: 5, ActivePrimaryShards: 3}}), + }} + statusFn := newRestoreStatusFn(client, log, expected, time.Nanosecond, 5) + + for range 3 { + status, isComplete, err := statusFn() + require.NoError(t, err) + assert.Equal(t, es.StatusInProgress, status) + assert.False(t, isComplete) + } + }) + + t.Run("progress resets the deadline", func(t *testing.T) { + client := &fakeHealthClient{results: []healthResult{ + ok(map[string]es.IndexHealth{}), + ok(map[string]es.IndexHealth{"sts_topology": restored(1)}), + }} + statusFn := newRestoreStatusFn(client, log, expected, time.Nanosecond, 5) + + _, _, err := statusFn() + require.NoError(t, err) + + status, isComplete, err := statusFn() + require.NoError(t, err) + assert.Equal(t, es.StatusInProgress, status) + assert.False(t, isComplete) + }) + + t.Run("zero waits indefinitely", func(t *testing.T) { + client := &fakeHealthClient{results: []healthResult{ + ok(map[string]es.IndexHealth{"sts_topology": restored(1)}), + }} + statusFn := newRestoreStatusFn(client, log, expected, 0, 5) + + for range 3 { + status, isComplete, err := statusFn() + require.NoError(t, err) + assert.Equal(t, es.StatusInProgress, status) + assert.False(t, isComplete) + } + }) +} + +func TestNewRestoreStatusFn_TransientErrors(t *testing.T) { + expected := statusFnExpected + log := logger.New(true, false) + + t.Run("a dropped port-forward does not end the restore", func(t *testing.T) { + client := &fakeHealthClient{results: []healthResult{ + {err: errors.New("connection refused")}, + {err: errors.New("connection refused")}, + ok(map[string]es.IndexHealth{"sts_topology": restored(1), "sts_events": restored(1)}), + }} + statusFn := newRestoreStatusFn(client, log, expected, time.Minute, 5) + + for range 2 { + status, isComplete, err := statusFn() + require.NoError(t, err) + assert.Equal(t, es.StatusInProgress, status) + assert.False(t, isComplete) + } + + status, isComplete, err := statusFn() + require.NoError(t, err) + assert.Equal(t, es.StatusSuccess, status) + assert.True(t, isComplete) + }) + + t.Run("errors surface once they stop being transient", func(t *testing.T) { + client := &fakeHealthClient{results: []healthResult{{err: errors.New("connection refused")}}} + statusFn := newRestoreStatusFn(client, log, expected, time.Minute, 3) + + for range 2 { + _, _, err := statusFn() + require.NoError(t, err) + } + + _, _, err := statusFn() + require.Error(t, err) + assert.ErrorContains(t, err, "connection refused") + }) + + t.Run("a successful check clears earlier errors", func(t *testing.T) { + client := &fakeHealthClient{results: []healthResult{ + {err: errors.New("blip")}, + ok(map[string]es.IndexHealth{"sts_topology": restored(1)}), + {err: errors.New("blip")}, + {err: errors.New("blip")}, + }} + statusFn := newRestoreStatusFn(client, log, expected, time.Minute, 3) + + for range 4 { + _, _, err := statusFn() + require.NoError(t, err) + } + }) +} + +func TestReconnectingHealthClient(t *testing.T) { + t.Run("passes through the caller's client while it works", func(t *testing.T) { + health := map[string]es.IndexHealth{"sts_topology": restored(1)} + seeded := &fakeHealthClient{results: []healthResult{ok(health)}} + client := &reconnectingHealthClient{client: seeded} + + got, err := client.GetIndicesHealth() + + require.NoError(t, err) + assert.Equal(t, health, got) + assert.Equal(t, 1, seeded.calls) + }) + + t.Run("drops a broken client but never closes a port-forward it did not open", func(t *testing.T) { + // The caller closes the port-forward it passed in, so adopting it here would double-close a + // channel and panic. pf must stay nil until connect() opens one. + seeded := &fakeHealthClient{results: []healthResult{{err: errors.New("connection refused")}}} + client := &reconnectingHealthClient{client: seeded} + + _, err := client.GetIndicesHealth() + + require.Error(t, err) + assert.Nil(t, client.client, "a broken client must be dropped so the next call reconnects") + assert.Nil(t, client.pf, "must not adopt a port-forward it did not open") + }) + + t.Run("disconnect is safe to call repeatedly", func(t *testing.T) { + client := &reconnectingHealthClient{client: &fakeHealthClient{}} + + assert.NotPanics(t, func() { + client.disconnect() + client.disconnect() + }) + }) +} diff --git a/cmd/elasticsearch/restore.go b/cmd/elasticsearch/restore.go index 675f980..77868d4 100644 --- a/cmd/elasticsearch/restore.go +++ b/cmd/elasticsearch/restore.go @@ -26,11 +26,12 @@ const ( // Restore command flags var ( - snapshotName string - useLatest bool - runBackground bool - skipConfirmation bool - allowPartial bool + snapshotName string + useLatest bool + runBackground bool + skipConfirmation bool + allowPartial bool + restoreNoProgressIn time.Duration ) func restoreCmd(globalFlags *config.CLIGlobalFlags) *cobra.Command { @@ -47,6 +48,7 @@ func restoreCmd(globalFlags *config.CLIGlobalFlags) *cobra.Command { cmd.Flags().BoolVar(&runBackground, "background", false, "Run restore in background without waiting for completion") cmd.Flags().BoolVarP(&skipConfirmation, "yes", "y", false, "Skip confirmation prompt") cmd.Flags().BoolVar(&allowPartial, "allow-partial", false, "Allow restoring from a PARTIAL snapshot without extra confirmation") + cmd.Flags().DurationVar(&restoreNoProgressIn, "no-progress-timeout", defaultNoProgressTimeout, noProgressTimeoutUsage) cmd.MarkFlagsMutuallyExclusive("snapshot", "latest") cmd.MarkFlagsOneRequired("snapshot", "latest") return cmd @@ -145,7 +147,7 @@ func runRestore(appCtx *app.Context) error { return nil } - return checkAndFinalize(esClient, appCtx, repository, selectedSnapshot, !runBackground) + return checkAndFinalize(esClient, appCtx, repository, selectedSnapshot, !runBackground, restoreNoProgressIn) } // getLatestSnapshot retrieves the most recent snapshot from the repository diff --git a/cmd/elasticsearch/restore_test.go b/cmd/elasticsearch/restore_test.go index e723b43..e611100 100644 --- a/cmd/elasticsearch/restore_test.go +++ b/cmd/elasticsearch/restore_test.go @@ -94,8 +94,8 @@ func (m *mockESClientForRestore) ConfigureSLMPolicy(_, _, _, _, _, _ string, _, return fmt.Errorf("not implemented") } -func (m *mockESClientForRestore) GetRestoreStatus(_, _ string) (string, bool, error) { - return "SUCCESS", true, nil +func (m *mockESClientForRestore) GetIndicesHealth() (map[string]elasticsearch.IndexHealth, error) { + return nil, fmt.Errorf("not implemented") } // TestRestoreCmd_Unit tests the command structure diff --git a/internal/clients/elasticsearch/client.go b/internal/clients/elasticsearch/client.go index 1619ab8..03db72d 100644 --- a/internal/clients/elasticsearch/client.go +++ b/internal/clients/elasticsearch/client.go @@ -390,52 +390,43 @@ func (c *Client) RestoreSnapshot(repository, snapshotName, indicesPattern string return nil } -// RecoveryInfo represents the recovery status of a shard from _cat/recovery API -type RecoveryInfo struct { - Index string `json:"index"` - Shard string `json:"shard"` - Type string `json:"type"` - Stage string `json:"stage"` - Repository string `json:"repository"` - Snapshot string `json:"snapshot"` +// IndexHealth is the per-index section of the _cluster/health response +type IndexHealth struct { + Status string `json:"status"` + NumberOfShards int `json:"number_of_shards"` + NumberOfReplicas int `json:"number_of_replicas"` + ActivePrimaryShards int `json:"active_primary_shards"` + ActiveShards int `json:"active_shards"` + InitializingShards int `json:"initializing_shards"` + UnassignedShards int `json:"unassigned_shards"` } -// GetRestoreStatus checks the status of a restore operation by examining active shard recoveries. -// When a snapshot is being restored, shards are recovered with type "snapshot". -// Returns: (statusMessage, isComplete, error) -// Status can be: "IN_PROGRESS", "SUCCESS" -func (c *Client) GetRestoreStatus(repository, snapshotName string) (string, bool, error) { - // Use _cat/recovery API to check for active snapshot recoveries. - // This shows shards that are currently being recovered from a snapshot. - res, err := c.es.Cat.Recovery( - c.es.Cat.Recovery.WithContext(context.Background()), - c.es.Cat.Recovery.WithFormat("json"), - c.es.Cat.Recovery.WithActiveOnly(true), - c.es.Cat.Recovery.WithH("index,shard,type,stage,repository,snapshot"), +// GetIndicesHealth returns per-index cluster health keyed by index name. +// Indices that do not exist are simply absent from the result. +func (c *Client) GetIndicesHealth() (map[string]IndexHealth, error) { + res, err := c.es.Cluster.Health( + c.es.Cluster.Health.WithContext(context.Background()), + c.es.Cluster.Health.WithLevel("indices"), ) if err != nil { - return "", false, fmt.Errorf("failed to get recovery status: %w", err) + return nil, fmt.Errorf("failed to get cluster health: %w", err) } defer res.Body.Close() if res.IsError() { - return "", false, fmt.Errorf("elasticsearch returned error: %s", res.String()) + return nil, fmt.Errorf("elasticsearch returned error: %s", res.String()) } - var recoveries []RecoveryInfo - if err := json.NewDecoder(res.Body).Decode(&recoveries); err != nil { - return "", false, fmt.Errorf("failed to decode response: %w", err) + var health struct { + Indices map[string]IndexHealth `json:"indices"` + } + if err := json.NewDecoder(res.Body).Decode(&health); err != nil { + return nil, fmt.Errorf("failed to decode response: %w", err) } - // Check if any active recovery is from the specified snapshot - for _, recovery := range recoveries { - if recovery.Type == "snapshot" && - recovery.Repository == repository && - recovery.Snapshot == snapshotName { - return StatusInProgress, false, nil - } + if health.Indices == nil { + health.Indices = map[string]IndexHealth{} } - // No active recoveries from this snapshot - restore is complete - return StatusSuccess, true, nil + return health.Indices, nil } diff --git a/internal/clients/elasticsearch/client_test.go b/internal/clients/elasticsearch/client_test.go index f8e0e8e..74671a5 100644 --- a/internal/clients/elasticsearch/client_test.go +++ b/internal/clients/elasticsearch/client_test.go @@ -412,6 +412,96 @@ func TestClient_RestoreSnapshot(t *testing.T) { } } +func TestClient_GetIndicesHealth(t *testing.T) { + tests := []struct { + name string + responseBody string + expectedCount int + assertContents func(t *testing.T, health map[string]IndexHealth) + }{ + { + name: "indices are keyed by name", + responseBody: `{ + "cluster_name": "test", + "status": "green", + "indices": { + "sts_topology": { + "status": "green", + "number_of_shards": 3, + "number_of_replicas": 1, + "active_primary_shards": 3, + "active_shards": 6, + "initializing_shards": 0, + "unassigned_shards": 0 + }, + "sts_events": { + "status": "yellow", + "number_of_shards": 1, + "number_of_replicas": 1, + "active_primary_shards": 1, + "active_shards": 1, + "initializing_shards": 0, + "unassigned_shards": 1 + } + } + }`, + expectedCount: 2, + assertContents: func(t *testing.T, health map[string]IndexHealth) { + assert.Equal(t, 3, health["sts_topology"].NumberOfShards) + assert.Equal(t, 3, health["sts_topology"].ActivePrimaryShards) + assert.Equal(t, "yellow", health["sts_events"].Status) + assert.Equal(t, 1, health["sts_events"].UnassignedShards) + }, + }, + { + name: "no indices section yields an empty map", + responseBody: `{"cluster_name": "test", "status": "green"}`, + expectedCount: 0, + assertContents: func(t *testing.T, health map[string]IndexHealth) { + assert.NotNil(t, health) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := mockESServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "/_cluster/health", r.URL.Path) + assert.Equal(t, "indices", r.URL.Query().Get("level")) + + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(tt.responseBody)) + })) + defer server.Close() + + client, err := NewClient(server.URL) + require.NoError(t, err) + + health, err := client.GetIndicesHealth() + + require.NoError(t, err) + assert.Len(t, health, tt.expectedCount) + tt.assertContents(t, health) + }) + } +} + +func TestClient_GetIndicesHealth_ElasticsearchError(t *testing.T) { + server := mockESServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error": "boom"}`)) + })) + defer server.Close() + + client, err := NewClient(server.URL) + require.NoError(t, err) + + health, err := client.GetIndicesHealth() + + require.Error(t, err) + assert.Nil(t, health) +} + func TestNewClient(t *testing.T) { client, err := NewClient("http://localhost:9200") require.NoError(t, err) diff --git a/internal/clients/elasticsearch/interface.go b/internal/clients/elasticsearch/interface.go index 5e1afd8..f8e0089 100644 --- a/internal/clients/elasticsearch/interface.go +++ b/internal/clients/elasticsearch/interface.go @@ -7,9 +7,9 @@ type Interface interface { ListSnapshots(repository string) ([]Snapshot, error) GetSnapshot(repository, snapshotName string) (*Snapshot, error) RestoreSnapshot(repository, snapshotName, indicesPattern string, partial bool) error - GetRestoreStatus(repository, snapshotName string) (string, bool, error) // Index operations + GetIndicesHealth() (map[string]IndexHealth, error) ListIndices(pattern string) ([]string, error) ListIndicesDetailed() ([]IndexInfo, error) DeleteIndex(index string) error diff --git a/internal/orchestration/restore/apirestore.go b/internal/orchestration/restore/apirestore.go index 1cd2d7c..5ffe37a 100644 --- a/internal/orchestration/restore/apirestore.go +++ b/internal/orchestration/restore/apirestore.go @@ -30,7 +30,8 @@ func WaitForAPIRestore( <-ticker.C statusMsg, isComplete, err := checkStatusFn() if err != nil { - return fmt.Errorf("failed to check restore status: %w", err) + // Not necessarily a failed check: a status function may also stop the wait deliberately. + return fmt.Errorf("stopped waiting for restore: %w", err) } log.Debugf("Restore status: %s (complete: %v)", statusMsg, isComplete) @@ -53,7 +54,7 @@ func PrintAPIWaitingMessage(serviceName, identifier, namespace string, log *logg log.Println() log.Infof("You can safely interrupt this command with Ctrl+C.") log.Infof("To check status and finalize later, run:") - log.Infof(" sts-backup %s check-and-finalize --operation-id %s -n %s", serviceName, identifier, namespace) + log.Infof(" sts-backup %s check-and-finalize --operation-id %s --wait -n %s", serviceName, identifier, namespace) } // PrintAPIRunningRestoreStatus prints status and instructions for a running restore