endpoint.rs 11.2 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
19
    storage::key_value_store,
    traits::DistributedRuntimeProvider,
20
    transports::nats,
21
};
22

Ryan Olson's avatar
Ryan Olson committed
23
24
25
26
27
28
29
30
31
32
#[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>,
33
34
35
36
37

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

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

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

    /// 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
53
54
55
56
57
58
59
}

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

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

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

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

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

83
84
85
        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());
86
        // Add metrics to the handler. The endpoint provides additional information to the handler.
87
        handler.add_metrics(&endpoint, metrics_labels.as_deref())?;
88

89
90
91
92
93
94
95
96
        // 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");
97
98
99
100
101
102
103
104
105
            // 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,
            );
106
107
        }

108
109
110
        // 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();
111

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

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

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

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

141
142
143
144
145
146
147
148
149
150
151
        // 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
152

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

        // Create clones for the async closure
161
162
163
        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();
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
209
        // 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"
                );
            }
210
211
212

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

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

220
221
222
223
224
        // 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();

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

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

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

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

        Ok(())
    }
}
252

253
/// Build transport type based on request plane mode
254
///
255
256
257
258
259
/// 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)
260
261
fn build_transport_type(
    mode: RequestPlaneMode,
262
263
    endpoint_id: &EndpointId,
    connection_id: u64,
264
265
266
267
268
269
270
271
272
273
274
275
) -> 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!(
276
277
                "http://{http_host}:{http_port}{rpc_root}/{}",
                endpoint_id.name
278
279
280
281
282
283
284
285
286
287
288
            );

            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);

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

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