From 643c2a310d437d6e30d78804ad521a33901e2763 Mon Sep 17 00:00:00 2001
From: Anthony Romano <anthony.romano@coreos.com>
Date: Wed, 24 May 2017 13:57:16 -0700
Subject: [PATCH] etcdctl: support exec on lock

The lock command is clumsy to use from the command line, needing mkfifo,
wait, etc. Instead, make like consul and support launching a command if
one is given.
---
 etcdctl/README.md                     | 13 ++++++++++-
 etcdctl/ctlv3/command/lock_command.go | 32 ++++++++++++++++++++++-----
 2 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/etcdctl/README.md b/etcdctl/README.md
index 2781177f2..22682370a 100644
--- a/etcdctl/README.md
+++ b/etcdctl/README.md
@@ -790,7 +790,7 @@ Prints a line of JSON encoding the database hash, revision, total keys, and size
 
 ## Concurrency commands
 
-### LOCK \<lockname\>
+### LOCK \<lockname\> [command arg1 arg2 ...]
 
 LOCK acquires a distributed named mutex with a given name. Once the lock is acquired, it will be held until etcdctl is terminated.
 
@@ -798,13 +798,24 @@ LOCK acquires a distributed named mutex with a given name. Once the lock is acqu
 
 Once the lock is acquired, the result for the GET on the unique lock holder key is displayed.
 
+If a command is given, it will be launched with environment variables `ETCD_LOCK_KEY` and `ETCD_LOCK_REV` set to the lock's holder key and revision.
+
 #### Example
 
+Acquire lock with standard output display:
+
 ```bash
 ./etcdctl lock mylock
 # mylock/1234534535445
 ```
 
+Acquire lock and execute `echo lock acquired`:
+
+```bash
+./etcdctl lock mylock echo lock acquired
+# lock acquired
+```
+
 #### Remarks
 
 LOCK returns a zero exit code only if it is terminated by a signal and releases the lock.
diff --git a/etcdctl/ctlv3/command/lock_command.go b/etcdctl/ctlv3/command/lock_command.go
index 2e55c49df..e130493f8 100644
--- a/etcdctl/ctlv3/command/lock_command.go
+++ b/etcdctl/ctlv3/command/lock_command.go
@@ -16,7 +16,9 @@ package command
 
 import (
 	"errors"
+	"fmt"
 	"os"
+	"os/exec"
 	"os/signal"
 	"syscall"
 
@@ -29,7 +31,7 @@ import (
 // NewLockCommand returns the cobra command for "lock".
 func NewLockCommand() *cobra.Command {
 	c := &cobra.Command{
-		Use:   "lock <lockname>",
+		Use:   "lock <lockname> [exec-command arg1 arg2 ...]",
 		Short: "Acquires a named lock",
 		Run:   lockCommandFunc,
 	}
@@ -37,16 +39,16 @@ func NewLockCommand() *cobra.Command {
 }
 
 func lockCommandFunc(cmd *cobra.Command, args []string) {
-	if len(args) != 1 {
-		ExitWithError(ExitBadArgs, errors.New("lock takes one lock name argument."))
+	if len(args) == 0 {
+		ExitWithError(ExitBadArgs, errors.New("lock takes a lock name argument and an optional command to execute."))
 	}
 	c := mustClientFromCmd(cmd)
-	if err := lockUntilSignal(c, args[0]); err != nil {
+	if err := lockUntilSignal(c, args[0], args[1:]); err != nil {
 		ExitWithError(ExitError, err)
 	}
 }
 
-func lockUntilSignal(c *clientv3.Client, lockname string) error {
+func lockUntilSignal(c *clientv3.Client, lockname string, cmdArgs []string) error {
 	s, err := concurrency.NewSession(c)
 	if err != nil {
 		return err
@@ -69,6 +71,18 @@ func lockUntilSignal(c *clientv3.Client, lockname string) error {
 		return err
 	}
 
+	if len(cmdArgs) > 0 {
+		cmd := exec.Command(cmdArgs[0], cmdArgs[1:]...)
+		cmd.Env = append(environLockResponse(m), os.Environ()...)
+		cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
+		err := cmd.Run()
+		unlockErr := m.Unlock(context.TODO())
+		if err != nil {
+			return err
+		}
+		return unlockErr
+	}
+
 	k, kerr := c.Get(ctx, m.Key())
 	if kerr != nil {
 		return kerr
@@ -76,7 +90,6 @@ func lockUntilSignal(c *clientv3.Client, lockname string) error {
 	if len(k.Kvs) == 0 {
 		return errors.New("lock lost on init")
 	}
-
 	display.Get(*k)
 
 	select {
@@ -87,3 +100,10 @@ func lockUntilSignal(c *clientv3.Client, lockname string) error {
 
 	return errors.New("session expired")
 }
+
+func environLockResponse(m *concurrency.Mutex) []string {
+	return []string{
+		"ETCD_LOCK_KEY=" + m.Key(),
+		fmt.Sprintf("ETCD_LOCK_REV=%d", m.Header().Revision),
+	}
+}