From 5b6b03d081da3f5913033d01f10132e0958e7d9d Mon Sep 17 00:00:00 2001 From: Jingyi Hu Date: Tue, 4 Dec 2018 18:55:15 -0800 Subject: [PATCH] etcdctl: add timeout to snapshot save command snapshot save command by default has no timeout, but respects user specified command timeout. --- etcdctl/ctlv3/command/snapshot_command.go | 9 ++++++++- etcdctl/ctlv3/command/util.go | 8 ++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/etcdctl/ctlv3/command/snapshot_command.go b/etcdctl/ctlv3/command/snapshot_command.go index 175ada4899e..32bcd7a61ed 100644 --- a/etcdctl/ctlv3/command/snapshot_command.go +++ b/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) diff --git a/etcdctl/ctlv3/command/util.go b/etcdctl/ctlv3/command/util.go index f536b2e5f5f..2a90ed9276a 100644 --- a/etcdctl/ctlv3/command/util.go +++ b/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 /metrics func endpointMemoryMetrics(host string) float64 { residentMemoryKey := "process_resident_memory_bytes"