dynamo_runtime/pipeline/network/ingress/
http_endpoint.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! HTTP endpoint for receiving requests via Axum/HTTP/2
5
6use super::*;
7use crate::SystemHealth;
8use crate::config::HealthStatus;
9use crate::logging::TraceParent;
10use anyhow::Result;
11use axum::{
12    Router,
13    body::Bytes,
14    extract::{Path, State as AxumState},
15    http::{HeaderMap, StatusCode},
16    response::IntoResponse,
17    routing::post,
18};
19use dashmap::DashMap;
20use hyper_util::rt::{TokioExecutor, TokioIo};
21use hyper_util::server::conn::auto::Builder as Http2Builder;
22use hyper_util::service::TowerToHyperService;
23use parking_lot::Mutex;
24use std::net::SocketAddr;
25use std::sync::atomic::{AtomicU64, Ordering};
26use tokio::sync::Notify;
27use tokio_util::sync::CancellationToken;
28use tower_http::trace::TraceLayer;
29use tracing::Instrument;
30
31/// Default root path for dynamo RPC endpoints
32const DEFAULT_RPC_ROOT_PATH: &str = "/v1/rpc";
33
34/// version of crate
35pub const VERSION: &str = env!("CARGO_PKG_VERSION");
36
37/// Shared HTTP server that handles multiple endpoints on a single port
38pub struct SharedHttpServer {
39    handlers: Arc<DashMap<String, Arc<EndpointHandler>>>,
40    bind_addr: SocketAddr,
41    cancellation_token: CancellationToken,
42}
43
44/// Handler for a specific endpoint
45struct EndpointHandler {
46    service_handler: Arc<dyn PushWorkHandler>,
47    instance_id: u64,
48    namespace: Arc<String>,
49    component_name: Arc<String>,
50    endpoint_name: Arc<String>,
51    system_health: Arc<Mutex<SystemHealth>>,
52    inflight: Arc<AtomicU64>,
53    notify: Arc<Notify>,
54}
55
56impl SharedHttpServer {
57    pub fn new(bind_addr: SocketAddr, cancellation_token: CancellationToken) -> Arc<Self> {
58        Arc::new(Self {
59            handlers: Arc::new(DashMap::new()),
60            bind_addr,
61            cancellation_token,
62        })
63    }
64
65    /// Register an endpoint handler with this server
66    #[allow(clippy::too_many_arguments)]
67    pub async fn register_endpoint(
68        &self,
69        subject: String,
70        service_handler: Arc<dyn PushWorkHandler>,
71        instance_id: u64,
72        namespace: String,
73        component_name: String,
74        endpoint_name: String,
75        system_health: Arc<Mutex<SystemHealth>>,
76    ) -> Result<()> {
77        let handler = Arc::new(EndpointHandler {
78            service_handler,
79            instance_id,
80            namespace: Arc::new(namespace),
81            component_name: Arc::new(component_name),
82            endpoint_name: Arc::new(endpoint_name.clone()),
83            system_health: system_health.clone(),
84            inflight: Arc::new(AtomicU64::new(0)),
85            notify: Arc::new(Notify::new()),
86        });
87
88        // Insert handler FIRST to ensure it's ready to receive requests
89        let subject_clone = subject.clone();
90        self.handlers.insert(subject, handler);
91
92        // THEN set health status to Ready (after handler is registered and ready)
93        system_health
94            .lock()
95            .set_endpoint_health_status(&endpoint_name, HealthStatus::Ready);
96
97        tracing::debug!("Registered endpoint handler for subject: {}", subject_clone);
98        Ok(())
99    }
100
101    /// Unregister an endpoint handler
102    pub async fn unregister_endpoint(&self, subject: &str, endpoint_name: &str) {
103        if let Some((_, handler)) = self.handlers.remove(subject) {
104            handler
105                .system_health
106                .lock()
107                .set_endpoint_health_status(endpoint_name, HealthStatus::NotReady);
108            tracing::debug!("Unregistered endpoint handler for subject: {}", subject);
109        }
110    }
111
112    /// Start the shared HTTP server
113    pub async fn start(self: Arc<Self>) -> Result<()> {
114        let rpc_root_path = std::env::var("DYN_HTTP_RPC_ROOT_PATH")
115            .unwrap_or_else(|_| DEFAULT_RPC_ROOT_PATH.to_string());
116        let route_pattern = format!("{}/{{*endpoint}}", rpc_root_path);
117
118        let app = Router::new()
119            .route(&route_pattern, post(handle_shared_request))
120            .layer(TraceLayer::new_for_http())
121            .with_state(self.clone());
122
123        tracing::info!(
124            "Starting shared HTTP/2 endpoint server on {} at path {}/:endpoint",
125            self.bind_addr,
126            rpc_root_path
127        );
128
129        let listener = tokio::net::TcpListener::bind(&self.bind_addr).await?;
130        let cancellation_token = self.cancellation_token.clone();
131
132        loop {
133            tokio::select! {
134                accept_result = listener.accept() => {
135                    match accept_result {
136                        Ok((stream, _addr)) => {
137                            let app_clone = app.clone();
138                            let cancel_clone = cancellation_token.clone();
139
140                            tokio::spawn(async move {
141                                // Create HTTP/2 connection builder with prior knowledge
142                                let http2_builder = Http2Builder::new(TokioExecutor::new());
143
144                                let io = TokioIo::new(stream);
145                                let tower_service = app_clone.into_service();
146
147                                // Wrap Tower service for Hyper compatibility
148                                let hyper_service = TowerToHyperService::new(tower_service);
149
150                                tokio::select! {
151                                    result = http2_builder.serve_connection(io, hyper_service) => {
152                                        if let Err(e) = result {
153                                            tracing::debug!("HTTP/2 connection error: {}", e);
154                                        }
155                                    }
156                                    _ = cancel_clone.cancelled() => {
157                                        tracing::trace!("Connection cancelled");
158                                    }
159                                }
160                            });
161                        }
162                        Err(e) => {
163                            tracing::error!("Failed to accept connection: {}", e);
164                        }
165                    }
166                }
167                _ = cancellation_token.cancelled() => {
168                    tracing::info!("SharedHttpServer received cancellation signal, shutting down");
169                    return Ok(());
170                }
171            }
172        }
173    }
174
175    /// Wait for all inflight requests across all endpoints
176    pub async fn wait_for_inflight(&self) {
177        for handler in self.handlers.iter() {
178            while handler.value().inflight.load(Ordering::SeqCst) > 0 {
179                tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
180            }
181        }
182    }
183}
184
185/// HTTP handler for the shared server
186async fn handle_shared_request(
187    AxumState(server): AxumState<Arc<SharedHttpServer>>,
188    Path(endpoint_path): Path<String>,
189    headers: HeaderMap,
190    body: Bytes,
191) -> impl IntoResponse {
192    // Look up the handler for this endpoint (lock-free read with DashMap)
193    let handler = match server.handlers.get(&endpoint_path) {
194        Some(h) => h.clone(),
195        None => {
196            tracing::warn!("No handler found for endpoint: {}", endpoint_path);
197            return (StatusCode::NOT_FOUND, "Endpoint not found");
198        }
199    };
200
201    // Increment inflight counter
202    handler.inflight.fetch_add(1, Ordering::SeqCst);
203
204    // Extract tracing headers
205    let traceparent = TraceParent::from_axum_headers(&headers);
206
207    // Spawn async handler
208    let service_handler = handler.service_handler.clone();
209    let inflight = handler.inflight.clone();
210    let notify = handler.notify.clone();
211    let namespace = handler.namespace.clone();
212    let component_name = handler.component_name.clone();
213    let endpoint_name = handler.endpoint_name.clone();
214    let instance_id = handler.instance_id;
215
216    tokio::spawn(async move {
217        tracing::trace!(instance_id, "handling new HTTP request");
218        let result = service_handler
219            .handle_payload(body)
220            .instrument(tracing::info_span!(
221                "handle_payload",
222                component = component_name.as_ref(),
223                endpoint = endpoint_name.as_ref(),
224                namespace = namespace.as_ref(),
225                instance_id = instance_id,
226                trace_id = traceparent.trace_id,
227                parent_id = traceparent.parent_id,
228                x_request_id = traceparent.x_request_id,
229                x_dynamo_request_id = traceparent.x_dynamo_request_id,
230                tracestate = traceparent.tracestate
231            ))
232            .await;
233        match result {
234            Ok(_) => {
235                tracing::trace!(instance_id, "request handled successfully");
236            }
237            Err(e) => {
238                tracing::warn!("Failed to handle request: {}", e.to_string());
239            }
240        }
241
242        // Decrease inflight counter
243        inflight.fetch_sub(1, Ordering::SeqCst);
244        notify.notify_one();
245    });
246
247    // Return 202 Accepted immediately (like NATS ack)
248    (StatusCode::ACCEPTED, "")
249}
250
251/// Extension trait for TraceParent to support Axum headers
252impl TraceParent {
253    pub fn from_axum_headers(headers: &HeaderMap) -> Self {
254        let mut traceparent = TraceParent::default();
255
256        if let Some(value) = headers.get("traceparent")
257            && let Ok(s) = value.to_str()
258        {
259            traceparent.trace_id = Some(s.to_string());
260        }
261
262        if let Some(value) = headers.get("tracestate")
263            && let Ok(s) = value.to_str()
264        {
265            traceparent.tracestate = Some(s.to_string());
266        }
267
268        if let Some(value) = headers.get("x-request-id")
269            && let Ok(s) = value.to_str()
270        {
271            traceparent.x_request_id = Some(s.to_string());
272        }
273
274        if let Some(value) = headers.get("x-dynamo-request-id")
275            && let Ok(s) = value.to_str()
276        {
277            traceparent.x_dynamo_request_id = Some(s.to_string());
278        }
279
280        traceparent
281    }
282}
283
284// Implement RequestPlaneServer trait for SharedHttpServer
285#[async_trait::async_trait]
286impl super::unified_server::RequestPlaneServer for SharedHttpServer {
287    async fn register_endpoint(
288        &self,
289        endpoint_name: String,
290        service_handler: Arc<dyn PushWorkHandler>,
291        instance_id: u64,
292        namespace: String,
293        component_name: String,
294        system_health: Arc<Mutex<SystemHealth>>,
295    ) -> Result<()> {
296        // For HTTP, we use endpoint_name as both the subject (routing key) and endpoint_name
297        self.register_endpoint(
298            endpoint_name.clone(),
299            service_handler,
300            instance_id,
301            namespace,
302            component_name,
303            endpoint_name,
304            system_health,
305        )
306        .await
307    }
308
309    async fn unregister_endpoint(&self, endpoint_name: &str) -> Result<()> {
310        self.unregister_endpoint(endpoint_name, endpoint_name).await;
311        Ok(())
312    }
313
314    fn address(&self) -> String {
315        format!("http://{}:{}", self.bind_addr.ip(), self.bind_addr.port())
316    }
317
318    fn transport_name(&self) -> &'static str {
319        "http"
320    }
321
322    fn is_healthy(&self) -> bool {
323        // Server is healthy if it has been created
324        // TODO: Add more sophisticated health checks (e.g., check if listener is active)
325        true
326    }
327}
328
329#[cfg(test)]
330mod tests {
331    use super::*;
332
333    #[test]
334    fn test_traceparent_from_axum_headers() {
335        let mut headers = HeaderMap::new();
336        headers.insert("traceparent", "test-trace-id".parse().unwrap());
337        headers.insert("tracestate", "test-state".parse().unwrap());
338        headers.insert("x-request-id", "req-123".parse().unwrap());
339        headers.insert("x-dynamo-request-id", "dyn-456".parse().unwrap());
340
341        let traceparent = TraceParent::from_axum_headers(&headers);
342        assert_eq!(traceparent.trace_id, Some("test-trace-id".to_string()));
343        assert_eq!(traceparent.tracestate, Some("test-state".to_string()));
344        assert_eq!(traceparent.x_request_id, Some("req-123".to_string()));
345        assert_eq!(traceparent.x_dynamo_request_id, Some("dyn-456".to_string()));
346    }
347
348    #[test]
349    fn test_shared_http_server_creation() {
350        use std::net::{IpAddr, Ipv4Addr};
351        let bind_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 0);
352        let token = CancellationToken::new();
353
354        let server = SharedHttpServer::new(bind_addr, token);
355        assert!(server.handlers.is_empty());
356    }
357}
dynamo_runtime/pipeline/network/ingress/http_endpoint.rs

dynamo_runtime/pipeline/network/ingress/
http_endpoint.rs
