feat: dispatch queued commands after reconciliation
This commit is contained in:
@@ -28,6 +28,7 @@ var (
|
||||
ErrUnexpectedOrigin = errors.New("agent connections must not send an Origin header")
|
||||
ErrUnexpectedMessage = errors.New("agent message is not a binary protobuf envelope")
|
||||
ErrStaleSession = errors.New("agent message is for an unknown or fenced session")
|
||||
ErrDispatchDataFull = errors.New("session dispatch data lane is full")
|
||||
)
|
||||
|
||||
// AgentServer is the transport edge for agent WebSocket sessions. Command
|
||||
@@ -137,6 +138,11 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
|
||||
|
||||
queue := NewWriterQueue(16, 64, 8)
|
||||
defer queue.Close()
|
||||
capacity := &CapacityShadow{}
|
||||
if !capacity.UpdateAdvertised(0, 0, hello.GetMaxRunningCommands(), hello.GetMaxQueuedCommands()) {
|
||||
server.close(connection, websocket.StatusPolicyViolation, "invalid initial client capacity")
|
||||
return
|
||||
}
|
||||
encodedSessionID := encodeSessionID(sessionID)
|
||||
welcome, err := proto.Marshal(&rvboxv1.AgentEnvelope{
|
||||
SessionId: encodedSessionID, SessionGeneration: registration.Generation,
|
||||
@@ -194,14 +200,80 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
|
||||
SessionId: encodedSessionID, SessionGeneration: registration.Generation,
|
||||
Payload: &rvboxv1.AgentEnvelope_ReconcileResult{ReconcileResult: result},
|
||||
})
|
||||
if err != nil || queue.EnqueueControl(Frame{Kind: FrameControl, Payload: encoded}) != nil {
|
||||
resultWritten := make(chan struct{})
|
||||
if err != nil || queue.EnqueueControl(Frame{Kind: FrameControl, Payload: encoded, Written: resultWritten}) != nil {
|
||||
server.close(connection, websocket.StatusInternalError, "could not queue reconciliation result")
|
||||
return
|
||||
}
|
||||
select {
|
||||
case <-resultWritten:
|
||||
case <-sessionContext.Done():
|
||||
return
|
||||
}
|
||||
lane := capacity.Reserve()
|
||||
dispatched, dispatchErr := false, error(nil)
|
||||
if lane != DispatchNone {
|
||||
dispatched, dispatchErr = server.enqueueNextDispatch(sessionContext, queue, hello.GetClientId(), hello.GetPlatform(), encodedSessionID, registration.Generation)
|
||||
if !dispatched {
|
||||
capacity.Release(lane)
|
||||
}
|
||||
}
|
||||
if dispatchErr != nil && !errors.Is(dispatchErr, ErrDispatchDataFull) {
|
||||
server.close(connection, websocket.StatusInternalError, "could not queue command dispatch")
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
if advertised := envelope.GetClientCapacity(); advertised != nil {
|
||||
if !capacity.UpdateAdvertised(advertised.GetRunningCommands(), advertised.GetQueuedCommands(), advertised.GetMaxRunningCommands(), advertised.GetMaxQueuedCommands()) {
|
||||
server.close(connection, websocket.StatusPolicyViolation, "invalid client capacity")
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// enqueueNextDispatch records the queued-to-dispatched transition before
|
||||
// exposing work to the network. A full data lane is a pre-write failure, so
|
||||
// only the owning generation can put the command back into the queue.
|
||||
func (server *AgentServer) enqueueNextDispatch(ctx context.Context, queue *WriterQueue, clientID string, platform rvboxv1.Platform, sessionID string, generation uint64) (bool, error) {
|
||||
candidate, err := server.Store.ClaimNextDispatch(ctx, clientID, generation, server.now())
|
||||
if err != nil || candidate == nil {
|
||||
return false, err
|
||||
}
|
||||
requeue := func(cause error) (bool, error) {
|
||||
_, rollbackErr := server.Store.RequeueDispatch(ctx, candidate.IssueUUID, clientID, generation)
|
||||
if rollbackErr != nil {
|
||||
return false, rollbackErr
|
||||
}
|
||||
return false, cause
|
||||
}
|
||||
spec := &rvboxv1.ExecutionSpec{}
|
||||
if err := proto.Unmarshal(candidate.ExecutionSpec, spec); err != nil {
|
||||
return requeue(fmt.Errorf("decode persisted execution spec: %w", err))
|
||||
}
|
||||
if err := agentproto.ValidateExecutionSpec(spec, server.limits(), platform); err != nil {
|
||||
return requeue(fmt.Errorf("validate persisted execution spec: %w", err))
|
||||
}
|
||||
dispatch := &rvboxv1.CommandDispatch{
|
||||
IssueUuid: candidate.IssueUUID.String(), CommandRevision: candidate.Revision, TargetSessionGeneration: generation,
|
||||
IssueTime: timestamppb.New(candidate.IssueTime), Spec: spec,
|
||||
}
|
||||
if candidate.QueueExpiryTime != nil {
|
||||
dispatch.QueueExpiryTime = timestamppb.New(*candidate.QueueExpiryTime)
|
||||
}
|
||||
encoded, err := proto.Marshal(&rvboxv1.AgentEnvelope{
|
||||
SessionId: sessionID, SessionGeneration: generation, Payload: &rvboxv1.AgentEnvelope_CommandDispatch{CommandDispatch: dispatch},
|
||||
})
|
||||
if err != nil {
|
||||
return requeue(err)
|
||||
}
|
||||
if !queue.EnqueueData(Frame{Kind: FrameData, Payload: encoded}) {
|
||||
return requeue(ErrDispatchDataFull)
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (server *AgentServer) writeLoop(ctx context.Context, connection *websocket.Conn, queue *WriterQueue, heartbeat *synchronizedHeartbeat, started time.Time) {
|
||||
for {
|
||||
frameContext, cancel := context.WithTimeout(ctx, server.heartbeatPollInterval())
|
||||
@@ -214,6 +286,9 @@ func (server *AgentServer) writeLoop(ctx context.Context, connection *websocket.
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if frame.Written != nil {
|
||||
close(frame.Written)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if !errors.Is(err, context.DeadlineExceeded) {
|
||||
|
||||
@@ -23,6 +23,9 @@ const (
|
||||
type Frame struct {
|
||||
Kind FrameKind
|
||||
Payload []byte
|
||||
// Written is closed by the sole socket writer after a successful write.
|
||||
// It is used only for protocol barriers such as reconciliation-before-work.
|
||||
Written chan<- struct{}
|
||||
}
|
||||
|
||||
// WriterQueue is owned by one socket writer. Data saturation leaves the work
|
||||
|
||||
Reference in New Issue
Block a user