From 46732b61d7399f9b067c1490d719bdb959d5de40 Mon Sep 17 00:00:00 2001
From: Dmitrii Stepanov <d.stepanov@yadro.com>
Date: Thu, 20 Jun 2024 16:04:55 +0300
Subject: [PATCH] [#60] cli: Add `await` flag to `control set-status`

Signed-off-by: Dmitrii Stepanov <d.stepanov@yadro.com>
---
 cmd/frostfs-cli/internal/commonflags/flags.go |  3 +
 .../modules/control/set_netmap_status.go      | 64 +++++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/cmd/frostfs-cli/internal/commonflags/flags.go b/cmd/frostfs-cli/internal/commonflags/flags.go
index 6915ef42b..23dbc8c2b 100644
--- a/cmd/frostfs-cli/internal/commonflags/flags.go
+++ b/cmd/frostfs-cli/internal/commonflags/flags.go
@@ -50,6 +50,9 @@ const (
 
 	TracingFlag      = "trace"
 	TracingFlagUsage = "Generate trace ID and print it."
+
+	AwaitFlag      = "await"
+	AwaitFlagUsage = "Wait for the operation to complete"
 )
 
 // Init adds common flags to the command:
diff --git a/cmd/frostfs-cli/modules/control/set_netmap_status.go b/cmd/frostfs-cli/modules/control/set_netmap_status.go
index 3aa741689..31ade1eb9 100644
--- a/cmd/frostfs-cli/modules/control/set_netmap_status.go
+++ b/cmd/frostfs-cli/modules/control/set_netmap_status.go
@@ -1,7 +1,10 @@
 package control
 
 import (
+	"crypto/ecdsa"
+	"errors"
 	"fmt"
+	"time"
 
 	rawclient "git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/rpc/client"
 	"git.frostfs.info/TrueCloudLab/frostfs-node/cmd/frostfs-cli/internal/common"
@@ -9,6 +12,7 @@ import (
 	"git.frostfs.info/TrueCloudLab/frostfs-node/cmd/frostfs-cli/internal/key"
 	commonCmd "git.frostfs.info/TrueCloudLab/frostfs-node/cmd/internal/common"
 	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/services/control"
+	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/client"
 	"github.com/spf13/cobra"
 )
 
@@ -18,8 +22,13 @@ const (
 	netmapStatusOnline      = "online"
 	netmapStatusOffline     = "offline"
 	netmapStatusMaintenance = "maintenance"
+
+	maxSetStatusMaxWaitTime = 30 * time.Minute
+	setStatusWaitTimeout    = 30 * time.Second
 )
 
+var errNetmapStatusAwaitFailed = errors.New("netmap status hasn't changed for 30 minutes")
+
 var setNetmapStatusCmd = &cobra.Command{
 	Use:   "set-status",
 	Short: "Set status of the storage node in FrostFS network map",
@@ -43,6 +52,8 @@ func initControlSetNetmapStatusCmd() {
 
 	flags.BoolP(commonflags.ForceFlag, commonflags.ForceFlagShorthand, false,
 		"Force turning to local maintenance")
+
+	flags.Bool(commonflags.AwaitFlag, false, commonflags.AwaitFlagUsage)
 }
 
 func setNetmapStatus(cmd *cobra.Command, _ []string) {
@@ -56,15 +67,19 @@ func setNetmapStatus(cmd *cobra.Command, _ []string) {
 		}
 	}
 
+	await, _ := cmd.Flags().GetBool(commonflags.AwaitFlag)
+	var targetStatus control.NetmapStatus
 	switch st, _ := cmd.Flags().GetString(netmapStatusFlag); st {
 	default:
 		commonCmd.ExitOnErr(cmd, "", fmt.Errorf("unsupported status %s", st))
 	case netmapStatusOnline:
 		body.SetStatus(control.NetmapStatus_ONLINE)
 		printIgnoreForce(control.NetmapStatus_ONLINE)
+		targetStatus = control.NetmapStatus_ONLINE
 	case netmapStatusOffline:
 		body.SetStatus(control.NetmapStatus_OFFLINE)
 		printIgnoreForce(control.NetmapStatus_OFFLINE)
+		targetStatus = control.NetmapStatus_OFFLINE
 	case netmapStatusMaintenance:
 		body.SetStatus(control.NetmapStatus_MAINTENANCE)
 
@@ -72,6 +87,7 @@ func setNetmapStatus(cmd *cobra.Command, _ []string) {
 			body.SetForceMaintenance()
 			common.PrintVerbose(cmd, "Local maintenance will be forced.")
 		}
+		targetStatus = control.NetmapStatus_MAINTENANCE
 	}
 
 	req := new(control.SetNetmapStatusRequest)
@@ -92,4 +108,52 @@ func setNetmapStatus(cmd *cobra.Command, _ []string) {
 	verifyResponse(cmd, resp.GetSignature(), resp.GetBody())
 
 	cmd.Println("Network status update request successfully sent.")
+
+	if await {
+		awaitSetNetmapStatus(cmd, pk, cli, targetStatus)
+	}
+}
+
+func awaitSetNetmapStatus(cmd *cobra.Command, pk *ecdsa.PrivateKey, cli *client.Client, targetStatus control.NetmapStatus) {
+	req := &control.GetNetmapStatusRequest{
+		Body: &control.GetNetmapStatusRequest_Body{},
+	}
+	signRequest(cmd, pk, req)
+	var epoch uint64
+	var status control.NetmapStatus
+	startTime := time.Now()
+	cmd.Println("Wait until epoch and netmap status change...")
+	for {
+		var resp *control.GetNetmapStatusResponse
+		var err error
+		err = cli.ExecRaw(func(client *rawclient.Client) error {
+			resp, err = control.GetNetmapStatus(client, req)
+			return err
+		})
+		commonCmd.ExitOnErr(cmd, "failed to get current netmap status: %w", err)
+
+		if epoch == 0 {
+			epoch = resp.GetBody().GetEpoch()
+		}
+
+		status = resp.GetBody().GetStatus()
+		if resp.GetBody().GetEpoch() > epoch {
+			epoch = resp.GetBody().GetEpoch()
+			cmd.Printf("Epoch changed to %d\n", resp.GetBody().GetEpoch())
+		}
+
+		if status == targetStatus {
+			break
+		}
+
+		if time.Since(startTime) > maxSetStatusMaxWaitTime {
+			commonCmd.ExitOnErr(cmd, "failed to wait netmap status: %w", errNetmapStatusAwaitFailed)
+			return
+		}
+
+		time.Sleep(setStatusWaitTimeout)
+
+		cmd.Printf("Current netmap status '%s', target status '%s'\n", status.String(), targetStatus.String())
+	}
+	cmd.Printf("Netmap status changed to '%s' successfully.\n", status.String())
 }