Skip to main content

witchcraft_server/
lib.rs

1// Copyright 2022 Palantir Technologies, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (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// http://www.apache.org/licenses/LICENSE-2.0
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//! A highly opinionated embedded application server for RESTy APIs.
15//!
16//! # Initialization
17//!
18//! The entrypoint of a Witchcraft server is an initialization function annotated with `#[witchcraft_server::main]`:
19//!
20//! ```ignore
21//! use conjure_error::Error;
22//! use refreshable::Refreshable;
23//! use witchcraft_server::config::install::InstallConfig;
24//! use witchcraft_server::config::runtime::RuntimeConfig;
25//!
26//! #[witchcraft_server::main]
27//! fn main(
28//!     install: InstallConfig,
29//!     runtime: Refreshable<RuntimeConfig>,
30//!     wc: &mut Witchcraft,
31//! ) -> Result<(), Error> {
32//!     wc.api(CustomApiEndpoints::new(CustomApiResource));
33//!
34//!     Ok(())
35//! }
36//! ```
37//!
38//! The function is provided with the server's install and runtime configuration, as well as the [`Witchcraft`] object
39//! which can be used to configure the server. Once the initialization function returns, the server will start.
40//!
41//! ## Note
42//!
43//! The initialization function is expected to return quickly - any long-running work required should happen in the
44//! background.
45//!
46//! # Configuration
47//!
48//! Witchcraft divides configuration into two categories:
49//!
50//! * *Install* - Configuration that is fixed for the lifetime of a service. For example, the port that the server
51//!   listens on is part of the server's install configuration.
52//! * *Runtime* - Configuration that can dynamically update while the service is running. For example, the logging
53//!   verbosity level is part of the server's runtime configuration.
54//!
55//! Configuration is loaded from the `var/conf/install.yml` and `var/conf/runtime.yml` files respectively. The
56//! `runtime.yml` file is automatically checked for updates every few seconds.
57//!
58//! ## Extension
59//!
60//! The configuration files are deserialized into Rust types via the [`serde::Deserialize`] trait. `witchcraft-server`'s
61//! own internal configuration is represented by the [`InstallConfig`] and [`RuntimeConfig`] types. Services that need
62//! their own configuration should embed the Witchcraft configuration within their own using the `#[serde(flatten)]`
63//! annotation and implement the [`AsRef`] trait:
64//!
65//! ```
66//! use serde::Deserialize;
67//! use witchcraft_server::config::install::InstallConfig;
68//!
69//! #[derive(Deserialize)]
70//! #[serde(rename_all = "kebab-case")]
71//! struct MyInstallConfig {
72//!     shave_yaks: bool,
73//!     #[serde(flatten)]
74//!     base: InstallConfig,
75//! }
76//!
77//! impl AsRef<InstallConfig> for MyInstallConfig {
78//!     fn as_ref(&self) -> &InstallConfig {
79//!         &self.base
80//!     }
81//! }
82//! ```
83//!
84//! The service's custom configuration will then sit next to the standard Witchcraft configuration:
85//!
86//! ```yml
87//! product-name: my-service
88//! product-version: 1.0.0
89//! port: 12345
90//! shave-yaks: true
91//! ```
92//!
93//! ## Sensitive values
94//!
95//! The server's configuration deserializer supports two methods to handle sensitive values:
96//!
97//! * `${enc:5BBfGvf90H6bApwfx...}` - inline in an encrypted form using [`serde_encrypted_value`] with the
98//!   key stored in `var/conf/encrypted-config-value.key`.
99//! * `${file:/mnt/secrets/foo}` - as a reference to a file containing the value using [`serde_file_value`].
100//!
101//! ## Refreshable runtime configuration
102//!
103//! The server's runtime configuration is wrapped in the [`Refreshable`] type to allow code to properly handle updates
104//! to the configuration. Depending on the use case, implementations can use the [`Refreshable::get`] to retrieve the
105//! current state of the configuration when needed or the [`Refreshable::subscribe`] method to be notified of changes
106//! to the configuration as they happen. See the documentation of the [`refreshable`] crate for more details.
107//!
108//! # HTTP APIs
109//!
110//! The server supports HTTP endpoints implementing the [`Service`] and [`AsyncService`]
111//! traits. These implementations can be generated from a [Conjure] YML [definition] with the [`conjure-codegen`] crate.
112//!
113//! While we strongly encourage the use of Conjure-generated APIs, some services may need to expose endpoints that can't
114//! be defined in Conjure. The [`conjure_http::conjure_endpoints`] macro can be used to define arbitrary HTTP endpoints.
115//!
116//! API endpoints should normally be registered with the [`Witchcraft::api`] and [`Witchcraft::blocking_api`] methods,
117//! which will place the endpoints under the `/api` route. If necessary, the [`Witchcraft::app`] and
118//! [`Witchcraft::blocking_app`] methods can be used to place the endpoints directly at the root route instead.
119//!
120//! [`Service`]: conjure_http::server::Service
121//! [Conjure]: https://github.com/palantir/conjure
122//! [definition]: https://palantir.github.io/conjure/#/docs/spec/conjure_definitions
123//! [`conjure-codegen`]: https://docs.rs/conjure-codegen
124//!
125//! # HTTP clients
126//!
127//! Remote services are configured in the `service-discovery` section of the runtime configuration, and clients can be
128//! created from the [`ClientFactory`] returned by the [`Witchcraft::client_factory`] method. The clients will
129//! automatically update based on changes to the runtime configuration. See the documentation of the [`conjure_runtime`]
130//! crate for more details.
131//!
132//! # Status endpoints
133//!
134//! The server exposes several "status" endpoints to report various aspects of the server.
135//!
136//! ## Liveness
137//!
138//! The `/status/liveness` endpoint returns a successful response to all requests, indicating that the server is alive.
139//!
140//! ## Readiness
141//!
142//! The `/status/readiness` endpoint returns a response indicating the server's readiness to handle requests to its
143//! endpoints. Deployment infrastructure uses the result of this endpoint to decide if requests should be routed to a
144//! given instance of the service. Custom readiness checks can be added to the server via the [`ReadinessCheckRegistry`]
145//! returned by the [`Witchcraft::readiness_checks`] method. Any long-running initialization logic should happen
146//! asynchronously and use a readiness check to indicate completion.
147//!
148//! ## Health
149//!
150//! The `/status/health` endpoint returns a response indicating the server's overall health. Deployment infrastructure
151//! uses the result of this endpoint to trigger alerts. Custom health checks can be added to the server via the
152//! [`HealthCheckRegistry`] returned by the [`Witchcraft::health_checks`] method. Requests to this endpoint must be
153//! authenticated with the `health-checks.shared-secret` bearer token in runtime configuration.
154//!
155//! The server registers several built-in health checks:
156//!
157//! * `CONFIG_RELOAD` - Reports an error state if the runtime configuration failed to reload properly.
158//! * `ENDPOINT_FIVE_HUNDREDS` - Reports a warning if an endpoint has a high rate of `500 Internal Server Error`
159//!   responses.
160//! * `SERVICE_DEPENDENCY` - Tracks the status of requests made with HTTP clients created via the server's client
161//!   factory, and reports a warning state of requests to a remote service have a high failure rate.
162//! * `PANICS` - Reports a warning if the server has panicked at any point.
163//!
164//! # Diagnostics
165//!
166//! The `/debug/diagnostic/{diagnosticType}` endpoint returns diagnostic information. Requests to this endpoint must be
167//! authenticated with the `diagnostics.debug-shared-secret` bearer token in the runtime configuration.
168//!
169//! Several diagnostic types are defined:
170//!
171//! * `diagnostic.types.v1` - Returns a JSON-encoded list of all valid diagnostic types.
172//! * `rust.heap.status.v1` - Returns detailed statistics about the state of the heap. Requires the `jemalloc` feature
173//!   (enabled by default).
174//! * `rust.heap.profile.v1` - Returns a profile of the source of a sample of live allocations. Use the `jeprof` tool
175//!   to analyze the profile. Requires the `jemalloc` feature (enabled by default).
176//! * `metric.names.v1` - Returns a JSON-encoded list of the names of all metrics registered with the server.
177//! * `rust.thread.dump.v1` - Returns a stack trace of every thread in the process. Only supported when running on
178//!   Linux.
179//!
180//! # Logging
181//!
182//! `witchcraft-server` emits JSON-encoded logs following the [witchcraft-api spec]. By default, logs will be written to
183//! a file in `var/log` corresponding to the type of log message (`service.log`, `request.log`, etc). These files are
184//! automatically rotated and compressed based on a non-configurable policy. If running in a Docker container or if the
185//! `use-console-log` setting is enabled in the install configuration, logs will instead be written to standard out.
186//!
187//! [witchcraft-api spec]: https://github.com/palantir/witchcraft-api
188//!
189//! ## Service
190//!
191//! The service log contains the messages emitted by invocations of the macros in the [`witchcraft_log`] crate. Messages
192//! emitted by the standard Rust [`log`] crate are additionally captured, but code that is written as part of a
193//! Witchcraft service should use [`witchcraft_log`] instead for better integration. See the documentation of that crate
194//! for more details.
195//!
196//! ## Request
197//!
198//! The request log records an entry for each HTTP request processed by the server. Parameters marked marked as safe by
199//! an endpoint's Conjure definition will be included as parameters in the log record.
200//!
201//! ## Trace
202//!
203//! The trace log records [Zipkin]-style trace spans. The server automatically creates spans for each incoming HTTP
204//! request based off of request's propagation metadata. Traces that have not alread had a sampling decision made will
205//! be sampled at the rate specified by the `logging.trace-rate` field in the server's runtime configuration, which
206//! defaults to 0.005%. Server logic can create additional spans with the [`zipkin`] crate. See the documentation of
207//! that crate for more details.
208//!
209//! [Zipkin]: https://zipkin.io/
210//!
211//! ## Metric
212//!
213//! The metric log contains the values of metrics reporting the state of various components of the server. Metrics are
214//! recorded every 30 seconds. Server logic can create additional metrics with the [`MetricRegistry`] returned by the
215//! [`Witchcraft::metrics`] method. See the documentation of the [`witchcraft_metrics`] crate for more details.
216//!
217//! # Metrics
218//!
219//! The server reports a variety of metrics by default:
220//!
221//! ## Thread Pool
222//!
223//! * `server.worker.max` (gauge) - The configured maximum size of the server's thread pool used for requests to
224//!   blocking endpoints.
225//! * `server.worker.active` (gauge) - The number of threads actively processing requests to blocking endpoints.
226//! * `server.worker.utilization-max` (gauge) - `server.worker.active` divided by `server.worker.max`. If this is 1, the
227//!   server will immediately reject calls to blocking endpoints with a `503 Service Unavailable` status code.
228//!
229//! ## Logging
230//!
231//! * `logging.queue (type: <log_type>)` (gauge) - The number of log messages queued for output.
232//!
233//! ## Process
234//!
235//! * `process.heap` (gauge) - The total number of bytes allocated from the heap. Requires the `jemalloc` feature
236//!   (enabled by default).
237//! * `process.heap.active` (gauge) - The total number of bytes in active pages. Requires the `jemalloc` feature
238//!   (enabled by default).
239//! * `process.heap.resident` (gauge) - The total number of bytes in physically resident pages. Requires the `jemalloc` feature
240//!   (enabled by default).
241//! * `process.uptime` (gauge) - The number of microseconds that have elapsed since the server started.
242//! * `process.panics` (counter) - The number of times the server has panicked.
243//! * `process.user-time` (gauge) - The number of microseconds the process has spent running in user-space.
244//! * `process.user-time.norm` (gauge) - `process.user-time` divided by the number of CPU cores.
245//! * `process.system-time` (gauge) - The number of microseconds the process has spent either running in kernel-space
246//!   or in uninterruptable IO wait.
247//! * `process.system-time.norm` (gauge) - `process.system-time` divided by the number of CPU cores.
248//! * `process.blocks-read` (gauge) - The number of filesystem blocks the server has read.
249//! * `process.blocks-written` (gauge) - The number of filesystem blocks the server has written.
250//! * `process.threads` (gauge) - The number of threads in the process.
251//! * `process.filedescriptor` (gauge) - The number of file descriptors held open by the process divided by the maximum
252//!   number of files the server may hold open.
253//!
254//! ## Connection
255//!
256//! * `server.connection.active` (counter) - The number of TCP sockets currently connected to the HTTP server.
257//! * `server.connection.utilization` (gauge) - `server.connection.active` divided by the maximum number of connections
258//!   the server will accept.
259//!
260//! ## TLS
261//!
262//! * `tls.handshake (context: server, protocol: <protocol>, cipher: <cipher>)` (meter) - The rate of TLS handshakes
263//!   completed by the HTTP server.
264//!
265//! ## Server
266//!
267//! * `server.request.active` (counter) - The number of requests being actively processed.
268//! * `server.request.unmatched` (meter) - The rate of `404 Not Found` responses returned by the server.
269//! * `server.response.all` (meter) - The rate of responses returned by the server.
270//! * `server.response.1xx` (meter) - The rate of `1xx` responses returned by the server.
271//! * `server.response.2xx` (meter) - The rate of `2xx` responses returned by the server.
272//! * `server.response.3xx` (meter) - The rate of `3xx` responses returned by the server.
273//! * `server.response.4xx` (meter) - The rate of `4xx` responses returned by the server.
274//! * `server.response.5xx` (meter) - The rate of `5xx` responses returned by the server.
275//! * `server.response.500` (meter) - The rate of `500 Internal Server Error` responses returned by the server.
276//!
277//! ## Endpoints
278//!
279//! * `server.response (service-name: <service_name>, endpoint: <endpoint>)` (timer) - The amount of time required to
280//!   process each request to the endpoint, including sending the entire response body.
281//! * `server.response.error (service-name: <service_name>, endpoint: <endpoint>)` (meter) - The rate of `5xx` errors
282//!   returned for requests to the endpoint.
283//!
284//! ## HTTP clients
285//!
286//! See the documentation of the [`conjure_runtime`] crate for the metrics reported by HTTP clients.
287#![warn(missing_docs)]
288
289use std::path::Path;
290use std::pin::pin;
291use std::process;
292use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
293use std::sync::Arc;
294use std::time::Duration;
295use std::{env, mem};
296
297use conjure_error::Error;
298use conjure_http::server::{AsyncService, ConjureRuntime};
299use conjure_runtime::{Agent, ClientFactory, HostMetricsRegistry, UserAgent};
300use debug::endpoint::DebugResource;
301use debug::endpoint::DebugServiceEndpoints;
302#[cfg(feature = "jemalloc")]
303use debug::heap_profile::{self, HeapProfileDiagnostic};
304use futures_util::{stream, Stream, StreamExt};
305use refreshable::Refreshable;
306use serde::de::DeserializeOwned;
307use status::StatusResource;
308use status::StatusServiceEndpoints;
309use tokio::runtime::{Handle, Runtime};
310use tokio::signal::unix::{self, SignalKind};
311use tokio::{runtime, select, time};
312use witchcraft_log::{error, fatal, info};
313use witchcraft_metrics::MetricRegistry;
314
315pub use body::{RequestBody, ResponseWriter};
316use config::install::InstallConfig;
317use config::runtime::RuntimeConfig;
318pub use witchcraft::Witchcraft;
319#[doc(inline)]
320pub use witchcraft_server_config as config;
321#[doc(inline)]
322pub use witchcraft_server_macros::main;
323
324use crate::debug::diagnostic_types::DiagnosticTypesDiagnostic;
325#[cfg(feature = "jemalloc")]
326use crate::debug::heap_stats::HeapStatsDiagnostic;
327use crate::debug::metric_names::MetricNamesDiagnostic;
328#[cfg(target_os = "linux")]
329use crate::debug::thread_dump::ThreadDumpDiagnostic;
330use crate::debug::DiagnosticRegistry;
331use crate::health::config_reload::ConfigReloadHealthCheck;
332use crate::health::endpoint_500s::Endpoint500sHealthCheck;
333use crate::health::minidump::MinidumpHealthCheck;
334use crate::health::panics::PanicsHealthCheck;
335use crate::health::service_dependency::ServiceDependencyHealthCheck;
336use crate::health::HealthCheckRegistry;
337use crate::readiness::ReadinessCheckRegistry;
338use crate::server::Listener;
339use crate::shutdown_hooks::ShutdownHooks;
340
341pub mod blocking;
342mod body;
343mod configs;
344pub mod debug;
345mod endpoint;
346pub mod extensions;
347pub mod health;
348pub mod logging;
349mod metrics;
350mod minidump;
351pub mod readiness;
352mod server;
353mod service;
354mod shutdown_hooks;
355mod status;
356pub mod tls;
357mod witchcraft;
358
359#[cfg(feature = "jemalloc")]
360#[global_allocator]
361static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
362
363/// Initializes a Witchcraft server.
364///
365/// `init` is invoked with the parsed install and runtime configs as well as the [`Witchcraft`] context object. It
366/// is expected to return quickly; any long running initialization should be spawned off into the background to run
367/// asynchronously.
368pub fn init<I, R, F>(init: F)
369where
370    I: AsRef<InstallConfig> + DeserializeOwned,
371    R: AsRef<RuntimeConfig> + DeserializeOwned + PartialEq + 'static + Sync + Send,
372    F: FnOnce(I, Refreshable<R, Error>, &mut Witchcraft) -> Result<(), Error>,
373{
374    init_with_configs(init, configs::load_install::<I>, configs::load_runtime::<R>)
375}
376
377/// Initializes a Witchcraft server with custom config loaders.
378///
379/// `init` is invoked with the install and runtime configs from the provided loaders as well as the [`Witchcraft`]
380/// context object. It is expected to return quickly; any long running initialization should be spawned off into
381/// the background to run asynchronously.
382pub fn init_with_configs<I, R, F, LI, LR>(init: F, load_install: LI, load_runtime: LR)
383where
384    I: AsRef<InstallConfig> + DeserializeOwned,
385    R: AsRef<RuntimeConfig> + DeserializeOwned + PartialEq + 'static + Sync + Send,
386    F: FnOnce(I, Refreshable<R, Error>, &mut Witchcraft) -> Result<(), Error>,
387    LI: FnOnce() -> Result<I, Error>,
388    LR: FnOnce(&Handle, &Arc<AtomicBool>) -> Result<Refreshable<R, Error>, Error>,
389{
390    logging::early_init();
391
392    let args = env::args_os().collect::<Vec<_>>();
393    if args.len() == 3 && args[1] == "minidump" {
394        let ret = match minidump::server(Path::new(&args[2])) {
395            Ok(()) => 0,
396            Err(e) => {
397                fatal!("error starting minidump server", error: e);
398                1
399            }
400        };
401        process::exit(ret);
402    } else {
403        let mut runtime_guard = None;
404
405        let ret = match init_inner(
406            init,
407            load_install,
408            load_runtime,
409            &mut runtime_guard,
410            InitOptions::default(),
411        ) {
412            Ok(witchcraft) => {
413                match witchcraft.handle.block_on(shutdown(
414                    witchcraft.shutdown_hooks,
415                    witchcraft.install_config.server().shutdown_timeout(),
416                )) {
417                    Ok(_) => 0,
418                    Err(e) => {
419                        fatal!("error after starting server", error: e);
420                        1
421                    }
422                }
423            }
424            Err(e) => {
425                fatal!("error starting server", error: e);
426                1
427            }
428        };
429        drop(runtime_guard);
430
431        process::exit(ret);
432    }
433}
434
435#[cfg(feature = "in-memory-testing")]
436/// Variation of the Witchcraft server initialization, which can be used for testing. Allows
437/// spawning multiple instances of the server in memory.
438pub mod in_memory_testing {
439    use crate::{drain_shutdown_hooks, init_inner, InitOptions, Witchcraft};
440    use conjure_error::Error;
441    use refreshable::Refreshable;
442    use serde::de::DeserializeOwned;
443    use std::sync::atomic::{AtomicBool, Ordering};
444    use std::sync::{mpsc, Arc, Once};
445    use std::thread::{self, JoinHandle};
446    use tokio::runtime::Handle;
447    use tokio::sync::oneshot;
448    use witchcraft_server_config::install::InstallConfig;
449    use witchcraft_server_config::runtime::RuntimeConfig;
450
451    static INIT: Once = Once::new();
452    static FIRST_INIT_SUCCESS: AtomicBool = AtomicBool::new(false);
453
454    /// Initializes a Witchcraft server for testing with custom config loaders. This variation of init
455    /// spawns a Witchcraft server in a thread, rather than as a separate process. Note: because
456    /// logging infrastructure is global to a process, the first initialized server will also
457    /// initialize logging appenders; all subsequent ones will reuse the same appenders. However,
458    /// the actual service logger implementation and log level must be set globally once by the
459    /// calling test code.
460    ///
461    /// The server runs on a dedicated thread, and the returned [`RunHandle`] shuts the
462    /// server down gracefully when dropped. Note that since the first call to this init function will
463    /// also initialize the global logging appenders; the corresponding returned handle will shut down the
464    /// appenders on drop, making logging unavailable to servers initialized after the first one.
465    /// Thus, ensure that first initialized server is the last one to go out of scope.
466    pub fn init_with_configs_for_tests<I, R, F, LI, LR>(
467        init: F,
468        load_install: LI,
469        load_runtime: LR,
470        thread_prefix: Option<String>,
471    ) -> Result<RunHandle, Error>
472    where
473        I: AsRef<InstallConfig> + DeserializeOwned,
474        R: AsRef<RuntimeConfig> + DeserializeOwned + PartialEq + 'static + Sync + Send,
475        F: FnOnce(I, Refreshable<R, Error>, &mut Witchcraft) -> Result<(), Error> + Send + 'static,
476        LI: FnOnce() -> Result<I, Error> + Send + 'static,
477        LR: FnOnce(&Handle, &Arc<AtomicBool>) -> Result<Refreshable<R, Error>, Error>
478            + Send
479            + 'static,
480    {
481        let init_fn = |init_globals: bool| {
482            let (startup_result_sender, startup_result_receiver) =
483                mpsc::channel::<Result<(), Error>>();
484            let (shutdown_signal_sender, shutdown_signal_receiver) = oneshot::channel::<()>();
485
486            // Dedicated thread for this server instance, allows it to create its own tokio runtime
487            let server_thread = thread::spawn(move || {
488                let mut runtime_guard = None;
489
490                let witchcraft = match init_inner(
491                    init,
492                    load_install,
493                    load_runtime,
494                    &mut runtime_guard,
495                    InitOptions {
496                        init_globals,
497                        thread_prefix,
498                    },
499                ) {
500                    Ok(witchcraft) => witchcraft,
501                    Err(e) => {
502                        let _ = startup_result_sender.send(Err(e));
503                        return;
504                    }
505                };
506                // Startup OK, notify receiver waiting below
507                let _ = startup_result_sender.send(Ok(()));
508
509                let handle = witchcraft.handle.clone();
510                let timeout = witchcraft.install_config.server().shutdown_timeout();
511
512                // Block the server until shutdown is called by the RunHandle's drop
513                handle.block_on(async move {
514                    let _ = shutdown_signal_receiver.await;
515                    drain_shutdown_hooks(witchcraft.shutdown_hooks, timeout).await;
516                });
517                drop(runtime_guard);
518            });
519
520            // Wait for startup to finish before handing back a handle. A receive error means the thread
521            // panicked before reporting.
522            match startup_result_receiver.recv() {
523                Ok(Ok(())) => Ok(RunHandle {
524                    shutdown_signal_sender: Some(shutdown_signal_sender),
525                    server_thread: Some(server_thread),
526                }),
527                Ok(Err(e)) => Err(e), // Error during init_inner()
528                Err(_) => Err(Error::internal_safe(
529                    "in-memory server thread panicked during initialization",
530                )),
531            }
532        };
533
534        let mut init_fn = Some(init_fn);
535
536        let mut first_init_result: Option<Result<RunHandle, Error>> = None;
537        INIT.call_once(|| {
538            // Only one branch of init_fn() is guaranteed to run, this take() lets us bypass
539            // borrow checker not detecting this branching.
540            let res = init_fn.take().unwrap()(true);
541            FIRST_INIT_SUCCESS.store(res.is_ok(), Ordering::Relaxed);
542            first_init_result = Some(res);
543        });
544
545        first_init_result.unwrap_or_else(|| {
546            if !FIRST_INIT_SUCCESS.load(Ordering::Relaxed) {
547                return Err(Error::internal_safe("Initial init failed, not continuing"));
548            }
549            init_fn.take().unwrap()(false)
550        })
551    }
552
553    /// Provides a handle for a running in-memory Witchcraft server. Gracefully shuts the server
554    /// down when it goes out of scope.
555    pub struct RunHandle {
556        shutdown_signal_sender: Option<oneshot::Sender<()>>,
557        server_thread: Option<JoinHandle<()>>,
558    }
559
560    impl Drop for RunHandle {
561        fn drop(&mut self) {
562            // Signal shutdown, then wait for the server thread to drain and tear down its runtime.
563            let _ = self.shutdown_signal_sender.take().unwrap().send(());
564            let _ = self.server_thread.take().unwrap().join();
565        }
566    }
567}
568
569struct InitOptions {
570    init_globals: bool,
571    thread_prefix: Option<String>,
572}
573
574impl Default for InitOptions {
575    fn default() -> Self {
576        Self {
577            init_globals: true,
578            thread_prefix: None,
579        }
580    }
581}
582
583fn init_inner<I, R, F, LI, LR>(
584    init: F,
585    load_install: LI,
586    load_runtime: LR,
587    runtime_guard: &mut Option<RuntimeGuard>,
588    init_options: InitOptions,
589) -> Result<Witchcraft, Error>
590where
591    I: AsRef<InstallConfig> + DeserializeOwned,
592    R: AsRef<RuntimeConfig> + DeserializeOwned + PartialEq + 'static + Sync + Send,
593    F: FnOnce(I, Refreshable<R, Error>, &mut Witchcraft) -> Result<(), Error>,
594    LI: FnOnce() -> Result<I, Error>,
595    LR: FnOnce(&Handle, &Arc<AtomicBool>) -> Result<Refreshable<R, Error>, Error>,
596{
597    let install_config = load_install()?;
598
599    let thread_id = AtomicUsize::new(0);
600    let thread_prefix = init_options
601        .thread_prefix
602        .clone()
603        .unwrap_or("runtime".to_string());
604    let runtime = runtime::Builder::new_multi_thread()
605        .enable_all()
606        .thread_name_fn(move || {
607            format!(
608                "{}-{}",
609                thread_prefix,
610                thread_id.fetch_add(1, Ordering::Relaxed)
611            )
612        })
613        .worker_threads(install_config.as_ref().server().io_threads())
614        .thread_keep_alive(install_config.as_ref().server().idle_thread_timeout())
615        .build()
616        .map_err(Error::internal_safe)?;
617
618    let handle = runtime.handle().clone();
619    let runtime = runtime_guard.insert(RuntimeGuard {
620        runtime: Some(runtime),
621        logger_shutdown: Some(ShutdownHooks::new()),
622    });
623
624    let runtime_config_ok = Arc::new(AtomicBool::new(true));
625    let runtime_config = load_runtime(&handle, &runtime_config_ok)?;
626
627    let metrics = Arc::new(MetricRegistry::new());
628    let host_metrics = Arc::new(HostMetricsRegistry::new());
629    let health_checks = Arc::new(HealthCheckRegistry::new(&handle));
630    let readiness_checks = Arc::new(ReadinessCheckRegistry::new());
631    let diagnostics = Arc::new(DiagnosticRegistry::new());
632
633    let loggers = if init_options.init_globals {
634        handle.block_on(logging::init(
635            &metrics,
636            install_config.as_ref(),
637            &runtime_config.map(|c| c.as_ref().logging().clone()),
638            runtime.logger_shutdown.as_mut().unwrap(),
639        ))
640    } else {
641        logging::get_existing()
642    }?;
643
644    info!("server starting");
645
646    if init_options.init_globals && install_config.as_ref().minidump().enabled() {
647        let minidump_ok = Arc::new(AtomicBool::new(true));
648        let socket_dir = install_config.as_ref().minidump().socket_dir();
649        handle.spawn({
650            let minidump_ok = minidump_ok.clone();
651            let socket_dir = socket_dir.to_owned();
652            async move {
653                let result = minidump::init(&socket_dir).await;
654                minidump_ok.store(result.is_ok(), Ordering::Relaxed);
655                if let Err(e) = result {
656                    error!("error during minidump init", error: e)
657                }
658            }
659        });
660
661        health_checks.register(MinidumpHealthCheck::new(minidump_ok));
662        #[cfg(target_os = "linux")]
663        diagnostics.register(ThreadDumpDiagnostic::new(socket_dir));
664    }
665
666    metrics::init(&metrics);
667
668    health_checks.register(ServiceDependencyHealthCheck::new(&host_metrics));
669    health_checks.register(PanicsHealthCheck::new());
670    health_checks.register(ConfigReloadHealthCheck::new(runtime_config_ok));
671
672    diagnostics.register(MetricNamesDiagnostic::new(&metrics));
673    #[cfg(feature = "jemalloc")]
674    {
675        if init_options.init_globals {
676            diagnostics.register(HeapStatsDiagnostic);
677            heap_profile::init(&runtime_config);
678            diagnostics.register(HeapProfileDiagnostic);
679        }
680    }
681    diagnostics.register(DiagnosticTypesDiagnostic::new(Arc::downgrade(&diagnostics)));
682
683    let client_factory = ClientFactory::builder()
684        .config(runtime_config.map(|c| c.as_ref().service_discovery().clone()))
685        .user_agent(UserAgent::new(Agent::new(
686            install_config.as_ref().product_name(),
687            install_config.as_ref().product_version(),
688        )))
689        .metrics(metrics.clone())
690        .host_metrics(host_metrics)
691        .blocking_handle(handle.clone());
692
693    let mut witchcraft = Witchcraft {
694        metrics,
695        health_checks,
696        readiness_checks,
697        client_factory,
698        diagnostics: diagnostics.clone(),
699        handle: handle.clone(),
700        install_config: install_config.as_ref().clone(),
701        thread_pool: None,
702        thread_prefix: init_options.thread_prefix,
703        endpoints: vec![],
704        shutdown_hooks: ShutdownHooks::new(),
705        conjure_runtime: Arc::new(ConjureRuntime::new()),
706    };
707
708    let status_endpoints = StatusServiceEndpoints::new(StatusResource::new(
709        &runtime_config,
710        &witchcraft.health_checks,
711        &witchcraft.readiness_checks,
712    ));
713    witchcraft.endpoints(
714        None,
715        status_endpoints.endpoints(&witchcraft.conjure_runtime),
716        false,
717    );
718
719    let debug_endpoints =
720        DebugServiceEndpoints::new(DebugResource::new(&runtime_config, &diagnostics));
721    witchcraft.app(debug_endpoints);
722
723    // A bit of a dance to avoid activating the management server before init returns, which would
724    // cause us to report ready too early.
725    let management_endpoints = mem::take(&mut witchcraft.endpoints);
726
727    init(install_config, runtime_config, &mut witchcraft)?;
728
729    witchcraft
730        .health_checks
731        .register(Endpoint500sHealthCheck::new(&witchcraft.endpoints));
732
733    let mut main_endpoints = mem::take(&mut witchcraft.endpoints);
734
735    match witchcraft.install_config.management_port() {
736        Some(management_port) if management_port != witchcraft.install_config.port() => {
737            handle.block_on(server::start(
738                &mut witchcraft,
739                management_endpoints,
740                &loggers,
741                Listener::Management,
742                management_port,
743            ))?;
744        }
745        _ => main_endpoints.extend(management_endpoints),
746    }
747
748    let port = witchcraft.install_config.port();
749    handle.block_on(server::start(
750        &mut witchcraft,
751        main_endpoints,
752        &loggers,
753        Listener::Service,
754        port,
755    ))?;
756
757    Ok(witchcraft)
758}
759
760async fn shutdown(shutdown_hooks: ShutdownHooks, timeout: Duration) -> Result<(), Error> {
761    let mut signals = pin!(signals()?);
762
763    signals.next().await;
764
765    select! {
766        _ = drain_shutdown_hooks(shutdown_hooks, timeout) => {}
767        _ = signals.next() => info!("graceful shutdown interrupted by signal"),
768    }
769
770    Ok(())
771}
772
773/// Awaits the server's shutdown hooks, giving up after `timeout` elapses.
774async fn drain_shutdown_hooks(shutdown_hooks: ShutdownHooks, timeout: Duration) {
775    info!("server shutting down");
776
777    select! {
778        _ = shutdown_hooks => {}
779        _ = time::sleep(timeout) => {
780            info!(
781                "graceful shutdown timed out",
782                safe: {
783                    timeout: format_args!("{timeout:?}"),
784                },
785            );
786        }
787    }
788}
789
790fn signals() -> Result<impl Stream<Item = ()>, Error> {
791    let sigint = signal(SignalKind::interrupt())?;
792    let sigterm = signal(SignalKind::terminate())?;
793    Ok(stream::select(sigint, sigterm))
794}
795
796fn signal(kind: SignalKind) -> Result<impl Stream<Item = ()>, Error> {
797    let mut signal = unix::signal(kind).map_err(Error::internal_safe)?;
798    Ok(stream::poll_fn(move |cx| signal.poll_recv(cx)))
799}
800
801struct RuntimeGuard {
802    runtime: Option<Runtime>,
803    logger_shutdown: Option<ShutdownHooks>,
804}
805
806impl Drop for RuntimeGuard {
807    fn drop(&mut self) {
808        let runtime = self.runtime.take().unwrap();
809        runtime.block_on(self.logger_shutdown.take().unwrap());
810        runtime.shutdown_background()
811    }
812}