endpoint.rs 11.1 KB
Newer Older
1
2
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
Ryan Olson's avatar
Ryan Olson committed
3

4
5
6
7
8
use std::sync::Arc;

use anyhow::Result;
pub use async_nats::service::endpoint::Stats as EndpointStats;
use derive_builder::Builder;
Ryan Olson's avatar
Ryan Olson committed
9
use derive_getters::Dissolve;
10
use educe::Educe;
11
use tokio_util::sync::CancellationToken;
Ryan Olson's avatar
Ryan Olson committed
12

13
14
use crate::{
    component::{Endpoint, Instance, TransportType, service::EndpointStatsHandler},
15
    distributed::RequestPlaneMode,
16
    pipeline::network::{PushWorkHandler, ingress::push_endpoint::PushEndpoint},
17
    protocols::EndpointId,
18
    traits::DistributedRuntimeProvider,
19
    transports::nats,
20
};
21

Ryan Olson's avatar
Ryan Olson committed
22
23
24
25
26
27
28
29
30
31
#[derive(Educe, Builder, Dissolve)]
#[educe(Debug)]
#[builder(pattern = "owned", build_fn(private, name = "build_internal"))]
pub struct EndpointConfig {
    #[builder(private)]
    endpoint: Endpoint,

    /// Endpoint handler
    #[educe(Debug(ignore))]
    handler: Arc<dyn PushWorkHandler>,
32
33
34
35
36

    /// Stats handler
    #[educe(Debug(ignore))]
    #[builder(default, private)]
    _stats_handler: Option<EndpointStatsHandler>,
37

38
39
40
41
    /// Additional labels for metrics
    #[builder(default, setter(into))]
    metrics_labels: Option<Vec<(String, String)>>,

42
43
44
    /// Whether to wait for inflight requests to complete during shutdown
    #[builder(default = "true")]
    graceful_shutdown: bool,
45
46
47
48
49
50
51

    /// Health check payload for this endpoint
    /// This payload will be sent to the endpoint during health checks
    /// to verify it's responding properly
    #[educe(Debug(ignore))]
    #[builder(default, setter(into, strip_option))]
    health_check_payload: Option<serde_json::Value>,
Ryan Olson's avatar
Ryan Olson committed
52
53
54
55
56
57
58
}

