Coverage Report

Created: 2026-11-13 23:29

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/build/source/src/lib.rs
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
}