iop/apps/node/internal/node/run_handler.go

277 lines
8.6 KiB
Go

package node
import (
"context"
"errors"
"fmt"
"time"
"go.uber.org/zap"
"iop/apps/node/internal/runtime"
"iop/apps/node/internal/store"
"iop/apps/node/internal/transport"
iop "iop/proto/gen/iop"
)
// OnRunRequest handles an incoming RunRequest from a transport Session.
func (n *Node) OnRunRequest(ctx context.Context, sess *transport.Session, req *iop.RunRequest) error {
n.logger.Info("run request received",
zap.String("run_id", req.GetRunId()),
zap.String("adapter", req.GetAdapter()),
zap.String("target", req.GetTarget()),
)
rr := runtime.RunRequest{
RunID: req.GetRunId(),
Adapter: req.GetAdapter(),
Target: req.GetTarget(),
SessionID: req.GetSessionId(),
SessionMode: sessionModeFromProto(req.GetSessionMode()),
Background: req.GetBackground(),
Workspace: req.GetWorkspace(),
Policy: structAsMap(req.GetPolicy()),
Input: structAsMap(req.GetInput()),
TimeoutSec: int(req.GetTimeoutSec()),
Metadata: req.GetMetadata(),
}
printEdgeMessage(n.out, rr.Input)
n.configSetMu.RLock()
configLocked := true
defer func() {
if configLocked {
n.configSetMu.RUnlock()
}
}()
spec, adapter, err := n.router.ResolveAdapter(ctx, rr)
if err != nil {
n.sendPreExecuteError(sess, req.GetRunId(), req.GetSessionId(), req.GetBackground(), n.nodeID, err.Error())
return fmt.Errorf("node: resolve: %w", err)
}
var caps runtime.Capabilities
if c, capsErr := adapter.Capabilities(ctx); capsErr == nil {
caps = c
}
admission := n.admissionFor(spec.Adapter, caps)
// Acquire safety capacity ticket. Since we no longer maintain a Node-local FIFO queue,
// if concurrency is full, we reject immediately.
ticket, err := admission.acquire()
if err != nil {
n.logger.Warn("run admission rejected",
zap.String("run_id", spec.RunID),
zap.String("adapter", spec.Adapter),
zap.String("reason", string(admissionRejectReason(err))),
)
n.rejectRun(ctx, sess, spec, err)
return fmt.Errorf("node: run %s: %w", spec.RunID, err)
}
// Record the request as running since it is admitted immediately without queueing.
if err := n.store.InsertRun(ctx, store.RunRecord{
RunID: spec.RunID,
Adapter: spec.Adapter,
Target: spec.Target,
SessionID: normalizeSessionID(spec.SessionID),
Background: spec.Background,
Status: "running",
CreatedAt: time.Now(),
}); err != nil {
n.logger.Warn("store: insert run", zap.String("run_id", spec.RunID), zap.Error(err))
}
execCtx, cancel := context.WithCancel(ctx)
if spec.TimeoutSec > 0 {
execCtx, cancel = context.WithTimeout(ctx, time.Duration(spec.TimeoutSec)*time.Second)
}
h := &runHandle{
runID: spec.RunID,
adapter: spec.Adapter,
target: spec.Target,
sessionID: normalizeSessionID(spec.SessionID),
cancel: cancel,
done: make(chan struct{}),
}
n.runs.register(h)
configLocked = false
n.configSetMu.RUnlock()
var sender protoSender = noopSender{}
if sess != nil && sess.IsAlive() {
sender = sess
}
sink := &sessionSink{
sess: sender,
out: n.out,
nodeID: n.nodeID,
sessionID: normalizeSessionID(spec.SessionID),
background: spec.Background,
}
runSink := &terminalDeferringSink{inner: sink}
run := func() error {
released := false
releaseTicket := func() {
if !released {
ticket.release()
released = true
}
}
defer releaseTicket()
defer cancel()
defer n.runs.deregister(spec.RunID)
defer close(h.done)
execErr := adapter.Execute(execCtx, spec, runSink)
releaseTicket()
if !runSink.hasTerminalObserved() {
if synthErr := n.synthAndEmitTerminal(ctx, runSink, spec, execErr); synthErr != nil {
if execErr == nil {
execErr = synthErr
}
}
}
n.completeRun(spec, execErr)
if flushErr := runSink.Flush(context.Background()); flushErr != nil {
n.logger.Warn("session: flush terminal events", zap.String("run_id", spec.RunID), zap.Error(flushErr))
if execErr == nil {
return flushErr
}
}
return execErr
}
if spec.Background {
go func() { _ = run() }()
return nil
}
return run()
}
// admissionFor returns an admissionManager with only the per-adapter gate for
// the given adapter instance key. Node-wide global admission gating has been
// removed; each adapter controls its own concurrency via Capabilities().MaxConcurrency.
func (n *Node) admissionFor(adapterKey string, caps runtime.Capabilities) *admissionManager {
return &admissionManager{
adapter: n.adapterGateFor(adapterKey, caps),
}
}
// adapterGateFor returns the shared gate for the given adapter instance key,
// creating one from caps if not yet seen. MaxConcurrency <= 0 means the adapter gate is unlimited.
func (n *Node) adapterGateFor(adapterKey string, caps runtime.Capabilities) *fifoGate {
n.adapterGatesMu.Lock()
defer n.adapterGatesMu.Unlock()
g, ok := n.adapterGates[adapterKey]
if !ok {
g = newFifoGate(caps.MaxConcurrency)
n.adapterGates[adapterKey] = g
}
return g
}
// rejectRun records a rejected run in the store and sends an error RunEvent to
// the session so Edge can observe the rejection through the normal event stream.
// The rejection message carries the concrete reason (concurrency_unavailable) extracted from the admission error.
func (n *Node) rejectRun(ctx context.Context, sess *transport.Session, spec runtime.ExecutionSpec, admitErr error) {
errMsg := admissionMessage(admitErr)
// Since we reject immediately without queueing, insert as a terminal rejected record.
if err := n.store.InsertRun(ctx, store.RunRecord{
RunID: spec.RunID,
Adapter: spec.Adapter,
Target: spec.Target,
SessionID: normalizeSessionID(spec.SessionID),
Background: spec.Background,
Status: "rejected",
CreatedAt: time.Now(),
}); err != nil {
n.logger.Warn("store: insert rejected run", zap.String("run_id", spec.RunID), zap.Error(err))
}
if err := n.store.CompleteRun(ctx, spec.RunID, "rejected", errMsg); err != nil {
n.logger.Warn("store: complete rejected run", zap.String("run_id", spec.RunID), zap.Error(err))
}
if sess != nil && sess.IsAlive() {
re := &iop.RunEvent{
RunId: spec.RunID,
Type: string(runtime.EventTypeError),
Error: errMsg,
Timestamp: time.Now().UnixNano(),
SessionId: normalizeSessionID(spec.SessionID),
Background: spec.Background,
NodeId: n.nodeID,
}
if err := sess.Send(re); err != nil {
n.logger.Warn("session: send reject event", zap.String("run_id", spec.RunID), zap.Error(err))
}
}
}
// sendPreExecuteError sends an error RunEvent when an error occurs before
// adapter execution (e.g. ResolveAdapter failure). No store record is needed
// because the run never reached execution. This ensures Edge can observe the
// failure and avoid inflight slot leaks.
func (n *Node) sendPreExecuteError(sess *transport.Session, runID, sessionID string, background bool, nodeID, errMsg string) {
if sess != nil && sess.IsAlive() {
re := &iop.RunEvent{
RunId: runID,
Type: string(runtime.EventTypeError),
Error: errMsg,
Timestamp: time.Now().UnixNano(),
SessionId: normalizeSessionID(sessionID),
Background: background,
NodeId: nodeID,
}
if err := sess.Send(re); err != nil {
n.logger.Warn("session: send pre-execute error event", zap.String("run_id", runID), zap.Error(err))
}
}
}
func (n *Node) completeRun(spec runtime.ExecutionSpec, execErr error) {
status := "completed"
errMsg := ""
if execErr != nil {
if errors.Is(execErr, runtime.ErrRunCancelled) {
status = "cancelled"
} else {
status = "failed"
errMsg = execErr.Error()
}
n.logger.Warn("run ended", zap.String("run_id", spec.RunID), zap.String("status", status), zap.Error(execErr))
}
if err := n.store.CompleteRun(context.Background(), spec.RunID, status, errMsg); err != nil {
n.logger.Warn("store: complete run", zap.String("run_id", spec.RunID), zap.Error(err))
}
}
// synthAndEmitTerminal queues a terminal event when the adapter returned
// without emitting one. The caller flushes it after local admission release, so
// Edge can observe run completion without over-dispatching back into Node.
func (n *Node) synthAndEmitTerminal(ctx context.Context, sink *terminalDeferringSink, spec runtime.ExecutionSpec, execErr error) error {
event := runtime.RuntimeEvent{
RunID: spec.RunID,
Timestamp: time.Now(),
}
switch {
case errors.Is(execErr, runtime.ErrRunCancelled):
event.Type = runtime.EventTypeCancelled
case execErr != nil:
event.Type = runtime.EventTypeError
event.Error = execErr.Error()
default:
event.Type = runtime.EventTypeComplete
event.Message = "adapter completed without terminal event"
}
if err := sink.Emit(ctx, event); err != nil {
return fmt.Errorf("synthesize terminal event: %w", err)
}
return nil
}