dynamo_runtime/pipeline/network/ingress/
http_endpoint.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 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, RwLock};
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    actual_addr: RwLock<Option<SocketAddr>>,
42    cancellation_token: CancellationToken,
43}
44
45/// Handler for a specific endpoint
46struct EndpointHandler {
47    service_handler: Arc<dyn PushWorkHandler>,
48    instance_id: u64,
49    namespace: Arc<String>,
50    component_name: Arc<String>,
51    endpoint_name: Arc<String>,
52    system_health: Arc<Mutex<SystemHealth>>,
53    inflight: Arc<AtomicU64>,
54    notify: Arc<Notify>,
55}
56
57impl SharedHttpServer {
58    pub fn new(bind_addr: SocketAddr, cancellation_token: CancellationToken) -> Arc<Self> {
59        Arc::new(Self {
60            handlers: Arc::new(DashMap::new()),
61            bind_addr,
62            actual_addr: RwLock::new(None),
63            cancellation_token,
64        })
65    }
66
67    /// Get the actual bound address (after `bind_and_start` resolves).
68    pub fn actual_address(&self) -> Option<SocketAddr> {
69        self.actual_addr.try_read().ok().and_then(|g| *g)
70    }
71
72    /// Register an endpoint handler with this server
73    #[allow(clippy::too_many_arguments)]
74    pub async fn register_endpoint(
75        &self,
76        subject: String,
77        service_handler: Arc<dyn PushWorkHandler>,
78        instance_id: u64,
79        namespace: String,
80        component_name: String,
81        endpoint_name: String,
82        system_health: Arc<Mutex<SystemHealth>>,
83    ) -> Result<()> {
84        let handler = Arc::new(EndpointHandler {
85            service_handler,
86            instance_id,
87            namespace: Arc::new(namespace),
88            component_name: Arc::new(component_name),
89            endpoint_name: Arc::new(endpoint_name.clone()),
90            system_health: system_health.clone(),
91            inflight: Arc::new(AtomicU64::new(0)),
92            notify: Arc::new(Notify::new()),
93        });
94
95        // Insert handler FIRST to ensure it's ready to receive requests
96        let subject_clone = subject.clone();
97        self.handlers.insert(subject, handler);
98
99        system_health.lock().set_endpoint_registered(&endpoint_name);
100
101        tracing::debug!("Registered endpoint handler for subject: {subject_clone}");
102        Ok(())
103    }
104
105    /// Unregister an endpoint handler
106    pub async fn unregister_endpoint(&self, subject: &str, endpoint_name: &str) {
107        if let Some((_, handler)) = self.handlers.remove(subject) {
108            handler
109                .system_health
110                .lock()
111                .set_endpoint_health_status(endpoint_name, HealthStatus::NotReady);
112            tracing::debug!(
113                endpoint_name = %endpoint_name,
114                subject = %subject,
115                "Unregistered HTTP endpoint handler"
116            );
117
118            let inflight_count = handler.inflight.load(Ordering::SeqCst);
119            if inflight_count > 0 {
120                tracing::info!(
121                    endpoint_name = %endpoint_name,
122                    inflight_count = inflight_count,
123                    "Waiting for inflight HTTP requests to complete"
124                );
125                while handler.inflight.load(Ordering::SeqCst) > 0 {
126                    handler.notify.notified().await;
127                }
128                tracing::info!(
129                    endpoint_name = %endpoint_name,
130                    "All inflight HTTP requests completed"
131                );
132            }
133        }
134    }
135
136    /// Bind the TCP listener and start the accept loop.
137    ///
138    /// Returns the actual bound `SocketAddr` (important when binding to port 0).
139    pub async fn bind_and_start(self: Arc<Self>) -> Result<SocketAddr> {
140        let rpc_root_path = std::env::var("DYN_HTTP_RPC_ROOT_PATH")
141            .unwrap_or_else(|_| DEFAULT_RPC_ROOT_PATH.to_string());
142        let route_pattern = format!("{}/{{*endpoint}}", rpc_root_path);
143
144        let app = Router::new()
145            .route(&route_pattern, post(handle_shared_request))
146            .layer(TraceLayer::new_for_http())
147            .with_state(self.clone());
148
149        let listener = tokio::net::TcpListener::bind(&self.bind_addr).await?;
150        let actual_addr = listener.local_addr()?;
151
152        tracing::info!(
153            requested = %self.bind_addr,
154            actual = %actual_addr,
155            rpc_root = %rpc_root_path,
156            "HTTP/2 endpoint server bound"
157        );
158
159        // Store the actual address so `address()` returns the real port.
160        *self.actual_addr.write().await = Some(actual_addr);
161
162        let cancellation_token = self.cancellation_token.clone();
163
164        // Spawn the accept loop in the background.
165        tokio::spawn(async move {
166            loop {
167                tokio::select! {
168                    accept_result = listener.accept() => {
169                        match accept_result {
170                            Ok((stream, _addr)) => {
171                                let app_clone = app.clone();
172                                let cancel_clone = cancellation_token.clone();
173
174                                tokio::spawn(async move {
175                                    let http2_builder = Http2Builder::new(TokioExecutor::new());
176
177                                    let io = TokioIo::new(stream);
178                                    let tower_service = app_clone.into_service();
179                                    let hyper_service = TowerToHyperService::new(tower_service);
180
181                                    tokio::select! {
182                                        result = http2_builder.serve_connection(io, hyper_service) => {
183                                            if let Err(e) = result {
184                                                tracing::debug!("HTTP/2 connection error: {e}");
185                                            }
186                                        }
187                                        _ = cancel_clone.cancelled() => {
188                                            tracing::trace!("Connection cancelled");
189                                        }
190                                    }
191                                });
192                            }
193                            Err(e) => {
194                                tracing::error!("Failed to accept connection: {e}");
195                            }
196                        }
197                    }
198                    _ = cancellation_token.cancelled() => {
199                        tracing::info!("SharedHttpServer received cancellation signal, shutting down");
200                        return;
201                    }
202                }
203            }
204        });
205
206        Ok(actual_addr)
207    }
208
209    /// Wait for all inflight requests across all endpoints
210    pub async fn wait_for_inflight(&self) {
211        for handler in self.handlers.iter() {
212            while handler.value().inflight.load(Ordering::SeqCst) > 0 {
213                tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
214            }
215        }
216    }
217}
218
219/// HTTP handler for the shared server
220async fn handle_shared_request(
221    AxumState(server): AxumState<Arc<SharedHttpServer>>,
222    Path(endpoint_path): Path<String>,
223    headers: HeaderMap,
224    body: Bytes,
225) -> impl IntoResponse {
226    // Look up the handler for this endpoint (lock-free read with DashMap)
227    let handler = match server.handlers.get(&endpoint_path) {
228        Some(h) => h.clone(),
229        None => {
230            tracing::warn!("No handler found for endpoint: {endpoint_path}");
231            return (StatusCode::NOT_FOUND, "Endpoint not found");
232        }
233    };
234
235    // Increment inflight counter
236    handler.inflight.fetch_add(1, Ordering::SeqCst);
237
238    // Extract tracing headers
239    let traceparent = TraceParent::from_axum_headers(&headers);
240
241    // Spawn async handler
242    let service_handler = handler.service_handler.clone();
243    let inflight = handler.inflight.clone();
244    let notify = handler.notify.clone();
245    let namespace = handler.namespace.clone();
246    let component_name = handler.component_name.clone();
247    let endpoint_name = handler.endpoint_name.clone();
248    let instance_id = handler.instance_id;
249
250    tokio::spawn(async move {
251        tracing::trace!(instance_id, "handling new HTTP request");
252        let result = service_handler
253            .handle_payload(body, traceparent.request_id.clone())
254            .instrument(tracing::info_span!(
255                "handle_payload",
256                component = component_name.as_ref(),
257                endpoint = endpoint_name.as_ref(),
258                namespace = namespace.as_ref(),
259                instance_id = instance_id,
260                trace_id = traceparent.trace_id,
261                parent_id = traceparent.parent_id,
262                x_request_id = traceparent.x_request_id,
263                request_id = traceparent.request_id,
264                tracestate = traceparent.tracestate
265            ))
266            .await;
267        match result {
268            Ok(_) => {
269                tracing::trace!(instance_id, "request handled successfully");
270            }
271            Err(e) => {
272                tracing::warn!("Failed to handle request: {}", e.to_string());
273            }
274        }
275
276        // Decrease inflight counter
277        inflight.fetch_sub(1, Ordering::SeqCst);
278        notify.notify_one();
279    });
280
281    // Return 202 Accepted immediately (like NATS ack)
282    (StatusCode::ACCEPTED, "")
283}
284
285/// Extension trait for TraceParent to support Axum headers
286impl TraceParent {
287    pub fn from_axum_headers(headers: &HeaderMap) -> Self {
288        let mut traceparent = TraceParent::default();
289
290        if let Some(value) = headers.get("traceparent")
291            && let Ok(s) = value.to_str()
292        {
293            traceparent.trace_id = Some(s.to_string());
294        }
295
296        if let Some(value) = headers.get("tracestate")
297            && let Ok(s) = value.to_str()
298        {
299            traceparent.tracestate = Some(s.to_string());
300        }
301
302        if let Some(value) = headers.get("x-request-id")
303            && let Ok(s) = value.to_str()
304        {
305            traceparent.x_request_id = Some(s.to_string());
306        }
307
308        // Read request-id from internal headers, with fallback to deprecated x-dynamo-request-id
309        if let Some(s) = headers
310            .get("request-id")
311            .and_then(|v| v.to_str().ok())
312            .filter(|s| uuid::Uuid::parse_str(s).is_ok())
313        {
314            traceparent.request_id = Some(s.to_string());
315        } else if let Some(s) = headers
316            .get("x-dynamo-request-id")
317            .and_then(|v| v.to_str().ok())
318            .filter(|s| uuid::Uuid::parse_str(s).is_ok())
319        {
320            traceparent.request_id = Some(s.to_string());
321        }
322
323        traceparent
324    }
325}
326
327// Implement RequestPlaneServer trait for SharedHttpServer
328#[async_trait::async_trait]
329impl super::unified_server::RequestPlaneServer for SharedHttpServer {
330    async fn register_endpoint(
331        &self,
332        endpoint_name: String,
333        service_handler: Arc<dyn PushWorkHandler>,
334        instance_id: u64,
335        namespace: String,
336        component_name: String,
337        system_health: Arc<Mutex<SystemHealth>>,
338    ) -> Result<()> {
339        // For HTTP, we use endpoint_name as both the subject (routing key) and endpoint_name
340        self.register_endpoint(
341            endpoint_name.clone(),
342            service_handler,
343            instance_id,
344            namespace,
345            component_name,
346            endpoint_name,
347            system_health,
348        )
349        .await
350    }
351
352    async fn unregister_endpoint(&self, endpoint_name: &str) -> Result<()> {
353        self.unregister_endpoint(endpoint_name, endpoint_name).await;
354        Ok(())
355    }
356
357    fn address(&self) -> String {
358        let addr = self.actual_address().unwrap_or(self.bind_addr);
359        format!("http://{}:{}", addr.ip(), addr.port())
360    }
361
362    fn transport_name(&self) -> &'static str {
363        "http"
364    }
365
366    fn is_healthy(&self) -> bool {
367        // Server is healthy if it has been created
368        // TODO: Add more sophisticated health checks (e.g., check if listener is active)
369        true
370    }
371}
372
373#[cfg(test)]
374mod tests {
375    use super::*;
376
377    #[test]
378    fn test_traceparent_from_axum_headers() {
379        let mut headers = HeaderMap::new();
380        headers.insert("traceparent", "test-trace-id".parse().unwrap());
381        headers.insert("tracestate", "test-state".parse().unwrap());
382        headers.insert("x-request-id", "req-123".parse().unwrap());
383        headers.insert(
384            "x-dynamo-request-id",
385            "550e8400-e29b-41d4-a716-446655440000".parse().unwrap(),
386        );
387
388        let traceparent = TraceParent::from_axum_headers(&headers);
389        assert_eq!(traceparent.trace_id, Some("test-trace-id".to_string()));
390        assert_eq!(traceparent.tracestate, Some("test-state".to_string()));
391        assert_eq!(traceparent.x_request_id, Some("req-123".to_string()));
392        assert_eq!(
393            traceparent.request_id,
394            Some("550e8400-e29b-41d4-a716-446655440000".to_string())
395        );
396    }
397
398    #[test]
399    fn test_shared_http_server_creation() {
400        use std::net::{IpAddr, Ipv4Addr};
401        let bind_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 0);
402        let token = CancellationToken::new();
403
404        let server = SharedHttpServer::new(bind_addr, token);
405        assert!(server.handlers.is_empty());
406    }
407
408    #[tokio::test]
409    async fn test_bind_and_start_assigns_os_port() {
410        use std::net::{IpAddr, Ipv4Addr};
411        let bind_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 0);
412        let token = CancellationToken::new();
413
414        let server = SharedHttpServer::new(bind_addr, token.clone());
415        let actual_addr = server.clone().bind_and_start().await.unwrap();
416
417        // OS should assign a non-zero port
418        assert_ne!(actual_addr.port(), 0);
419
420        // actual_address() should return the real bound address
421        assert_eq!(server.actual_address(), Some(actual_addr));
422
423        // address() should contain the real port
424        let addr_str =
425            <SharedHttpServer as super::unified_server::RequestPlaneServer>::address(&*server);
426        assert!(addr_str.contains(&actual_addr.port().to_string()));
427
428        token.cancel();
429    }
430
431    #[tokio::test]
432    async fn test_two_servers_get_different_ports() {
433        use std::net::{IpAddr, Ipv4Addr};
434        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 0);
435
436        let token1 = CancellationToken::new();
437        let token2 = CancellationToken::new();
438
439        let server1 = SharedHttpServer::new(addr, token1.clone());
440        let server2 = SharedHttpServer::new(addr, token2.clone());
441
442        let actual1 = server1.clone().bind_and_start().await.unwrap();
443        let actual2 = server2.clone().bind_and_start().await.unwrap();
444
445        // Two servers binding to port 0 must get different ports
446        assert_ne!(actual1.port(), actual2.port());
447
448        token1.cancel();
449        token2.cancel();
450    }
451
452    #[tokio::test]
453    async fn test_bind_and_start_with_explicit_port() {
454        use std::net::{IpAddr, Ipv4Addr};
455
456        // First bind to port 0 to get a free port
457        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
458        let free_port = listener.local_addr().unwrap().port();
459        drop(listener); // Release the port
460
461        let bind_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), free_port);
462        let token = CancellationToken::new();
463
464        let server = SharedHttpServer::new(bind_addr, token.clone());
465        let actual_addr = server.clone().bind_and_start().await.unwrap();
466
467        // When binding to an explicit port, actual should match
468        assert_eq!(actual_addr.port(), free_port);
469
470        token.cancel();
471    }
472}
dynamo_runtime/pipeline/network/ingress/http_endpoint.rs

dynamo_runtime/pipeline/network/ingress/
http_endpoint.rs
