Centralize command-line parsing

This commit is contained in:
Seth Vargo 2016-10-09 12:10:40 +08:00
parent 41f1764434
commit de1718a8a6
No known key found for this signature in database
GPG Key ID: 905A90C2949E8787
4 changed files with 329 additions and 101 deletions

View File

@ -1,7 +1,6 @@
package command
import (
"flag"
"fmt"
"os"
"path"
@ -12,7 +11,6 @@ import (
"github.com/hashicorp/consul/api"
"github.com/hashicorp/consul/command/agent"
"github.com/mitchellh/cli"
)
const (
@ -37,54 +35,42 @@ const (
// LockCommand is a Command implementation that is used to setup
// a "lock" which manages lock acquisition and invokes a sub-process
type LockCommand struct {
Meta
ShutdownCh <-chan struct{}
Ui cli.Ui
child *os.Process
childLock sync.Mutex
verbose bool
limit int
monitorRetry int
name string
passStdin bool
timeout time.Duration
verbose bool
}
func (c *LockCommand) Help() string {
helpText := `
Usage: consul lock [options] prefix child...
Acquires a lock or semaphore at a given path, and invokes a child
process when successful. The child process can assume the lock is
held while it executes. If the lock is lost or communication is
disrupted the child process will be sent a SIGTERM signal and given
time to gracefully exit. After the grace period expires the process
will be hard terminated.
Acquires a lock or semaphore at a given path, and invokes a child process
when successful. The child process can assume the lock is held while it
executes. If the lock is lost or communication is disrupted the child
process will be sent a SIGTERM signal and given time to gracefully exit.
After the grace period expires the process will be hard terminated.
For Consul agents on Windows, the child process is always hard
terminated with a SIGKILL, since Windows has no POSIX compatible
notion for SIGTERM.
For Consul agents on Windows, the child process is always hard terminated
with a SIGKILL, since Windows has no POSIX compatible notion for SIGTERM.
When -n=1, only a single lock holder or leader exists providing
mutual exclusion. Setting a higher value switches to a semaphore
allowing multiple holders to coordinate.
When -n=1, only a single lock holder or leader exists providing mutual
exclusion. Setting a higher value switches to a semaphore allowing multiple
holders to coordinate.
The prefix provided must have write privileges.
Options:
` + c.Meta.Help()
-http-addr=127.0.0.1:8500 HTTP address of the Consul agent.
-n=1 Maximum number of allowed lock holders. If this
value is one, it operates as a lock, otherwise
a semaphore is used.
-name="" Optional name to associate with lock session.
-token="" ACL token to use. Defaults to that of agent.
-pass-stdin Pass stdin to child process.
-try=timeout Attempt to acquire the lock up to the given
timeout (eg. "15s").
-monitor-retry=n Retry up to n times if Consul returns a 500 error
while monitoring the lock. This allows riding out brief
periods of unavailability without causing leader
elections, but increases the amount of time required
to detect a lost lock in some cases. Defaults to 3,
with a 1s wait between retries. Set to 0 to disable.
-verbose Enables verbose output
`
return strings.TrimSpace(helpText)
}
@ -93,40 +79,49 @@ func (c *LockCommand) Run(args []string) int {
return c.run(args, &lu)
}
// run exposes the underlying lock for testing.
func (c *LockCommand) run(args []string, lu **LockUnlock) int {
var childDone chan struct{}
var name, token string
var limit int
var passStdin bool
var try string
var retry int
cmdFlags := flag.NewFlagSet("watch", flag.ContinueOnError)
cmdFlags.Usage = func() { c.Ui.Output(c.Help()) }
cmdFlags.IntVar(&limit, "n", 1, "")
cmdFlags.StringVar(&name, "name", "", "")
cmdFlags.StringVar(&token, "token", "", "")
cmdFlags.BoolVar(&passStdin, "pass-stdin", false, "")
cmdFlags.StringVar(&try, "try", "", "")
cmdFlags.IntVar(&retry, "monitor-retry", defaultMonitorRetry, "")
cmdFlags.BoolVar(&c.verbose, "verbose", false, "")
httpAddr := HTTPAddrFlag(cmdFlags)
if err := cmdFlags.Parse(args); err != nil {
f := c.Meta.NewFlagSet(c)
f.IntVar(&c.limit, "limit", 1,
"Optional limit on the number of concurrent lock holders. The underlying "+
"implementation switches from a lock to a semaphore when the value is "+
"greater than 1. The default value is 1.")
f.IntVar(&c.monitorRetry, "monitor-retry", defaultMonitorRetry,
"Number of times to retry Consul returns a 500 error while monitoring "+
"the lock. This allows riding out brief periods of unavailability "+
"without causing leader elections, but increases the amount of time "+
"required to detect a lost lock in some cases. The default value is 3, "+
"with a 1s wait between retries. Set this value to 0 to disable retires.")
f.StringVar(&c.name, "name", "",
"Optional name to associate with the lock session. It not provided, one "+
"is generated based on the provided child command.")
f.BoolVar(&c.passStdin, "pass-stdin", false,
"Pass stdin to the child process.")
f.DurationVar(&c.timeout, "timeout", 0,
"Maximum amount of time to wait to acquire the lock, specified as a "+
"timestamp like \"1s\" or \"3h\". The default value is 0.")
f.BoolVar(&c.verbose, "verbose", false,
"Enable verbose (debugging) output.")
// Deprecations
f.DurationVar(&c.timeout, "try", 0,
"DEPRECATED. Use -timeout instead.")
if err := c.Meta.Parse(args); err != nil {
return 1
}
// Check the limit
if limit <= 0 {
if c.limit <= 0 {
c.Ui.Error(fmt.Sprintf("Lock holder limit must be positive"))
return 1
}
// Verify the prefix and child are provided
extra := cmdFlags.Args()
extra := f.Args()
if len(extra) < 2 {
c.Ui.Error("Key prefix and child command must be specified")
c.Ui.Error("")
c.Ui.Error(c.Help())
return 1
}
prefix := extra[0]
@ -134,40 +129,21 @@ func (c *LockCommand) run(args []string, lu **LockUnlock) int {
script := strings.Join(extra[1:], " ")
// Calculate a session name if none provided
if name == "" {
name = fmt.Sprintf("Consul lock for '%s' at '%s'", script, prefix)
if c.name == "" {
c.name = fmt.Sprintf("Consul lock for '%s' at '%s'", script, prefix)
}
// Verify the duration if given.
oneshot := false
var wait time.Duration
if try != "" {
var err error
wait, err = time.ParseDuration(try)
if err != nil {
c.Ui.Error(fmt.Sprintf("Error parsing try timeout: %s", err))
return 1
}
if wait <= 0 {
c.Ui.Error("Try timeout must be positive")
return 1
}
oneshot = true
}
// Calculate oneshot
oneshot := c.timeout > 0
// Check the retry parameter
if retry < 0 {
if c.monitorRetry < 0 {
c.Ui.Error("Number for 'monitor-retry' must be >= 0")
return 1
}
// Create and test the HTTP client
conf := api.DefaultConfig()
conf.Address = *httpAddr
conf.Token = token
client, err := api.NewClient(conf)
client, err := c.Meta.HTTPClient()
if err != nil {
c.Ui.Error(fmt.Sprintf("Error connecting to Consul agent: %s", err))
return 1
@ -179,10 +155,10 @@ func (c *LockCommand) run(args []string, lu **LockUnlock) int {
}
// Setup the lock or semaphore
if limit == 1 {
*lu, err = c.setupLock(client, prefix, name, oneshot, wait, retry)
if c.limit == 1 {
*lu, err = c.setupLock(client, prefix, c.name, oneshot, c.timeout, c.monitorRetry)
} else {
*lu, err = c.setupSemaphore(client, limit, prefix, name, oneshot, wait, retry)
*lu, err = c.setupSemaphore(client, c.limit, prefix, c.name, oneshot, c.timeout, c.monitorRetry)
}
if err != nil {
c.Ui.Error(fmt.Sprintf("Lock setup failed: %s", err))
@ -214,7 +190,7 @@ func (c *LockCommand) run(args []string, lu **LockUnlock) int {
// Start the child process
childDone = make(chan struct{})
go func() {
if err := c.startChild(script, childDone, passStdin); err != nil {
if err := c.startChild(script, childDone, c.passStdin); err != nil {
c.Ui.Error(fmt.Sprintf("%s", err))
}
}()

View File

@ -12,13 +12,22 @@ import (
"github.com/mitchellh/cli"
)
func testLockCommand(t *testing.T) (*cli.MockUi, *LockCommand) {
ui := new(cli.MockUi)
return ui, &LockCommand{
Meta: Meta{
Ui: ui,
Flags: FlagSetHTTP,
},
}
}
func TestLockCommand_implements(t *testing.T) {
var _ cli.Command = &LockCommand{}
}
func argFail(t *testing.T, args []string, expected string) {
ui := new(cli.MockUi)
c := &LockCommand{Ui: ui}
ui, c := testLockCommand(t)
if code := c.Run(args); code != 1 {
t.Fatalf("expected return code 1, got %d", code)
}
@ -40,8 +49,7 @@ func TestLockCommand_Run(t *testing.T) {
defer a1.Shutdown()
waitForLeader(t, a1.httpAddr)
ui := new(cli.MockUi)
c := &LockCommand{Ui: ui}
ui, c := testLockCommand(t)
filePath := filepath.Join(a1.dir, "test_touch")
touchCmd := fmt.Sprintf("touch '%s'", filePath)
args := []string{"-http-addr=" + a1.httpAddr, "test/prefix", touchCmd}
@ -63,8 +71,7 @@ func TestLockCommand_Try_Lock(t *testing.T) {
defer a1.Shutdown()
waitForLeader(t, a1.httpAddr)
ui := new(cli.MockUi)
c := &LockCommand{Ui: ui}
ui, c := testLockCommand(t)
filePath := filepath.Join(a1.dir, "test_touch")
touchCmd := fmt.Sprintf("touch '%s'", filePath)
args := []string{"-http-addr=" + a1.httpAddr, "-try=10s", "test/prefix", touchCmd}
@ -95,8 +102,7 @@ func TestLockCommand_Try_Semaphore(t *testing.T) {
defer a1.Shutdown()
waitForLeader(t, a1.httpAddr)
ui := new(cli.MockUi)
c := &LockCommand{Ui: ui}
ui, c := testLockCommand(t)
filePath := filepath.Join(a1.dir, "test_touch")
touchCmd := fmt.Sprintf("touch '%s'", filePath)
args := []string{"-http-addr=" + a1.httpAddr, "-n=3", "-try=10s", "test/prefix", touchCmd}
@ -127,8 +133,7 @@ func TestLockCommand_MonitorRetry_Lock_Default(t *testing.T) {
defer a1.Shutdown()
waitForLeader(t, a1.httpAddr)
ui := new(cli.MockUi)
c := &LockCommand{Ui: ui}
ui, c := testLockCommand(t)
filePath := filepath.Join(a1.dir, "test_touch")
touchCmd := fmt.Sprintf("touch '%s'", filePath)
args := []string{"-http-addr=" + a1.httpAddr, "test/prefix", touchCmd}
@ -160,8 +165,7 @@ func TestLockCommand_MonitorRetry_Semaphore_Default(t *testing.T) {
defer a1.Shutdown()
waitForLeader(t, a1.httpAddr)
ui := new(cli.MockUi)
c := &LockCommand{Ui: ui}
ui, c := testLockCommand(t)
filePath := filepath.Join(a1.dir, "test_touch")
touchCmd := fmt.Sprintf("touch '%s'", filePath)
args := []string{"-http-addr=" + a1.httpAddr, "-n=3", "test/prefix", touchCmd}
@ -193,8 +197,7 @@ func TestLockCommand_MonitorRetry_Lock_Arg(t *testing.T) {
defer a1.Shutdown()
waitForLeader(t, a1.httpAddr)
ui := new(cli.MockUi)
c := &LockCommand{Ui: ui}
ui, c := testLockCommand(t)
filePath := filepath.Join(a1.dir, "test_touch")
touchCmd := fmt.Sprintf("touch '%s'", filePath)
args := []string{"-http-addr=" + a1.httpAddr, "-monitor-retry=9", "test/prefix", touchCmd}
@ -226,8 +229,7 @@ func TestLockCommand_MonitorRetry_Semaphore_Arg(t *testing.T) {
defer a1.Shutdown()
waitForLeader(t, a1.httpAddr)
ui := new(cli.MockUi)
c := &LockCommand{Ui: ui}
ui, c := testLockCommand(t)
filePath := filepath.Join(a1.dir, "test_touch")
touchCmd := fmt.Sprintf("touch '%s'", filePath)
args := []string{"-http-addr=" + a1.httpAddr, "-n=3", "-monitor-retry=9", "test/prefix", touchCmd}

247
command/meta.go Normal file
View File

@ -0,0 +1,247 @@
package command
import (
"bufio"
"bytes"
"flag"
"fmt"
"io"
"strings"
"github.com/hashicorp/consul/api"
"github.com/mitchellh/cli"
text "github.com/tonnerre/golang-text"
)
// maxLineLength is the maximum width of any line.
const maxLineLength int = 72
// FlagSetFlags is an enum to define what flags are present in the
// default FlagSet returned.
type FlagSetFlags uint
const (
FlagSetNone FlagSetFlags = iota << 1
FlagSetHTTP FlagSetFlags = iota << 1
FlagSetRPC FlagSetFlags = iota << 1
)
type Meta struct {
Ui cli.Ui
Flags FlagSetFlags
flagSet *flag.FlagSet
// These are the options which correspond to the HTTP API options
httpAddr string
datacenter string
token string
stale bool
rpcAddr string
}
// HTTPClient returns a client with the parsed flags. It panics if the command
// does not accept HTTP flags or if the flags have not been parsed.
func (m *Meta) HTTPClient() (*api.Client, error) {
if !m.hasHTTP() {
panic("no http flags defined")
}
if !m.flagSet.Parsed() {
panic("flags have not been parsed")
}
return api.NewClient(&api.Config{
Datacenter: m.datacenter,
Address: m.httpAddr,
Token: m.token,
})
}
// httpFlags is the list of flags that apply to HTTP connections.
func (m *Meta) httpFlags(f *flag.FlagSet) *flag.FlagSet {
if f == nil {
f = flag.NewFlagSet("", flag.ContinueOnError)
}
f.StringVar(&m.datacenter, "datacenter", "",
"Name of the datacenter to query. If unspecified, this will default to "+
"the datacenter of the queried agent.")
f.StringVar(&m.httpAddr, "http-addr", "",
"Address and port to the Consul HTTP agent. The value can be an IP "+
"address or DNS address, but it must also include the port. This can "+
"also be specified via the CONSUL_HTTP_ADDR environment variable. The "+
"default value is 127.0.0.1:8500.")
f.StringVar(&m.token, "token", "",
"ACL token to use in the request. This can also be specified via the "+
"CONSUL_HTTP_TOKEN environment variable. If unspecified, the query will "+
"default to the token of the Consul agent at the HTTP address.")
f.BoolVar(&m.stale, "stale", false,
"Permit any Consul server (non-leader) to respond to this request. This "+
"allows for lower latency and higher throughput, but can result in "+
"stale data. This option has no effect on non-read operations. The "+
"default value is false.")
return f
}
// RPCClient returns a client with the parsed flags. It panics if the command
// does not accept RPC flags or if the flags have not been parsed.
func (m *Meta) RPCClient() (*api.Client, error) {
if !m.hasRPC() {
panic("no rpc flags defined")
}
if !m.flagSet.Parsed() {
panic("flags have not been parsed")
}
// TODO
return nil, nil
}
// rpcFlags is the list of flags that apply to RPC connections.
func (m *Meta) rpcFlags(f *flag.FlagSet) *flag.FlagSet {
if f == nil {
f = flag.NewFlagSet("", flag.ContinueOnError)
}
f.StringVar(&m.rpcAddr, "rpc-addr", "",
"Address and port to the Consul RPC agent. The value can be an IP "+
"address or DNS address, but it must also include the port. This can "+
"also be specified via the CONSUL_RPC_ADDR environment variable. The "+
"default value is 127.0.0.1:8400.")
return f
}
// NewFlagSet creates a new flag set for the given command. It automatically
// generates help output and adds the appropriate API flags.
func (m *Meta) NewFlagSet(c cli.Command) *flag.FlagSet {
f := flag.NewFlagSet("", flag.ContinueOnError)
f.Usage = func() { m.Ui.Error(c.Help()) }
if m.hasHTTP() {
m.httpFlags(f)
}
if m.hasRPC() {
m.rpcFlags(f)
}
errR, errW := io.Pipe()
errScanner := bufio.NewScanner(errR)
go func() {
for errScanner.Scan() {
m.Ui.Error(errScanner.Text())
}
}()
f.SetOutput(errW)
m.flagSet = f
return f
}
// Parse is used to parse the underlying flag set.
func (m *Meta) Parse(args []string) error {
return m.flagSet.Parse(args)
}
// Help returns the help for this flagSet.
func (m *Meta) Help() string {
return m.helpFlagsFor(m.flagSet)
}
// hasHTTP returns true if this meta command contains HTTP flags.
func (m *Meta) hasHTTP() bool {
return m.Flags&FlagSetHTTP != 0
}
// hasRPC returns true if this meta command contains RPC flags.
func (m *Meta) hasRPC() bool {
return m.Flags&FlagSetRPC != 0
}
// helpFlagsFor visits all flags in the given flag set and prints formatted
// help output. This function is sad because there's no "merging" of command
// line flags. We explicitly pull out our "common" options into another section
// by doing string comparisons :(.
func (m *Meta) helpFlagsFor(f *flag.FlagSet) string {
httpFlags := m.httpFlags(nil)
rpcFlags := m.rpcFlags(nil)
var out bytes.Buffer
printTitle(&out, "Command Options")
f.VisitAll(func(f *flag.Flag) {
// Skip HTTP and RPC flags as they will be grouped separately
if flagContains(httpFlags, f) || flagContains(rpcFlags, f) {
return
}
printFlag(&out, f)
})
if m.hasHTTP() {
printTitle(&out, "HTTP API Options")
httpFlags.VisitAll(func(f *flag.Flag) {
printFlag(&out, f)
})
}
if m.hasRPC() {
printTitle(&out, "RPC API Options")
rpcFlags.VisitAll(func(f *flag.Flag) {
printFlag(&out, f)
})
}
return strings.TrimRight(out.String(), "\n")
}
// printTitle prints a consistently-formatted title to the given writer.
func printTitle(w io.Writer, s string) {
fmt.Fprintf(w, "%s\n\n", s)
}
// printFlag prints a single flag to the given writer.
func printFlag(w io.Writer, f *flag.Flag) {
example, _ := flag.UnquoteUsage(f)
if example != "" {
fmt.Fprintf(w, " -%s=<%s>\n", f.Name, example)
} else {
fmt.Fprintf(w, " -%s\n", f.Name)
}
indented := wrapAtLength(f.Usage, 5)
fmt.Fprintf(w, "%s\n\n", indented)
}
// flagContains returns true if the given flag is contained in the given flag
// set or false otherwise.
func flagContains(fs *flag.FlagSet, f *flag.Flag) bool {
var skip bool
fs.VisitAll(func(hf *flag.Flag) {
if skip {
return
}
if f.Name == hf.Name && f.Usage == hf.Usage {
skip = true
return
}
})
return skip
}
// wrapAtLength wraps the given text at the maxLineLength, taxing into account
// any provided left padding.
func wrapAtLength(s string, pad int) string {
wrapped := text.Wrap(s, maxLineLength-pad)
lines := strings.Split(wrapped, "\n")
for i, line := range lines {
lines[i] = strings.Repeat(" ", pad) + line
}
return strings.Join(lines, "\n")
}

View File

@ -104,7 +104,10 @@ func init() {
"lock": func() (cli.Command, error) {
return &command.LockCommand{
ShutdownCh: makeShutdownCh(),
Ui: ui,
Meta: command.Meta{
Flags: command.FlagSetHTTP,
Ui: ui,
},
}, nil
},