Add persistent Run control stream

2026-08-26 23:05:40 +08:00
parent 3b4e857ef7
commit ac14f80306
8 changed files with 245 additions and 9 deletions
@@ -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
@@ -905,6 +905,16 @@ func (client *fakeWorkerClient) Heartbeat(_ context.Context, request protocol.Ru
 	return protocol.RunHeartbeatResponse{Accepted: true, RunEndpointID: request.RunEndpointID, NextHeartbeatSeconds: 15, ServerTime: workerTestTime()}, nil
 }

+func (client *fakeWorkerClient) StreamControlEvents(ctx context.Context, request protocol.RunControlStreamRequest, handle func(protocol.RunControlEvent) error) error {
+	if handle != nil {
+		if err := handle(protocol.RunControlEvent{RunEndpointID: request.RunEndpointID, Type: protocol.RunControlEventTypeReady, Sequence: request.LastEventSeq, ServerTime: workerTestTime()}); err != nil {
+			return err
+		}
+	}
+	<-ctx.Done()
+	return ctx.Err()
+}
+
 func (client *fakeWorkerClient) ReportLifecycle(_ context.Context, request protocol.RunLifecycleReportRequest) (protocol.RunLifecycleReportResponse, error) {
 	client.lifecycleReports = append(client.lifecycleReports, request)
 	if client.lifecycleReportErr != nil {