Merge pull request #4817 from mqliang/time-out

etcdctlv3: add timeout support

etcdctlv3/command/auth_command.go

@@ -18,7 +18,6 @@ import (
 	"fmt"
 
 	"github.com/spf13/cobra"
-	"golang.org/x/net/context"
 )
 
 // NewAuthCommand returns the cobra command for "auth".
@@ -47,7 +46,9 @@ func authEnableCommandFunc(cmd *cobra.Command, args []string) {
 		ExitWithError(ExitBadArgs, fmt.Errorf("auth enable command does not accept argument."))
 	}
 
-	_, err := mustClientFromCmd(cmd).Auth.AuthEnable(context.TODO())
+	ctx, cancel := commandCtx(cmd)
+	_, err := mustClientFromCmd(cmd).Auth.AuthEnable(ctx)
+	cancel()
 	if err != nil {
 		ExitWithError(ExitError, err)
 	}

etcdctlv3/command/compaction_command.go

@@ -19,7 +19,6 @@ import (
 	"strconv"
 
 	"github.com/spf13/cobra"
-	"golang.org/x/net/context"
 )
 
 // NewCompactionCommand returns the cobra command for "compaction".
@@ -43,7 +42,10 @@ func compactionCommandFunc(cmd *cobra.Command, args []string) {
 	}
 
 	c := mustClientFromCmd(cmd)
