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
cmd/rvbox-server/main.go
+3 -1
View File
@@ -69,8 +69,10 @@ func run(configPath string) error {
return err
}
defer func() { _ = cleanupControl() }()
registry := session.NewRegistry()
controlService, err := control.NewService(control.Options{
Store: persistence, DefaultQueueTTL: configured.Queue.DefaultTTL, DefaultQueueTTLSet: true,
WakeClient: registry.Wake,
Limits: agentproto.Limits{MaxEnvelopeBytes: configured.Protocol.MaxAgentEnvelopeBytes, MaxExecutionSpecBytes: configured.Protocol.MaxExecutionSpecBytes, MaxRawChunkBytes: configured.Protocol.MaxRawChunkBytes, MaxScriptBytes: configured.Protocol.MaxScriptBytes, MaxDetailBytes: configured.Storage.ProtocolDetailMaxBytes},
})
if err != nil {
@@ -79,7 +81,7 @@ func run(configPath string) error {
grpcServer := grpc.NewServer(grpc.MaxRecvMsgSize(int(configured.Protocol.MaxControlRequestBytes)), grpc.MaxSendMsgSize(int(configured.Protocol.MaxControlRequestBytes)))
rvboxv1.RegisterControlServer(grpcServer, controlService)
agent := &session.AgentServer{
Store: persistence, Registry: session.NewRegistry(), Path: configured.Server.AgentPath,
Store: persistence, Registry: registry, Path: configured.Server.AgentPath,
Limits: agentproto.Limits{
MaxEnvelopeBytes: configured.Protocol.MaxAgentEnvelopeBytes,
MaxExecutionSpecBytes: configured.Protocol.MaxExecutionSpecBytes,
internal/server/control/service.go
+6 -1
View File
@@ -44,6 +44,7 @@ type Options struct {
Limits agentproto.Limits
Now func() time.Time
CursorKey []byte
WakeClient func(string) bool
}
type Service struct {
@@ -53,6 +54,7 @@ type Service struct {
limits agentproto.Limits
now func() time.Time
cursors *domain.CursorCodec
wakeClient func(string) bool
}
func NewService(options Options) (*Service, error) {
@@ -82,7 +84,7 @@ func NewService(options Options) (*Service, error) {
if err != nil {
return nil, err
}
return &Service{store: options.Store, defaultQueueTTL: options.DefaultQueueTTL, limits: options.Limits, now: options.Now, cursors: codec}, nil
return &Service{store: options.Store, defaultQueueTTL: options.DefaultQueueTTL, limits: options.Limits, now: options.Now, cursors: codec, wakeClient: options.WakeClient}, nil
}
func (service *Service) ListClients(ctx context.Context, request *rvboxv1.ListClientsRequest) (*rvboxv1.ListClientsResponse, error) {
@@ -264,6 +266,9 @@ func (service *Service) RunCommand(ctx context.Context, request *rvboxv1.RunComm
if err != nil {
return nil, mapStoreError(err)
}
if service.wakeClient != nil {
_ = service.wakeClient(request.GetTargetClientId())
}
return &rvboxv1.RunCommandResponse{IssueUuid: issue.String(), Lifecycle: rvboxv1.CommandLifecycle(queued.Lifecycle)}, nil
}
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")
}
test/coverage.toml
+6
View File
@@ -254,6 +254,12 @@ layer = "unit"
status = "implemented"
tests = ["internal/server/store/command_test.go:TestRecordCommandAcceptanceFencesGeneration_HP_DISPATCH_05"]
[[requirements]]
id = "HP-DISPATCH-06"
layer = "integration"
status = "implemented"
tests = ["test/integration/clientagent/clientagent_integration_test.go:TestControlQueueWakesReconciledSession_HP_DISPATCH_06"]
[[requirements]]
id = "HP-EVENT-01"
layer = "unit"
test/integration/clientagent/clientagent_integration_test.go
+61
View File
@@ -13,8 +13,11 @@ import (
"github.com/rvbox/rvbox/internal/agentproto"
"github.com/rvbox/rvbox/internal/client/agent"
"github.com/rvbox/rvbox/internal/domain"
"github.com/rvbox/rvbox/internal/server/control"
"github.com/rvbox/rvbox/internal/server/session"
"github.com/rvbox/rvbox/internal/server/store"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/known/timestamppb"
)
@@ -137,3 +140,61 @@ func TestWebSocketDispatchAfterReconciliation_HP_DISPATCH_03(t *testing.T) {
t.Fatalf("command lifecycle = %d, %v", lifecycle, err)
}
}
func TestControlQueueWakesReconciledSession_HP_DISPATCH_06(t *testing.T) {
ctx := context.Background()
persistence, err := store.Open(ctx, store.Options{DataDir: filepath.Join(t.TempDir(), "state"), BusyTimeout: time.Second})
if err != nil {
t.Fatal(err)
}
defer persistence.Close()
registry := session.NewRegistry()
server := httptest.NewServer(&session.AgentServer{Store: persistence, Registry: registry, Path: "/v1/agent", Limits: agentproto.DefaultLimits()})
defer server.Close()
controlService, err := control.NewService(control.Options{Store: persistence, WakeClient: registry.Wake, CursorKey: []byte("0123456789abcdef0123456789abcdef")})
if err != nil {
t.Fatal(err)
}
address := "ws" + strings.TrimPrefix(server.URL, "http") + "/v1/agent"
transport, err := agent.DialWebSocket(ctx, address, nil)
if err != nil {
t.Fatal(err)
}
defer transport.Close()
hello := &rvboxv1.ClientHello{ClientId: "win-wake-client", SupportedProtocol: &rvboxv1.ProtocolRange{Major: 1, MinMinor: 0, MaxMinor: 0}, DaemonVersion: "test", Platform: rvboxv1.Platform_PLATFORM_WINDOWS, Architecture: "amd64", DaemonCwd: `C:\`, SupportedShells: []rvboxv1.ShellType{rvboxv1.ShellType_SHELL_POWERSHELL}, ClientInstanceId: "019c46f1-1d02-7000-8000-000000000076", MaxRunningCommands: 1, MaxQueuedCommands: 1, SentAt: timestamppb.Now()}
accepted, err := agent.Handshake(ctx, transport, hello, agentproto.DefaultLimits())
if err != nil {
t.Fatal(err)
}
if _, err := agent.Reconcile(ctx, transport, accepted, &rvboxv1.ReconcileSnapshot{}, agentproto.DefaultLimits()); err != nil {
t.Fatal(err)
}
issue, err := domain.ParseUUIDv7("019c46f1-1d02-7000-8000-000000000077")
if err != nil {
t.Fatal(err)
}
response, err := controlService.RunCommand(ctx, &rvboxv1.RunCommandRequest{TargetClientId: hello.GetClientId(), RequestId: issue.String(), Spec: &rvboxv1.ExecutionSpec{ShellType: rvboxv1.ShellType_SHELL_POWERSHELL, Source: &rvboxv1.ExecutionSpec_CommandText{CommandText: "Write-Output wake"}}})
if err != nil || response.GetIssueUuid() != issue.String() {
t.Fatalf("control queue response = %#v, %v", response, err)
}
readContext, cancel := context.WithTimeout(ctx, time.Second)
defer cancel()
encoded, err := transport.Read(readContext)
if err != nil {
t.Fatal(err)
}
envelope, err := agentproto.DecodeEnvelope(encoded, agentproto.DefaultLimits(), rvboxv1.Platform_PLATFORM_WINDOWS)
if err != nil || envelope.GetCommandDispatch() == nil || envelope.GetCommandDispatch().GetIssueUuid() != issue.String() {
t.Fatalf("woken dispatch = %#v, %v", envelope, err)
}
_, missingErr := controlService.RunCommand(ctx, &rvboxv1.RunCommandRequest{TargetClientId: "missing", RequestId: fixedIssueID(0x78), Spec: &rvboxv1.ExecutionSpec{ShellType: rvboxv1.ShellType_SHELL_POWERSHELL, Source: &rvboxv1.ExecutionSpec_CommandText{CommandText: "x"}}})
if status.Code(missingErr) != codes.NotFound {
t.Fatal("missing target did not return not found")
}
}
func fixedIssueID(last byte) string {
issue, _ := domain.ParseUUIDv7("019c46f1-1d02-7000-8000-000000000001")
issue[15] = last
return issue.String()
}