godaemon/godaemon.go

package godaemon

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"

	"git.pyer.club/kingecg/gologger"
)

const (
	daemon_env_key  = "_go_daemon"
	daemon_process  = "g_daemon"
	daemon_task     = "g_dtask"
	daemon_taskargs = "g_args"
)

type GoDaemon struct {
	pidFile     string
	taskPidFile string
	flag        *string
	sigChan     chan os.Signal
	state       string
	*gologger.Logger
	Running *exec.Cmd

	StartFn func(*GoDaemon)
	StopFn  func(*GoDaemon)
}

func (g *GoDaemon) GetPid() int {
	pids, ferr := os.ReadFile(g.pidFile)
	pid, err := strconv.Atoi(string(pids))
	if err != nil || ferr != nil {
		return 0
	}
	return pid
}
func (g *GoDaemon) GetTaskPid() int {
	pids, ferr := os.ReadFile(g.taskPidFile)
	pid, err := strconv.Atoi(string(pids))
	if err != nil || ferr != nil {
		return 0
	}
	return pid
}

func (g *GoDaemon) Start() {
	if g.flag == nil {
		g.flag = flag.String("s", "", "send signal to daemon. support: reload and quit")
	}

	if IsMaster() {
		flag.Parse()
		if *g.flag == "" {
			g.startDaemon()
			return
		}

		var sig syscall.Signal
		if *g.flag == "reload" {
			sig = syscall.SIGHUP
		} else if *g.flag == "quit" {
			sig = syscall.SIGTERM
		} else {
			fmt.Println("Not supported signal")
			return
		}

		p := g.getDaemonProcess()
		if p == nil {
			fmt.Println("Daemon process not found")
			return
		}
		g.Debug("Send signal:", p.Pid, sig)
		p.Signal(sig)
	} else if IsDaemon() {
		pid := os.Getpid()

		os.WriteFile(g.pidFile, []byte(strconv.Itoa(pid)), 0644)
		g.sigChan = make(chan os.Signal)

		signal.Notify(g.sigChan, syscall.SIGTERM, syscall.SIGHUP)
		go g.serveSignal()

		for {

			g.Debug("Starting new task")
			g.Running = g.startTask()
			g.state = "running"
			g.Running.Process.Wait()
			if g.state == "stopped" {
				g.Debug("daemon is stopped, exit now")
				break
			}
		}

	} else {
		waiter := make(chan os.Signal, 1)
		g.StartFn(g)
		g.Info("daemon task is started")
		<-waiter
		g.Info("daemon task will be stopped")
		g.StopFn(g)
	}
}

func (g *GoDaemon) serveSignal() {
	sig := <-g.sigChan
	if sig == syscall.SIGTERM {
		g.state = "stopped"
	} else {
		g.state = "restart"
	}

	g.Running.Process.Signal(syscall.SIGTERM)
}

func (g *GoDaemon) getDaemonProcess() *os.Process {
	pid := g.GetPid()
	if pid == 0 {
		return nil
	}
	p, err := os.FindProcess(pid)
	if err != nil {
		g.Debug(err)
	}
	serr := p.Signal(syscall.Signal(0))
	if serr != nil {
		return nil
	}
	return p
}

func (g *GoDaemon) startDaemon() {
	dp := g.getDaemonProcess()

	if dp != nil {
		fmt.Println("daemon is running with pid:", dp.Pid)
		return
	}

	execName, _ := os.Executable()
	cmd := exec.Command(execName)
	cmd.Env = append(cmd.Env, daemon_env_key+"="+daemon_process)
	pargs := os.Args[1:]
	cmd.Env = append(cmd.Env, daemon_taskargs+"="+strings.Join(pargs, ";"))
	cmd.Start()
}

func (g *GoDaemon) startTask() *exec.Cmd {
	extraArgs, _ := os.LookupEnv(daemon_taskargs)
	var cmd *exec.Cmd
	execName, _ := os.Executable()
	if extraArgs != "" {
		eargs := strings.Split(extraArgs, ";")
		cmd = exec.Command(execName, eargs...)
	} else {
		cmd = exec.Command(execName)
	}
	cmd.Env = append(cmd.Env, daemon_env_key+"="+daemon_task)
	cmd.Start()
	return cmd
}
func IsMaster() bool {
	goDaemonEnv, _ := os.LookupEnv(daemon_env_key)
	return goDaemonEnv == ""
}

func IsDaemon() bool {
	goDaemonEnv, _ := os.LookupEnv(daemon_env_key)
	return goDaemonEnv == daemon_process
}

func IsDaemonTask() bool {
	goDaemonEnv, _ := os.LookupEnv(daemon_env_key)
	return goDaemonEnv == daemon_task
}

func NewGoDaemon(start, stop func(*GoDaemon)) *GoDaemon {
	godaemon := &GoDaemon{
		Logger: gologger.GetLogger("daemon"),
	}

	execName, _ := os.Executable()
	if filepath.Ext(execName) != "" {
		execName = strings.TrimSuffix(execName, filepath.Ext(execName))
	}
	godaemon.pidFile = execName + ".pid"
	godaemon.taskPidFile = execName + ".task.pid"
	godaemon.StartFn = start
	godaemon.StopFn = stop
	return godaemon
}