godaemon/godaemon.go

// Package godaemon provides a simple daemon library for Go applications.
// It allows applications to run as daemon processes with start/stop/restart capabilities.
package godaemon

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

	"git.kingecg.top/kingecg/gologger"
)

// Constants defining environment variable keys and values used for process identification
const (
	daemon_env_key  = "_go_daemon" // Environment variable key for process role identification
	daemon_process  = "g_daemon"   // Value indicating daemon process role
	daemon_task     = "g_dtask"    // Value indicating task process role
	daemon_taskargs = "g_args"     // Environment variable key for storing task arguments
)

// GoDaemon represents a daemon process manager
type GoDaemon struct {
	pidFile          string         // Path to file storing daemon process PID
	taskPidFile      string         // Path to file storing task process PID
	flag             *string        // Command line flag for signal control
	sigChan          chan os.Signal // Channel for receiving OS signals
	state            string         // Current state of the daemon ("running", "stopped", etc.)
	*gologger.Logger                // Embedded logger for logging
	Running          *exec.Cmd      // Currently running task process

	StartFn func(*GoDaemon) // Function called when task starts
	StopFn  func(*GoDaemon) // Function called when task stops
}

// GetPid retrieves the daemon process ID from the PID file
//
// Returns:
//   - int: process ID if found and valid, 0 otherwise
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
}

// GetTaskPid retrieves the task process ID from the task PID file
//
// Returns:
//   - int: process ID if found and valid, 0 otherwise
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
}

// Start begins the daemon process management based on the current process role
//
// Behavior depends on process role:
//   - Master: starts the daemon process or sends signals to running daemon
//   - Daemon: manages task process lifecycle
//   - Task: executes the user-provided StartFn
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, 1)

		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)
	}
}

// serveSignal handles incoming OS signals for the daemon process
//
// Signals handled:
//   - SIGTERM: stops the daemon
//   - SIGHUP: restarts the task process
func (g *GoDaemon) serveSignal() {
	sig := <-g.sigChan
	if sig == syscall.SIGTERM {
		g.state = "stopped"
	} else {
		g.state = "restart"
	}

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

// getDaemonProcess retrieves the daemon process instance if it's running
//
// Returns:
//   - *os.Process: running daemon process if found, nil otherwise
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
}

// startDaemon starts a new daemon process in the background
//
// Checks if daemon is already running before starting a new instance
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()
}

// startTask starts a new task process with the configured arguments
//
// Returns:
//   - *exec.Cmd: the started command representing the task process
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.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Env = append(cmd.Env, daemon_env_key+"="+daemon_task)
	cmd.Start()
	return cmd
}

// IsMaster checks if the current process is the master process
//
// Returns:
//   - bool: true if this is the master process, false otherwise
func IsMaster() bool {
	goDaemonEnv, _ := os.LookupEnv(daemon_env_key)
	return goDaemonEnv == ""
}

// IsDaemon checks if the current process is the daemon process
//
// Returns:
//   - bool: true if this is the daemon process, false otherwise
func IsDaemon() bool {
	goDaemonEnv, _ := os.LookupEnv(daemon_env_key)
	return goDaemonEnv == daemon_process
}

// IsDaemonTask checks if the current process is the task process
//
// Returns:
//   - bool: true if this is the task process, false otherwise
func IsDaemonTask() bool {
	goDaemonEnv, _ := os.LookupEnv(daemon_env_key)
	return goDaemonEnv == daemon_task
}

// NewGoDaemon creates a new GoDaemon instance
//
// Parameters:
//   - start: function to be called when the task process starts
//   - stop: function to be called when the task process stops
//
// Returns:
//   - *GoDaemon: initialized GoDaemon instance
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
}