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