STAC-25639: fix false-success on elasticsearch restore

cmd/elasticsearch/check_and_finalize.go

@@ -2,21 +2,25 @@ 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"
 )
 
 // 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
-}