impl EndpointConfigBuilder {
    pub(crate) fn from_endpoint(endpoint: Endpoint) -> Self {
        Self::default().endpoint(endpoint)
    }

59
60
    pub fn stats_handler<F>(self, handler: F) -> Self
    where
61
        F: FnMut(EndpointStats) -> serde_json::Value + Send + Sync + 'static,
62
63
64
65
    {
        self._stats_handler(Some(Box::new(handler)))
    }

Ryan Olson's avatar
Ryan Olson committed
66
    pub async fn start(self) -> Result<()> {
67
        let (
68
            endpoint,
69
70
71
72
73
74
            handler,
            stats_handler,
            metrics_labels,
            graceful_shutdown,
            health_check_payload,
        ) = self.build_internal()?.dissolve();
75
        let connection_id = endpoint.drt().connection_id();
76
        let endpoint_id = endpoint.id();
Ryan Olson's avatar
Ryan Olson committed
77

78
        tracing::debug!("Starting endpoint: {endpoint_id}");
Ryan Olson's avatar
Ryan Olson committed
79

80
81
        let service_name = endpoint.component.service_name();

82
83
84
        let metrics_labels: Option<Vec<(&str, &str)>> = metrics_labels
            .as_ref()
            .map(|v| v.iter().map(|(k, v)| (k.as_str(), v.as_str())).collect());
85
        // Add metrics to the handler. The endpoint provides additional information to the handler.
86
        handler.add_metrics(&endpoint, metrics_labels.as_deref())?;
87

88
89
90
91
92
93
94
95
        // Insert the stats handler. depends on NATS.
        if let Some(stats_handler) = stats_handler {
            let registry = endpoint.drt().component_registry().inner.lock().await;
            let handler_map = registry
                .stats_handlers
                .get(&service_name)
                .cloned()
                .expect("no stats handler registry; this is unexpected");
96
97
98
99
100
101
102
103
104
            // There is something wrong with the stats handler map I think.
            // Here the connection_id is included, but in component/service.rs add_stats_service it uses service_name,
            // no connection id so it's per-endpoint not per-instance. Doesn't match.
            // To not block current refactor I am keeping previous behavior, but I think needs
            // investigation.
            handler_map.lock().insert(
                nats::instance_subject(&endpoint_id, connection_id),
                stats_handler,
            );
105
106
        }

107
108
109
        // This creates a child token of the runtime's endpoint_shutdown_token. That token is
        // cancelled first as part of graceful shutdown. See Runtime::shutdown.
        let endpoint_shutdown_token = endpoint.drt().child_token();
110

111
        let system_health = endpoint.drt().system_health();
112

113
114
115
        let request_plane_mode = endpoint.drt().request_plane();
        tracing::info!("Endpoint starting with request plane mode: {request_plane_mode}",);

116
117
        // Register health check target in SystemHealth if provided
        if let Some(health_check_payload) = &health_check_payload {
118
            // Build transport based on request plane mode
119
            let transport = build_transport_type(request_plane_mode, &endpoint_id, connection_id);
120

121
            let instance = Instance {
122
123
124
                component: endpoint_id.component.clone(),
                endpoint: endpoint_id.name.clone(),
                namespace: endpoint_id.namespace.clone(),
125
                instance_id: connection_id,
126
                transport,
127
            };
128
            tracing::debug!(endpoint_name = %endpoint.name, "Registering endpoint health check target");
129
            let guard = system_health.lock();
130
            guard.register_health_check_target(
131
                &endpoint.name,
132
133
134
                instance,
                health_check_payload.clone(),
            );
135
            if let Some(notifier) = guard.get_endpoint_health_check_notifier(&endpoint.name) {
136
137
138
139
                handler.set_endpoint_health_check_notifier(notifier)?;
            }
        }

140
141
142
143
144
145
146
147
148
149
150
        // Register with graceful shutdown tracker if needed
        if graceful_shutdown {
            tracing::debug!(
                "Registering endpoint '{}' with graceful shutdown tracker",
                endpoint.name
            );
            let tracker = endpoint.drt().graceful_shutdown_tracker();
            tracker.register_endpoint();
        } else {
            tracing::debug!("Endpoint '{}' has graceful_shutdown=false", endpoint.name);
        }
Ryan Olson's avatar
Ryan Olson committed
151

152
        // Launch endpoint based on request plane mode
153
154
155
156
157
158
159
        let tracker_clone = if graceful_shutdown {
            Some(endpoint.drt().graceful_shutdown_tracker())
        } else {
            None
        };

        // Create clones for the async closure
160
161
162
        let namespace_name_for_task = endpoint_id.namespace.clone();
        let component_name_for_task = endpoint_id.component.clone();
        let endpoint_name_for_task = endpoint_id.name.clone();
163

164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
        // Get the unified request plane server (works for all transport types)
        let server = endpoint.drt().request_plane_server().await?;

        tracing::info!(
            endpoint = %endpoint_name_for_task,
            transport = server.transport_name(),
            "Registering endpoint with request plane server"
        );

        // Register endpoint with the server (unified interface)
        server
            .register_endpoint(
                endpoint_name_for_task.clone(),
                handler,
                connection_id,
                namespace_name_for_task.clone(),
                component_name_for_task.clone(),
                system_health.clone(),
            )
            .await?;

        // Create cleanup task that unregisters on cancellation
        let endpoint_name_for_cleanup = endpoint_name_for_task.clone();
        let server_for_cleanup = server.clone();
        let cancel_token_for_cleanup = endpoint_shutdown_token.clone();

        let task: tokio::task::JoinHandle<anyhow::Result<()>> = tokio::spawn(async move {
            cancel_token_for_cleanup.cancelled().await;

            tracing::debug!(
                endpoint = %endpoint_name_for_cleanup,
                "Unregistering endpoint from request plane server"
            );

            // Unregister from server
            if let Err(e) = server_for_cleanup
                .unregister_endpoint(&endpoint_name_for_cleanup)
                .await
            {
                tracing::warn!(
                    endpoint = %endpoint_name_for_cleanup,
                    error = %e,
                    "Failed to unregister endpoint"
                );
            }
209
210
211

            // Unregister from graceful shutdown tracker
            if let Some(tracker) = tracker_clone {
212
                tracing::debug!("Unregister endpoint from graceful shutdown tracker");
213
214
215
                tracker.unregister_endpoint();
            }

216
            anyhow::Ok(())
217
        });
Ryan Olson's avatar
Ryan Olson committed
218

219
220
221
222
223
        // Register this endpoint instance in the discovery plane
        // The discovery interface abstracts storage backend (etcd, k8s, etc) and provides
        // consistent registration/discovery across the system.
        let discovery = endpoint.drt().discovery();

224
        // Build transport for discovery service based on request plane mode
225
        let transport = build_transport_type(request_plane_mode, &endpoint_id, connection_id);
226

227
        let discovery_spec = crate::discovery::DiscoverySpec::Endpoint {
228
229
230
            namespace: endpoint_id.namespace.clone(),
            component: endpoint_id.component.clone(),
            endpoint: endpoint_id.name.clone(),
231
            transport,
Ryan Olson's avatar
Ryan Olson committed
232
233
        };

234
        if let Err(e) = discovery.register(discovery_spec).await {
235
            tracing::error!(
236
                %endpoint_id,
237
                error = %e,
238
239
                "Unable to register service for discovery"
            );
240
            endpoint_shutdown_token.cancel();
241
            anyhow::bail!(
242
                "Unable to register service for discovery. Check discovery service status"
243
            );
Ryan Olson's avatar
Ryan Olson committed
244
        }
245

Ryan Olson's avatar
Ryan Olson committed
246
247
248
249
250
        task.await??;

        Ok(())
    }
}
251

