Add persistent Run control stream

This commit is contained in:
npc0-hue
2026-08-26 23:05:40 +08:00
parent 3b4e857ef7
commit ac14f80306
8 changed files with 245 additions and 9 deletions
+87 -6
View File
@@ -24,6 +24,7 @@ import (
type WorkerClient interface {
Hello(context.Context, protocol.RunHelloRequest) (protocol.RunHelloResponse, error)
Heartbeat(context.Context, protocol.RunHeartbeatRequest) (protocol.RunHeartbeatResponse, error)
StreamControlEvents(context.Context, protocol.RunControlStreamRequest, func(protocol.RunControlEvent) error) error
ReportLifecycle(context.Context, protocol.RunLifecycleReportRequest) (protocol.RunLifecycleReportResponse, error)
GetRunLogStreamProgress(context.Context, protocol.RunLogStreamProgressRequest) (protocol.RunLogStreamProgressResponse, error)
ClaimJob(context.Context, protocol.RunJobClaimRequest) (protocol.RunJobClaimResponse, error)
@@ -430,6 +431,10 @@ func (worker *Worker) reregisterAndReconcile(ctx context.Context, reason string,
}
func (worker *Worker) ClaimAndRunOnce(ctx context.Context) (bool, error) {
return worker.claimAndRunOnce(ctx, jobClaimWaitSeconds)
}
func (worker *Worker) claimAndRunOnce(ctx context.Context, waitSeconds int) (bool, error) {
state, err := worker.registeredState()
if err != nil {
return false, err
@@ -440,7 +445,7 @@ func (worker *Worker) ClaimAndRunOnce(ctx context.Context) (bool, error) {
SessionToken: state.SessionToken,
Capabilities: state.Capabilities,
Capacity: worker.capacityReportFor(state),
WaitSeconds: jobClaimWaitSeconds,
WaitSeconds: waitSeconds,
})
if err != nil {
if sessionInvalidError(err) {
@@ -973,7 +978,10 @@ func (worker *Worker) Run(ctx context.Context) error {
uploaderDone := make(chan struct{})
go worker.runDurableUploaders(workerCtx, uploaderDone)
jobDone := make(chan error, 1)
go func() { jobDone <- worker.runJobLoop(workerCtx, jobInterval) }()
controlWake := make(chan struct{}, 1)
controlDone := make(chan error, 1)
go func() { controlDone <- worker.runControlStreamLoop(workerCtx, controlWake) }()
go func() { jobDone <- worker.runJobLoop(workerCtx, jobInterval, controlWake) }()
defer func() {
cancelWorker()
<-uploaderDone
@@ -986,6 +994,12 @@ func (worker *Worker) Run(ctx context.Context) error {
case err := <-jobDone:
log.Printf("RUN phase=run status=job_loop_done error=%s", errorSummary(err))
return err
case err := <-controlDone:
if err != nil && err != context.Canceled {
log.Printf("RUN phase=run status=control_stream_done error=%s", errorSummary(err))
return err
}
return err
case <-heartbeatTicker.C:
if err := worker.HeartbeatOnce(ctx); err != nil {
log.Printf("RUN phase=heartbeat status=retry_scheduled backoffMs=%d", boundedRetryBackoff(worker.cfg.RetryBackoff).Milliseconds())
@@ -1036,8 +1050,59 @@ func (worker *Worker) reportRunUpdateHealth(ctx context.Context) error {
return nil
}
func (worker *Worker) runJobLoop(ctx context.Context, interval time.Duration) error {
log.Printf("RUN phase=job_loop status=starting fallbackPollMs=%d claimWaitSeconds=%d", interval.Milliseconds(), jobClaimWaitSeconds)
func (worker *Worker) runControlStreamLoop(ctx context.Context, wake chan<- struct{}) error {
log.Printf("RUN phase=control_stream status=starting endpoint=%s", worker.cfg.RunEndpointID)
var lastSeq uint64
for {
state, err := worker.registeredState()
if err != nil {
if waitErr := waitWorkerLoop(ctx, boundedRetryBackoff(worker.cfg.RetryBackoff)); waitErr != nil {
return waitErr
}
continue
}
err = worker.client.StreamControlEvents(ctx, protocol.RunControlStreamRequest{RunEndpointID: state.RunEndpointID, SessionToken: state.SessionToken, LastEventSeq: lastSeq}, func(event protocol.RunControlEvent) error {
if event.RunEndpointID != "" && event.RunEndpointID != state.RunEndpointID {
return fmt.Errorf("control event endpoint mismatch")
}
if event.Sequence > lastSeq {
lastSeq = event.Sequence
}
log.Printf("RUN phase=control_stream status=event endpoint=%s type=%s seq=%d", state.RunEndpointID, safeOptional(event.Type), event.Sequence)
signalControlWake(wake)
return nil
})
if ctx.Err() != nil {
log.Printf("RUN phase=control_stream status=context_done error=%s", RedactText(ctx.Err().Error()))
return ctx.Err()
}
if err != nil {
if sessionInvalidError(err) {
log.Printf("RUN phase=control_stream status=session_invalid endpoint=%s error=%s", state.RunEndpointID, RedactText(err.Error()))
if refreshErr := worker.reregisterAndReconcile(ctx, "control_stream_session_invalid", state.SessionToken); refreshErr != nil {
log.Printf("RUN phase=control_stream status=reregister_failed endpoint=%s error=%s", state.RunEndpointID, RedactText(refreshErr.Error()))
}
lastSeq = 0
signalControlWake(wake)
} else {
log.Printf("RUN phase=control_stream status=failed endpoint=%s error=%s", state.RunEndpointID, RedactText(err.Error()))
}
}
if waitErr := waitWorkerLoop(ctx, boundedRetryBackoff(worker.cfg.RetryBackoff)); waitErr != nil {
return waitErr
}
}
}
func signalControlWake(wake chan<- struct{}) {
select {
case wake <- struct{}{}:
default:
}
}
func (worker *Worker) runJobLoop(ctx context.Context, interval time.Duration, wake <-chan struct{}) error {
log.Printf("RUN phase=job_loop status=starting fallbackPollMs=%d controlStream=true", interval.Milliseconds())
for {
select {
case <-ctx.Done():
@@ -1063,7 +1128,7 @@ func (worker *Worker) runJobLoop(ctx context.Context, interval time.Duration) er
}
}
claimStartedAt := time.Now()
handled, err := worker.ClaimAndRunOnce(ctx)
handled, err := worker.claimAndRunOnce(ctx, 0)
if err != nil {
log.Printf("RUN phase=job_loop status=claim_failed error=%s", RedactText(err.Error()))
if err := waitWorkerLoop(ctx, boundedRetryBackoff(worker.cfg.RetryBackoff)); err != nil {
@@ -1079,13 +1144,29 @@ func (worker *Worker) runJobLoop(ctx context.Context, interval time.Duration) er
return ErrSelfUpdateRestartRequested
}
if !handled && time.Since(claimStartedAt) < interval {
if err := waitWorkerLoop(ctx, interval-time.Since(claimStartedAt)); err != nil {
if err := waitWorkerJobWake(ctx, interval-time.Since(claimStartedAt), wake); err != nil {
return err
}
}
}
}
func waitWorkerJobWake(ctx context.Context, delay time.Duration, wake <-chan struct{}) error {
if delay <= 0 {
return nil
}
timer := time.NewTimer(delay)
defer timer.Stop()
select {
case <-ctx.Done():
return ctx.Err()
case <-wake:
return nil
case <-timer.C:
return nil
}
}
func waitWorkerLoop(ctx context.Context, delay time.Duration) error {
if delay <= 0 {
return nil