1use aion_proto::{
4 ProtoActivityResult, ProtoRegisterWorker,
5 generated::{
6 self,
7 worker_protocol_server::{WorkerProtocol, WorkerProtocolServer},
8 },
9};
10use tokio::sync::mpsc;
11use tokio_stream::wrappers::ReceiverStream;
12use tonic::{Request, Response, Status, Streaming};
13
14use crate::worker::PendingActivities;
15use crate::worker::dispatch::{ActivityCompletion, ActivityCompletionSink};
16use crate::worker::registry::{WorkerId, WorkerMessage};
17use crate::{CallerIdentity, ServerState};
18
19#[derive(Clone)]
21pub struct WorkerGrpcService {
22 state: ServerState,
23}
24
25impl WorkerGrpcService {
26 #[must_use]
28 pub const fn new(state: ServerState) -> Self {
29 Self { state }
30 }
31}
32
33#[must_use]
35pub fn worker_service(state: ServerState) -> WorkerProtocolServer<WorkerGrpcService> {
36 WorkerProtocolServer::new(WorkerGrpcService::new(state))
37}
38
39#[tonic::async_trait]
40impl WorkerProtocol for WorkerGrpcService {
41 type StreamWorkerStream = ReceiverStream<Result<generated::ServerToWorker, Status>>;
42
43 async fn stream_worker(
44 &self,
45 request: Request<Streaming<generated::WorkerToServer>>,
46 ) -> Result<Response<Self::StreamWorkerStream>, Status> {
47 let metadata = request.metadata().clone();
48 let caller = worker_caller_from_metadata(&metadata, &self.state).await?;
49 let token_expires_at = token_expiration_from_metadata(&metadata, &self.state).await?;
50 let heartbeat_grace = self.state.runtime_config().worker.heartbeat_window;
51 let mut inbound = request.into_inner();
52
53 let first = inbound
54 .message()
55 .await?
56 .and_then(|msg| msg.message)
57 .ok_or_else(|| Status::invalid_argument("first message must be RegisterWorker"))?;
58
59 let register = match first {
60 generated::worker_to_server::Message::Register(r) => decode_register(r),
61 _ => {
62 return Err(Status::invalid_argument(
63 "first message must be RegisterWorker",
64 ));
65 }
66 };
67
68 let (task_tx, task_rx) = mpsc::channel::<Result<generated::ServerToWorker, Status>>(32);
69 let (worker_tx, worker_rx) = mpsc::channel(32);
70
71 let registration = self
72 .state
73 .worker_registry()
74 .accept_registration(self.state.namespace_guard(), &caller, ®ister, worker_tx)
75 .await
76 .map_err(|error| status_from_server_error(&error))?;
77
78 let pending = self.state.pending_activities().clone();
79 let heartbeat = self.state.heartbeat_tracker().clone();
80 let drain = self.state.drain_state().clone();
81 let registry = self.state.worker_registry().clone();
82 let worker_id = registration
83 .worker_id()
84 .ok_or_else(|| Status::internal("worker registration missing id"))?;
85 let authorized_namespace = registration
89 .namespaces()
90 .filter(|namespaces| !namespaces.is_empty())
91 .ok_or_else(|| Status::internal("worker registration missing namespace"))?
92 .iter()
93 .cloned()
94 .collect::<Vec<_>>()
95 .join(",");
96
97 task_tx
102 .try_send(Ok(register_ack_frame(
103 worker_id,
104 &authorized_namespace,
105 heartbeat_grace,
106 )))
107 .map_err(|_| Status::internal("worker response channel closed before RegisterAck"))?;
108
109 tokio::spawn(async move {
110 let write_handle = spawn_write_forwarder(worker_rx, task_tx.clone());
111
112 let teardown = StreamTeardown {
117 worker_id,
118 heartbeat: &heartbeat,
119 registry: ®istry,
120 pending: &pending,
121 drain: &drain,
122 };
123 let session = WorkerSession {
124 worker_id,
125 pending: &pending,
126 heartbeat: &heartbeat,
127 drain: &drain,
128 token_expires_at,
129 heartbeat_grace,
130 task_tx: task_tx.clone(),
131 };
132 if let Err(status) = process_inbound(inbound, session).await {
133 tracing::info!(
134 worker_id = ?worker_id,
135 %status,
136 "worker stream closed with status"
137 );
138 }
139
140 write_handle.abort();
141 drop(task_tx);
142 drop(teardown);
143 if let Err(error) = registration.deregister() {
147 tracing::error!(
148 worker_id = ?worker_id,
149 %error,
150 "worker deregistration failed during stream teardown"
151 );
152 }
153 });
154
155 Ok(Response::new(ReceiverStream::new(task_rx)))
156 }
157}
158
159fn spawn_write_forwarder(
176 mut worker_rx: mpsc::Receiver<WorkerMessage>,
177 task_tx: mpsc::Sender<Result<generated::ServerToWorker, Status>>,
178) -> tokio::task::JoinHandle<()> {
179 tokio::spawn(async move {
180 while let Some(message) = worker_rx.recv().await {
181 let msg = encode_server_to_worker(message);
182 if task_tx.send(Ok(msg)).await.is_err() {
183 return;
184 }
185 }
186 let _ = task_tx
187 .send(Err(Status::unavailable(
188 "worker was deregistered by the server (heartbeat window expired); \
189 reconnect and re-register",
190 )))
191 .await;
192 })
193}
194
195struct StreamTeardown<'a> {
204 worker_id: WorkerId,
205 heartbeat: &'a crate::worker::HeartbeatTracker,
206 registry: &'a crate::worker::ConnectedWorkerRegistry,
207 pending: &'a PendingActivities,
208 drain: &'a crate::shutdown::DrainState,
209}
210
211impl Drop for StreamTeardown<'_> {
212 fn drop(&mut self) {
213 teardown_worker_stream(
214 self.worker_id,
215 self.heartbeat,
216 self.registry,
217 self.pending,
218 self.drain,
219 );
220 }
221}
222
223fn teardown_worker_stream(
243 worker_id: WorkerId,
244 heartbeat: &crate::worker::HeartbeatTracker,
245 registry: &crate::worker::ConnectedWorkerRegistry,
246 pending: &PendingActivities,
247 drain: &crate::shutdown::DrainState,
248) {
249 if drain.is_draining() {
250 match heartbeat.park_disconnected_worker(worker_id, registry, pending) {
251 Ok(report) if report.tasks.is_empty() => {}
252 Ok(report) => {
253 tracing::info!(
254 worker_id = ?worker_id,
255 parked_tasks = report.tasks.len(),
256 "worker stream ended during drain; in-flight activities \
257 parked for restart recovery"
258 );
259 }
260 Err(error) => {
261 tracing::error!(
262 worker_id = ?worker_id,
263 %error,
264 "failed to park draining worker's in-flight activities"
265 );
266 }
267 }
268 } else {
269 match heartbeat.fail_disconnected_worker(worker_id, registry, pending) {
270 Ok(report) if report.tasks.is_empty() => {}
271 Ok(report) => {
272 tracing::warn!(
273 worker_id = ?worker_id,
274 failed_tasks = report.tasks.len(),
275 "worker disconnected with in-flight activities; \
276 surfaced as retryable lost-worker failures"
277 );
278 }
279 Err(error) => {
280 tracing::error!(
281 worker_id = ?worker_id,
282 %error,
283 "failed to sweep disconnected worker's in-flight activities"
284 );
285 }
286 }
287 }
288 drain.notify_activity_drained();
291}
292
293struct WorkerSession<'a> {
294 worker_id: WorkerId,
295 pending: &'a PendingActivities,
296 heartbeat: &'a crate::worker::HeartbeatTracker,
297 drain: &'a crate::shutdown::DrainState,
298 token_expires_at: Option<u64>,
299 heartbeat_grace: std::time::Duration,
300 task_tx: mpsc::Sender<Result<generated::ServerToWorker, Status>>,
301}
302
303async fn process_inbound(
304 mut inbound: Streaming<generated::WorkerToServer>,
305 session: WorkerSession<'_>,
306) -> Result<(), Status> {
307 let mut expired_since: Option<std::time::Instant> = None;
308 while let Some(msg) = inbound.message().await? {
309 let Some(inner) = msg.message else {
310 continue;
311 };
312 match inner {
313 generated::worker_to_server::Message::Result(result) => {
314 let proto_result = decode_activity_result(result);
315 match ActivityCompletion::try_from(proto_result) {
316 Ok(completion) => {
317 let workflow_id = completion.workflow_id.clone();
318 let activity_id = completion.activity_id.clone();
319 if let Err(error) = session.heartbeat.complete_task(
320 session.worker_id,
321 &workflow_id,
322 &activity_id,
323 ) {
324 tracing::error!(
328 worker_id = ?session.worker_id,
329 workflow_id = %workflow_id,
330 activity_id = %activity_id,
331 %error,
332 "failed to clear in-flight tracking for completed activity"
333 );
334 }
335 session.drain.notify_activity_drained();
336 if let Err(error) = session.pending.complete_activity(completion) {
337 tracing::error!(
338 worker_id = ?session.worker_id,
339 workflow_id = %workflow_id,
340 activity_id = %activity_id,
341 %error,
342 "activity completion handoff failed"
343 );
344 }
345 let ack = result_ack_frame(&workflow_id, &activity_id);
352 if let Err(error) = session.task_tx.try_send(Ok(ack)) {
353 tracing::warn!(
354 worker_id = ?session.worker_id,
355 workflow_id = %workflow_id,
356 activity_id = %activity_id,
357 %error,
358 "result ack dropped: worker stream channel unavailable"
359 );
360 }
361 }
362 Err(error) => {
363 tracing::error!(
367 worker_id = ?session.worker_id,
368 %error,
369 "malformed activity result frame; no ack sent"
370 );
371 }
372 }
373 }
374 generated::worker_to_server::Message::Register(_) => {
375 tracing::warn!(
376 worker_id = ?session.worker_id,
377 "ignoring subsequent RegisterWorker message; \
378 only the first registration is accepted per stream"
379 );
380 }
381 generated::worker_to_server::Message::Heartbeat(heartbeat_msg) => {
382 if let Err(error) = session.heartbeat.record_heartbeat(
383 session.worker_id,
384 decode_heartbeat(heartbeat_msg),
385 std::time::Instant::now(),
386 ) {
387 if matches!(error, crate::ServerError::LockPoisoned { .. }) {
392 tracing::error!(
393 worker_id = ?session.worker_id,
394 %error,
395 "heartbeat tracker lock poisoned; liveness state untrustworthy"
396 );
397 } else {
398 tracing::warn!(
399 worker_id = ?session.worker_id,
400 %error,
401 "worker heartbeat rejected"
402 );
403 }
404 }
405 if token_expired(session.token_expires_at) {
406 let first_expired = *expired_since.get_or_insert_with(std::time::Instant::now);
407 let _ = session
408 .task_tx
409 .send(Err(Status::unauthenticated(
410 "worker token expired; re-authentication required",
411 )))
412 .await;
413 if first_expired.elapsed() >= session.heartbeat_grace {
414 return Err(Status::unauthenticated("worker token expired"));
415 }
416 }
417 }
418 }
419 }
420 Ok(())
421}
422
423async fn worker_caller_from_metadata(
424 metadata: &tonic::metadata::MetadataMap,
425 state: &ServerState,
426) -> Result<CallerIdentity, Status> {
427 crate::api::grpc::caller_from_metadata(metadata, state).await
428}
429
430async fn token_expiration_from_metadata(
431 metadata: &tonic::metadata::MetadataMap,
432 state: &ServerState,
433) -> Result<Option<u64>, Status> {
434 if !state.runtime_config().auth.enabled {
435 return Ok(None);
436 }
437 #[cfg(feature = "auth")]
438 {
439 let bearer = metadata
440 .get("authorization")
441 .and_then(|value| value.to_str().ok())
442 .and_then(parse_bearer)
443 .ok_or_else(|| Status::unauthenticated("missing bearer token"))?;
444 let Some(cache) = state.jwks_cache() else {
445 return Err(Status::unauthenticated("invalid bearer token"));
446 };
447 return cache
448 .validate(&bearer)
449 .await
450 .map(|claims| Some(claims.expires_at()))
451 .map_err(|_error| Status::unauthenticated("invalid bearer token"));
452 }
453 #[cfg(not(feature = "auth"))]
454 {
455 let _ = metadata;
456 tokio::task::yield_now().await;
458 Ok(None)
459 }
460}
461
462#[cfg(feature = "auth")]
463fn parse_bearer(value: &str) -> Option<String> {
464 let token = value.strip_prefix("Bearer ")?.trim();
465 if token.is_empty() {
466 return None;
467 }
468 Some(token.to_owned())
469}
470
471fn token_expired(expires_at: Option<u64>) -> bool {
472 expires_at.is_some_and(|expires_at| {
473 #[cfg(feature = "auth")]
474 {
475 crate::auth::jwks::is_expired(expires_at)
476 }
477 #[cfg(not(feature = "auth"))]
478 {
479 let _ = expires_at;
480 false
481 }
482 })
483}
484
485fn status_from_server_error(error: &crate::ServerError) -> Status {
486 let wire = error.to_wire_error();
487 if wire.code == aion_proto::WireErrorCode::NamespaceDenied {
488 Status::permission_denied(wire.message)
489 } else {
490 Status::internal(wire.message)
491 }
492}
493
494fn register_ack_frame(
497 worker_id: WorkerId,
498 namespace: &str,
499 heartbeat_window: std::time::Duration,
500) -> generated::ServerToWorker {
501 generated::ServerToWorker {
502 message: Some(generated::server_to_worker::Message::RegisterAck(
503 generated::RegisterAck {
504 worker_id: worker_id.value(),
505 namespace: namespace.to_owned(),
506 heartbeat_window_ms: u64::try_from(heartbeat_window.as_millis())
507 .unwrap_or(u64::MAX),
508 },
509 )),
510 }
511}
512
513fn result_ack_frame(
515 workflow_id: &aion_core::WorkflowId,
516 activity_id: &aion_core::ActivityId,
517) -> generated::ServerToWorker {
518 generated::ServerToWorker {
519 message: Some(generated::server_to_worker::Message::ResultAck(
520 generated::ResultAck {
521 workflow_id: Some(generated::WorkflowId {
522 uuid: workflow_id.to_string(),
523 }),
524 activity_id: Some(generated::ActivityId {
525 sequence_position: activity_id.sequence_position(),
526 }),
527 },
528 )),
529 }
530}
531
532fn decode_register(r: generated::RegisterWorker) -> ProtoRegisterWorker {
533 ProtoRegisterWorker {
534 namespaces: r.namespaces,
535 activity_types: r.activity_types,
536 task_queue: r.task_queue,
537 node: r.node,
538 }
539}
540
541fn encode_server_to_worker(message: WorkerMessage) -> generated::ServerToWorker {
542 let message = match message {
543 WorkerMessage::ActivityTask(task) => {
544 generated::server_to_worker::Message::Task(encode_task(task))
545 }
546 WorkerMessage::DrainRequest => {
547 generated::server_to_worker::Message::Drain(generated::DrainRequest {})
548 }
549 };
550 generated::ServerToWorker {
551 message: Some(message),
552 }
553}
554
555fn encode_task(task: aion_proto::ProtoActivityTask) -> generated::ActivityTask {
556 generated::ActivityTask {
557 workflow_id: task
558 .workflow_id
559 .map(|id| generated::WorkflowId { uuid: id.uuid }),
560 activity_id: task.activity_id.map(|id| generated::ActivityId {
561 sequence_position: id.sequence_position,
562 }),
563 activity_type: task.activity_type,
564 input: task.input.map(|p| generated::Payload {
565 content_type: p.content_type,
566 bytes: p.bytes,
567 }),
568 attempt: task.attempt,
569 labels: task.labels,
570 run_id: task.run_id.map(|id| generated::RunId { uuid: id.uuid }),
571 }
572}
573
574fn decode_activity_result(r: generated::ActivityResult) -> ProtoActivityResult {
575 ProtoActivityResult {
576 workflow_id: r
577 .workflow_id
578 .map(|id| aion_proto::ProtoWorkflowId { uuid: id.uuid }),
579 activity_id: r.activity_id.map(|id| aion_proto::ProtoActivityId {
580 sequence_position: id.sequence_position,
581 }),
582 outcome: r.outcome.map(decode_outcome),
583 run_id: r.run_id.map(|id| aion_proto::ProtoRunId { uuid: id.uuid }),
584 }
585}
586
587fn decode_heartbeat(r: generated::Heartbeat) -> aion_proto::ProtoHeartbeat {
588 aion_proto::ProtoHeartbeat {
589 workflow_id: r
590 .workflow_id
591 .map(|id| aion_proto::ProtoWorkflowId { uuid: id.uuid }),
592 activity_id: r.activity_id.map(|id| aion_proto::ProtoActivityId {
593 sequence_position: id.sequence_position,
594 }),
595 progress: r.progress.map(|p| aion_proto::ProtoPayload {
596 content_type: p.content_type,
597 bytes: p.bytes,
598 }),
599 }
600}
601
602fn decode_outcome(
603 outcome: generated::activity_result::Outcome,
604) -> aion_proto::proto_activity_result::Outcome {
605 match outcome {
606 generated::activity_result::Outcome::Result(p) => {
607 aion_proto::proto_activity_result::Outcome::Result(aion_proto::ProtoPayload {
608 content_type: p.content_type,
609 bytes: p.bytes,
610 })
611 }
612 generated::activity_result::Outcome::Error(e) => {
613 aion_proto::proto_activity_result::Outcome::Error(aion_proto::ProtoActivityError {
614 kind: e.kind,
615 message: e.message,
616 details: e.details.map(|p| aion_proto::ProtoPayload {
617 content_type: p.content_type,
618 bytes: p.bytes,
619 }),
620 })
621 }
622 }
623}
624
625#[cfg(test)]
626mod tests {
627 use std::time::{Duration, Instant};
628
629 use aion_core::{ActivityId, WorkflowId};
630
631 use crate::shutdown::DrainState;
632 use crate::worker::heartbeat::InFlightActivity;
633 use crate::worker::registry::ConnectedWorkerRegistry;
634 use crate::worker::{HeartbeatTracker, PendingActivities};
635
636 use super::teardown_worker_stream;
637
638 type TestError = Box<dyn std::error::Error>;
639
640 struct TeardownFixture {
644 registry: ConnectedWorkerRegistry,
645 tracker: HeartbeatTracker,
646 pending: PendingActivities,
647 drain: DrainState,
648 worker_id: crate::worker::registry::WorkerId,
649 workflow_id: WorkflowId,
650 activity_id: ActivityId,
651 rx: std::sync::mpsc::Receiver<Result<String, String>>,
652 _registration: crate::worker::registry::WorkerRegistration,
655 }
656
657 fn fixture() -> Result<TeardownFixture, TestError> {
658 let registry = ConnectedWorkerRegistry::default();
659 let (tx, _rx) = tokio::sync::mpsc::channel(1);
660 let activity_types = [String::from("greet")];
661 let registration = registry.register("default", activity_types.iter(), tx)?;
662 let worker_id = registration
663 .worker_id()
664 .ok_or("test worker registration missing id")?;
665 let tracker = HeartbeatTracker::new(Duration::from_secs(5));
666 let pending = PendingActivities::default();
667 let workflow_id = WorkflowId::new_v4();
668 let activity_id = ActivityId::from_sequence_position(0);
669 tracker.track_task(
670 worker_id,
671 InFlightActivity {
672 workflow_id: workflow_id.clone(),
673 activity_id: activity_id.clone(),
674 },
675 Instant::now(),
676 )?;
677 let rx = pending.insert_for_test(workflow_id.clone(), activity_id.clone());
678 Ok(TeardownFixture {
679 registry,
680 tracker,
681 pending,
682 drain: DrainState::default(),
683 worker_id,
684 workflow_id,
685 activity_id,
686 rx,
687 _registration: registration,
688 })
689 }
690
691 #[test]
695 fn teardown_under_drain_parks_instead_of_failing() -> Result<(), TestError> {
696 let fixture = fixture()?;
697 assert!(fixture.drain.begin());
698
699 teardown_worker_stream(
700 fixture.worker_id,
701 &fixture.tracker,
702 &fixture.registry,
703 &fixture.pending,
704 &fixture.drain,
705 );
706
707 let resolved = fixture.rx.recv_timeout(Duration::from_millis(200))?;
708 assert_eq!(
709 resolved,
710 Err(aion::PARKED_ACTIVITY_REASON.to_owned()),
711 "a drain teardown must resolve the waiter with the parked sentinel"
712 );
713 assert_eq!(fixture.tracker.in_flight_count()?, 0);
714 assert!(
715 !fixture.tracker.is_tracked(
716 fixture.worker_id,
717 &fixture.workflow_id,
718 &fixture.activity_id
719 )?,
720 "parking must retire the tracked entry"
721 );
722 Ok(())
723 }
724
725 #[test]
729 fn teardown_without_drain_still_fails_with_retryable_lost_worker() -> Result<(), TestError> {
730 let fixture = fixture()?;
731
732 teardown_worker_stream(
733 fixture.worker_id,
734 &fixture.tracker,
735 &fixture.registry,
736 &fixture.pending,
737 &fixture.drain,
738 );
739
740 let resolved = fixture.rx.recv_timeout(Duration::from_millis(200))?;
741 let reason = resolved.err().ok_or("expected a lost-worker failure")?;
742 assert!(
743 reason.starts_with("retryable:"),
744 "a mid-run teardown must stay a retryable failure: {reason}"
745 );
746 assert!(
747 reason.contains("lost before reporting activity result"),
748 "the failure must name worker loss: {reason}"
749 );
750 assert_eq!(fixture.tracker.in_flight_count()?, 0);
751 Ok(())
752 }
753}