-	if cerr := c.Compact(context.TODO(), rev); cerr != nil {
+	ctx, cancel := commandCtx(cmd)
+	cerr := c.Compact(ctx, rev)
+	cancel()
+	if cerr != nil {
 		ExitWithError(ExitError, cerr)
 		return
 	}

etcdctlv3/command/defrag_command.go

@@ -19,7 +19,6 @@ import (
 	"os"
 
 	"github.com/spf13/cobra"
-	"golang.org/x/net/context"
 )
 
 // NewDefragCommand returns the cobra command for "Defrag".
@@ -34,7 +33,9 @@ func NewDefragCommand() *cobra.Command {
 func defragCommandFunc(cmd *cobra.Command, args []string) {
 	c := mustClientFromCmd(cmd)
 	for _, ep := range c.Endpoints() {
-		_, err := c.Defragment(context.TODO(), ep)
+		ctx, cancel := commandCtx(cmd)
+		_, err := c.Defragment(ctx, ep)
+		cancel()
 		if err != nil {
 			fmt.Fprintf(os.Stderr, "Failed to defragment etcd member[%s] (%v)\n", ep, err)
 		} else {

etcdctlv3/command/del_command.go

@@ -19,7 +19,6 @@ import (
 
 	"github.com/coreos/etcd/clientv3"
 	"github.com/spf13/cobra"
-	"golang.org/x/net/context"
 )
 
 // NewDelCommand returns the cobra command for "del".
@@ -34,7 +33,9 @@ func NewDelCommand() *cobra.Command {
 // delCommandFunc executes the "del" command.
 func delCommandFunc(cmd *cobra.Command, args []string) {
 	key, opts := getDelOp(cmd, args)
-	resp, err := mustClientFromCmd(cmd).Delete(context.TODO(), key, opts...)
+	ctx, cancel := commandCtx(cmd)
+	resp, err := mustClientFromCmd(cmd).Delete(ctx, key, opts...)
+	cancel()
 	if err != nil {
 		ExitWithError(ExitError, err)
 	}

etcdctlv3/command/ep_health_command.go

@@ -21,7 +21,6 @@ import (
 
 	"github.com/coreos/etcd/clientv3"
 	"github.com/spf13/cobra"
-	"golang.org/x/net/context"
 )
 
 // NewEpHealthCommand returns the cobra command for "endpoint-health".
@@ -67,7 +66,9 @@ func epHealthCommandFunc(cmd *cobra.Command, args []string) {
 			st := time.Now()
 			// get a random key. As long as we can get the response without an error, the
 			// endpoint is health.
-			_, err = cli.Get(context.TODO(), "health")
+			ctx, cancel := commandCtx(cmd)
+			_, err = cli.Get(ctx, "health")
+			cancel()
 			if err != nil {
 				fmt.Printf("%s is unhealthy: failed to commit proposal: %v\n", ep, err)
 			} else {

etcdctlv3/command/get_command.go

@@ -20,7 +20,6 @@ import (
 
 	"github.com/coreos/etcd/clientv3"
 	"github.com/spf13/cobra"
-	"golang.org/x/net/context"
 )
 
 var (
@@ -52,7 +51,9 @@ func NewGetCommand() *cobra.Command {
 // getCommandFunc executes the "get" command.
 func getCommandFunc(cmd *cobra.Command, args []string) {
 	key, opts := getGetOp(cmd, args)
-	resp, err := mustClientFromCmd(cmd).Get(context.TODO(), key, opts...)
+	ctx, cancel := commandCtx(cmd)
+	resp, err := mustClientFromCmd(cmd).Get(ctx, key, opts...)
+	cancel()
 	if err != nil {
 		ExitWithError(ExitError, err)
 	}

etcdctlv3/command/global.go

@@ -29,9 +29,10 @@ import (
 // GlobalFlags are flags that defined globally
 // and are inherited to all sub-commands.
 type GlobalFlags struct {
-	Insecure    bool
-	Endpoints   []string
-	DialTimeout time.Duration
+	Insecure       bool
+	Endpoints      []string
+	DialTimeout    time.Duration
+	CommandTimeOut time.Duration
 
 	TLS transport.TLSInfo
 

etcdctlv3/command/lease_command.go

@@ -61,7 +61,9 @@ func leaseCreateCommandFunc(cmd *cobra.Command, args []string) {
 		ExitWithError(ExitBadArgs, fmt.Errorf("bad TTL (%v)", err))
 	}
 
-	resp, err := mustClientFromCmd(cmd).Create(context.TODO(), ttl)
+	ctx, cancel := commandCtx(cmd)
+	resp, err := mustClientFromCmd(cmd).Create(ctx, ttl)
+	cancel()
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "failed to create lease (%v)\n", err)
 		return
@@ -92,7 +94,9 @@ func leaseRevokeCommandFunc(cmd *cobra.Command, args []string) {
 		ExitWithError(ExitBadArgs, fmt.Errorf("bad lease ID arg (%v), expecting ID in Hex", err))
 	}
 
-	_, err = mustClientFromCmd(cmd).Revoke(context.TODO(), v3.LeaseID(id))
+	ctx, cancel := commandCtx(cmd)
+	_, err = mustClientFromCmd(cmd).Revoke(ctx, v3.LeaseID(id))
+	cancel()
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "failed to revoke lease (%v)\n", err)
 		return

etcdctlv3/command/member_command.go

@@ -20,7 +20,6 @@ import (
 	"strings"
 
 	"github.com/spf13/cobra"
-	"golang.org/x/net/context"
 )
 
 var (
@@ -106,8 +105,9 @@ func memberAddCommandFunc(cmd *cobra.Command, args []string) {
 	}
 
 	urls := strings.Split(memberPeerURLs, ",")
-
-	resp, err := mustClientFromCmd(cmd).MemberAdd(context.TODO(), urls)
+	ctx, cancel := commandCtx(cmd)
+	resp, err := mustClientFromCmd(cmd).MemberAdd(ctx, urls)
+	cancel()
 	if err != nil {
 		ExitWithError(ExitError, err)
 	}
@@ -126,7 +126,9 @@ func memberRemoveCommandFunc(cmd *cobra.Command, args []string) {
 		ExitWithError(ExitBadArgs, fmt.Errorf("bad member ID arg (%v), expecting ID in Hex", err))
 	}
 
-	resp, err := mustClientFromCmd(cmd).MemberRemove(context.TODO(), id)
+	ctx, cancel := commandCtx(cmd)
+	resp, err := mustClientFromCmd(cmd).MemberRemove(ctx, id)
+	cancel()
 	if err != nil {
 		ExitWithError(ExitError, err)
 	}
@@ -151,7 +153,9 @@ func memberUpdateCommandFunc(cmd *cobra.Command, args []string) {
 
 	urls := strings.Split(memberPeerURLs, ",")
 
-	resp, err := mustClientFromCmd(cmd).MemberUpdate(context.TODO(), id, urls)
+	ctx, cancel := commandCtx(cmd)
+	resp, err := mustClientFromCmd(cmd).MemberUpdate(ctx, id, urls)
+	cancel()
 	if err != nil {
 		ExitWithError(ExitError, err)
 	}
@@ -161,7 +165,9 @@ func memberUpdateCommandFunc(cmd *cobra.Command, args []string) {
 
 // memberListCommandFunc executes the "member list" command.
 func memberListCommandFunc(cmd *cobra.Command, args []string) {
-	resp, err := mustClientFromCmd(cmd).MemberList(context.TODO())
+	ctx, cancel := commandCtx(cmd)
+	resp, err := mustClientFromCmd(cmd).MemberList(ctx)
+	cancel()
 	if err != nil {
 		ExitWithError(ExitError, err)
 	}

etcdctlv3/command/put_command.go

@@ -21,7 +21,6 @@ import (
 
 	"github.com/coreos/etcd/clientv3"
 	"github.com/spf13/cobra"
-	"golang.org/x/net/context"
 )
 
 var (
@@ -57,7 +56,9 @@ will store the content of the file to <key>.
 func putCommandFunc(cmd *cobra.Command, args []string) {
 	key, value, opts := getPutOp(cmd, args)
 
-	resp, err := mustClientFromCmd(cmd).Put(context.TODO(), key, value, opts...)
+	ctx, cancel := commandCtx(cmd)
+	resp, err := mustClientFromCmd(cmd).Put(ctx, key, value, opts...)
+	cancel()
 	if err != nil {
 		ExitWithError(ExitError, err)
 	}

etcdctlv3/command/util.go

@@ -20,6 +20,8 @@ import (
 	"regexp"
 
 	pb "github.com/coreos/etcd/storage/storagepb"
+	"github.com/spf13/cobra"
+	"golang.org/x/net/context"
 )
 
 func printKV(isHex bool, kv *pb.KeyValue) {
@@ -47,3 +49,11 @@ func argify(s string) []string {
 	r := regexp.MustCompile("'.+'|\".+\"|\\S+")
 	return r.FindAllString(s, -1)
 }
+
+func commandCtx(cmd *cobra.Command) (context.Context, context.CancelFunc) {
+	timeOut, err := cmd.Flags().GetDuration("command-timeout")
+	if err != nil {
+		ExitWithError(ExitError, err)
+	}
+	return context.WithTimeout(context.Background(), timeOut)
+}

etcdctlv3/main.go

@@ -27,7 +27,8 @@ const (
 	cliName        = "etcdctlv3"
 	cliDescription = "A simple command line client for etcd3."
 
-	defaultDialTimeout = 2 * time.Second
+	defaultDialTimeout    = 2 * time.Second
+	defaultCommandTimeOut = 5 * time.Second
 )
 
 var (
@@ -50,6 +51,7 @@ func init() {
 	rootCmd.PersistentFlags().BoolVar(&globalFlags.IsHex, "hex", false, "print byte strings as hex encoded strings")
 
 	rootCmd.PersistentFlags().DurationVar(&globalFlags.DialTimeout, "dial-timeout", defaultDialTimeout, "dial timeout for client connections")
+	rootCmd.PersistentFlags().DurationVar(&globalFlags.CommandTimeOut, "command-timeout", defaultCommandTimeOut, "timeout for short running command (excluding dial timeout)")
 
 	// TODO: secure by default when etcd enables secure gRPC by default.
 	rootCmd.PersistentFlags().BoolVar(&globalFlags.Insecure, "insecure-transport", true, "disable transport security for client connections")