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
	flag     *string
	Running  *exec.Cmd
	sigChan  chan os.Signal
	doneChan chan int
	state    string
	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) 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
		}
		fmt.Println("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)
		g.doneChan = make(chan int)
		signal.Notify(g.sigChan, syscall.SIGTERM, syscall.SIGHUP)
		go g.serveSignal()
		l := gologger.GetLogger("daemon")

		for {
			if g.state == "running" {
				l.Debug("Due to for loop ")
				break
			}
			l.Debug("Starting task new")
			g.Running = g.startTask()
			g.state = "running"
			g.Running.Process.Wait()
			if g.state == "stopped" {
				break
			}
		}

	} else {

		g.StartFn(g)
	}
}

func (g *GoDaemon) serveSignal() {
	sig := <-g.sigChan
	if sig == syscall.SIGTERM {
		g.state = "stopped"
	} else {
		g.state = "restart"
	}
	l := gologger.GetLogger("daemon")
	l.Debug("Stop it")
	g.StopFn(g)
}

func (g *GoDaemon) getDaemonProcess() *os.Process {
	pid := g.GetPid()
	if pid == 0 {
		return nil
	}
	p, err := os.FindProcess(pid)
	if err != nil {
		l := gologger.GetLogger("default")
		l.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{}

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