golang-github-hashicorp-serf-dev 0.7.0~ds1-1 / usr / share / gocode / src / github.com / hashicorp / serf / command / agent / invoke.go

package agent

import (
	"fmt"
	"github.com/armon/circbuf"
	"github.com/armon/go-metrics"
	"github.com/hashicorp/serf/serf"
	"io"
	"log"
	"os"
	"os/exec"
	"regexp"
	"runtime"
	"strings"
	"time"
)

const (
	windows = "windows"

	// maxBufSize limits how much data we collect from a handler.
	// This is to prevent Serf's memory from growing to an enormous
	// amount due to a faulty handler.
	maxBufSize = 8 * 1024

	// warnSlow is used to warn about a slow handler invocation
	warnSlow = time.Second
)

var sanitizeTagRegexp = regexp.MustCompile(`[^A-Z0-9_]`)

// invokeEventScript will execute the given event script with the given
// event. Depending on the event, the semantics of how data are passed
// are a bit different. For all events, the SERF_EVENT environmental
// variable is the type of the event. For user events, the SERF_USER_EVENT
// environmental variable is also set, containing the name of the user
// event that was fired.
//
// In all events, data is passed in via stdin to facilitate piping. See
// the various stdin functions below for more information.
func invokeEventScript(logger *log.Logger, script string, self serf.Member, event serf.Event) error {
	defer metrics.MeasureSince([]string{"agent", "invoke", script}, time.Now())
	output, _ := circbuf.NewBuffer(maxBufSize)

	// 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"
	}

	cmd := exec.Command(shell, flag, script)
	cmd.Env = append(os.Environ(),
		"SERF_EVENT="+event.EventType().String(),
		"SERF_SELF_NAME="+self.Name,
		"SERF_SELF_ROLE="+self.Tags["role"],
	)
	cmd.Stderr = output
	cmd.Stdout = output

	// Add all the tags
	for name, val := range self.Tags {
		//http://stackoverflow.com/questions/2821043/allowed-characters-in-linux-environment-variable-names
		//(http://pubs.opengroup.org/onlinepubs/000095399/basedefs/xbd_chap08.html for the long version)
		//says that env var names must be in [A-Z0-9_] and not start with [0-9].
		//we only care about the first part, so convert all chars not in [A-Z0-9_] to _
		sanitizedName := sanitizeTagRegexp.ReplaceAllString(strings.ToUpper(name), "_")
		tag_env := fmt.Sprintf("SERF_TAG_%s=%s", sanitizedName, val)
		cmd.Env = append(cmd.Env, tag_env)
	}

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return err
	}

	switch e := event.(type) {
	case serf.MemberEvent:
		go memberEventStdin(logger, stdin, &e)
	case serf.UserEvent:
		cmd.Env = append(cmd.Env, "SERF_USER_EVENT="+e.Name)
		cmd.Env = append(cmd.Env, fmt.Sprintf("SERF_USER_LTIME=%d", e.LTime))
		go streamPayload(logger, stdin, e.Payload)
	case *serf.Query:
		cmd.Env = append(cmd.Env, "SERF_QUERY_NAME="+e.Name)
		cmd.Env = append(cmd.Env, fmt.Sprintf("SERF_QUERY_LTIME=%d", e.LTime))
		go streamPayload(logger, stdin, e.Payload)
	default:
		return fmt.Errorf("Unknown event type: %s", event.EventType().String())
	}

	// Start a timer to warn about slow handlers
	slowTimer := time.AfterFunc(warnSlow, func() {
		logger.Printf("[WARN] agent: Script '%s' slow, execution exceeding %v",
			script, warnSlow)
	})

	if err := cmd.Start(); err != nil {
		return err
	}

	// Warn if buffer is overritten
	if output.TotalWritten() > output.Size() {
		logger.Printf("[WARN] agent: Script '%s' generated %d bytes of output, truncated to %d",
			script, output.TotalWritten(), output.Size())
	}

	err = cmd.Wait()
	slowTimer.Stop()
	logger.Printf("[DEBUG] agent: Event '%s' script output: %s",
		event.EventType().String(), output.String())
	if err != nil {
		return err
	}

	// If this is a query and we have output, respond
	if query, ok := event.(*serf.Query); ok && output.TotalWritten() > 0 {
		if err := query.Respond(output.Bytes()); err != nil {
			logger.Printf("[WARN] agent: Failed to respond to query '%s': %s",
				event.String(), err)
		}
	}

	return nil
}

// eventClean cleans a value to be a parameter in an event line.
func eventClean(v string) string {
	v = strings.Replace(v, "\t", "\\t", -1)
	v = strings.Replace(v, "\n", "\\n", -1)
	return v
}

// Sends data on stdin for a member event.
//
// The format for the data is unix tool friendly, separated by whitespace
// and newlines. The structure of each line for any member event is:
// "NAME    ADDRESS    ROLE    TAGS" where the whitespace is actually tabs.
// The name and role are cleaned so that newlines and tabs are replaced
// with "\n" and "\t" respectively.
func memberEventStdin(logger *log.Logger, stdin io.WriteCloser, e *serf.MemberEvent) {
	defer stdin.Close()
	for _, member := range e.Members {
		// Format the tags as tag1=v1,tag2=v2,...
		var tagPairs []string
		for name, value := range member.Tags {
			tagPairs = append(tagPairs, fmt.Sprintf("%s=%s", name, value))
		}
		tags := strings.Join(tagPairs, ",")

		// Send the entire line
		_, err := stdin.Write([]byte(fmt.Sprintf(
			"%s\t%s\t%s\t%s\n",
			eventClean(member.Name),
			member.Addr.String(),
			eventClean(member.Tags["role"]),
			eventClean(tags))))
		if err != nil {
			return
		}
	}
}

// Sends data on stdin for an event. The stdin simply contains the
// payload (if any).
// Most shells read implementations need a newline, force it to be there
func streamPayload(logger *log.Logger, stdin io.WriteCloser, buf []byte) {
	defer stdin.Close()

	// Append a newline to payload if missing
	payload := buf
	if len(payload) > 0 && payload[len(payload)-1] != '\n' {
		payload = append(payload, '\n')
	}

	if _, err := stdin.Write(payload); err != nil {
		logger.Printf("[ERR] Error writing payload: %s", err)
		return
	}
}