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    /// An already-authorized namespace mint to relay to the namespace registry
52    /// shard's owner, which applies its own auto-create policy to it.
53    MintNamespace(generated::MintNamespaceRequest),
54}
55
56/// The owner's reply, relayed back to the original caller unchanged.
57#[derive(Clone, Debug)]
58pub enum ForwardReply {
59    /// The owner's `start` reply (R-4 steered start).
60    Start(generated::StartWorkflowResponse),
61    /// The owner's `signal` reply.
62    Signal(generated::SignalResponse),
63    /// The owner's `query` reply.
64    Query(generated::QueryResponse),
65    /// The owner's `cancel` reply.
66    Cancel(generated::CancelResponse),
67    /// The owner's `reopen` reply.
68    Reopen(generated::ReopenResponse),
69    /// The owner's `pause` reply (#204).
70    Pause(generated::PauseResponse),
71    /// The owner's `resume` reply (#204).
72    Resume(generated::ResumeResponse),
73    /// The owner's namespace-mint ack.
74    MintNamespace(generated::MintNamespaceResponse),
75}
76
77/// Relays a client request to a remote shard owner and returns its reply.
78///
79/// The transport is abstracted (Decision C): gRPC now, liminal later, behind one
80/// trait. Errors are returned as a tonic [`Status`] so the edge can relay them
81/// verbatim (an owner `NotOwner` stays `NotOwner`; a transport failure surfaces
82/// as `Unavailable`).
83#[async_trait]
84pub trait RequestForwarder: Send + Sync {
85    /// Forward `request` to `target` (the owner's gRPC address), copying the
86    /// caller `metadata` and stamping the next hop count, then relay the reply.
87    async fn forward(
88        &self,
89        target: SocketAddr,
90        metadata: tonic::metadata::MetadataMap,
91        request: ForwardRequest,
92    ) -> Result<ForwardReply, Status>;
93}
94
95/// gRPC forwarder: dials the owner's `grpc_address` with a tonic
96/// `WorkflowService` client and re-issues the RPC.
97#[derive(Clone, Default)]
98pub struct GrpcRequestForwarder;
99
100impl GrpcRequestForwarder {
101    /// A new gRPC forwarder.
102    #[must_use]
103    pub const fn new() -> Self {
104        Self
105    }
106}
107
108/// Read the current hop count from request metadata (`0` when absent/malformed).
109#[must_use]
110pub fn current_hops(metadata: &tonic::metadata::MetadataMap) -> u32 {
111    metadata
112        .get(FORWARD_HOPS_METADATA)
113        .and_then(|value| value.to_str().ok())
114        .and_then(|value| value.parse().ok())
115        .unwrap_or(0)
116}
117
118/// Stamp `metadata` with the incremented hop count for the outbound forward.
119fn stamp_next_hop(metadata: &mut tonic::metadata::MetadataMap) -> Result<(), Status> {
120    let next = current_hops(metadata)
121        .checked_add(1)
122        .ok_or_else(|| Status::internal("forward hop counter overflow"))?;
123    let value = tonic::metadata::MetadataValue::try_from(next.to_string())
124        .map_err(|_| Status::internal("invalid forward hop metadata value"))?;
125    metadata.insert(FORWARD_HOPS_METADATA, value);
126    Ok(())
127}
128
129async fn connect(target: SocketAddr) -> Result<WorkflowServiceClient<Channel>, Status> {
130    let uri = format!("http://{target}");
131    let endpoint = Endpoint::try_from(uri)
132        .map_err(|error| Status::unavailable(format!("invalid forward target: {error}")))?;
133    let channel = endpoint
134        .connect()
135        .await
136        .map_err(|error| Status::unavailable(format!("forward dial failed: {error}")))?;
137    Ok(WorkflowServiceClient::new(channel))
138}
139
140#[async_trait]
141impl RequestForwarder for GrpcRequestForwarder {
142    async fn forward(
143        &self,
144        target: SocketAddr,
145        mut metadata: tonic::metadata::MetadataMap,
146        request: ForwardRequest,
147    ) -> Result<ForwardReply, Status> {
148        stamp_next_hop(&mut metadata)?;
149        let mut client = connect(target).await?;
150        match request {
151            ForwardRequest::Start(message) => {
152                let mut outbound = Request::new(message);
153                *outbound.metadata_mut() = metadata;
154                client
155                    .start_workflow(outbound)
156                    .await
157                    .map(|response| ForwardReply::Start(response.into_inner()))
158            }
159            ForwardRequest::Signal(message) => {
160                let mut outbound = Request::new(message);
161                *outbound.metadata_mut() = metadata;
162                client
163                    .signal(outbound)
164                    .await
165                    .map(|response| ForwardReply::Signal(response.into_inner()))
166            }
167            ForwardRequest::Query(message) => {
168                let mut outbound = Request::new(message);
169                *outbound.metadata_mut() = metadata;
170                client
171                    .query(outbound)
172                    .await
173                    .map(|response| ForwardReply::Query(response.into_inner()))
174            }
175            ForwardRequest::Cancel(message) => {
176                let mut outbound = Request::new(message);
177                *outbound.metadata_mut() = metadata;
178                client
179                    .cancel(outbound)
180                    .await
181                    .map(|response| ForwardReply::Cancel(response.into_inner()))
182            }
183            ForwardRequest::Reopen(message) => {
184                let mut outbound = Request::new(message);
185                *outbound.metadata_mut() = metadata;
186                client
187                    .reopen(outbound)
188                    .await
189                    .map(|response| ForwardReply::Reopen(response.into_inner()))
190            }
191            ForwardRequest::Pause(message) => {
192                let mut outbound = Request::new(message);
193                *outbound.metadata_mut() = metadata;
194                client
195                    .pause(outbound)
196                    .await
197                    .map(|response| ForwardReply::Pause(response.into_inner()))
198            }
199            ForwardRequest::Resume(message) => {
200                let mut outbound = Request::new(message);
201                *outbound.metadata_mut() = metadata;
202                client
203                    .resume(outbound)
204                    .await
205                    .map(|response| ForwardReply::Resume(response.into_inner()))
206            }
207            ForwardRequest::MintNamespace(message) => {
208                let mut outbound = Request::new(message);
209                *outbound.metadata_mut() = metadata;
210                client
211                    .mint_namespace(outbound)
212                    .await
213                    .map(|response| ForwardReply::MintNamespace(response.into_inner()))
214            }
215        }
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::{FORWARD_HOPS_METADATA, MAX_FORWARD_HOPS, current_hops, stamp_next_hop};
222
223    #[test]
224    fn current_hops_defaults_to_zero_when_absent() {
225        let metadata = tonic::metadata::MetadataMap::new();
226        assert_eq!(current_hops(&metadata), 0);
227    }
228
229    #[test]
230    fn stamp_next_hop_increments_from_zero() -> Result<(), tonic::Status> {
231        let mut metadata = tonic::metadata::MetadataMap::new();
232        stamp_next_hop(&mut metadata)?;
233        assert_eq!(current_hops(&metadata), 1);
234        stamp_next_hop(&mut metadata)?;
235        assert_eq!(current_hops(&metadata), 2);
236        Ok(())
237    }
238
239    #[test]
240    fn malformed_hop_value_reads_as_zero() -> Result<(), tonic::Status> {
241        let mut metadata = tonic::metadata::MetadataMap::new();
242        metadata.insert(
243            FORWARD_HOPS_METADATA,
244            tonic::metadata::MetadataValue::try_from("not-a-number")
245                .map_err(|_| tonic::Status::internal("bad fixture"))?,
246        );
247        assert_eq!(current_hops(&metadata), 0);
248        Ok(())
249    }
250
251    #[test]
252    fn hop_cap_is_two() {
253        assert_eq!(MAX_FORWARD_HOPS, 2);
254    }
255}
aion_server/routing/forwarder.rs

aion_server/routing/
forwarder.rs
