webui4frpc/internal/process/process.go

// Package process supervises worker frpc processes, one per remote.
package process

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"sync"
	"syscall"
	"time"
)

// Status describes the current state of a worker process.
type Status struct {
	State        string `json:"state"` // stopped | starting | running | restarting | crashed
	Pid          int    `json:"pid,omitempty"`
	StartTime    int64  `json:"startTime,omitempty"`
	RestartCount int    `json:"restartCount"`
	ExitCode     int    `json:"exitCode,omitempty"`
	Err          string `json:"err,omitempty"`
}

const (
	stopGraceTimeout = 10 * time.Second
	maxRestartDelay  = 60 * time.Second
)

// Options configures a Manager.
type Options struct {
	// ConfigsDir is where rendered JSON configs are written.
	ConfigsDir string
	// LogsDir is where worker output is captured.
	LogsDir string
	// BinaryPath returns the frpc binary to spawn (resolved dynamically).
	BinaryPath func() string
	// Render returns the config bytes for a remote.
	Render func(remoteName string) ([]byte, error)
	// AutoRestart returns whether to auto-restart a remote's worker.
	AutoRestart func(remoteName string) bool
	// RestartInterval returns the base restart interval in seconds.
	RestartInterval func() int
}

type worker struct {
	name    string
	proc    *exec.Cmd
	logFile *rotatingFile

	stopOnce  sync.Once
	closeOnce sync.Once
	stopCh    chan struct{}
	doneCh    chan struct{}

	mu          sync.Mutex
	status      Status
	supervising bool // guarded by Manager.mu
}

// Manager supervises all workers.
type Manager struct {
	opts    Options
	mu      sync.Mutex
	workers map[string]*worker
}

// NewManager builds a worker supervisor.
func NewManager(opts Options) *Manager {
	return &Manager{opts: opts, workers: make(map[string]*worker)}
}

func (m *Manager) getStatus(w *worker) Status {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.status
}

func (m *Manager) setState(w *worker, s Status) {
	w.mu.Lock()
	w.status = s
	w.mu.Unlock()
}

// Start renders and spawns a remote's worker. Idempotent if already running.
func (m *Manager) Start(name string) error {
	m.mu.Lock()
	w := m.workers[name]
	m.mu.Unlock()
	if w != nil && statusRunning(w) {
		return nil
	}

	w = &worker{
		name:   name,
		stopCh: make(chan struct{}),
		doneCh: make(chan struct{}),
	}
	m.mu.Lock()
	m.workers[name] = w
	m.mu.Unlock()

	if err := m.spawn(w); err != nil {
		m.mu.Lock()
		delete(m.workers, name)
		m.mu.Unlock()
		return err
	}
	m.mu.Lock()
	if !w.supervising {
		w.supervising = true
		go m.supervise(w)
	}
	m.mu.Unlock()
	return nil
}

func statusRunning(w *worker) bool {
	select {
	case <-w.doneCh:
		return false
	default:
		return w.proc != nil && w.proc.Process != nil
	}
}

func (m *Manager) spawn(w *worker) error {
	data, err := m.opts.Render(w.name)
	if err != nil {
		return fmt.Errorf("render config for %q: %w", w.name, err)
	}
	cfgPath := filepath.Join(m.opts.ConfigsDir, w.name+".json")
	if err := os.MkdirAll(m.opts.ConfigsDir, 0o755); err != nil {
		return err
	}
	if err := os.WriteFile(cfgPath, data, 0o600); err != nil {
		return err
	}

	logWriter, err := newRotatingFile(filepath.Join(m.opts.LogsDir, w.name+".log"), 10*1024*1024)
	if err != nil {
		return err
	}
	cmd := exec.Command(m.opts.BinaryPath(), "-c", cfgPath)
	cmd.Stdout = logWriter
	cmd.Stderr = logWriter
	setSysProcAttr(cmd)
	if err := cmd.Start(); err != nil {
		_ = logWriter.Close()
		return fmt.Errorf("start worker %q: %w", w.name, err)
	}

	restartCount := m.getStatus(w).RestartCount
	w.proc = cmd
	w.logFile = logWriter
	m.setState(w, Status{State: "running", Pid: cmd.Process.Pid, StartTime: time.Now().Unix(), RestartCount: restartCount})
	// NOTE: no new supervise goroutine here. The first Start()'s supervise loop
	// re-spawns on restart; spawning here would race on doneCh close.
	return nil
}

func (m *Manager) supervise(w *worker) {
	defer w.closeOnce.Do(func() { close(w.doneCh) })
	for {
		err := w.proc.Wait()
		exitCode := 0
		if ee, ok := err.(*exec.ExitError); ok {
			exitCode = ee.ExitCode()
		}

		m.mu.Lock()
		if w.logFile != nil {
			_ = w.logFile.Close()
			w.logFile = nil
		}
		w.proc = nil
		stopRequested := false
		select {
		case <-w.stopCh:
			stopRequested = true
		default:
		}
		m.mu.Unlock()

		m.setState(w, Status{State: "stopped", ExitCode: exitCode, RestartCount: m.getStatus(w).RestartCount})
		if stopRequested {
			return
		}
		if !m.opts.AutoRestart(w.name) {
			return
		}

		rc := m.getStatus(w).RestartCount + 1
		m.setState(w, Status{State: "restarting", RestartCount: rc})
		select {
		case <-time.After(m.restartDelay(rc)):
		case <-w.stopCh:
			m.setState(w, Status{State: "stopped", RestartCount: rc})
			return
		}
		if err := m.spawn(w); err != nil {
			m.setState(w, Status{State: "crashed", Err: err.Error(), RestartCount: rc})
		}
	}
}

