my-rook-config/pkg/util/proc/procmanager.go

package proc

import (
	"fmt"
	"log"
	"os"
	"sync"
	"syscall"

	"github.com/rook/rook/pkg/util/exec"
)

type ProcStartPolicy int

const (
	StartAction = "start"
	RunAction   = "run"
	StopAction  = "stop"

	RestartExisting ProcStartPolicy = iota
	ReuseExisting
)

type ProcManager struct {
	sync.RWMutex
	procs    []*MonitoredProc
	executor exec.Executor
}

// Create a new proc manager
func New(executor exec.Executor) *ProcManager {
	return &ProcManager{executor: executor}
}

// Start a child process and wait for its completion
func (p *ProcManager) Run(daemon string, args ...string) error {

	log.Printf("Running process %s with args: %v", daemon, args)
	err := p.executor.ExecuteCommand("daemon "+daemon, os.Args[0], createDaemonArgs(daemon, args...)...)
	if err != nil {
		return fmt.Errorf("failed to run %s: %+v", daemon, err)
	}

	return nil
}

// Start the given daemon and provided arguments.  Handling of any matching existing process will be in accordance
// with the given ProcStartPolicy.  The search pattern will be used to search through the cmdline args of existing
// processes to find any matching existing process.  Therefore, it should be a regex pattern that can uniquely
// identify the process (e.g., --id=1)
func (p *ProcManager) Start(daemon, procSearchPattern string, policy ProcStartPolicy, args ...string) (*MonitoredProc, error) {
	// look for an existing process first
	shouldStart, err := p.checkProcessExists(os.Args[0], procSearchPattern, policy)
	if err != nil {
		return nil, err
	}

	if !shouldStart {
		// based on the presence of an existing process and the given policy, we should not start a process, bail out.
		return nil, nil
	}

	log.Printf("Starting process %s with args: %v", daemon, args)
	cmd, err := p.executor.StartExecuteCommand("daemon "+daemon, os.Args[0], createDaemonArgs(daemon, args...)...)
	if err != nil {
		return nil, fmt.Errorf("failed to start daemon %s: %+v", daemon, err)
	}

	proc := newMonitoredProc(p, cmd)

	// monitor the process if it was not mocked
	if cmd != nil && cmd.Process != nil {
		go proc.Monitor()
	}

	p.Lock()
	p.procs = append(p.procs, proc)
	p.Unlock()
	return proc, nil
}

func (p *ProcManager) Shutdown() {
	p.RLock()
	for _, proc := range p.procs {
		proc.Stop()
	}
	p.procs = nil
	p.RUnlock()
}

func (p *ProcManager) checkProcessExists(binary, procSearchPattern string, policy ProcStartPolicy) (bool, error) {
	existingProc, err := findProcessSearch(binary, procSearchPattern)
	if err != nil {
		return false, fmt.Errorf("failed to search for process %s with pattern '%s': %+v", binary, procSearchPattern, err)
	}

	if existingProc == nil {
		log.Printf("no existing process found for binary %s and pattern '%s'", binary, procSearchPattern)
		return true, nil
	}

	log.Printf("existing process found for binary %s with pid %d, cmdline '%s'.", binary, existingProc.pid, existingProc.cmdline)
	if policy == ReuseExisting {
		log.Printf("Policy is 'reuse', reusing existing process.")
		return false, nil
	}

	log.Printf("Policy is 'restart', restarting existing process.")
	// find our managed instance and try to stop the process through that
	stopped := false
	p.RLock()
	for i := range p.procs {
		proc := p.procs[i].cmd
		if proc != nil && proc.Process != nil && proc.Process.Pid == existingProc.pid {
			if err := p.procs[i].Stop(); err != nil {
				log.Printf("failed to stop child process %d: %v", existingProc.pid, err)
				break
			}

			stopped = true
			break
		}
	}
	p.RUnlock()

	if !stopped {
		// we could't stop the existing process through our own managed proces set, try to stop the process
		// via a direct signal to its PID
		if err := syscall.Kill(existingProc.pid, syscall.SIGKILL); err != nil {
			return false, fmt.Errorf("failed to stop child process %d: %v", existingProc.pid, err)
		}
		stopped = true
	}

	if stopped {
		// the child process was stopped, remove it from our managed list (if it's in there)
		p.removeManagedProc(existingProc.pid)
	}

	return true, nil
}

func (p *ProcManager) removeManagedProc(pid int) {
	procNum := -1

	p.Lock()
	for i := range p.procs {
		proc := p.procs[i].cmd
		if proc != nil && proc.Process != nil && proc.Process.Pid == pid {
			procNum = i
			break
		}
	}

	if procNum >= 0 {
		p.procs[procNum] = p.procs[len(p.procs)-1]
		p.procs[len(p.procs)-1] = nil
		p.procs = p.procs[:len(p.procs)-1]
	}
	p.Unlock()
}

func createDaemonArgs(daemon string, args ...string) []string {
	return append(
		[]string{"daemon", fmt.Sprintf("--type=%s", daemon), "--"},
		args...)
}