etcdctl: add timeout to snapshot save command

snapshot save command by default has no timeout, but respects user
specified command timeout.

Author: jingyih

etcdctl/ctlv3/command/snapshot_command.go

@@ -102,8 +102,15 @@ func snapshotSaveCommandFunc(cmd *cobra.Command, args []string) {
 	sp := snapshot.NewV3(lg)
 	cfg := mustClientCfgFromCmd(cmd)
 
+	// if user does not specify "--command-timeout" flag, there will be no timeout for snapshot save command
+	ctx, cancel := context.WithCancel(context.Background())
+	if isCommandTimeoutFlagSet(cmd) {
+		ctx, cancel = commandCtx(cmd)
+	}
+	defer cancel()
+
 	path := args[0]
-	if err := sp.Save(context.TODO(), *cfg, path); err != nil {
+	if err := sp.Save(ctx, *cfg, path); err != nil {
 		ExitWithError(ExitInterrupted, err)
 	}
 	fmt.Printf("Snapshot saved at %s\n", path)

etcdctl/ctlv3/command/util.go

@@ -82,6 +82,14 @@ func commandCtx(cmd *cobra.Command) (context.Context, context.CancelFunc) {
 	return context.WithTimeout(context.Background(), timeOut)
 }
 
+func isCommandTimeoutFlagSet(cmd *cobra.Command) bool {
+	commandTimeoutFlag := cmd.Flags().Lookup("command-timeout")
+	if commandTimeoutFlag == nil {
+		panic("expect command-timeout flag to exist")
+	}
+	return commandTimeoutFlag.Changed
+}
+
 // get the process_resident_memory_bytes from <server:2379>/metrics
 func endpointMemoryMetrics(host string) float64 {
 	residentMemoryKey := "process_resident_memory_bytes"