252
/// Build transport type based on request plane mode
253
///
254
255
256
257
258
/// This function handles both health check and discovery transport building.
/// All transport modes use consistent addressing:
/// - HTTP: Uses full URL path including endpoint name (e.g., http://host:port/v1/rpc/endpoint_name)
/// - TCP: Includes endpoint name for routing (e.g., host:port/endpoint_name)
/// - NATS: Uses subject-based addressing (unique per endpoint)
259
260
fn build_transport_type(
    mode: RequestPlaneMode,
261
262
    endpoint_id: &EndpointId,
    connection_id: u64,
263
264
265
266
267
268
269
270
271
272
273
274
) -> TransportType {
    match mode {
        RequestPlaneMode::Http => {
            let http_host = crate::utils::get_http_rpc_host_from_env();
            let http_port = std::env::var("DYN_HTTP_RPC_PORT")
                .ok()
                .and_then(|p| p.parse::<u16>().ok())
                .unwrap_or(8888);
            let rpc_root =
                std::env::var("DYN_HTTP_RPC_ROOT_PATH").unwrap_or_else(|_| "/v1/rpc".to_string());

            let http_endpoint = format!(
275
276
                "http://{http_host}:{http_port}{rpc_root}/{}",
                endpoint_id.name
277
278
279
280
281
282
283
284
285
286
287
            );

            TransportType::Http(http_endpoint)
        }
        RequestPlaneMode::Tcp => {
            let tcp_host = crate::utils::get_tcp_rpc_host_from_env();
            let tcp_port = std::env::var("DYN_TCP_RPC_PORT")
                .ok()
                .and_then(|p| p.parse::<u16>().ok())
                .unwrap_or(9999);

288
289
            // Include endpoint name for proper TCP routing
            // TCP client parses this format and adds x-endpoint-path header for server-side routing
290
            let tcp_endpoint = format!("{}:{}/{}", tcp_host, tcp_port, endpoint_id.name);
291
292
293

            TransportType::Tcp(tcp_endpoint)
        }
294
295
296
        RequestPlaneMode::Nats => {
            TransportType::Nats(nats::instance_subject(endpoint_id, connection_id))
        }
297
298
    }
}