orca/internal/runtime/process.go

package runtime

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

// StopGrace is the default grace period between SIGTERM and SIGKILL
// for ProcessRuntime.Stop (10s, matching the systemd default TimeoutStopSec).
const StopGrace = 10 * time.Second

// ProcessRuntime implements Runtime for the "process" one_of using
// os/exec. It is the in-process equivalent of the systemd unit the CLI
// emits in production: the CLI emits a .service file and the peer's
// systemd runs the process; ProcessRuntime starts the process directly
// in the current Go process. It is intended for LOCAL testing and
// hermetic CI — NOT for production (production uses the systemd emitter
// + the peer's systemd, not this in-process path).
//
// It tracks started PIDs in an in-memory map; restarts of the CLI lose
// that state (acceptable for the local-test use case).
type ProcessRuntime struct {
	mu      sync.Mutex
	pids    map[string]int      // alloc.ID -> PID
	procs   map[int]*os.Process // PID -> process handle
	stopped map[string]bool     // alloc.ID -> reported stopped after Stop
}

// NewProcessRuntime returns a ProcessRuntime.
func NewProcessRuntime() *ProcessRuntime {
	return &ProcessRuntime{
		pids:    make(map[string]int),
		procs:   make(map[int]*os.Process),
		stopped: make(map[string]bool),
	}
}

// Prepare is a no-op for the process runtime: the systemd unit is
// emitted by the systemd emitter (internal/emitter), not by the runtime.
func (p *ProcessRuntime) Prepare(ctx context.Context, alloc *Alloc) error {
	_ = ctx
	_ = alloc
	return nil
}

// Start execs the alloc's command and returns the PID. The process is
// left running in the background; Stop terminates it.
func (p *ProcessRuntime) Start(ctx context.Context, alloc *Alloc) (int, error) {
	cmdStr, err := commandFor(alloc)
	if err != nil {
		return 0, err
	}
	// Parse the command string into argv. A leading "exec" form
	// (shell-style) is NOT supported — the command must be a direct
	// argv[0] + args. Split on whitespace (simple, matches the existing
	// executor.go behaviour which takes Command + Args separately).
	parts := strings.Fields(cmdStr)
	if len(parts) == 0 {
		return 0, fmt.Errorf("process: empty command for alloc %s", alloc.ID)
	}

	// Use a detached context so the process survives the request
	// context cancellation (the request ends; the workload keeps
	// running until Stop). We apply our own timeout for Start only.
	startCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	cmd := exec.CommandContext(startCtx, parts[0], parts[1:]...)
	// Detach the child from the parent's process group so it survives.
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	// Discard output for the runtime; the systemd unit captures logs
	// in production. For tests, callers that need output run their
	// own exec.Command.
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Start(); err != nil {
		return 0, fmt.Errorf("process: start: %w", err)
	}
	pid := cmd.Process.Pid

	// Background-reap the process so it doesn't become a zombie; we
	// only need the PID for Stop/Status. When the process exits
	// naturally, mark the alloc as stopped.
	go func() {
		_ = cmd.Wait()
		p.mu.Lock()
		delete(p.procs, pid)
		// Only mark stopped if the alloc is still associated with
		// this PID (Stop may have already removed the mapping).
		if cur, ok := p.pids[alloc.ID]; ok && cur == pid {
			delete(p.pids, alloc.ID)
			p.stopped[alloc.ID] = true
		}
		p.mu.Unlock()
	}()

	p.mu.Lock()
	p.pids[alloc.ID] = pid
	p.procs[pid] = cmd.Process
	delete(p.stopped, alloc.ID)
	p.mu.Unlock()

	return pid, nil
}

// Stop sends SIGTERM, waits the grace period, then SIGKILL.
func (p *ProcessRuntime) Stop(ctx context.Context, alloc *Alloc) error {
	p.mu.Lock()
	pid, ok := p.pids[alloc.ID]
	proc := p.procs[pid]
	p.mu.Unlock()
	if !ok || proc == nil {
		return nil // not running; idempotent
	}
	// SIGTERM the process group (negative PID).
	_ = syscall.Kill(-pid, syscall.SIGTERM)

	grace := StopGrace
	if dl, ok := ctx.Deadline(); ok {
		if remaining := time.Until(dl); remaining > 0 && remaining < grace {
			grace = remaining
		}
	}
	deadline := time.Now().Add(grace)
	for time.Now().Before(deadline) {
		if !p.alive(pid) {
			p.forget(alloc.ID, pid)
			return nil
		}
		select {
		case <-ctx.Done():
			p.forget(alloc.ID, pid)
			return ctx.Err()
		case <-time.After(100 * time.Millisecond):
		}
	}
	// SIGKILL the group.
	_ = syscall.Kill(-pid, syscall.SIGKILL)
	p.forget(alloc.ID, pid)
	return nil
}

// Status reports the alloc's state by checking if the process is alive.
func (p *ProcessRuntime) Status(ctx context.Context, alloc *Alloc) (State, error) {
	_ = ctx
	p.mu.Lock()
	pid, ok := p.pids[alloc.ID]
	wasStopped := p.stopped[alloc.ID]
	p.mu.Unlock()
	if !ok {
		if wasStopped {
			return StateStopped, nil
		}
		return StatePending, nil
	}
	if !p.alive(pid) {
		// Process exited but the reaper hasn't run yet; mark it
		// stopped and clean up.
		p.forget(alloc.ID, pid)
		return StateStopped, nil
	}
	return StateRunning, nil
}

// alive reports whether the process with the given PID is still running.
func (p *ProcessRuntime) alive(pid int) bool {
	proc, err := os.FindProcess(pid)
	if err != nil {
		return false
	}
	if err := proc.Signal(syscall.Signal(0)); err != nil {
		// ESRCH means the process is gone.
		return false
	}
	return true
}

// forget removes the alloc/PID mapping and marks the alloc stopped.
func (p *ProcessRuntime) forget(allocID string, pid int) {
	p.mu.Lock()
	delete(p.pids, allocID)
	delete(p.procs, pid)
	p.stopped[allocID] = true
	p.mu.Unlock()
}

// PID returns the recorded PID for alloc (for tests/inspection).
func (p *ProcessRuntime) PID(allocID string) (int, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	pid, ok := p.pids[allocID]
	if !ok {
		return 0, errors.New("process: no pid for alloc " + strconv.Quote(allocID))
	}
	return pid, nil
}