agent: Refactor script invoke

This commit is contained in:
Armon Dadgar 2014-08-21 14:28:16 -07:00
parent 4b547a43d0
commit 46a96d9c42
4 changed files with 72 additions and 26 deletions

View File

@ -6,7 +6,6 @@ import (
"github.com/hashicorp/consul/consul/structs" "github.com/hashicorp/consul/consul/structs"
"log" "log"
"os/exec" "os/exec"
"runtime"
"sync" "sync"
"syscall" "syscall"
"time" "time"
@ -106,18 +105,13 @@ func (c *CheckMonitor) run() {
// check is invoked periodically to perform the script check // check is invoked periodically to perform the script check
func (c *CheckMonitor) check() { func (c *CheckMonitor) check() {
// Determine the shell invocation based on OS
var shell, flag string
if runtime.GOOS == "windows" {
shell = "cmd"
flag = "/C"
} else {
shell = "/bin/sh"
flag = "-c"
}
// Create the command // Create the command
cmd := exec.Command(shell, flag, c.Script) cmd, err := ExecScript(c.Script)
if err != nil {
c.Logger.Printf("[ERR] agent: failed to setup invoke '%s': %s", c.Script, err)
c.Notify.UpdateCheck(c.CheckID, structs.HealthUnknown, err.Error())
return
}
// Collect the output // Collect the output
output, _ := circbuf.NewBuffer(CheckBufSize) output, _ := circbuf.NewBuffer(CheckBufSize)
@ -140,7 +134,7 @@ func (c *CheckMonitor) check() {
time.Sleep(30 * time.Second) time.Sleep(30 * time.Second)
errCh <- fmt.Errorf("Timed out running check '%s'", c.Script) errCh <- fmt.Errorf("Timed out running check '%s'", c.Script)
}() }()
err := <-errCh err = <-errCh
// Get the output, add a message about truncation // Get the output, add a message about truncation
outputStr := string(output.Bytes()) outputStr := string(output.Bytes())

View File

@ -3,6 +3,8 @@ package agent
import ( import (
"math" "math"
"math/rand" "math/rand"
"os/exec"
"runtime"
"time" "time"
) )
@ -39,3 +41,17 @@ func strContains(l []string, s string) bool {
} }
return false return false
} }
// ExecScript returns a command to execute a script
func ExecScript(script string) (*exec.Cmd, error) {
var shell, flag string
if runtime.GOOS == "windows" {
shell = "cmd"
flag = "/C"
} else {
shell = "/bin/sh"
flag = "-c"
}
cmd := exec.Command(shell, flag, script)
return cmd, nil
}

View File

@ -7,8 +7,6 @@ import (
"io" "io"
"log" "log"
"os" "os"
"os/exec"
"runtime"
"strconv" "strconv"
"github.com/armon/circbuf" "github.com/armon/circbuf"
@ -39,18 +37,12 @@ func makeWatchHandler(logOutput io.Writer, params interface{}) watch.HandlerFunc
script := params.(string) script := params.(string)
logger := log.New(logOutput, "", log.LstdFlags) logger := log.New(logOutput, "", log.LstdFlags)
fn := func(idx uint64, data interface{}) { fn := func(idx uint64, data interface{}) {
// Determine the shell invocation based on OS
var shell, flag string
if runtime.GOOS == "windows" {
shell = "cmd"
flag = "/C"
} else {
shell = "/bin/sh"
flag = "-c"
}
// Create the command // Create the command
cmd := exec.Command(shell, flag, script) cmd, err := ExecScript(script)
if err != nil {
logger.Printf("[ERR] agent: Failed to setup watch: %v", err)
return
}
cmd.Env = append(os.Environ(), cmd.Env = append(os.Environ(),
"CONSUL_INDEX="+strconv.FormatUint(idx, 10), "CONSUL_INDEX="+strconv.FormatUint(idx, 10),
) )

View File

@ -0,0 +1,44 @@
package agent
import (
"io/ioutil"
"os"
"testing"
)
func TestVerifyWatchHandler(t *testing.T) {
if err := verifyWatchHandler(nil); err == nil {
t.Fatalf("should err")
}
if err := verifyWatchHandler(123); err == nil {
t.Fatalf("should err")
}
if err := verifyWatchHandler([]string{"foo"}); err == nil {
t.Fatalf("should err")
}
if err := verifyWatchHandler("foo"); err != nil {
t.Fatalf("err: %v", err)
}
}
func TestMakeWatchHandler(t *testing.T) {
defer os.Remove("handler_out")
defer os.Remove("handler_index_out")
script := "echo $CONSUL_INDEX >> handler_index_out && cat >> handler_out"
handler := makeWatchHandler(os.Stderr, script)
handler(100, []string{"foo", "bar", "baz"})
raw, err := ioutil.ReadFile("handler_out")
if err != nil {
t.Fatalf("err: %v", err)
}
if string(raw) != "[\"foo\",\"bar\",\"baz\"]\n" {
t.Fatalf("bad: %s", raw)
}
raw, err = ioutil.ReadFile("handler_index_out")
if err != nil {
t.Fatalf("err: %v", err)
}
if string(raw) != "100\n" {
t.Fatalf("bad: %s", raw)
}
}