391 lines
14 KiB
Go
391 lines
14 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"sync"
|
|
"time"
|
|
|
|
rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
|
|
"github.com/rvbox/rvbox/internal/client/spool"
|
|
"github.com/rvbox/rvbox/internal/client/supervisor"
|
|
"github.com/rvbox/rvbox/internal/domain"
|
|
"google.golang.org/protobuf/proto"
|
|
"google.golang.org/protobuf/types/known/timestamppb"
|
|
)
|
|
|
|
// Executor bridges durable dispatch records to the platform supervisor. It
|
|
// owns no network state: a command's process, pipes, and spool writes continue
|
|
// after the WebSocket session is replaced.
|
|
type Executor struct {
|
|
Store *spool.Store
|
|
Supervisor supervisor.Supervisor
|
|
WorkDir string
|
|
Now func() time.Time
|
|
Notify func(domain.UUID)
|
|
|
|
mu sync.Mutex
|
|
active map[domain.UUID]supervisor.Process
|
|
cancel map[domain.UUID]context.CancelFunc
|
|
forced map[domain.UUID]bool
|
|
}
|
|
|
|
type ExecutorOptions struct {
|
|
Store *spool.Store
|
|
Supervisor supervisor.Supervisor
|
|
WorkDir string
|
|
Now func() time.Time
|
|
Notify func(domain.UUID)
|
|
}
|
|
|
|
func NewExecutor(options ExecutorOptions) (*Executor, error) {
|
|
if options.Store == nil || options.Supervisor == nil || options.WorkDir == "" {
|
|
return nil, errors.New("executor requires durable store, supervisor, and work directory")
|
|
}
|
|
if options.Now == nil {
|
|
options.Now = func() time.Time { return time.Now().UTC() }
|
|
}
|
|
return &Executor{Store: options.Store, Supervisor: options.Supervisor, WorkDir: options.WorkDir, Now: options.Now, Notify: options.Notify, active: make(map[domain.UUID]supervisor.Process), cancel: make(map[domain.UUID]context.CancelFunc), forced: make(map[domain.UUID]bool)}, nil
|
|
}
|
|
|
|
// Dispatch is safe to invoke after CommandAccepted has been sent. Script
|
|
// commands intentionally wait for ScriptCommit; the server may deliver the
|
|
// descriptor and body in separate frames.
|
|
func (executor *Executor) Dispatch(ctx context.Context, _ Session, dispatch *rvboxv1.CommandDispatch) error {
|
|
if executor == nil || dispatch == nil {
|
|
return errors.New("invalid command dispatch")
|
|
}
|
|
issue, err := domain.ParseUUIDv7(dispatch.GetIssueUuid())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if dispatch.GetSpec().GetScript() != nil {
|
|
if _, err := executor.Store.ScriptBody(ctx, issue); errors.Is(err, spool.ErrScriptNotReady) {
|
|
return nil
|
|
} else if err != nil {
|
|
return executor.reject(ctx, issue, dispatch.GetCommandRevision(), err)
|
|
}
|
|
}
|
|
return executor.launch(ctx, issue, dispatch.GetCommandRevision(), dispatch.GetSpec())
|
|
}
|
|
|
|
// ScriptReady launches a script after the durable commit barrier. Repeated
|
|
// progress/commit frames are harmless because the active map and spool phase
|
|
// make launch at-most-once.
|
|
func (executor *Executor) ScriptReady(ctx context.Context, _ Session, issue domain.UUID) error {
|
|
command, err := executor.Store.GetCommand(ctx, issue)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var spec rvboxv1.ExecutionSpec
|
|
if err := proto.Unmarshal(command.ExecutionSpec, &spec); err != nil {
|
|
return executor.reject(ctx, issue, command.Revision, err)
|
|
}
|
|
if spec.GetScript() == nil {
|
|
return nil
|
|
}
|
|
return executor.launch(ctx, issue, command.Revision, &spec)
|
|
}
|
|
|
|
func (executor *Executor) launch(ctx context.Context, issue domain.UUID, revision uint64, spec *rvboxv1.ExecutionSpec) error {
|
|
if spec == nil || revision == 0 {
|
|
return executor.reject(ctx, issue, revision, errors.New("missing execution specification"))
|
|
}
|
|
executor.mu.Lock()
|
|
if _, exists := executor.active[issue]; exists {
|
|
executor.mu.Unlock()
|
|
return nil
|
|
}
|
|
executor.mu.Unlock()
|
|
var scriptBody []byte
|
|
if spec.GetScript() != nil {
|
|
body, err := executor.Store.ScriptBody(ctx, issue)
|
|
if err != nil {
|
|
if errors.Is(err, spool.ErrScriptNotReady) {
|
|
return nil
|
|
}
|
|
return executor.reject(ctx, issue, revision, err)
|
|
}
|
|
scriptBody = body
|
|
}
|
|
workingDirectory := spec.GetCwd()
|
|
if workingDirectory == "" {
|
|
workingDirectory = executor.WorkDir
|
|
}
|
|
if err := executor.Store.SetLaunchPhase(ctx, issue, domain.LaunchPhasePrepared, "", 0); err != nil {
|
|
return err
|
|
}
|
|
// The native Windows supervisor creates/assigns the Job while the child is
|
|
// suspended. It does not cross the durable authorization barrier until the
|
|
// process object has been fully prepared and its immutable identity is
|
|
// available. A daemon crash before this write leaves only launch_prepared;
|
|
// the kill-on-close Job prevents a suspended child from surviving restart.
|
|
runContext, cancel := context.WithCancel(context.Background())
|
|
process, err := executor.Supervisor.Start(runContext, supervisor.StartSpec{IssueUUID: issue, CommandRevision: revision, Execution: proto.Clone(spec).(*rvboxv1.ExecutionSpec), ScriptBody: scriptBody, WorkingDirectory: workingDirectory, Environment: cloneEnvironment(spec.GetEnvOverrides()), ExecutionProfiles: executionProfileNames(spec.GetExecutionProfiles())})
|
|
if err != nil {
|
|
cancel()
|
|
_ = executor.Store.SetLaunchPhase(context.Background(), issue, domain.LaunchPhaseNone, "", 0)
|
|
var startErr *supervisor.StartError
|
|
if errors.As(err, &startErr) {
|
|
return executor.rejectWithIdentity(ctx, issue, revision, err, startErr.WindowsIdentity)
|
|
}
|
|
return executor.reject(ctx, issue, revision, err)
|
|
}
|
|
identity := process.Identity()
|
|
if err := executor.Store.SetLaunchPhase(ctx, issue, domain.LaunchPhaseAuthorized, identity.Context, 0); err != nil {
|
|
_, _ = executor.Supervisor.Signal(context.Background(), process, supervisor.SignalKill)
|
|
cancel()
|
|
return err
|
|
}
|
|
if err := process.Release(runContext); err != nil {
|
|
_, _ = executor.Supervisor.Signal(context.Background(), process, supervisor.SignalKill)
|
|
cancel()
|
|
return err
|
|
}
|
|
executor.mu.Lock()
|
|
executor.active[issue] = process
|
|
executor.cancel[issue] = cancel
|
|
executor.mu.Unlock()
|
|
if _, err := executor.Store.AppendLifecycleWithIdentity(ctx, issue, uint32(rvboxv1.CommandLifecycle_COMMAND_RUNNING), revision, "process started", executor.Now(), identity.WindowsIdentity()); err != nil {
|
|
_, _ = executor.Supervisor.Signal(context.Background(), process, supervisor.SignalKill)
|
|
cancel()
|
|
executor.remove(issue)
|
|
return err
|
|
}
|
|
executor.notify(issue)
|
|
go executor.watch(runContext, issue, revision, process, cancel)
|
|
return nil
|
|
}
|
|
|
|
func (executor *Executor) watch(ctx context.Context, issue domain.UUID, revision uint64, process supervisor.Process, cancel context.CancelFunc) {
|
|
defer cancel()
|
|
for {
|
|
chunk, err := process.ReadOutput(ctx)
|
|
if errors.Is(err, io.EOF) {
|
|
break
|
|
}
|
|
if err != nil {
|
|
_ = executor.appendIncomplete(context.Background(), issue, revision, err.Error())
|
|
break
|
|
}
|
|
if len(chunk.Data) == 0 {
|
|
continue
|
|
}
|
|
if _, err := executor.Store.AppendOutput(context.Background(), issue, spool.OutputInput{Stream: chunk.Stream, Raw: chunk.Data, ObservedAt: executor.Now()}); err != nil {
|
|
_ = executor.appendIncomplete(context.Background(), issue, revision, err.Error())
|
|
_, _ = executor.Supervisor.Signal(context.Background(), process, supervisor.SignalKill)
|
|
break
|
|
}
|
|
executor.notify(issue)
|
|
}
|
|
status, waitErr := process.Wait(context.Background())
|
|
phase := rvboxv1.CommandLifecycle_COMMAND_FAILED
|
|
if waitErr == nil && status.Code == 0 {
|
|
phase = rvboxv1.CommandLifecycle_COMMAND_SUCCEEDED
|
|
}
|
|
executor.mu.Lock()
|
|
if executor.forced[issue] {
|
|
phase = rvboxv1.CommandLifecycle_COMMAND_TERMINATED
|
|
}
|
|
executor.mu.Unlock()
|
|
detail := fmt.Sprintf("exit code %d", status.Code)
|
|
if waitErr != nil {
|
|
detail = boundedError(waitErr)
|
|
}
|
|
if err := executor.appendLifecycleWithIdentity(context.Background(), issue, revision, phase, detail, process.Identity().WindowsIdentity()); err == nil {
|
|
executor.notify(issue)
|
|
}
|
|
executor.remove(issue)
|
|
}
|
|
|
|
func (executor *Executor) appendLifecycle(ctx context.Context, issue domain.UUID, revision uint64, phase rvboxv1.CommandLifecycle, detail string) error {
|
|
return executor.appendLifecycleWithIdentity(ctx, issue, revision, phase, detail, nil)
|
|
}
|
|
|
|
func (executor *Executor) appendLifecycleWithIdentity(ctx context.Context, issue domain.UUID, revision uint64, phase rvboxv1.CommandLifecycle, detail string, identity *rvboxv1.WindowsExecutionIdentity) error {
|
|
_, err := executor.Store.AppendLifecycleWithIdentity(ctx, issue, uint32(phase), revision, detail, executor.Now(), identity)
|
|
return err
|
|
}
|
|
|
|
func (executor *Executor) reject(ctx context.Context, issue domain.UUID, revision uint64, cause error) error {
|
|
return executor.rejectWithIdentity(ctx, issue, revision, cause, nil)
|
|
}
|
|
|
|
func (executor *Executor) rejectWithIdentity(ctx context.Context, issue domain.UUID, revision uint64, cause error, identity *rvboxv1.WindowsExecutionIdentity) error {
|
|
if revision == 0 {
|
|
return cause
|
|
}
|
|
if err := executor.appendLifecycleWithIdentity(ctx, issue, revision, rvboxv1.CommandLifecycle_COMMAND_REJECTED, boundedError(cause), identity); err != nil {
|
|
return err
|
|
}
|
|
executor.notify(issue)
|
|
return nil
|
|
}
|
|
|
|
func (executor *Executor) appendIncomplete(ctx context.Context, issue domain.UUID, _ uint64, detail string) error {
|
|
payload, err := proto.MarshalOptions{Deterministic: true}.Marshal(&rvboxv1.CommandEvent{IssueUuid: issue.String(), ObservedAt: timestamppb.New(executor.Now()), Payload: &rvboxv1.CommandEvent_OutputIncomplete{OutputIncomplete: &rvboxv1.OutputIncomplete{Reason: boundedError(errors.New(detail))}}})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = executor.Store.AppendEvent(ctx, issue, spool.EventInput{Kind: 11, Compression: 1, RawBytes: uint64(len(payload)), Payload: payload, CreatedAt: executor.Now()})
|
|
if err == nil {
|
|
executor.notify(issue)
|
|
}
|
|
return err
|
|
}
|
|
|
|
func (executor *Executor) Stdin(ctx context.Context, _ Session, input *rvboxv1.StdinWrite) error {
|
|
if input == nil {
|
|
return errors.New("missing stdin request")
|
|
}
|
|
issue, err := domain.ParseUUIDv7(input.GetIssueUuid())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
executor.mu.Lock()
|
|
process := executor.active[issue]
|
|
executor.mu.Unlock()
|
|
if process == nil {
|
|
return executor.appendStdinAck(ctx, issue, input.GetWriteSeq(), false, "command is not running")
|
|
}
|
|
err = process.WriteStdin(ctx, input.GetData(), input.GetAppendNewline())
|
|
detail := ""
|
|
if err != nil {
|
|
detail = boundedError(err)
|
|
}
|
|
return executor.appendStdinAck(ctx, issue, input.GetWriteSeq(), err == nil, detail)
|
|
}
|
|
|
|
func (executor *Executor) CloseStdin(ctx context.Context, _ Session, input *rvboxv1.CloseStdin) error {
|
|
if input == nil {
|
|
return errors.New("missing close-stdin request")
|
|
}
|
|
issue, err := domain.ParseUUIDv7(input.GetIssueUuid())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
executor.mu.Lock()
|
|
process := executor.active[issue]
|
|
executor.mu.Unlock()
|
|
if process == nil {
|
|
return executor.appendStdinAck(ctx, issue, input.GetWriteSeq(), false, "command is not running")
|
|
}
|
|
err = process.CloseStdin(ctx)
|
|
detail := ""
|
|
if err != nil {
|
|
detail = boundedError(err)
|
|
}
|
|
return executor.appendStdinAck(ctx, issue, input.GetWriteSeq(), err == nil, detail)
|
|
}
|
|
|
|
func (executor *Executor) appendStdinAck(ctx context.Context, issue domain.UUID, writeSeq uint64, accepted bool, detail string) error {
|
|
payload, err := proto.MarshalOptions{Deterministic: true}.Marshal(&rvboxv1.CommandEvent{IssueUuid: issue.String(), ObservedAt: timestamppb.New(executor.Now()), Payload: &rvboxv1.CommandEvent_StdinAck{StdinAck: &rvboxv1.StdinAcknowledgement{WriteSeq: writeSeq, Detail: detail, StdinClosed: !accepted}}})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = executor.Store.AppendEvent(ctx, issue, spool.EventInput{Kind: 7, Compression: 1, RawBytes: uint64(len(payload)), Payload: payload, CreatedAt: executor.Now(), UseCloseout: false})
|
|
if err == nil {
|
|
executor.notify(issue)
|
|
}
|
|
return err
|
|
}
|
|
|
|
func (executor *Executor) Signal(ctx context.Context, _ Session, input *rvboxv1.SignalCommand) error {
|
|
if input == nil {
|
|
return errors.New("missing signal request")
|
|
}
|
|
issue, err := domain.ParseUUIDv7(input.GetIssueUuid())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
executor.mu.Lock()
|
|
process := executor.active[issue]
|
|
executor.mu.Unlock()
|
|
if process == nil {
|
|
return errors.New("command is not running")
|
|
}
|
|
kind := supervisor.SignalTerm
|
|
if input.GetSignal() == rvboxv1.SignalKind_SIGNAL_KILL {
|
|
kind = supervisor.SignalKill
|
|
} else if input.GetSignal() != rvboxv1.SignalKind_SIGNAL_TERM {
|
|
return errors.New("unsupported Windows signal")
|
|
}
|
|
outcome, err := executor.Supervisor.Signal(ctx, process, kind)
|
|
if kind == supervisor.SignalKill || outcome.Escalated {
|
|
executor.mu.Lock()
|
|
executor.forced[issue] = true
|
|
executor.mu.Unlock()
|
|
}
|
|
return executor.appendSignalResult(ctx, issue, input.GetCommandRevision(), input.GetSignal(), err == nil && outcome.Delivered, outcome, err)
|
|
}
|
|
|
|
func (executor *Executor) appendSignalResult(ctx context.Context, issue domain.UUID, revision uint64, signal rvboxv1.SignalKind, accepted bool, outcome supervisor.SignalOutcome, cause error) error {
|
|
detail := outcome.Detail
|
|
if cause != nil {
|
|
detail = boundedError(cause)
|
|
}
|
|
payload, err := proto.MarshalOptions{Deterministic: true}.Marshal(&rvboxv1.CommandEvent{IssueUuid: issue.String(), ObservedAt: timestamppb.New(executor.Now()), Payload: &rvboxv1.CommandEvent_SignalResult{SignalResult: &rvboxv1.SignalResult{Signal: signal, Accepted: accepted, GracefulDeliveryAttempted: signal == rvboxv1.SignalKind_SIGNAL_TERM, ForcedTerminationUsed: outcome.Escalated || signal == rvboxv1.SignalKind_SIGNAL_KILL, Detail: detail, CommandRevision: revision}}})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = executor.Store.AppendEvent(ctx, issue, spool.EventInput{Kind: 8, Compression: 1, RawBytes: uint64(len(payload)), Payload: payload, CreatedAt: executor.Now()})
|
|
if err == nil {
|
|
executor.notify(issue)
|
|
}
|
|
return err
|
|
}
|
|
|
|
func (executor *Executor) Terminate(ctx context.Context, issue domain.UUID) error {
|
|
executor.mu.Lock()
|
|
process := executor.active[issue]
|
|
executor.forced[issue] = true
|
|
executor.mu.Unlock()
|
|
if process == nil {
|
|
return nil
|
|
}
|
|
returnError := error(nil)
|
|
if _, err := executor.Supervisor.Signal(ctx, process, supervisor.SignalKill); err != nil {
|
|
returnError = err
|
|
}
|
|
return returnError
|
|
}
|
|
|
|
func (executor *Executor) StopAll(ctx context.Context) error {
|
|
return executor.Supervisor.StopAll(ctx)
|
|
}
|
|
|
|
func (executor *Executor) remove(issue domain.UUID) {
|
|
executor.mu.Lock()
|
|
delete(executor.active, issue)
|
|
delete(executor.cancel, issue)
|
|
delete(executor.forced, issue)
|
|
executor.mu.Unlock()
|
|
}
|
|
|
|
func (executor *Executor) notify(issue domain.UUID) {
|
|
if executor.Notify != nil {
|
|
executor.Notify(issue)
|
|
}
|
|
}
|
|
|
|
func cloneEnvironment(input map[string]string) map[string]string {
|
|
if input == nil {
|
|
return nil
|
|
}
|
|
result := make(map[string]string, len(input))
|
|
for key, value := range input {
|
|
result[key] = value
|
|
}
|
|
return result
|
|
}
|
|
|
|
func executionProfileNames(input []rvboxv1.ExecutionProfile) []string {
|
|
result := make([]string, 0, len(input))
|
|
for _, profile := range input {
|
|
result = append(result, profile.String())
|
|
}
|
|
return result
|
|
}
|