aion_server/routing/
forwarder.rs

1//! R-3 request forwarding: the `RequestForwarder` trait + a gRPC implementation
2//! that relays a non-local signal/query/cancel to the shard's current owner and
3//! returns its reply (DISTRIBUTED-ROUTING-DESIGN §2.2 / §2.6 / Decision C).
4//!
5//! The trait mirrors the `OutboxRowDispatch` gRPC/liminal seam: a gRPC forwarder
6//! ships now (a tonic `WorkflowService` client to the owner's `grpc_address`); a
7//! liminal forwarder (`request_reply_conversation()`) drops in behind the same
8//! trait when liminal 13-L0/L1 land (R-6). Loop prevention is a hop cap carried
9//! in request metadata; stale-target handling is bounded re-resolution at the
10//! edge (§2.5).
11
12use std::net::SocketAddr;
13
14use async_trait::async_trait;
15use tonic::transport::{Channel, Endpoint};
16use tonic::{Request, Status};
17
18use aion_proto::generated::{self, workflow_service_client::WorkflowServiceClient};
19
20/// The metadata key carrying the number of cluster hops a forwarded request has
21/// already taken. A request arriving with a value at or above the hop cap is NOT
22/// forwarded again (loop prevention) — the receiver returns `NotOwner` so the
23/// original caller re-resolves and retries with backoff.
24pub const FORWARD_HOPS_METADATA: &str = "x-aion-forward-hops";
25
26/// The maximum number of intra-cluster forward hops a single client request may
27/// take before the chain is broken with `NotOwner` (§2.5: "a hop counter caps
28/// forwards (e.g. 2)").
29pub const MAX_FORWARD_HOPS: u32 = 2;
30
31/// The forwardable client RPCs. For `signal`/`query`/`cancel` the target
32/// `workflow_id` is known at the edge; `start` is forwardable once the caller
33/// supplies an R-4 `routing_key` so the edge can resolve the target shard's owner
34/// before placing the start.
35#[derive(Clone, Debug)]
36pub enum ForwardRequest {
37    /// A steered `start` to relay verbatim to the routing key's shard owner (R-4).
38    Start(generated::StartWorkflowRequest),
39    /// A `signal` to relay verbatim to the owner.
40    Signal(generated::SignalRequest),
41    /// A `query` to relay verbatim to the owner.
42    Query(generated::QueryRequest),
43    /// A `cancel` to relay verbatim to the owner.
44    Cancel(generated::CancelRequest),
45    /// A `reopen` to relay verbatim to the owner.
46    Reopen(generated::ReopenRequest),
47    /// A `pause` to relay verbatim to the owner (#204).
48    Pause(generated::PauseRequest),
49    /// A `resume` to relay verbatim to the owner (#204).
50    Resume(generated::ResumeRequest),
51}
52
53/// The owner's reply, relayed back to the original caller unchanged.
54#[derive(Clone, Debug)]
55pub enum ForwardReply {
56    /// The owner's `start` reply (R-4 steered start).
57    Start(generated::StartWorkflowResponse),
58    /// The owner's `signal` reply.
59    Signal(generated::SignalResponse),
60    /// The owner's `query` reply.
61    Query(generated::QueryResponse),
62    /// The owner's `cancel` reply.
63    Cancel(generated::CancelResponse),
64    /// The owner's `reopen` reply.
65    Reopen(generated::ReopenResponse),
66    /// The owner's `pause` reply (#204).
67    Pause(generated::PauseResponse),
68    /// The owner's `resume` reply (#204).
69    Resume(generated::ResumeResponse),
70}
71
72/// Relays a client request to a remote shard owner and returns its reply.
73///
74/// The transport is abstracted (Decision C): gRPC now, liminal later, behind one
75/// trait. Errors are returned as a tonic [`Status`] so the edge can relay them
76/// verbatim (an owner `NotOwner` stays `NotOwner`; a transport failure surfaces
77/// as `Unavailable`).
78#[async_trait]
79pub trait RequestForwarder: Send + Sync {
80    /// Forward `request` to `target` (the owner's gRPC address), copying the
81    /// caller `metadata` and stamping the next hop count, then relay the reply.
82    async fn forward(
83        &self,
84        target: SocketAddr,
85        metadata: tonic::metadata::MetadataMap,
86        request: ForwardRequest,
87    ) -> Result<ForwardReply, Status>;
88}
89
90/// gRPC forwarder: dials the owner's `grpc_address` with a tonic
91/// `WorkflowService` client and re-issues the RPC.
92#[derive(Clone, Default)]
93pub struct GrpcRequestForwarder;
94
95impl GrpcRequestForwarder {
96    /// A new gRPC forwarder.
97    #[must_use]
98    pub const fn new() -> Self {
99        Self
100    }
101}
102
103/// Read the current hop count from request metadata (`0` when absent/malformed).
104#[must_use]
105pub fn current_hops(metadata: &tonic::metadata::MetadataMap) -> u32 {
106    metadata
107        .get(FORWARD_HOPS_METADATA)
108        .and_then(|value| value.to_str().ok())
109        .and_then(|value| value.parse().ok())
110        .unwrap_or(0)
111}
112
113/// Stamp `metadata` with the incremented hop count for the outbound forward.
114fn stamp_next_hop(metadata: &mut tonic::metadata::MetadataMap) -> Result<(), Status> {
115    let next = current_hops(metadata)
116        .checked_add(1)
117        .ok_or_else(|| Status::internal("forward hop counter overflow"))?;
118    let value = tonic::metadata::MetadataValue::try_from(next.to_string())
119        .map_err(|_| Status::internal("invalid forward hop metadata value"))?;
120    metadata.insert(FORWARD_HOPS_METADATA, value);
121    Ok(())
122}
123
124async fn connect(target: SocketAddr) -> Result<WorkflowServiceClient<Channel>, Status> {
125    let uri = format!("http://{target}");
126    let endpoint = Endpoint::try_from(uri)
127        .map_err(|error| Status::unavailable(format!("invalid forward target: {error}")))?;
128    let channel = endpoint
129        .connect()
130        .await
131        .map_err(|error| Status::unavailable(format!("forward dial failed: {error}")))?;
132    Ok(WorkflowServiceClient::new(channel))
133}
134
135#[async_trait]
136impl RequestForwarder for GrpcRequestForwarder {
137    async fn forward(
138        &self,
139        target: SocketAddr,
140        mut metadata: tonic::metadata::MetadataMap,
141        request: ForwardRequest,
142    ) -> Result<ForwardReply, Status> {
143        stamp_next_hop(&mut metadata)?;
144        let mut client = connect(target).await?;
145        match request {
146            ForwardRequest::Start(message) => {
147                let mut outbound = Request::new(message);
148                *outbound.metadata_mut() = metadata;
149                client
150                    .start_workflow(outbound)
151                    .await
152                    .map(|response| ForwardReply::Start(response.into_inner()))
153            }
154            ForwardRequest::Signal(message) => {
155                let mut outbound = Request::new(message);
156                *outbound.metadata_mut() = metadata;
157                client
158                    .signal(outbound)
159                    .await
160                    .map(|response| ForwardReply::Signal(response.into_inner()))
161            }
162            ForwardRequest::Query(message) => {
163                let mut outbound = Request::new(message);
164                *outbound.metadata_mut() = metadata;
165                client
166                    .query(outbound)
167                    .await
168                    .map(|response| ForwardReply::Query(response.into_inner()))
169            }
170            ForwardRequest::Cancel(message) => {
171                let mut outbound = Request::new(message);
172                *outbound.metadata_mut() = metadata;
173                client
174                    .cancel(outbound)
175                    .await
176                    .map(|response| ForwardReply::Cancel(response.into_inner()))
177            }
178            ForwardRequest::Reopen(message) => {
179                let mut outbound = Request::new(message);
180                *outbound.metadata_mut() = metadata;
181                client
182                    .reopen(outbound)
183                    .await
184                    .map(|response| ForwardReply::Reopen(response.into_inner()))
185            }
186            ForwardRequest::Pause(message) => {
187                let mut outbound = Request::new(message);
188                *outbound.metadata_mut() = metadata;
189                client
190                    .pause(outbound)
191                    .await
192                    .map(|response| ForwardReply::Pause(response.into_inner()))
193            }
194            ForwardRequest::Resume(message) => {
195                let mut outbound = Request::new(message);
196                *outbound.metadata_mut() = metadata;
197                client
198                    .resume(outbound)
199                    .await
200                    .map(|response| ForwardReply::Resume(response.into_inner()))
201            }
202        }
203    }
204}
205
206#[cfg(test)]
207mod tests {
208    use super::{FORWARD_HOPS_METADATA, MAX_FORWARD_HOPS, current_hops, stamp_next_hop};
209
210    #[test]
211    fn current_hops_defaults_to_zero_when_absent() {
212        let metadata = tonic::metadata::MetadataMap::new();
213        assert_eq!(current_hops(&metadata), 0);
214    }
215
216    #[test]
217    fn stamp_next_hop_increments_from_zero() -> Result<(), tonic::Status> {
218        let mut metadata = tonic::metadata::MetadataMap::new();
219        stamp_next_hop(&mut metadata)?;
220        assert_eq!(current_hops(&metadata), 1);
221        stamp_next_hop(&mut metadata)?;
222        assert_eq!(current_hops(&metadata), 2);
223        Ok(())
224    }
225
226    #[test]
227    fn malformed_hop_value_reads_as_zero() -> Result<(), tonic::Status> {
228        let mut metadata = tonic::metadata::MetadataMap::new();
229        metadata.insert(
230            FORWARD_HOPS_METADATA,
231            tonic::metadata::MetadataValue::try_from("not-a-number")
232                .map_err(|_| tonic::Status::internal("bad fixture"))?,
233        );
234        assert_eq!(current_hops(&metadata), 0);
235        Ok(())
236    }
237
238    #[test]
239    fn hop_cap_is_two() {
240        assert_eq!(MAX_FORWARD_HOPS, 2);
241    }
242}
aion_server/routing/forwarder.rs

aion_server/routing/
forwarder.rs
