support daemon

pkg/util/util.go

@@ -9,6 +9,7 @@ import (
 	"path"
 	"path/filepath"
 	"strings"
+	"syscall"
 )
 
 // IsFileExist returns true if the file exists.
@@ -17,6 +18,15 @@ func IsFileExist(name string) bool {
 	return err == nil
 }
 
+// IsProcessExist returns true if the porcess still exists.
+func IsProcessExist(pid int) bool {
+	p, err := os.FindProcess(pid)
+	if err != nil {
+		return false
+	}
+	return p.Signal(syscall.Signal(0)) == nil
+}
+
 // Wget downloads a string URL to the dest directory and returns the file path.
 // SKips if the file already exists.
 func Wget(ctx context.Context, rawURL string, dest string) (string, error) {
@@ -92,3 +102,102 @@ func InstallArchive(ctx context.Context, rawURL string, dest string) error {
 
 	return os.Rename(tmpDir, dest)
 }
+
+// DaemonOptions is the options to start a command in daemon mode.
+type DaemonOptions struct {
+	Background       bool
+	ChDir            string
+	LogFile          string
+	MakePidFile      bool
+	MatchExecutable  bool
+	MatchProcessName bool
+	PidFile          string
+	ProcessName      string
+}
+
+// NewDaemonOptions returns a default daemon options.
+func NewDaemonOptions(chDir string, pidFile string, logFile string) DaemonOptions {
+	return DaemonOptions{
+		Background:       true,
+		MakePidFile:      true,
+		MatchExecutable:  true,
+		MatchProcessName: false,
+		ChDir:            chDir,
+		PidFile:          pidFile,
+		LogFile:          logFile,
+	}
+}
+
+// StartDaemon starts a daemon process with options
+func StartDaemon(ctx context.Context, opts DaemonOptions, cmd string, cmdArgs ...string) error {
+	var args []string
+	args = append(args, "--start")
+
+	if opts.Background {
+		args = append(args, "--background", "--no-close")
+	}
+
+	if opts.MakePidFile {
+		args = append(args, "--make-pidfile")
+	}
+
+	if opts.MatchExecutable {
+		args = append(args, "--exec", cmd)
+	}
+
+	if opts.MatchProcessName {
+		processName := opts.ProcessName
+		if len(processName) == 0 {
+			processName = path.Base(cmd)
+		}
+		args = append(args, "--name", processName)
+	}
+
+	args = append(args, "--pidfile", opts.PidFile)
+	args = append(args, "--chdir", opts.ChDir)
+	args = append(args, "--oknodo", "--startas", cmd)
+	args = append(args, "--")
+	args = append(args, cmdArgs...)
+
+	c := exec.CommandContext(ctx, "start-stop-daemon", args...)
+	logFile, err := os.Create(opts.LogFile)
+	if err != nil {
+		return err
+	}
+	defer logFile.Close()
+	c.Stdout = logFile
+	c.Stderr = logFile
+
+	return c.Run()
+}
+
+func parsePID(pidFile string) string {
+	data, err := ioutil.ReadFile(pidFile)
+	if err != nil {
+		return ""
+	}
+
+	return strings.TrimSpace(string(data))
+}
+
+// StopDaemon kills the daemon process by pid file or by the command name.
+func StopDaemon(ctx context.Context, cmd string, pidFile string) error {
+	var err error
+	if len(cmd) > 0 {
+		err = exec.CommandContext(ctx, "killall", "-9", "-w", cmd).Run()
+	} else {
+		if pid := parsePID(pidFile); len(pid) > 0 {
+			err = exec.CommandContext(ctx, "kill", "-9", pid).Run()
+		}
+	}
+
+	if err != nil {
+		return err
+	}
+
+	if err = os.Remove(pidFile); os.IsNotExist(err) {
+		err = nil
+	}
+
+	return nil
+}

pkg/util/util_test.go

@@ -2,7 +2,10 @@ package util
 
 import (
 	"context"
+	"io/ioutil"
 	"os"
+	"path"
+	"strconv"
 	"testing"
 )
 
@@ -26,17 +29,50 @@ func TestWget(t *testing.T) {
 }
 
 func TestInstallArchive(t *testing.T) {
-	os.RemoveAll("./var")
-	defer os.RemoveAll("./var")
+	tmpDir, _ := ioutil.TempDir(".", "var")
+	defer os.RemoveAll(tmpDir)
 
-	err := InstallArchive(context.Background(), "https://github.com/siddontang/chaos/archive/master.zip", "./var/1")
+	err := InstallArchive(context.Background(), "https://github.com/siddontang/chaos/archive/master.zip", path.Join(tmpDir, "1"))
 	if err != nil {
 		t.Fatalf("install archive failed %v", err)
 	}
 
-	err = InstallArchive(context.Background(), "https://github.com/siddontang/chaos/archive/master.tar.gz", "./var/2")
+	err = InstallArchive(context.Background(), "https://github.com/siddontang/chaos/archive/master.tar.gz", path.Join(tmpDir, "2"))
 	if err != nil {
 		t.Fatalf("install archive failed %v", err)
 	}
+}
+
+func TestDaemon(t *testing.T) {
+	t.Log("test may only be run in the chaos docker")
+
+	tmpDir, _ := ioutil.TempDir(".", "var")
+	defer os.RemoveAll(tmpDir)
+
+	pidFile := path.Join(tmpDir, "sleep.pid")
+	opts := NewDaemonOptions(tmpDir, pidFile, path.Join(tmpDir, "sleep.log"))
+	err := StartDaemon(context.Background(), opts, "/bin/sleep", "100")
+	if err != nil {
+		t.Fatalf("start daemon failed %v", err)
+	}
+
+	pidStr := parsePID(pidFile)
+	if pidStr == "" {
+		t.Fatal("must have a pid file")
+	}
+
+	pid, _ := strconv.Atoi(pidStr)
+	if !IsProcessExist(pid) {
+		t.Fatalf("pid %d must exist", pid)
+	}
+
+	err = StopDaemon(context.Background(), "", pidFile)
+	if err != nil {
+		t.Fatalf("stop daemon failed %v", err)
+	}
+
+	if IsProcessExist(pid) {
+		t.Fatalf("pid %d must not exist", pid)
+	}
 
 }