hyperactor/
init.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9use std::sync::OnceLock;
10
11use crate::clock::ClockKind;
12use crate::panic_handler;
13
14/// A global runtime handle used for spawning tasks. Do not use for executing long running or
15/// compute intensive tasks.
16static RUNTIME: OnceLock<tokio::runtime::Handle> = OnceLock::new();
17
18/// Get a handle to the global runtime.
19///
20/// Panics if the runtime has not been initialized *and* the caller is not in an
21/// async context.
22pub(crate) fn get_runtime() -> tokio::runtime::Handle {
23    match RUNTIME.get() {
24        Some(handle) => handle.clone(),
25        None => tokio::runtime::Handle::current(),
26    }
27}
28
29/// Initialize the Hyperactor runtime. Specifically:
30/// - Set up panic handling, so that we get consistent panic stack traces in Actors.
31/// - Initialize logging defaults.
32/// - Store the provided tokio runtime handle for use by the hyperactor system.
33pub fn initialize(handle: tokio::runtime::Handle) {
34    initialize_with_log_prefix(handle, Option::None);
35}
36
37/// Initialize the Hyperactor runtime. Specifically:
38/// - Set up panic handling, so that we get consistent panic stack traces in Actors.
39/// - Initialize logging defaults.
40/// - Store the provided tokio runtime handle for use by the hyperactor system.
41/// - Set the env var whose value should be used to prefix log messages.
42pub fn initialize_with_log_prefix(
43    handle: tokio::runtime::Handle,
44    env_var_log_prefix: Option<String>,
45) {
46    RUNTIME
47        .set(handle)
48        .expect("hyperactor::initialize must only be called once");
49
50    panic_handler::set_panic_hook();
51    hyperactor_telemetry::initialize_logging_with_log_prefix(
52        ClockKind::default(),
53        env_var_log_prefix,
54    );
55}
56
57/// Initialize the Hyperactor runtime using the current tokio runtime handle.
58pub fn initialize_with_current_runtime() {
59    let handle = tokio::runtime::Handle::current();
60    initialize(handle);
61}