sarulabs/hydre

package main

import (
    "fmt"
    "io/ioutil"
    "log"
    "os"
    "os/exec"
    "strconv"
    "strings"
    "sync"
    "syscall"
    "time"

    "github.com/hpcloud/tail"
)

// Service is the interface used to mock the Daemon structure.
type Service interface {
    Start(stopNotifier chan<- string)
    Stop()
    Kill()
}

// Daemon is a program executed in the background.
// Command is the UNIX command that starts the Daemon.
// The program will be stopped with a SIGTERM signal,
// but if StopCommand is defined, it will be used to stop the Daemon instead.
// The program running in the background is Command except if PidFile is defined.
// In this case the program is the one that has the pid specified in the PidFile.
// If LogFiles is not empty, the files will be used to stream
// the logs of the Daemon in addition to the Command output.
// You should probably use LogFiles if PidFile is defined.
type Daemon struct {
    Name        string
    Command     []string `yaml:"command"`
    StopCommand []string `yaml:"stopCommand"`
    PidFile     string   `yaml:"pidFile"`
    LogFiles    []string `yaml:"logFiles"`
    mu          sync.Mutex
    process     *os.Process
}

// Start executes the Daemon Command.
// It sets the Daemon process that is either the command process
// or the process defined by the PidFile.
// It also streams the Daemon logs on the standard output.
// When the daemon dies or when its logs are not  streamed anymore,
// it sends a message in the stopNotifier channel
// to report that the Daemon is not working properly anymore.
func (d *Daemon) Start(stopNotifier chan<- string) {
    log.Printf("starting `%s` daemon with command : %s\n", d.Name, d.Command)

    cmd := exec.Command(d.Command[0], d.Command[1:]...)
    cmd.Stdout = os.Stdout
    cmd.Stderr = os.Stderr

    for _, logFile := range d.LogFiles {
        go d.printLogs(logFile, stopNotifier)
    }

    if d.PidFile == "" {
        d.startForeground(cmd, stopNotifier)
        return
    }

    d.startBackground(cmd, stopNotifier)
}

func (d *Daemon) startForeground(cmd *exec.Cmd, stopNotifier chan<- string) {
    if err := cmd.Start(); err != nil {
        stopNotifier <- fmt.Sprintf("could not start daemon `%s` : %s", d.Name, err)
        return
    }

    d.setProcess(cmd.Process)

    // send notification when the process dies
    go d.waitChild(stopNotifier)
}

func (d *Daemon) startBackground(cmd *exec.Cmd, stopNotifier chan<- string) {
    // remove pidfile content to be sure to read the correct value
    ioutil.WriteFile(d.PidFile, []byte{}, 0777)

    cmd.Run()

    pid := d.readPidFileRetry(500, 10*time.Millisecond)
    process, _ := os.FindProcess(pid)
    d.setProcess(process)

    // send notification when the process dies
    go d.wait(stopNotifier)
}

// waitChild waits for the Daemon process to finish and sends a message
// in the stopNotifier channel when that happens.
// The Daemon process need to be a child of this process.
func (d *Daemon) waitChild(stopNotifier chan<- string) {
    if p := d.getProcess(); p != nil {
        p.Wait()
    }

    stopNotifier <- fmt.Sprintf("daemon `%s` has stopped", d.Name)
}

// wait waits for the Daemon process to finish and sends a message
// in the stopNotifier channel when that happens.
// It checks every second that the process is still running.
func (d *Daemon) wait(stopNotifier chan<- string) {
    if p := d.getProcess(); p != nil {
        for isAlive(p.Pid) {
            time.Sleep(time.Second)
        }
    }

    stopNotifier <- fmt.Sprintf("daemon `%s` has stopped", d.Name)
}

// readPidFile returns the pid from the Daemon PidFile.
func (d *Daemon) readPidFile() (int, error) {
    b, err := ioutil.ReadFile(d.PidFile)
    if err != nil {
        return 0, err
    }

    pid := strings.Trim(string(b), " \n")

    return strconv.Atoi(pid)
}

// readPidFileRetry tries to read the pid from the Daemon PidFile.
// If it could not read it, it will retry maxAttempts times
// and will wait before each new attempt.
// If the pid file could not be red, it returns 0.
func (d *Daemon) readPidFileRetry(maxAttempts int, wait time.Duration) int {
    for i := 0; i < maxAttempts; i++ {
        if pid, _ := d.readPidFile(); pid != 0 {
            return pid
        }
    }

    return 0
}

// printLogs streams the content of a file on the standard output.
// When it stops working, it sends a message in the stopNotifier channel.
func (d *Daemon) printLogs(filename string, stopNotifier chan<- string) {
    t, _ := tail.TailFile(filename, tail.Config{Follow: true})

    for line := range t.Lines {
        fmt.Println(line.Text)
    }

    stopNotifier <- fmt.Sprintf("could not stream logs from `%s`", filename)
}

// Stop tries to gracefully stop the Daemon.
// If StopCommand is defined, it executes the command.
// If not it sends a SIGTERM signal to the Daemon process.
func (d *Daemon) Stop() {
    // there is a stop command to stop the process, execute it.
    if len(d.StopCommand) > 0 {
        cmd := exec.Command(d.StopCommand[0], d.StopCommand[1:]...)
        cmd.Stdout = os.Stdout
        cmd.Stderr = os.Stderr
        cmd.Run()
        return
    }

    // there is no defined command to stop the process,
    // just send a SIGTERM signal.
    if p := d.getProcess(); len(d.StopCommand) == 0 && p != nil {
        p.Signal(syscall.SIGTERM)
        return
    }
}

// Kill sends a SIGKILL signal to the Daemon process.
func (d *Daemon) Kill() {
    if p := d.getProcess(); p != nil {
        p.Signal(syscall.SIGKILL)
    }
}

// getProcess is the thread safe process getter.
func (d *Daemon) getProcess() *os.Process {
    d.mu.Lock()
    defer d.mu.Unlock()
    return d.process
}

// getProcess is the thread safe process setter.
func (d *Daemon) setProcess(process *os.Process) {
    d.mu.Lock()
    defer d.mu.Unlock()
    d.process = process
}