feature(worker): added exec_on_success and exec_on_failure option and hooks

worker/config.go

@@ -76,6 +76,9 @@ type mirrorConfig struct {
 	LogDir    string            `toml:"log_dir"`
 	Env       map[string]string `toml:"env"`
 
+	ExecOnSuccess string `toml:"exec_on_success"`
+	ExecOnFailure string `toml:"exec_on_failure"`
+
 	Command       string `toml:"command"`
 	UseIPv6       bool   `toml:"use_ipv6"`
 	ExcludeFile   string `toml:"exclude_file"`

worker/config_test.go

@@ -34,6 +34,7 @@ provider = "command"
 upstream = "https://aosp.google.com/"
 interval = 720
 mirror_dir = "/data/git/AOSP"
+exec_on_success = "bash -c 'echo ${TUNASYNC_JOB_EXIT_STATUS} > ${TUNASYNC_WORKING_DIR}/exit_status'"
 	[mirrors.env]
 	REPO = "/usr/local/bin/aosp-repo"
 
@@ -51,6 +52,7 @@ provider = "rsync"
 upstream = "rsync://ftp.fedoraproject.org/fedora/"
 use_ipv6 = true
 exclude_file = "/etc/tunasync.d/fedora-exclude.txt"
+exec_on_failure = "bash -c 'echo ${TUNASYNC_JOB_EXIT_STATUS} > ${TUNASYNC_WORKING_DIR}/exit_status'"
 	`
 
 	Convey("When giving invalid file", t, func() {
@@ -123,6 +125,12 @@ exclude_file = "/etc/tunasync.d/fedora-exclude.txt"
 		So(p.LogFile(), ShouldEqual, "/var/log/tunasync/AOSP/latest.log")
 		_, ok := p.(*cmdProvider)
 		So(ok, ShouldBeTrue)
+		for _, hook := range p.Hooks() {
+			switch h := hook.(type) {
+			case *execPostHook:
+				So(h.command, ShouldResemble, []string{"bash", "-c", `echo ${TUNASYNC_JOB_EXIT_STATUS} > ${TUNASYNC_WORKING_DIR}/exit_status`})
+			}
+		}
 
 		p = w.providers["debian"]
 		So(p.Name(), ShouldEqual, "debian")

worker/exec_post_hook.go

@@ -0,0 +1,96 @@
+package worker
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/anmitsu/go-shlex"
+	"github.com/codeskyblue/go-sh"
+)
+
+// hook to execute command after syncing
+// typically setting timestamp, etc.
+
+const (
+	execOnSuccess uint8 = iota
+	execOnFailure
+)
+
+type execPostHook struct {
+	emptyHook
+	provider mirrorProvider
+
+	// exec on success or on failure
+	execOn uint8
+	// command
+	command []string
+}
+
+func newExecPostHook(provider mirrorProvider, execOn uint8, command string) (*execPostHook, error) {
+	cmd, err := shlex.Split(command, true)
+	if err != nil {
+		// logger.Errorf("Failed to create exec-post-hook for command: %s", command)
+		return nil, err
+	}
+	if execOn != execOnSuccess && execOn != execOnFailure {
+		return nil, fmt.Errorf("Invalid option for exec-on: %d", execOn)
+	}
+
+	return &execPostHook{
+		provider: provider,
+		execOn:   execOn,
+		command:  cmd,
+	}, nil
+}
+
+func (h *execPostHook) postSuccess() error {
+	if h.execOn == execOnSuccess {
+		return h.Do()
+	}
+	return nil
+}
+
+func (h *execPostHook) postFail() error {
+	if h.execOn == execOnFailure {
+		return h.Do()
+	}
+	return nil
+}
+
+func (h *execPostHook) Do() error {
+	p := h.provider
+
+	exitStatus := ""
+	if h.execOn == execOnSuccess {
+		exitStatus = "success"
+	} else {
+		exitStatus = "failure"
+	}
+
+	env := map[string]string{
+		"TUNASYNC_MIRROR_NAME":     p.Name(),
+		"TUNASYNC_WORKING_DIR":     p.WorkingDir(),
+		"TUNASYNC_UPSTREAM_URL":    p.Upstream(),
+		"TUNASYNC_LOG_FILE":        p.LogFile(),
+		"TUNASYNC_JOB_EXIT_STATUS": exitStatus,
+	}
+
+	session := sh.NewSession()
+	for k, v := range env {
+		session.SetEnv(k, v)
+	}
+
+	var cmd string
+	args := []interface{}{}
+	if len(h.command) == 1 {
+		cmd = h.command[0]
+	} else if len(h.command) > 1 {
+		cmd = h.command[0]
+		for _, arg := range h.command[1:] {
+			args = append(args, arg)
+		}
+	} else {
+		return errors.New("Invalid Command")
+	}
+	return session.Command(cmd, args...).Run()
+}

worker/exec_post_test.go

@@ -0,0 +1,112 @@
+package worker
+
+import (
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	. "github.com/smartystreets/goconvey/convey"
+	. "github.com/tuna/tunasync/internal"
+)
+
+func TestExecPost(t *testing.T) {
+	Convey("ExecPost should work", t, func(ctx C) {
+		tmpDir, err := ioutil.TempDir("", "tunasync")
+		defer os.RemoveAll(tmpDir)
+		So(err, ShouldBeNil)
+		scriptFile := filepath.Join(tmpDir, "cmd.sh")
+
+		c := cmdConfig{
+			name:        "tuna-exec-post",
+			upstreamURL: "http://mirrors.tuna.moe/",
+			command:     scriptFile,
+			workingDir:  tmpDir,
+			logDir:      tmpDir,
+			logFile:     filepath.Join(tmpDir, "latest.log"),
+			interval:    600 * time.Second,
+		}
+
+		provider, err := newCmdProvider(c)
+		So(err, ShouldBeNil)
+
+		Convey("On success", func() {
+			hook, err := newExecPostHook(provider, execOnSuccess, "bash -c 'echo ${TUNASYNC_JOB_EXIT_STATUS} > ${TUNASYNC_WORKING_DIR}/exit_status'")
+			So(err, ShouldBeNil)
+			provider.AddHook(hook)
+			managerChan := make(chan jobMessage)
+			semaphore := make(chan empty, 1)
+			job := newMirrorJob(provider)
+
+			scriptContent := `#!/bin/bash
+echo $TUNASYNC_WORKING_DIR
+echo $TUNASYNC_MIRROR_NAME
+echo $TUNASYNC_UPSTREAM_URL
+echo $TUNASYNC_LOG_FILE
+			`
+
+			err = ioutil.WriteFile(scriptFile, []byte(scriptContent), 0755)
+			So(err, ShouldBeNil)
+
+			go job.Run(managerChan, semaphore)
+			job.ctrlChan <- jobStart
+			msg := <-managerChan
+			So(msg.status, ShouldEqual, PreSyncing)
+			msg = <-managerChan
+			So(msg.status, ShouldEqual, Syncing)
+			msg = <-managerChan
+			So(msg.status, ShouldEqual, Success)
+
+			time.Sleep(200 * time.Millisecond)
+			job.ctrlChan <- jobDisable
+			<-job.disabled
+
+			expectedOutput := "success\n"
+
+			outputContent, err := ioutil.ReadFile(filepath.Join(provider.WorkingDir(), "exit_status"))
+			So(err, ShouldBeNil)
+			So(string(outputContent), ShouldEqual, expectedOutput)
+		})
+
+		Convey("On failure", func() {
+			hook, err := newExecPostHook(provider, execOnFailure, "bash -c 'echo ${TUNASYNC_JOB_EXIT_STATUS} > ${TUNASYNC_WORKING_DIR}/exit_status'")
+			So(err, ShouldBeNil)
+			provider.AddHook(hook)
+			managerChan := make(chan jobMessage)
+			semaphore := make(chan empty, 1)
+			job := newMirrorJob(provider)
+
+			scriptContent := `#!/bin/bash
+echo $TUNASYNC_WORKING_DIR
+echo $TUNASYNC_MIRROR_NAME
+echo $TUNASYNC_UPSTREAM_URL
+echo $TUNASYNC_LOG_FILE
+sleep 5
+exit 1
+			`
+
+			err = ioutil.WriteFile(scriptFile, []byte(scriptContent), 0755)
+			So(err, ShouldBeNil)
+
+			go job.Run(managerChan, semaphore)
+			job.ctrlChan <- jobStart
+			msg := <-managerChan
+			So(msg.status, ShouldEqual, PreSyncing)
+			msg = <-managerChan
+			So(msg.status, ShouldEqual, Syncing)
+			msg = <-managerChan
+			So(msg.status, ShouldEqual, Failed)
+
+			time.Sleep(200 * time.Millisecond)
+			job.ctrlChan <- jobDisable
+			<-job.disabled
+
+			expectedOutput := "failure\n"
+
+			outputContent, err := ioutil.ReadFile(filepath.Join(provider.WorkingDir(), "exit_status"))
+			So(err, ShouldBeNil)
+			So(string(outputContent), ShouldEqual, expectedOutput)
+		})
+	})
+}

worker/worker.go

@@ -160,6 +160,25 @@ func (w *Worker) initProviders() {
 			)
 		}
 
+		// ExecOnSuccess hook
+		if mirror.ExecOnSuccess != "" {
+			h, err := newExecPostHook(provider, execOnSuccess, mirror.ExecOnSuccess)
+			if err != nil {
+				logger.Errorf("Error initializing mirror %s: %s", mirror.Name, err.Error())
+				panic(err)
+			}
+			provider.AddHook(h)
+		}
+		// ExecOnFailure hook
+		if mirror.ExecOnFailure != "" {
+			h, err := newExecPostHook(provider, execOnFailure, mirror.ExecOnFailure)
+			if err != nil {
+				logger.Errorf("Error initializing mirror %s: %s", mirror.Name, err.Error())
+				panic(err)
+			}
+			provider.AddHook(h)
+		}
+
 		w.providers[provider.Name()] = provider
 
 	}