cabbage/rvbox
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: wake reconciled sessions for queued dispatch

Browse Source
This commit is contained in:
cabbage
2026-09-06 10:03:11 +00:00
parent 3f6fb89a41
commit ef43e9592c
7 changed files with 210 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/server/session/agent_server.go
+86 -20
View File
@@ -139,10 +139,15 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
queue := NewWriterQueue(16, 64, 8)
defer queue.Close()
capacity := &CapacityShadow{}
var capacityMu sync.Mutex
reservations := make(map[string]DispatchLane)
reconciled := make(chan struct{})
var reconcileOnce sync.Once
if !capacity.UpdateAdvertised(0, 0, hello.GetMaxRunningCommands(), hello.GetMaxQueuedCommands()) {
server.close(connection, websocket.StatusPolicyViolation, "invalid initial client capacity")
return
}
go server.dispatchLoop(sessionContext, queue, handle.DispatchWake(), reconciled, &capacityMu, capacity, reservations, hello.GetClientId(), hello.GetPlatform(), encodeSessionID(sessionID), registration.Generation, cancel)
encodedSessionID := encodeSessionID(sessionID)
welcome, err := proto.Marshal(&rvboxv1.AgentEnvelope{
SessionId: encodedSessionID, SessionGeneration: registration.Generation,
@@ -210,25 +215,19 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
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
}
reconcileOnce.Do(func() { close(reconciled) })
handle.SignalDispatch()
continue
}
if advertised := envelope.GetClientCapacity(); advertised != nil {
if !capacity.UpdateAdvertised(advertised.GetRunningCommands(), advertised.GetQueuedCommands(), advertised.GetMaxRunningCommands(), advertised.GetMaxQueuedCommands()) {
capacityMu.Lock()
validCapacity := capacity.UpdateAdvertised(advertised.GetRunningCommands(), advertised.GetQueuedCommands(), advertised.GetMaxRunningCommands(), advertised.GetMaxQueuedCommands())
capacityMu.Unlock()
if !validCapacity {
server.close(connection, websocket.StatusPolicyViolation, "invalid client capacity")
return
}
handle.SignalDispatch()
continue
}
if acknowledgement := envelope.GetCommandAccepted(); acknowledgement != nil {
@@ -241,6 +240,13 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
server.close(connection, websocket.StatusPolicyViolation, "invalid command acceptance")
return
}
capacityMu.Lock()
if lane, reserved := reservations[issue.String()]; reserved {
capacity.Release(lane)
delete(reservations, issue.String())
}
capacityMu.Unlock()
handle.SignalDispatch()
continue
}
if event := envelope.GetCommandEvent(); event != nil {
@@ -315,10 +321,10 @@ func eventType(event *rvboxv1.CommandEvent) uint16 {
// 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) {
func (server *AgentServer) enqueueNextDispatch(ctx context.Context, queue *WriterQueue, clientID string, platform rvboxv1.Platform, sessionID string, generation uint64, beforeEnqueue func(domain.UUID)) (domain.UUID, bool, error) {
candidate, err := server.Store.ClaimNextDispatch(ctx, clientID, generation, server.now())
if err != nil || candidate == nil {
return false, err
return domain.UUID{}, false, err
}
requeue := func(cause error) (bool, error) {
_, rollbackErr := server.Store.RequeueDispatch(ctx, candidate.IssueUUID, clientID, generation)
@@ -329,10 +335,12 @@ func (server *AgentServer) enqueueNextDispatch(ctx context.Context, queue *Write
}
spec := &rvboxv1.ExecutionSpec{}
if err := proto.Unmarshal(candidate.ExecutionSpec, spec); err != nil {
return requeue(fmt.Errorf("decode persisted execution spec: %w", err))
sent, requeueErr := requeue(fmt.Errorf("decode persisted execution spec: %w", err))
return candidate.IssueUUID, sent, requeueErr
}
if err := agentproto.ValidateExecutionSpec(spec, server.limits(), platform); err != nil {
return requeue(fmt.Errorf("validate persisted execution spec: %w", err))
sent, requeueErr := requeue(fmt.Errorf("validate persisted execution spec: %w", err))
return candidate.IssueUUID, sent, requeueErr
}
dispatch := &rvboxv1.CommandDispatch{
IssueUuid: candidate.IssueUUID.String(), CommandRevision: candidate.Revision, TargetSessionGeneration: generation,
@@ -345,12 +353,70 @@ func (server *AgentServer) enqueueNextDispatch(ctx context.Context, queue *Write
SessionId: sessionID, SessionGeneration: generation, Payload: &rvboxv1.AgentEnvelope_CommandDispatch{CommandDispatch: dispatch},
})
if err != nil {
return requeue(err)
sent, requeueErr := requeue(err)
return candidate.IssueUUID, sent, requeueErr
}
if beforeEnqueue != nil {
beforeEnqueue(candidate.IssueUUID)
}
if !queue.EnqueueData(Frame{Kind: FrameData, Payload: encoded}) {
return requeue(ErrDispatchDataFull)
sent, requeueErr := requeue(ErrDispatchDataFull)
return candidate.IssueUUID, sent, requeueErr
}
return candidate.IssueUUID, true, nil
}
// dispatchLoop is the per-session serialized dispatcher. It waits for a
// complete reconciliation result before consuming queued work, then coalesces
// wakeups from local control RPCs, capacity advertisements, and acceptances.
func (server *AgentServer) dispatchLoop(ctx context.Context, queue *WriterQueue, wake <-chan struct{}, reconciled <-chan struct{}, capacityMu *sync.Mutex, capacity *CapacityShadow, reservations map[string]DispatchLane, clientID string, platform rvboxv1.Platform, sessionID string, generation uint64, cancel context.CancelFunc) {
select {
case <-reconciled:
case <-ctx.Done():
return
}
for {
for {
capacityMu.Lock()
lane := capacity.Reserve()
capacityMu.Unlock()
if lane == DispatchNone {
break
}
inserted := false
issue, sent, err := server.enqueueNextDispatch(ctx, queue, clientID, platform, sessionID, generation, func(issue domain.UUID) {
capacityMu.Lock()
reservations[issue.String()] = lane
inserted = true
capacityMu.Unlock()
})
if !sent {
capacityMu.Lock()
if inserted {
delete(reservations, issue.String())
}
capacity.Release(lane)
capacityMu.Unlock()
if err != nil && !errors.Is(err, ErrDispatchDataFull) {
server.closeForDispatchFailure(cancel)
return
}
break
}
}
select {
case <-wake:
continue
case <-ctx.Done():
return
}
}
}
func (server *AgentServer) closeForDispatchFailure(cancel context.CancelFunc) {
if cancel != nil {
cancel()
}
return true, nil
}
func (server *AgentServer) writeLoop(ctx context.Context, connection *websocket.Conn, queue *WriterQueue, heartbeat *synchronizedHeartbeat, started time.Time) {
internal/server/session/session.go
+36 -1
View File
@@ -231,6 +231,7 @@ type Handle struct {
SessionID [16]byte
Generation uint64
Context context.Context
wake chan struct{}
cancel context.CancelFunc
}
@@ -253,11 +254,45 @@ func (registry *Registry) Install(clientID string, sessionID [16]byte, generatio
old.cancel()
}
ctx, cancel := context.WithCancel(context.Background())
handle := &Handle{ClientID: clientID, SessionID: sessionID, Generation: generation, Context: ctx, cancel: cancel}
handle := &Handle{ClientID: clientID, SessionID: sessionID, Generation: generation, Context: ctx, wake: make(chan struct{}, 1), cancel: cancel}
registry.clients[clientID] = handle
return handle, nil
}
// Wake requests a best-effort dispatch pass for the current generation. The
// signal is edge-triggered and coalesced; durable queue state is the source of
// truth, so dropping a redundant wake cannot lose work.
func (registry *Registry) Wake(clientID string) bool {
if registry == nil || clientID == "" {
return false
}
registry.mu.Lock()
handle := registry.clients[clientID]
registry.mu.Unlock()
if handle == nil {
return false
}
handle.SignalDispatch()
return true
}
func (handle *Handle) DispatchWake() <-chan struct{} {
if handle == nil {
return nil
}
return handle.wake
}
func (handle *Handle) SignalDispatch() {
if handle == nil || handle.wake == nil {
return
}
select {
case handle.wake <- struct{}{}:
default:
}
}
func (registry *Registry) Remove(handle *Handle) bool {
if handle == nil {
return false
internal/server/session/session_test.go
+11
View File
@@ -100,6 +100,17 @@ func TestCapacityShadowAndRegistry_BH_SES_03(t *testing.T) {
if registry.Remove(first) || registry.Get("client-a") != second {
t.Fatal("stale handle removed current session")
}
if !registry.Wake("client-a") {
t.Fatal("wake did not find current session")
}
select {
case <-second.DispatchWake():
default:
t.Fatal("wake was not delivered to current session")
}
if registry.Wake("missing") {
t.Fatal("wake reported a missing session")
}
if !registry.Remove(second) {
t.Fatal("current handle removal failed")
}