func (m *Manager) restartDelay(count int) time.Duration {
	base := 5 * time.Second
	if m.opts.RestartInterval != nil {
		if s := m.opts.RestartInterval(); s > 0 {
			base = time.Duration(s) * time.Second
		}
	}
	d := base
	for i := 1; i < count; i++ {
		d *= 2
		if d >= maxRestartDelay {
			return maxRestartDelay
		}
	}
	return d
}

// Stop gracefully stops a worker and waits for exit.
func (m *Manager) Stop(name string) error {
	m.mu.Lock()
	w := m.workers[name]
	m.mu.Unlock()
	if w == nil {
		return nil
	}
	w.stopOnce.Do(func() { close(w.stopCh) })

	m.mu.Lock()
	proc := w.proc
	m.mu.Unlock()
	if proc != nil {
		signalGroup(proc, syscall.SIGTERM)
	}
	select {
	case <-w.doneCh:
	case <-time.After(stopGraceTimeout):
		m.mu.Lock()
		proc := w.proc
		m.mu.Unlock()
		if proc != nil {
			signalGroup(proc, syscall.SIGKILL)
		}
		<-w.doneCh
	}
	m.mu.Lock()
	delete(m.workers, name)
	m.mu.Unlock()
	return nil
}

// Restart stops and starts a worker.
func (m *Manager) Restart(name string) error {
	if err := m.Stop(name); err != nil {
		return err
	}
	return m.Start(name)
}

// Status returns the current status of a remote's worker.
func (m *Manager) Status(name string) (Status, bool) {
	m.mu.Lock()
	w := m.workers[name]
	m.mu.Unlock()
	if w == nil {
		return Status{State: "stopped"}, false
	}
	return m.getStatus(w), true
}

// ListWorkers returns the names of all workers currently under management
// (running, starting, crashed-but-supervised, etc.), sorted for stable output.
// Stopped workers are removed from the map and not listed.
func (m *Manager) ListWorkers() []string {
	m.mu.Lock()
	out := make([]string, 0, len(m.workers))
	for name := range m.workers {
		out = append(out, name)
	}
	m.mu.Unlock()
	sort.Strings(out)
	return out
}

// RemoteWorkers returns the worker keys of all per-forward workers whose key
// parses to the given remote name. In the per-forward model every link has
// its own process keyed "local~remote~port"; this aggregates them so remote-
// centric callers (status page, profile endpoints) can report an overall
// view without knowing the key format.
func (m *Manager) RemoteWorkers(remote string) []string {
	var out []string
	for _, key := range m.ListWorkers() {
		if _, r, _, ok := ParseWorkerKey(key); ok && r == remote {
			out = append(out, key)
		}
	}
	return out
}

// RemoteStatus summarizes a remote's per-forward workers into one Status:
// running when ANY forward worker runs; crashed/starting/restarting take
// priority in that order; stopped (has=false) when no worker exists.
// This keeps the legacy per-remote status shape working on top of the
// per-forward worker model.
func (m *Manager) RemoteStatus(remote string) (Status, bool) {
	keys := m.RemoteWorkers(remote)
	if len(keys) == 0 {
		return Status{State: "stopped"}, false
	}
	summary := Status{State: "stopped"}
	worst := -1 // rank: running=0 < others=1 < crashed=2
	for _, k := range keys {
		st, _ := m.Status(k)
		rank := 1
		switch st.State {
		case "running":
			rank = 0
		case "crashed":
			rank = 2
		}
		if rank > worst || (rank == worst && summary.State == "stopped") {
			if rank != worst || st.State != "stopped" {
				summary = st
				worst = rank
			}
		}
	}
	return summary, true
}

// StartRemote starts every per-forward worker of one remote. Missing workers
// are skipped; the first error wins.
func (m *Manager) StartRemote(remote string) error {
	var firstErr error
	for _, key := range m.RemoteWorkers(remote) {
		if err := m.Start(key); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	return firstErr
}

// RestartRemote restarts every per-forward worker belonging to one remote
// (used where the old model had a single per-remote worker). Missing workers
// are skipped; the first error wins.
func (m *Manager) RestartRemote(remote string) error {
	var firstErr error
	for _, key := range m.RemoteWorkers(remote) {
		if err := m.Restart(key); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	return firstErr
}

// StopRemote stops every per-forward worker of one remote.
func (m *Manager) StopRemote(remote string) error {
	var firstErr error
	for _, key := range m.RemoteWorkers(remote) {
		if err := m.Stop(key); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	return firstErr
}

// StopAll stops all workers (used on manager shutdown).
func (m *Manager) StopAll() {
	m.mu.Lock()
	names := make([]string, 0, len(m.workers))
	for n := range m.workers {
		names = append(names, n)
	}
	m.mu.Unlock()
	var wg sync.WaitGroup
	for _, n := range names {
		wg.Add(1)
		go func(name string) {
			defer wg.Done()
			_ = m.Stop(name)
		}(n)
	}
	wg.Wait()
}

// LogPath returns the log file path for a remote.
func (m *Manager) LogPath(name string) string {
	return filepath.Join(m.opts.LogsDir, name+".log")
}

// ConfigPath returns the rendered config path for a remote.
func (m *Manager) ConfigPath(name string) string {
	return filepath.Join(m.opts.ConfigsDir, name+".json")
}