Line | Count | Source |
1 | | // Copyright 2026 Trace Machina, Inc. All rights reserved. |
2 | | // |
3 | | // Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); |
4 | | // you may not use this file except in compliance with the License. |
5 | | // You may obtain a copy of the License at |
6 | | // |
7 | | // See LICENSE file for details |
8 | | // |
9 | | // Unless required by applicable law or agreed to in writing, software |
10 | | // distributed under the License is distributed on an "AS IS" BASIS, |
11 | | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | | // See the License for the specific language governing permissions and |
13 | | // limitations under the License. |
14 | | |
15 | | // The worker's action future is a long chain of combinators; proving it |
16 | | // Send for the spawn walks the whole chain, and the default limit of 128 is |
17 | | // a few steps short of it on the sanitizer toolchain. |
18 | | #![recursion_limit = "256"] |
19 | | |
20 | | use core::net::SocketAddr; |
21 | | use core::time::Duration; |
22 | | use std::collections::{HashMap, HashSet}; |
23 | | use std::io::ErrorKind; |
24 | | use std::pin::Pin; |
25 | | use std::sync::Arc; |
26 | | |
27 | | use async_lock::Mutex as AsyncMutex; |
28 | | use axum::Router; |
29 | | use axum::http::Uri; |
30 | | use futures::FutureExt; |
31 | | use futures::future::{BoxFuture, Either, OptionFuture, TryFutureExt}; |
32 | | use hyper::StatusCode; |
33 | | use hyper_util::rt::TokioTimer; |
34 | | use hyper_util::rt::tokio::TokioIo; |
35 | | use hyper_util::server::conn::auto; |
36 | | use hyper_util::service::TowerToHyperService; |
37 | | use nativelink_config::cas_server::{ |
38 | | CasConfig, CasStoreConfig, HttpCompressionAlgorithm, ListenerConfig, SchedulerConfig, |
39 | | ServerConfig, StoreConfig, WithInstanceName, WorkerConfig, |
40 | | }; |
41 | | use nativelink_error::{Code, Error, ResultExt, make_err, make_input_err}; |
42 | | use nativelink_scheduler::default_scheduler_factory::scheduler_factory; |
43 | | use nativelink_service::ac_server::AcServer; |
44 | | use nativelink_service::bep_server::BepServer; |
45 | | use nativelink_service::bytestream_server::ByteStreamServer; |
46 | | use nativelink_service::capabilities_server::CapabilitiesServer; |
47 | | use nativelink_service::cas_server::CasServer; |
48 | | use nativelink_service::execution_server::ExecutionServer; |
49 | | use nativelink_service::fetch_server::FetchServer; |
50 | | use nativelink_service::health_server::{HealthServer, health_paths}; |
51 | | use nativelink_service::push_server::PushServer; |
52 | | use nativelink_service::wire_compression::RemoteCacheCompressionInstances; |
53 | | use nativelink_service::worker_api_server::WorkerApiServer; |
54 | | use nativelink_store::default_store_factory::store_factory; |
55 | | use nativelink_store::store_manager::StoreManager; |
56 | | use nativelink_util::health_utils::HealthRegistryBuilder; |
57 | | use nativelink_util::origin_event_publisher::OriginEventPublisher; |
58 | | use nativelink_util::shutdown_guard::ShutdownGuard; |
59 | | use nativelink_util::task::TaskExecutor; |
60 | | use nativelink_util::{background_spawn, fs, spawn}; |
61 | | use nativelink_worker::local_worker::{WorkerRegistration, new_local_worker}; |
62 | | use rustls_pki_types::pem::PemObject; |
63 | | use rustls_pki_types::{CertificateRevocationListDer, PrivateKeyDer}; |
64 | | use tokio::net::{TcpListener, TcpSocket, TcpStream}; |
65 | | use tokio::select; |
66 | | use tokio::sync::oneshot::Sender; |
67 | | use tokio::sync::{broadcast, mpsc, oneshot}; |
68 | | use tokio_rustls::TlsAcceptor; |
69 | | use tokio_rustls::rustls::pki_types::CertificateDer; |
70 | | use tokio_rustls::rustls::server::WebPkiClientVerifier; |
71 | | use tokio_rustls::rustls::{RootCertStore, ServerConfig as TlsServerConfig}; |
72 | | use tonic::codec::CompressionEncoding; |
73 | | use tonic::service::Routes; |
74 | | use tracing::{error, error_span, info, trace_span, warn}; |
75 | | |
76 | | // Note: This must be kept in sync with the documentation in |
77 | | // `OriginEventsConfig::max_event_queue_size`. |
78 | | const DEFAULT_MAX_QUEUE_EVENTS: usize = 0x0001_0000; |
79 | | |
80 | | /// If this value changes update the documentation in the config definition. |
81 | | const DEFAULT_MAX_DECODING_MESSAGE_SIZE: usize = 4 * 1024 * 1024; |
82 | | |
83 | | /// Note: This must be kept in sync with the documentation in `AdminConfig::path`. |
84 | | const DEFAULT_ADMIN_API_PATH: &str = "/admin"; |
85 | | |
86 | | /// Bind a [`TcpListener`] with `IP_FREEBIND` set. |
87 | 0 | fn bind_freebind(socket_addr: SocketAddr) -> Result<TcpListener, std::io::Error> { |
88 | 0 | let socket = match socket_addr { |
89 | 0 | SocketAddr::V4(_) => TcpSocket::new_v4(), |
90 | 0 | SocketAddr::V6(_) => TcpSocket::new_v6(), |
91 | 0 | }?; |
92 | 0 | fs::set_freebind(&socket)?; |
93 | 0 | socket.bind(socket_addr)?; |
94 | 0 | socket.listen(1024) |
95 | 0 | } |
96 | | |
97 | | /// Configure accepted sockets before either the TLS handshake or HTTP serving. |
98 | 8 | async fn accept_tcp_connection( |
99 | 8 | listener: &TcpListener, |
100 | 8 | tcp_nodelay: bool, |
101 | 8 | ) -> Result<(TcpStream, SocketAddr), Error> { |
102 | 8 | let (stream7 , remote_addr7 ) = listener |
103 | 8 | .accept() |
104 | 8 | .await |
105 | 7 | .err_tip(|| "Failed to accept TCP connection")?0 ; |
106 | 7 | stream |
107 | 7 | .set_nodelay(tcp_nodelay) |
108 | 7 | .err_tip(|| format!0 ("Failed to set TCP_NODELAY to {tcp_nodelay} for {remote_addr}"))?0 ; |
109 | 7 | Ok((stream, remote_addr)) |
110 | 7 | } |
111 | | |
112 | | /// A JSON body with its content type, for the admin API's listings. |
113 | 0 | const fn json_response( |
114 | 0 | body: String, |
115 | 0 | ) -> ([(axum::http::header::HeaderName, &'static str); 1], String) { |
116 | 0 | ( |
117 | 0 | [(axum::http::header::CONTENT_TYPE, "application/json")], |
118 | 0 | body, |
119 | 0 | ) |
120 | 0 | } |
121 | | |
122 | | macro_rules! service_setup { |
123 | | ($service: expr, $http_config: ident) => {{ |
124 | | let mut service = $service; |
125 | | let max_decoding_message_size = if $http_config.max_decoding_message_size == 0 { |
126 | | DEFAULT_MAX_DECODING_MESSAGE_SIZE |
127 | | } else { |
128 | | $http_config.max_decoding_message_size |
129 | | }; |
130 | | service = service.max_decoding_message_size(max_decoding_message_size); |
131 | | let send_algo = &$http_config.compression.send_compression_algorithm; |
132 | | if let Some(encoding) = into_encoding(send_algo.unwrap_or(HttpCompressionAlgorithm::None)) { |
133 | | service = service.send_compressed(encoding); |
134 | | } |
135 | | for encoding in $http_config |
136 | | .compression |
137 | | .accepted_compression_algorithms |
138 | | .iter() |
139 | | // Filter None values. |
140 | 0 | .filter_map(|from: &HttpCompressionAlgorithm| into_encoding(*from)) |
141 | | { |
142 | | service = service.accept_compressed(encoding); |
143 | | } |
144 | | service |
145 | | }}; |
146 | | } |
147 | | |
148 | | trait RoutesExt { |
149 | | fn add_optional_service<S>(self, svc: Option<S>) -> Self |
150 | | where |
151 | | S: tower::Service< |
152 | | axum::http::Request<tonic::body::Body>, |
153 | | Error = core::convert::Infallible, |
154 | | > + tonic::server::NamedService |
155 | | + Clone |
156 | | + Send |
157 | | + Sync |
158 | | + 'static, |
159 | | S::Response: axum::response::IntoResponse, |
160 | | S::Future: Send + 'static; |
161 | | } |
162 | | |
163 | | impl RoutesExt for Routes { |
164 | 10 | fn add_optional_service<S>(mut self, svc: Option<S>) -> Self |
165 | 10 | where |
166 | 10 | S: tower::Service< |
167 | 10 | axum::http::Request<tonic::body::Body>, |
168 | 10 | Error = core::convert::Infallible, |
169 | 10 | > + tonic::server::NamedService |
170 | 10 | + Clone |
171 | 10 | + Send |
172 | 10 | + Sync |
173 | 10 | + 'static, |
174 | 10 | S::Response: axum::response::IntoResponse, |
175 | 10 | S::Future: Send + 'static, |
176 | | { |
177 | 10 | if let Some(svc0 ) = svc { |
178 | 0 | self = self.add_service(svc); |
179 | 10 | } |
180 | 10 | self |
181 | 10 | } |
182 | | } |
183 | | |
184 | 0 | const fn into_encoding(from: HttpCompressionAlgorithm) -> Option<CompressionEncoding> { |
185 | 0 | match from { |
186 | 0 | HttpCompressionAlgorithm::Gzip => Some(CompressionEncoding::Gzip), |
187 | 0 | HttpCompressionAlgorithm::None => None, |
188 | | } |
189 | 0 | } |
190 | | |
191 | | pub struct Server { |
192 | | #[allow(clippy::type_complexity)] |
193 | | pub root_futures: Vec<Pin<Box<dyn Future<Output = Result<(), Error>> + Send>>>, |
194 | | pub single_use_enabled: bool, |
195 | | pub single_use_complete_rx: oneshot::Receiver<()>, |
196 | | } |
197 | | |
198 | | impl std::fmt::Debug for Server { |
199 | 0 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> { |
200 | 0 | f.debug_struct("Server") |
201 | 0 | .field("single_use_enabled", &self.single_use_enabled) |
202 | 0 | .field("single_use_complete_rx", &self.single_use_complete_rx) |
203 | 0 | .finish_non_exhaustive() |
204 | 0 | } |
205 | | } |
206 | | |
207 | 1 | pub async fn build_server( |
208 | 1 | cfg: CasConfig, |
209 | 1 | shutdown_tx: broadcast::Sender<ShutdownGuard>, |
210 | 1 | scheduler_shutdown_tx: Sender<()>, |
211 | 1 | ) -> Result<Server, Error> { |
212 | 1 | let health_registry_builder = |
213 | 1 | Arc::new(AsyncMutex::new(HealthRegistryBuilder::new("nativelink"))); |
214 | 1 | let mut worker_registrations: Vec<Arc<WorkerRegistration>> = Vec::new(); |
215 | | |
216 | 1 | let store_manager = Arc::new(StoreManager::new()); |
217 | | { |
218 | 1 | let mut health_registry_lock = health_registry_builder.lock().await; |
219 | | |
220 | 1 | for StoreConfig { name0 , spec0 } in cfg.stores { |
221 | 0 | let health_component_name = format!("stores/{name}"); |
222 | 0 | let mut health_register_store = |
223 | 0 | health_registry_lock.sub_builder(&health_component_name); |
224 | 0 | let store = store_factory(&spec, &store_manager, Some(&mut health_register_store)) |
225 | 0 | .await |
226 | 0 | .err_tip(|| format!("Failed to create store '{name}'"))?; |
227 | 0 | store_manager |
228 | 0 | .add_store(&name, store) |
229 | 0 | .err_tip(|| format!("Failed to add store '{name}'"))?; |
230 | | } |
231 | 1 | store_manager.run_post_init().await?0 ; |
232 | | |
233 | | // Workers start after the listeners are up, but the health registry |
234 | | // is built with the listeners, so their registration flags are |
235 | | // made and registered here and handed to the workers later. |
236 | 1 | let mut workers_health = health_registry_lock.sub_builder("workers"); |
237 | 1 | for (i0 , worker_cfg0 ) in cfg.workers.iter().flatten().enumerate() { |
238 | 0 | let WorkerConfig::Local(local_worker_cfg) = worker_cfg; |
239 | 0 | let name = if local_worker_cfg.name.is_empty() { |
240 | 0 | format!("worker_{i}") |
241 | | } else { |
242 | 0 | local_worker_cfg.name.clone() |
243 | | }; |
244 | 0 | let registration = WorkerRegistration::new(&name); |
245 | 0 | workers_health |
246 | 0 | .sub_builder(&name) |
247 | 0 | .register_indicator(registration.clone()); |
248 | 0 | worker_registrations.push(registration); |
249 | | } |
250 | | } |
251 | | |
252 | 1 | let mut root_futures: Vec<BoxFuture<Result<(), Error>>> = Vec::new(); |
253 | 1 | let (single_use_complete_tx, single_use_complete_rx) = oneshot::channel(); |
254 | 1 | let mut single_use_complete_tx = Some(single_use_complete_tx); |
255 | 1 | let single_use_enabled = cfg.workers.as_ref().is_some_and(|workers| {0 |
256 | 0 | workers.iter().any(|worker| match worker { |
257 | 0 | WorkerConfig::Local(worker) => worker.single_use, |
258 | 0 | }) |
259 | 0 | }); |
260 | | |
261 | 1 | let maybe_origin_event_tx = cfg |
262 | 1 | .experimental_origin_events |
263 | 1 | .as_ref() |
264 | 1 | .map(|origin_events_cfg| {0 |
265 | 0 | let mut max_queued_events = origin_events_cfg.max_event_queue_size; |
266 | 0 | if max_queued_events == 0 { |
267 | 0 | max_queued_events = DEFAULT_MAX_QUEUE_EVENTS; |
268 | 0 | } |
269 | 0 | let (tx, rx) = mpsc::channel(max_queued_events); |
270 | 0 | let store_name = origin_events_cfg.publisher.store.as_str(); |
271 | 0 | let store = store_manager.get_store(store_name).err_tip(|| { |
272 | 0 | format!("Could not get store {store_name} for origin event publisher") |
273 | 0 | })?; |
274 | | |
275 | 0 | root_futures.push(Box::pin( |
276 | 0 | OriginEventPublisher::new(store, rx, shutdown_tx.clone()) |
277 | 0 | .run() |
278 | 0 | .map(Ok), |
279 | 0 | )); |
280 | | |
281 | 0 | Ok::<_, Error>(tx) |
282 | 0 | }) |
283 | 1 | .transpose()?0 ; |
284 | | |
285 | 1 | let mut action_schedulers = HashMap::new(); |
286 | 1 | let mut worker_schedulers = HashMap::new(); |
287 | 1 | let mut resource_sketches = HashMap::new(); |
288 | 1 | for SchedulerConfig { name0 , spec0 } in cfg.schedulers.iter().flatten() { |
289 | 0 | let (maybe_action_scheduler, maybe_worker_scheduler, maybe_resource_sketch) = |
290 | 0 | scheduler_factory(spec, &store_manager, maybe_origin_event_tx.as_ref()) |
291 | 0 | .await |
292 | 0 | .err_tip(|| format!("Failed to create scheduler '{name}'"))?; |
293 | 0 | if let Some(action_scheduler) = maybe_action_scheduler { |
294 | 0 | action_schedulers.insert(name.clone(), action_scheduler.clone()); |
295 | 0 | } |
296 | 0 | if let Some(worker_scheduler) = maybe_worker_scheduler { |
297 | 0 | worker_schedulers.insert(name.clone(), worker_scheduler.clone()); |
298 | 0 | } |
299 | 0 | if let Some(resource_sketch) = maybe_resource_sketch { |
300 | 0 | resource_sketches.insert(name.clone(), resource_sketch); |
301 | 0 | } |
302 | | } |
303 | | |
304 | 1 | let server_cfgs: Vec<ServerConfig> = cfg.servers.into_iter().collect(); |
305 | | |
306 | | // The capabilities service advertises chunking support for CAS instances |
307 | | // that may be served from a different server block (e.g. behind an L7 |
308 | | // router), so collect the CAS configs across all blocks. |
309 | 1 | let all_cas_configs: Vec<WithInstanceName<CasStoreConfig>> = server_cfgs |
310 | 1 | .iter() |
311 | 1 | .filter_map(|server_cfg| server_cfg.services.as_ref()) |
312 | 1 | .filter_map(|services| services.cas.as_deref()) |
313 | 1 | .flatten() |
314 | 1 | .cloned() |
315 | 1 | .collect(); |
316 | | |
317 | 1 | for server_cfg in server_cfgs { |
318 | 1 | let services = server_cfg |
319 | 1 | .services |
320 | 1 | .err_tip(|| "'services' must be configured")?0 ; |
321 | | |
322 | | // Currently we only support http as our socket type. |
323 | 1 | let ListenerConfig::Http(http_config) = server_cfg.listener; |
324 | | |
325 | 1 | let execution_server = services |
326 | 1 | .execution |
327 | 1 | .as_ref() |
328 | 1 | .map(|cfg| ExecutionServer::new0 (cfg0 , &action_schedulers0 , &store_manager0 )) |
329 | 1 | .transpose() |
330 | 1 | .err_tip(|| "Could not create Execution service")?0 ; |
331 | | |
332 | 1 | let capabilities_configs = services.capabilities.as_deref().unwrap_or_default(); |
333 | 1 | let remote_cache_compression_instances = |
334 | 1 | RemoteCacheCompressionInstances::from_capabilities_configs(capabilities_configs); |
335 | | |
336 | 1 | let tonic_services = Routes::builder() |
337 | 1 | .routes() |
338 | 1 | .add_optional_service( |
339 | 1 | services |
340 | 1 | .ac |
341 | 1 | .map_or(Ok(None), |cfg| {0 |
342 | 0 | AcServer::new(&cfg, &store_manager) |
343 | 0 | .map(|v| Some(service_setup!(v.into_service(), http_config))) |
344 | 0 | }) |
345 | 1 | .err_tip(|| "Could not create AC service")?0 , |
346 | | ) |
347 | 1 | .add_optional_service( |
348 | 1 | services |
349 | 1 | .cas |
350 | 1 | .as_deref() |
351 | 1 | .map_or(Ok(None), |cfg| {0 |
352 | 0 | CasServer::new(cfg, &store_manager, &remote_cache_compression_instances) |
353 | 0 | .map(|v| Some(service_setup!(v.into_service(), http_config))) |
354 | 0 | }) |
355 | 1 | .err_tip(|| "Could not create CAS service")?0 , |
356 | | ) |
357 | 1 | .add_optional_service( |
358 | 1 | execution_server |
359 | 1 | .clone() |
360 | 1 | .map(|v| service_setup!0 (v0 .into_service0 (), http_config)), |
361 | | ) |
362 | 1 | .add_optional_service( |
363 | 1 | execution_server.map(|v| service_setup!0 (v0 .into_operations_service0 (), http_config)), |
364 | | ) |
365 | 1 | .add_optional_service( |
366 | 1 | services |
367 | 1 | .fetch |
368 | 1 | .map_or(Ok(None), |cfg| {0 |
369 | 0 | FetchServer::new(&cfg, &store_manager) |
370 | 0 | .map(|v| Some(service_setup!(v.into_service(), http_config))) |
371 | 0 | }) |
372 | 1 | .err_tip(|| "Could not create Fetch service")?0 , |
373 | | ) |
374 | 1 | .add_optional_service( |
375 | 1 | services |
376 | 1 | .push |
377 | 1 | .map_or(Ok(None), |cfg| {0 |
378 | 0 | PushServer::new(&cfg, &store_manager) |
379 | 0 | .map(|v| Some(service_setup!(v.into_service(), http_config))) |
380 | 0 | }) |
381 | 1 | .err_tip(|| "Could not create Push service")?0 , |
382 | | ) |
383 | 1 | .add_optional_service( |
384 | 1 | services |
385 | 1 | .bytestream |
386 | 1 | .map_or(Ok(None), |cfg| {0 |
387 | 0 | ByteStreamServer::new( |
388 | 0 | &cfg, |
389 | 0 | &store_manager, |
390 | 0 | &remote_cache_compression_instances, |
391 | | ) |
392 | 0 | .map(|v| Some(service_setup!(v.into_service(), http_config))) |
393 | 0 | }) |
394 | 1 | .err_tip(|| "Could not create ByteStream service")?0 , |
395 | | ) |
396 | 1 | .add_optional_service( |
397 | 1 | OptionFuture::from(services.capabilities.as_ref().map(|cfg| {0 |
398 | 0 | CapabilitiesServer::new( |
399 | 0 | cfg, |
400 | 0 | &action_schedulers, |
401 | 0 | &remote_cache_compression_instances, |
402 | 0 | &all_cas_configs, |
403 | | ) |
404 | 0 | })) |
405 | 1 | .await |
406 | 1 | .map_or(Ok::<Option<CapabilitiesServer>, Error>(None), |server| {0 |
407 | 0 | Ok(Some(server?)) |
408 | 0 | }) |
409 | 1 | .err_tip(|| "Could not create Capabilities service")?0 |
410 | 1 | .map(|v| service_setup!0 (v0 .into_service0 (), http_config)), |
411 | | ) |
412 | 1 | .add_optional_service( |
413 | 1 | services |
414 | 1 | .worker_api |
415 | 1 | .map_or(Ok(None), |cfg| {0 |
416 | 0 | WorkerApiServer::new(&cfg, &worker_schedulers) |
417 | 0 | .map(|v| Some(service_setup!(v.into_service(), http_config))) |
418 | 0 | }) |
419 | 1 | .err_tip(|| "Could not create WorkerApi service")?0 , |
420 | | ) |
421 | 1 | .add_optional_service( |
422 | 1 | services |
423 | 1 | .experimental_bep |
424 | 1 | .map_or(Ok(None), |cfg| {0 |
425 | 0 | BepServer::new(&cfg, &store_manager) |
426 | 0 | .map(|v| Some(service_setup!(v.into_service(), http_config))) |
427 | 0 | }) |
428 | 1 | .err_tip(|| "Could not create BEP service")?0 , |
429 | | ); |
430 | | |
431 | 1 | let health_registry = health_registry_builder.lock().await.build(); |
432 | | |
433 | 1 | let mut svc = |
434 | 1 | tonic_services |
435 | 1 | .into_axum_router() |
436 | 1 | .layer(nativelink_util::telemetry::OtlpLayer::new( |
437 | 1 | server_cfg.experimental_identity_header.required, |
438 | | )); |
439 | | |
440 | 1 | if let Some(health_cfg) = services.health { |
441 | 1 | let (path, readiness_path) = health_paths(&health_cfg)?0 ; |
442 | 1 | svc = svc.route_service( |
443 | 1 | &path, |
444 | 1 | HealthServer::new(health_registry.clone(), &health_cfg), |
445 | | ); |
446 | 1 | svc = svc.route_service( |
447 | 1 | &readiness_path, |
448 | 1 | HealthServer::readiness(health_registry, &health_cfg), |
449 | | ); |
450 | 0 | } |
451 | | |
452 | 1 | if let Some(admin_config0 ) = services.admin { |
453 | 0 | let path = if admin_config.path.is_empty() { |
454 | 0 | DEFAULT_ADMIN_API_PATH |
455 | | } else { |
456 | 0 | &admin_config.path |
457 | | }; |
458 | 0 | let worker_schedulers = Arc::new(worker_schedulers.clone()); |
459 | 0 | let listing_worker_schedulers = worker_schedulers.clone(); |
460 | 0 | let listing_action_schedulers = Arc::new(action_schedulers.clone()); |
461 | 0 | let listing_resource_sketches = Arc::new(resource_sketches.clone()); |
462 | | // Read-only views for the provisioner and operators: what the |
463 | | // scheduler has connected and what it has queued, with the |
464 | | // reservations the queued actions carry. |
465 | 0 | let admin_router = Router::new() |
466 | 0 | .route( |
467 | 0 | "/scheduler/{instance_name}/workers", |
468 | 0 | axum::routing::get(move |params: axum::extract::Path<String>| async move { |
469 | 0 | let instance_name = params.0; |
470 | 0 | let Some(scheduler) = listing_worker_schedulers.get(&instance_name) else { |
471 | 0 | return Err(( |
472 | 0 | StatusCode::NOT_FOUND, |
473 | 0 | format!("No scheduler named '{instance_name}'"), |
474 | 0 | )); |
475 | | }; |
476 | 0 | let workers = scheduler.worker_snapshot().await; |
477 | 0 | nativelink_scheduler::admin::workers_json(&workers) |
478 | 0 | .map(json_response) |
479 | 0 | .map_err(|e| { |
480 | 0 | (StatusCode::INTERNAL_SERVER_ERROR, format!("Error: {e:?}")) |
481 | 0 | }) |
482 | 0 | }), |
483 | | ) |
484 | 0 | .route( |
485 | 0 | "/scheduler/{instance_name}/demand", |
486 | 0 | axum::routing::get(move |params: axum::extract::Path<String>| async move { |
487 | 0 | let instance_name = params.0; |
488 | 0 | let Some(scheduler) = listing_action_schedulers.get(&instance_name) else { |
489 | 0 | return Err(( |
490 | 0 | StatusCode::NOT_FOUND, |
491 | 0 | format!("No scheduler named '{instance_name}'"), |
492 | 0 | )); |
493 | | }; |
494 | 0 | nativelink_scheduler::admin::queued_demand_json(scheduler.as_ref()) |
495 | 0 | .await |
496 | 0 | .map(json_response) |
497 | 0 | .map_err(|e| { |
498 | 0 | (StatusCode::INTERNAL_SERVER_ERROR, format!("Error: {e:?}")) |
499 | 0 | }) |
500 | 0 | }), |
501 | | ) |
502 | | // What the scheduler reserves for each kind of action it has |
503 | | // measured, in the hints-file shape, so the UI shows what the |
504 | | // scheduler uses and a new fleet can be seeded from an old one. |
505 | | // `family=target|mnemonic|command` and `limit=N` narrow it. |
506 | 0 | .route( |
507 | 0 | "/scheduler/{instance_name}/hints", |
508 | 0 | axum::routing::get( |
509 | | move |params: axum::extract::Path<String>, |
510 | 0 | axum::extract::RawQuery(query): axum::extract::RawQuery| async move { |
511 | 0 | let instance_name = params.0; |
512 | | // Two plain parameters; no need for the query feature. |
513 | 0 | let param = |name: &str| -> Option<String> { |
514 | 0 | query.as_deref()?.split('&').find_map(|pair| { |
515 | 0 | let (key, value) = pair.split_once('=')?; |
516 | 0 | (key == name).then(|| value.to_string()) |
517 | 0 | }) |
518 | 0 | }; |
519 | 0 | let Some(sketch) = listing_resource_sketches.get(&instance_name) |
520 | | else { |
521 | 0 | return Err(( |
522 | 0 | StatusCode::NOT_FOUND, |
523 | 0 | format!("No historical_resource scheduler named '{instance_name}'"), |
524 | 0 | )); |
525 | | }; |
526 | 0 | nativelink_scheduler::admin::hints_json( |
527 | 0 | sketch, |
528 | 0 | param("family").as_deref(), |
529 | 0 | param("limit").and_then(|limit| limit.parse().ok()), |
530 | | ) |
531 | 0 | .map(json_response) |
532 | 0 | .map_err(|e| (StatusCode::BAD_REQUEST, format!("Error: {e:?}"))) |
533 | 0 | }, |
534 | | ), |
535 | | ); |
536 | 0 | svc = svc.nest_service( |
537 | 0 | path, |
538 | 0 | admin_router.route( |
539 | 0 | "/scheduler/{instance_name}/set_drain_worker/{worker_id}/{is_draining}", |
540 | 0 | axum::routing::post( |
541 | 0 | move |params: axum::extract::Path<(String, String, String)>| async move { |
542 | 0 | let (instance_name, worker_id, is_draining) = params.0; |
543 | 0 | (async move { |
544 | 0 | let is_draining = match is_draining.as_str() { |
545 | 0 | "0" => false, |
546 | 0 | "1" => true, |
547 | | _ => { |
548 | 0 | return Err(make_err!( |
549 | 0 | Code::Internal, |
550 | 0 | "{} is neither 0 nor 1", |
551 | 0 | is_draining |
552 | 0 | )); |
553 | | } |
554 | | }; |
555 | 0 | worker_schedulers |
556 | 0 | .get(&instance_name) |
557 | 0 | .err_tip(|| { |
558 | 0 | format!( |
559 | | "Can not get an instance with the name of '{instance_name}'", |
560 | | ) |
561 | 0 | })? |
562 | 0 | .clone() |
563 | 0 | .set_drain_worker(&worker_id.clone().into(), is_draining) |
564 | 0 | .await?; |
565 | 0 | Ok::<_, Error>(format!("Draining worker {worker_id}")) |
566 | | }) |
567 | 0 | .await |
568 | 0 | .map_err(|e| { |
569 | 0 | Err::<String, _>(( |
570 | 0 | StatusCode::INTERNAL_SERVER_ERROR, |
571 | 0 | format!("Error: {e:?}"), |
572 | 0 | )) |
573 | 0 | }) |
574 | 0 | }, |
575 | | ), |
576 | | ), |
577 | | ); |
578 | 1 | } |
579 | | |
580 | | // This is the default service that executes if no other endpoint matches. |
581 | 1 | svc = svc.fallback(|uri: Uri| async move {0 |
582 | 0 | warn!("No route for {uri}"); |
583 | 0 | (StatusCode::NOT_FOUND, format!("No route for {uri}")) |
584 | 0 | }); |
585 | | |
586 | | // Configure our TLS acceptor if we have TLS configured. |
587 | 1 | let maybe_tls_acceptor = http_config.tls.map_or(Ok(None), |tls_config| {0 |
588 | 0 | fn read_cert(cert_file: &str) -> Result<Vec<CertificateDer<'static>>, Error> { |
589 | 0 | let mut cert_reader = std::io::BufReader::new( |
590 | 0 | std::fs::File::open(cert_file) |
591 | 0 | .err_tip(|| format!("Could not open cert file {cert_file}"))?, |
592 | | ); |
593 | 0 | let certs = CertificateDer::pem_reader_iter(&mut cert_reader) |
594 | 0 | .collect::<Result<Vec<CertificateDer<'_>>, _>>() |
595 | 0 | .err_tip(|| format!("Could not extract certs from file {cert_file}"))?; |
596 | 0 | Ok(certs) |
597 | 0 | } |
598 | 0 | let certs = read_cert(&tls_config.cert_file)?; |
599 | 0 | let mut key_reader = std::io::BufReader::new( |
600 | 0 | std::fs::File::open(&tls_config.key_file) |
601 | 0 | .err_tip(|| format!("Could not open key file {}", tls_config.key_file))?, |
602 | | ); |
603 | 0 | let key = match PrivateKeyDer::from_pem_reader(&mut key_reader) |
604 | 0 | .err_tip(|| format!("Could not extract key(s) from file {}", tls_config.key_file))? |
605 | | { |
606 | 0 | PrivateKeyDer::Pkcs8(key) => key.into(), |
607 | 0 | PrivateKeyDer::Sec1(key) => key.into(), |
608 | 0 | PrivateKeyDer::Pkcs1(key) => key.into(), |
609 | | _ => { |
610 | 0 | return Err(make_err!( |
611 | 0 | Code::Internal, |
612 | 0 | "No keys found in file {}", |
613 | 0 | tls_config.key_file |
614 | 0 | )); |
615 | | } |
616 | | }; |
617 | 0 | if PrivateKeyDer::from_pem_reader(&mut key_reader).is_ok() { |
618 | 0 | return Err(make_err!( |
619 | 0 | Code::InvalidArgument, |
620 | 0 | "Expected 1 key in file {}", |
621 | 0 | tls_config.key_file |
622 | 0 | )); |
623 | 0 | } |
624 | 0 | let verifier = if let Some(client_ca_file) = &tls_config.client_ca_file { |
625 | 0 | let mut client_auth_roots = RootCertStore::empty(); |
626 | 0 | for cert in read_cert(client_ca_file)? { |
627 | 0 | client_auth_roots.add(cert).map_err(|e| { |
628 | 0 | Error::from_std_err(Code::Internal, &e).append("Could not read client CA") |
629 | 0 | })?; |
630 | | } |
631 | 0 | let crls = if let Some(client_crl_file) = &tls_config.client_crl_file { |
632 | 0 | let mut crl_reader = std::io::BufReader::new( |
633 | 0 | std::fs::File::open(client_crl_file) |
634 | 0 | .err_tip(|| format!("Could not open CRL file {client_crl_file}"))?, |
635 | | ); |
636 | 0 | CertificateRevocationListDer::pem_reader_iter(&mut crl_reader) |
637 | 0 | .collect::<Result<_, _>>() |
638 | 0 | .err_tip(|| format!("Could not extract CRLs from file {client_crl_file}"))? |
639 | | } else { |
640 | 0 | Vec::new() |
641 | | }; |
642 | 0 | WebPkiClientVerifier::builder(Arc::new(client_auth_roots)) |
643 | 0 | .with_crls(crls) |
644 | 0 | .build() |
645 | 0 | .map_err(|e| { |
646 | 0 | Error::from_std_err(Code::Internal, &e) |
647 | 0 | .append("Could not create WebPkiClientVerifier") |
648 | 0 | })? |
649 | | } else { |
650 | 0 | WebPkiClientVerifier::no_client_auth() |
651 | | }; |
652 | 0 | let mut config = TlsServerConfig::builder() |
653 | 0 | .with_client_cert_verifier(verifier) |
654 | 0 | .with_single_cert(certs, key) |
655 | 0 | .map_err(|e| { |
656 | 0 | Error::from_std_err(Code::Internal, &e) |
657 | 0 | .append("Could not create TlsServerConfig") |
658 | 0 | })?; |
659 | | |
660 | 0 | config.alpn_protocols.push("h2".into()); |
661 | 0 | Ok(Some(TlsAcceptor::from(Arc::new(config)))) |
662 | 0 | })?; |
663 | | |
664 | 1 | let socket_addr = http_config |
665 | 1 | .socket_address |
666 | 1 | .parse::<SocketAddr>() |
667 | 1 | .map_err(|e| {0 |
668 | 0 | Error::from_std_err(Code::InvalidArgument, &e) |
669 | 0 | .append(format!("Invalid address '{}'", http_config.socket_address)) |
670 | 0 | })?; |
671 | 1 | let tcp_listener = if http_config.freebind { |
672 | 0 | bind_freebind(socket_addr) |
673 | | } else { |
674 | 1 | TcpListener::bind(&socket_addr).await |
675 | | } |
676 | 1 | .map_err(|e| match e.kind()0 { |
677 | 0 | ErrorKind::AddrInUse => make_err!( |
678 | 0 | Code::AlreadyExists, |
679 | | "Address '{socket_addr}' is already in use by another process.", |
680 | | ), |
681 | 0 | ErrorKind::PermissionDenied => make_err!( |
682 | 0 | Code::PermissionDenied, |
683 | | "Permission denied. You may need root privileges to bind to address '{socket_addr}'.", |
684 | | ), |
685 | | ErrorKind::InvalidInput => { |
686 | 0 | make_input_err!("The provided address '{socket_addr}' is invalid.") |
687 | | } |
688 | 0 | _ => Error::from_std_err(Code::Internal, &e) |
689 | 0 | .append(format!("Failed to bind to socket address '{socket_addr}'")), |
690 | 0 | })?; |
691 | 1 | let mut http = auto::Builder::new(TaskExecutor::default()); |
692 | 1 | http.http2().timer(TokioTimer::new()); |
693 | | |
694 | 1 | let tcp_nodelay = http_config.tcp_nodelay; |
695 | 1 | let http_config = &http_config.advanced_http; |
696 | 1 | if let Some(value0 ) = http_config.http2_keep_alive_interval { |
697 | 0 | http.http2() |
698 | 0 | .keep_alive_interval(Duration::from_secs(u64::from(value))); |
699 | 1 | } |
700 | | |
701 | 1 | if let Some(value0 ) = http_config.experimental_http2_max_pending_accept_reset_streams { |
702 | 0 | http.http2() |
703 | 0 | .max_pending_accept_reset_streams(usize::try_from(value).err_tip( |
704 | | || "Could not convert experimental_http2_max_pending_accept_reset_streams", |
705 | 0 | )?); |
706 | 1 | } |
707 | 1 | if let Some(value0 ) = http_config.experimental_http2_initial_stream_window_size { |
708 | 0 | http.http2().initial_stream_window_size(value); |
709 | 1 | } |
710 | 1 | if let Some(value0 ) = http_config.experimental_http2_initial_connection_window_size { |
711 | 0 | http.http2().initial_connection_window_size(value); |
712 | 1 | } |
713 | 1 | if let Some(value0 ) = http_config.experimental_http2_adaptive_window { |
714 | 0 | http.http2().adaptive_window(value); |
715 | 1 | } |
716 | 1 | if let Some(value0 ) = http_config.experimental_http2_max_frame_size { |
717 | 0 | http.http2().max_frame_size(value); |
718 | 1 | } |
719 | 1 | if let Some(value0 ) = http_config.experimental_http2_max_concurrent_streams { |
720 | 0 | http.http2().max_concurrent_streams(value); |
721 | 1 | } |
722 | 1 | if let Some(value0 ) = http_config.experimental_http2_keep_alive_timeout_s { |
723 | 0 | http.http2() |
724 | 0 | .keep_alive_timeout(Duration::from_secs(u64::from(value))); |
725 | 1 | } |
726 | 1 | if let Some(value0 ) = http_config.experimental_http2_max_send_buf_size { |
727 | 0 | http.http2().max_send_buf_size( |
728 | 0 | usize::try_from(value).err_tip(|| "Could not convert http2_max_send_buf_size")?, |
729 | | ); |
730 | 1 | } |
731 | 1 | if http_config.experimental_http2_enable_connect_protocol == Some(true) { |
732 | 0 | http.http2().enable_connect_protocol(); |
733 | 1 | } |
734 | 1 | if let Some(value0 ) = http_config.experimental_http2_max_header_list_size { |
735 | 0 | http.http2().max_header_list_size(value); |
736 | 1 | } |
737 | 1 | info!("Ready, listening on {socket_addr}",); |
738 | 1 | root_futures.push(Box::pin(async move { |
739 | | loop { |
740 | 2 | select! { |
741 | 2 | accept_result1 = accept_tcp_connection(&tcp_listener, tcp_nodelay) => { |
742 | 1 | match accept_result { |
743 | 1 | Ok((tcp_stream, remote_addr)) => { |
744 | 1 | info!( |
745 | | target: "nativelink::services", |
746 | | ?remote_addr, |
747 | | ?socket_addr, |
748 | | "Client connected" |
749 | | ); |
750 | | |
751 | 1 | let (http, svc, maybe_tls_acceptor) = |
752 | 1 | (http.clone(), svc.clone(), maybe_tls_acceptor.clone()); |
753 | | |
754 | 1 | background_spawn!( |
755 | | name: "http_connection", |
756 | 1 | fut: error_span!( |
757 | | "http_connection", |
758 | | remote_addr = %remote_addr, |
759 | | socket_addr = %socket_addr, |
760 | 1 | ).in_scope(|| async move { |
761 | 1 | let connected_at = std::time::Instant::now(); |
762 | 1 | let serve_connection = if let Some(tls_acceptor0 ) = maybe_tls_acceptor { |
763 | 0 | match tls_acceptor.accept(tcp_stream).await { |
764 | 0 | Ok(tls_stream) => Either::Left(http.serve_connection( |
765 | 0 | TokioIo::new(tls_stream), |
766 | 0 | TowerToHyperService::new(svc), |
767 | 0 | )), |
768 | 0 | Err(err) => { |
769 | 0 | error!(?err, "Failed to accept tls stream"); |
770 | 0 | return; |
771 | | } |
772 | | } |
773 | | } else { |
774 | 1 | Either::Right(http.serve_connection( |
775 | 1 | TokioIo::new(tcp_stream), |
776 | 1 | TowerToHyperService::new(svc), |
777 | 1 | )) |
778 | | }; |
779 | | |
780 | 1 | if let Err(err0 ) = serve_connection.await { |
781 | | // The HTTP/2 layer closes a connection with |
782 | | // ENHANCE_YOUR_CALM when one of its internal |
783 | | // limits trips, taking every multiplexed |
784 | | // stream with it. The connection's age tells |
785 | | // a load-shaped limit (any age, during load) |
786 | | // from a lifetime cap (old connections), which |
787 | | // the GOAWAY's debug data alone cannot. |
788 | 0 | let local_err: &(dyn core::error::Error + 'static) = err.as_ref(); |
789 | 0 | let enhance_your_calm = core::iter::successors( |
790 | 0 | Some(local_err), |
791 | 0 | |err| err.source(), |
792 | | ) |
793 | 0 | .filter_map(|err| err.downcast_ref::<h2::Error>()) |
794 | 0 | .any(|err| { |
795 | 0 | err.is_go_away() |
796 | 0 | && err.reason() == Some(h2::Reason::ENHANCE_YOUR_CALM) |
797 | 0 | }); |
798 | 0 | if enhance_your_calm { |
799 | 0 | warn!( |
800 | | target: "nativelink::services", |
801 | | ?err, |
802 | | %remote_addr, |
803 | 0 | connection_age_s = connected_at.elapsed().as_secs(), |
804 | | "Connection closed with an ENHANCE_YOUR_CALM GOAWAY" |
805 | | ); |
806 | 0 | } |
807 | 0 | error!( |
808 | | target: "nativelink::services", |
809 | | ?err, |
810 | | "Failed running service" |
811 | | ); |
812 | 0 | } |
813 | 1 | }), |
814 | | target: "nativelink::services", |
815 | | ?remote_addr, |
816 | | ?socket_addr, |
817 | | ); |
818 | | }, |
819 | 0 | Err(err) => { |
820 | 0 | error!(?err, "Failed to initialize TCP connection"); |
821 | | } |
822 | | } |
823 | | }, |
824 | | } |
825 | | } |
826 | | // Unreachable |
827 | | })); |
828 | | } |
829 | | |
830 | | { |
831 | | // We start workers after our TcpListener is setup so if our worker connects to one |
832 | | // of these services it will be able to connect. |
833 | 1 | let worker_cfgs = cfg.workers.unwrap_or_default(); |
834 | 1 | let mut worker_names = HashSet::with_capacity(worker_cfgs.len()); |
835 | 1 | for (i0 , worker_cfg0 ) in worker_cfgs.into_iter().enumerate() { |
836 | 0 | let spawn_fut = match worker_cfg { |
837 | 0 | WorkerConfig::Local(local_worker_cfg) => { |
838 | 0 | let completion_tx = if local_worker_cfg.single_use { |
839 | 0 | single_use_complete_tx.take() |
840 | | } else { |
841 | 0 | None |
842 | | }; |
843 | 0 | let fast_slow_store = store_manager |
844 | 0 | .get_store(&local_worker_cfg.cas_fast_slow_store) |
845 | 0 | .err_tip(|| { |
846 | 0 | format!( |
847 | | "Failed to find store for cas_store_ref in worker config : {}", |
848 | | local_worker_cfg.cas_fast_slow_store |
849 | | ) |
850 | 0 | })?; |
851 | | |
852 | 0 | let maybe_ac_store = if let Some(ac_store_ref) = |
853 | 0 | &local_worker_cfg.upload_action_result.ac_store |
854 | | { |
855 | 0 | Some(store_manager.get_store(ac_store_ref).err_tip(|| { |
856 | 0 | format!("Failed to find store for ac_store in worker config : {ac_store_ref}") |
857 | 0 | })?) |
858 | | } else { |
859 | 0 | None |
860 | | }; |
861 | | // Note: Defaults to fast_slow_store if not specified. If this ever changes it must |
862 | | // be updated in config documentation for the `historical_results_store` the field. |
863 | 0 | let historical_store = if let Some(cas_store_ref) = &local_worker_cfg |
864 | 0 | .upload_action_result |
865 | 0 | .historical_results_store |
866 | | { |
867 | 0 | store_manager.get_store(cas_store_ref).err_tip(|| { |
868 | 0 | format!( |
869 | | "Failed to find store for historical_results_store in worker config : {cas_store_ref}" |
870 | | ) |
871 | 0 | })? |
872 | | } else { |
873 | 0 | fast_slow_store.clone() |
874 | | }; |
875 | 0 | let local_worker = new_local_worker( |
876 | 0 | Arc::new(local_worker_cfg), |
877 | 0 | fast_slow_store, |
878 | 0 | maybe_ac_store, |
879 | 0 | historical_store, |
880 | 0 | ) |
881 | 0 | .await |
882 | 0 | .err_tip(|| "Could not make LocalWorker")? |
883 | 0 | .with_registration(worker_registrations[i].clone()); |
884 | | |
885 | 0 | let name = if local_worker.name().is_empty() { |
886 | 0 | format!("worker_{i}") |
887 | | } else { |
888 | 0 | local_worker.name().clone() |
889 | | }; |
890 | | |
891 | 0 | if worker_names.contains(&name) { |
892 | 0 | Err(make_input_err!( |
893 | 0 | "Duplicate worker name '{}' found in config", |
894 | 0 | name |
895 | 0 | ))?; |
896 | 0 | } |
897 | 0 | worker_names.insert(name.clone()); |
898 | 0 | let shutdown_rx = shutdown_tx.subscribe(); |
899 | 0 | let fut = trace_span!("worker_ctx", worker_name = %name) |
900 | 0 | .in_scope(|| local_worker.run(shutdown_rx)); |
901 | 0 | spawn!( |
902 | | "worker", |
903 | 0 | async move { |
904 | 0 | fut.await?; |
905 | 0 | if let Some(completion_tx) = completion_tx { |
906 | 0 | let _ = completion_tx.send(()); |
907 | 0 | } |
908 | 0 | Ok::<(), Error>(()) |
909 | 0 | }, |
910 | | ?name |
911 | | ) |
912 | | } |
913 | | }; |
914 | 0 | root_futures.push(Box::pin(spawn_fut.map_ok_or_else(|e| Err(e.into()), |v| v))); |
915 | | } |
916 | | } |
917 | | |
918 | | // Set up a shutdown handler for the worker schedulers. |
919 | 1 | let mut shutdown_rx = shutdown_tx.subscribe(); |
920 | 1 | root_futures.push(Box::pin(async move { |
921 | 1 | if let Ok(shutdown_guard0 ) = shutdown_rx.recv().await { |
922 | 0 | let _ = scheduler_shutdown_tx.send(()); |
923 | 0 | for (_name, scheduler) in worker_schedulers { |
924 | 0 | scheduler.shutdown(shutdown_guard.clone()).await; |
925 | | } |
926 | 1 | } |
927 | 1 | Ok(()) |
928 | 1 | })); |
929 | 1 | Ok(Server { |
930 | 1 | root_futures, |
931 | 1 | single_use_enabled, |
932 | 1 | single_use_complete_rx, |
933 | 1 | }) |
934 | 1 | } |
935 | | |
936 | | #[cfg(test)] |
937 | | mod tests { |
938 | | use futures::future::try_join_all; |
939 | | use nativelink_config::cas_server::HttpListener; |
940 | | use nativelink_macro::nativelink_test; |
941 | | |
942 | | use super::*; |
943 | | |
944 | | #[nativelink_test] |
945 | | async fn accepted_connections_honor_tcp_nodelay() -> Result<(), Error> { |
946 | | let listener = TcpListener::bind("127.0.0.1:0").await?; |
947 | | for (config, expected) in [ |
948 | | (HttpListener::default(), false), |
949 | | ( |
950 | | HttpListener { |
951 | | tcp_nodelay: false, |
952 | | ..Default::default() |
953 | | }, |
954 | | false, |
955 | | ), |
956 | | ( |
957 | | HttpListener { |
958 | | tcp_nodelay: true, |
959 | | ..Default::default() |
960 | | }, |
961 | | true, |
962 | | ), |
963 | | ] { |
964 | | // Every accepted connection must receive the configured socket option. |
965 | | for _ in 0..2 { |
966 | | let (_client, (server, _)) = tokio::try_join!( |
967 | | TcpStream::connect(listener.local_addr()?).map_err(Error::from), |
968 | | accept_tcp_connection(&listener, config.tcp_nodelay), |
969 | | )?; |
970 | | assert_eq!(server.nodelay()?, expected); |
971 | | } |
972 | | } |
973 | | Ok(()) |
974 | | } |
975 | | |
976 | | /// Ask the OS for a currently-free TCP port. |
977 | 1 | fn get_ephemeral_port() -> Result<u16, Error> { |
978 | 1 | let listener = std::net::TcpListener::bind(("127.0.0.1", 0)) |
979 | 1 | .err_tip(|| "Binding an ephemeral port for test server")?0 ; |
980 | 1 | Ok(listener |
981 | 1 | .local_addr() |
982 | 1 | .err_tip(|| "Getting ephemeral port for test server")?0 |
983 | 1 | .port()) |
984 | 1 | } |
985 | | |
986 | | #[nativelink_test] |
987 | | async fn build_server_boots() -> Result<(), Error> { |
988 | | let (scheduler_shutdown_tx, _scheduler_shutdown_rx) = oneshot::channel(); |
989 | | let (shutdown_tx, _) = broadcast::channel::<ShutdownGuard>(1); |
990 | | let port = get_ephemeral_port()?; |
991 | | let config = CasConfig::try_from_json5_str(&format!( |
992 | | r#"{{ stores: [], servers: [ |
993 | | {{ |
994 | | name: "test_server", |
995 | | listener: {{ |
996 | | http: {{ |
997 | | socket_address: "0.0.0.0:{port}", |
998 | | }}, |
999 | | }}, |
1000 | | services: {{ |
1001 | | health: {{}}, |
1002 | | }}, |
1003 | | }}, |
1004 | | ] }}"# |
1005 | | )) |
1006 | | .unwrap(); |
1007 | | let server = build_server(config, shutdown_tx, scheduler_shutdown_tx).await?; |
1008 | | let client = reqwest::Client::builder() |
1009 | | .timeout(Duration::from_secs(2)) |
1010 | | .build()?; |
1011 | | |
1012 | | drop(select! { |
1013 | | result = try_join_all(server.root_futures) => result.map(|_| ()), |
1014 | | result = client.get(format!("http://localhost:{port}/status")).send() => { |
1015 | | let ok_res = result?; |
1016 | | assert_eq!(ok_res.status(), 200); |
1017 | | assert_eq!(ok_res.text().await?, "[]"); |
1018 | | Ok(()) |
1019 | | }, |
1020 | | }); |
1021 | | |
1022 | | Ok(()) |
1023 | | } |
1024 | | } |