solar_cli/
utils.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
//! Utility functions used by the Solar CLI.

use solar_interface::{diagnostics::DiagCtxt, SessionGlobals};

#[cfg(all(feature = "jemalloc", unix))]
use tikv_jemallocator as _;

// We use jemalloc for performance reasons.
// Except in tests, where we spawn a ton of processes and jemalloc has a higher startup cost.
cfg_if::cfg_if! {
    if #[cfg(all(feature = "jemalloc", unix, not(debug_assertions)))] {
        type AllocatorInner = tikv_jemallocator::Jemalloc;
    } else {
        type AllocatorInner = std::alloc::System;
    }
}

cfg_if::cfg_if! {
    if #[cfg(feature = "tracy-allocator")] {
        pub(super) type WrappedAllocator = tracing_tracy::client::ProfiledAllocator<AllocatorInner>;
        pub(super) const fn new_wrapped_allocator() -> WrappedAllocator {
            Allocator::new(AllocatorInner {}, 100)
        }
    } else {
        pub(super) type WrappedAllocator = AllocatorInner;
        pub(super) const fn new_wrapped_allocator() -> WrappedAllocator {
            AllocatorInner {}
        }
    }
}

/// The global allocator used by the compiler.
pub type Allocator = WrappedAllocator;

/// Create a new instance of the global allocator.
pub const fn new_allocator() -> Allocator {
    new_wrapped_allocator()
}

/// Initialize the tracing logger.
pub fn init_logger() -> impl Sized {
    match try_init_logger() {
        Ok(guard) => guard,
        Err(e) => DiagCtxt::new_early().fatal(e).emit(),
    }
}

fn try_init_logger() -> Result<impl Sized, String> {
    use tracing_subscriber::prelude::*;

    let (profile_layer, guard) = match std::env::var("SOLAR_PROFILE").as_deref() {
        Ok("chrome") => {
            if !cfg!(feature = "tracing-chrome") {
                return Err("chrome profiler support is not compiled in".to_string());
            }
            let (layer, guard) = chrome_layer();
            (Some(layer.boxed()), Some(guard))
        }
        Ok("tracy") => {
            if !cfg!(feature = "tracy") {
                return Err("tracy profiler support is not compiled in".to_string());
            }
            (Some(tracy_layer().boxed()), Default::default())
        }
        Ok(s) => return Err(format!("unknown profiler '{s}'; valid values: 'chrome', 'tracy'")),
        Err(_) => Default::default(),
    };
    tracing_subscriber::Registry::default()
        .with(tracing_subscriber::EnvFilter::from_default_env())
        .with(profile_layer)
        .with(tracing_subscriber::fmt::layer())
        .try_init()
        .map(|()| guard)
        .map_err(|e| e.to_string())
}

#[cfg(feature = "tracy")]
fn tracy_layer() -> tracing_tracy::TracyLayer<impl tracing_tracy::Config> {
    struct Config(tracing_subscriber::fmt::format::DefaultFields);
    impl tracing_tracy::Config for Config {
        type Formatter = tracing_subscriber::fmt::format::DefaultFields;
        fn formatter(&self) -> &Self::Formatter {
            &self.0
        }
        fn format_fields_in_zone_name(&self) -> bool {
            false
        }
    }

    tracing_tracy::client::register_demangler!();

    tracing_tracy::TracyLayer::new(Config(Default::default()))
}

#[cfg(not(feature = "tracy"))]
fn tracy_layer() -> tracing_subscriber::layer::Identity {
    tracing_subscriber::layer::Identity::new()
}

#[cfg(feature = "tracing-chrome")]
fn chrome_layer<S>() -> (tracing_chrome::ChromeLayer<S>, tracing_chrome::FlushGuard)
where
    S: tracing::Subscriber
        + for<'span> tracing_subscriber::registry::LookupSpan<'span>
        + Send
        + Sync,
{
    tracing_chrome::ChromeLayerBuilder::new().include_args(true).build()
}

#[cfg(not(feature = "tracing-chrome"))]
fn chrome_layer() -> (tracing_subscriber::layer::Identity, ()) {
    (tracing_subscriber::layer::Identity::new(), ())
}

#[allow(dead_code)]
pub(crate) fn env_to_bool(value: Option<&std::ffi::OsStr>) -> bool {
    value.is_some_and(|value| value == "1" || value == "true")
}

/// Runs the given closure in a thread pool with the given number of threads.
pub fn run_in_thread_pool_with_globals<R: Send>(
    threads: usize,
    f: impl FnOnce(usize) -> R + Send,
) -> R {
    let mut builder =
        rayon::ThreadPoolBuilder::new().thread_name(|i| format!("solar-{i}")).num_threads(threads);
    // We still want to use a rayon thread pool with 1 thread so that `ParallelIterator` don't
    // install their own thread pool.
    if threads == 1 {
        builder = builder.use_current_thread();
    }

    // We create the session globals on the main thread, then create the thread pool. Upon creation,
    // each worker thread created gets a copy of the session globals in TLS. This is possible
    // because `SessionGlobals` impls `Send`.
    SessionGlobals::new().set(|| {
        SessionGlobals::with(|session_globals| {
            builder
                .build_scoped(
                    // Initialize each new worker thread when created.
                    move |thread| session_globals.set(|| thread.run()),
                    // Run `f` on the first thread in the thread pool.
                    move |pool| pool.install(|| f(pool.current_num_threads())),
                )
                .unwrap()
        })
    })
}