Skip to main content

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