etcdctl: use user specified timeout value for entire command execution

etcdctl should be capable to use a user specified timeout value for
total command execution, not only per request timeout. This commit
adds a new option --total-timeout to the command. The value passed via
this option is used as a timeout value of entire command execution.

This commit introduces a new rule of timeout for the commands:
 - If a command issues only one API call, the call is executed in a
   total context. The total context is created with
   contextWithTotalTimeout() and its timeout is passed to etcdctl with
   --total-timeout. It means that the API call can use whole timeout.
 - if a command issues more than twice API calls, the calls is
   executed in a per request context. The per request context is
   created with contextWithPerRequestTimeout() and its timeout is
   passed passed to etcdctl with --timeout. The context should be a
   child of the total context.

Fixes coreos#3517

etcdctl/command/auth_commands.go

@@ -20,7 +20,6 @@ import (
 	"strings"
 
 	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/codegangsta/cli"
-	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
 	"github.com/coreos/etcd/client"
 )
 
@@ -67,7 +66,7 @@ func authEnableDisable(c *cli.Context, enable bool) {
 		os.Exit(1)
 	}
 	s := mustNewAuthAPI(c)
-	ctx, cancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel := contextWithTotalTimeout(c)
 	var err error
 	if enable {
 		err = s.Enable(ctx)

etcdctl/command/member_commands.go

@@ -59,9 +59,9 @@ func actionMemberList(c *cli.Context) {
 		os.Exit(1)
 	}
 	mAPI := mustNewMembersAPI(c)
-	ctx, cancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel := contextWithTotalTimeout(c)
+	defer cancel()
 	members, err := mAPI.List(ctx)
-	cancel()
 	if err != nil {
 		fmt.Fprintln(os.Stderr, err.Error())
 		os.Exit(1)
@@ -86,7 +86,10 @@ func actionMemberAdd(c *cli.Context) {
 	mAPI := mustNewMembersAPI(c)
 
 	url := args[1]
-	ctx, cancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	totalCtx, totalCancel := contextWithTotalTimeout(c)
+	defer totalCancel()
+
+	ctx, cancel := contextWithPerRequestTimeout(c, totalCtx)
 	m, err := mAPI.Add(ctx, url)
 	cancel()
 	if err != nil {
@@ -98,7 +101,7 @@ func actionMemberAdd(c *cli.Context) {
 	newName := args[0]
 	fmt.Printf("Added member named %s with ID %s to cluster\n", newName, newID)
 
-	ctx, cancel = context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel = contextWithPerRequestTimeout(c, totalCtx)
 	members, err := mAPI.List(ctx)
 	cancel()
 	if err != nil {
@@ -132,10 +135,15 @@ func actionMemberRemove(c *cli.Context) {
 	removalID := args[0]
 
 	mAPI := mustNewMembersAPI(c)
+
+	totalCtx, _ := contextWithTotalTimeout(c)
 	// Get the list of members.
-	listctx, listCancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
-	members, err := mAPI.List(listctx)
-	listCancel()
+	listCtx, listCancel := context.WithTimeout(totalCtx, client.DefaultRequestTimeout)
+	defer listCancel()
+
+	ctx, cancel := contextWithPerRequestTimeout(c, listCtx)
+	members, err := mAPI.List(ctx)
+	cancel()
 	if err != nil {
 		fmt.Fprintln(os.Stderr, "Error while verifying ID against known members:", err.Error())
 		os.Exit(1)
@@ -158,9 +166,9 @@ func actionMemberRemove(c *cli.Context) {
 	}
 
 	// Actually attempt to remove the member.
-	ctx, removeCancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel = contextWithPerRequestTimeout(c, listCtx)
 	err = mAPI.Remove(ctx, removalID)
-	removeCancel()
+	cancel()
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "Received an error trying to remove member %s: %s", removalID, err.Error())
 		os.Exit(1)
@@ -180,7 +188,7 @@ func actionMemberUpdate(c *cli.Context) {
 
 	mid := args[0]
 	urls := args[1]
-	ctx, cancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel := contextWithTotalTimeout(c)
 	err := mAPI.Update(ctx, mid, strings.Split(urls, ","))
 	cancel()
 	if err != nil {

etcdctl/command/role_commands.go

@@ -21,7 +21,6 @@ import (
 	"strings"
 
 	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/codegangsta/cli"
-	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
 	"github.com/coreos/etcd/client"
 	"github.com/coreos/etcd/pkg/pathutil"
 )
@@ -93,7 +92,7 @@ func actionRoleList(c *cli.Context) {
 		os.Exit(1)
 	}
 	r := mustNewAuthRoleAPI(c)
-	ctx, cancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel := contextWithTotalTimeout(c)
 	roles, err := r.ListRoles(ctx)
 	cancel()
 	if err != nil {
@@ -108,14 +107,13 @@ func actionRoleList(c *cli.Context) {
 
 func actionRoleAdd(c *cli.Context) {
 	api, role := mustRoleAPIAndName(c)
-	ctx, cancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel := contextWithTotalTimeout(c)
 	currentRole, err := api.GetRole(ctx, role)
-	cancel()
 	if currentRole != nil {
 		fmt.Fprintf(os.Stderr, "Role %s already exists\n", role)
 		os.Exit(1)
 	}
-	ctx, cancel = context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+
 	err = api.AddRole(ctx, role)
 	cancel()
 	if err != nil {
@@ -128,7 +126,7 @@ func actionRoleAdd(c *cli.Context) {
 
 func actionRoleRemove(c *cli.Context) {
 	api, role := mustRoleAPIAndName(c)
-	ctx, cancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel := contextWithTotalTimeout(c)
 	err := api.RemoveRole(ctx, role)
 	cancel()
 	if err != nil {
@@ -182,14 +180,16 @@ func roleGrantRevoke(c *cli.Context, grant bool) {
 	}
 
 	api, role := mustRoleAPIAndName(c)
-	ctx, cancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	totalCtx, totalCancel := contextWithTotalTimeout(c)
+	defer totalCancel()
+	ctx, cancel := contextWithPerRequestTimeout(c, totalCtx)
 	currentRole, err := api.GetRole(ctx, role)
 	cancel()
 	if err != nil {
 		fmt.Fprintln(os.Stderr, err.Error())
 		os.Exit(1)
 	}
-	ctx, cancel = context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel = contextWithPerRequestTimeout(c, totalCtx)
 	var newRole *client.Role
 	if grant {
 		newRole, err = api.GrantRoleKV(ctx, role, []string{path}, permType)
@@ -215,7 +215,7 @@ func roleGrantRevoke(c *cli.Context, grant bool) {
 func actionRoleGet(c *cli.Context) {
 	api, rolename := mustRoleAPIAndName(c)
 
-	ctx, cancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel := contextWithTotalTimeout(c)
 	role, err := api.GetRole(ctx, rolename)
 	cancel()
 	if err != nil {

etcdctl/command/user_commands.go

@@ -23,7 +23,6 @@ import (
 
 	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/bgentry/speakeasy"
 	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/codegangsta/cli"
-	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
 	"github.com/coreos/etcd/client"
 )
 
@@ -89,7 +88,7 @@ func actionUserList(c *cli.Context) {
 		os.Exit(1)
 	}
 	u := mustNewAuthUserAPI(c)
-	ctx, cancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel := contextWithTotalTimeout(c)
 	users, err := u.ListUsers(ctx)
 	cancel()
 	if err != nil {
@@ -104,7 +103,9 @@ func actionUserList(c *cli.Context) {
 
 func actionUserAdd(c *cli.Context) {
 	api, user := mustUserAPIAndName(c)
-	ctx, cancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	totalCtx, totalCancel := contextWithTotalTimeout(c)
+	defer totalCancel()
+	ctx, cancel := contextWithPerRequestTimeout(c, totalCtx)
 	currentUser, err := api.GetUser(ctx, user)
 	cancel()
 	if currentUser != nil {
@@ -116,7 +117,7 @@ func actionUserAdd(c *cli.Context) {
 		fmt.Fprintln(os.Stderr, "Error reading password:", err)
 		os.Exit(1)
 	}
-	ctx, cancel = context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel = contextWithPerRequestTimeout(c, totalCtx)
 	err = api.AddUser(ctx, user, pass)
 	cancel()
 	if err != nil {
@@ -129,7 +130,7 @@ func actionUserAdd(c *cli.Context) {
 
 func actionUserRemove(c *cli.Context) {
 	api, user := mustUserAPIAndName(c)
-	ctx, cancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel := contextWithTotalTimeout(c)
 	err := api.RemoveUser(ctx, user)
 	cancel()
 	if err != nil {
@@ -142,7 +143,9 @@ func actionUserRemove(c *cli.Context) {
 
 func actionUserPasswd(c *cli.Context) {
 	api, user := mustUserAPIAndName(c)
-	ctx, cancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	totalCtx, totalCancel := contextWithTotalTimeout(c)
+	defer totalCancel()
+	ctx, cancel := contextWithPerRequestTimeout(c, totalCtx)
 	currentUser, err := api.GetUser(ctx, user)
 	cancel()
 	if currentUser == nil {
@@ -155,7 +158,7 @@ func actionUserPasswd(c *cli.Context) {
 		os.Exit(1)
 	}
 
-	ctx, cancel = context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel = contextWithPerRequestTimeout(c, totalCtx)
 	_, err = api.ChangePassword(ctx, user, pass)
 	cancel()
 	if err != nil {
@@ -181,16 +184,19 @@ func userGrantRevoke(c *cli.Context, grant bool) {
 		os.Exit(1)
 	}
 
+	totalCtx, totalCancel := contextWithTotalTimeout(c)
+	defer totalCancel()
+
 	api, user := mustUserAPIAndName(c)
-	ctx, cancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel := contextWithPerRequestTimeout(c, totalCtx)
 	currentUser, err := api.GetUser(ctx, user)
 	cancel()
 	if currentUser == nil {
 		fmt.Fprintln(os.Stderr, err.Error())
 		os.Exit(1)
 	}
 
-	ctx, cancel = context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel = contextWithPerRequestTimeout(c, totalCtx)
 	var newUser *client.User
 	if grant {
 		newUser, err = api.GrantUser(ctx, user, roles)
@@ -217,7 +223,7 @@ func userGrantRevoke(c *cli.Context, grant bool) {
 
 func actionUserGet(c *cli.Context) {
 	api, username := mustUserAPIAndName(c)
-	ctx, cancel := context.WithTimeout(context.Background(), client.DefaultRequestTimeout)
+	ctx, cancel := contextWithTotalTimeout(c)
 	user, err := api.GetUser(ctx, username)
 	cancel()
 	if err != nil {

etcdctl/command/util.go

@@ -263,3 +263,11 @@ func newClient(c *cli.Context) (client.Client, error) {
 
 	return client.New(cfg)
 }
+
+func contextWithTotalTimeout(c *cli.Context) (context.Context, context.CancelFunc) {
+	return context.WithTimeout(context.Background(), c.GlobalDuration("total-timeout"))
+}
+
+func contextWithPerRequestTimeout(c *cli.Context, parent context.Context) (context.Context, context.CancelFunc) {
+	return context.WithTimeout(parent, c.GlobalDuration("timeout"))
+}

etcdctl/main.go

@@ -40,6 +40,7 @@ func main() {
 		cli.StringFlag{Name: "ca-file", Value: "", Usage: "verify certificates of HTTPS-enabled servers using this CA bundle"},
 		cli.StringFlag{Name: "username, u", Value: "", Usage: "provide username[:password] and prompt if password is not supplied."},
 		cli.DurationFlag{Name: "timeout", Value: time.Second, Usage: "connection timeout per request"},
+		cli.DurationFlag{Name: "total-timeout", Value: 5 * time.Second, Usage: "timeout for the command execution (except watch)"},
 	}
 	app.Commands = []cli.Command{
 		command.NewBackupCommand(),