1use solar_interface::diagnostics::DiagCtxt;
4use std::io;
5
6#[cfg(feature = "tracing")]
7use solar_sema::ast::Either;
8
9#[cfg(feature = "mimalloc")]
10use mimalloc as _;
11#[cfg(all(feature = "jemalloc", unix))]
12use tikv_jemallocator as _;
13
14cfg_if::cfg_if! {
17 if #[cfg(debug_assertions)] {
18 type AllocatorInner = std::alloc::System;
19 } else if #[cfg(feature = "mimalloc")] {
20 type AllocatorInner = mimalloc::MiMalloc;
21 } else if #[cfg(all(feature = "jemalloc", unix))] {
22 type AllocatorInner = tikv_jemallocator::Jemalloc;
23 } else {
24 type AllocatorInner = std::alloc::System;
25 }
26}
27
28cfg_if::cfg_if! {
29 if #[cfg(feature = "tracy-allocator")] {
30 pub(super) type WrappedAllocator = tracing_tracy::client::ProfiledAllocator<AllocatorInner>;
31 pub(super) const fn new_wrapped_allocator() -> WrappedAllocator {
32 Allocator::new(AllocatorInner {}, 100)
33 }
34 } else {
35 pub(super) type WrappedAllocator = AllocatorInner;
36 pub(super) const fn new_wrapped_allocator() -> WrappedAllocator {
37 AllocatorInner {}
38 }
39 }
40}
41
42pub type Allocator = WrappedAllocator;
44
45pub const fn new_allocator() -> Allocator {
47 new_wrapped_allocator()
48}
49
50#[derive(Default)]
52pub enum LogDestination {
53 #[default]
55 Stdout,
56 Stderr,
58}
59
60#[cfg(feature = "tracing")]
61impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for LogDestination {
62 type Writer = Either<io::Stdout, io::Stderr>;
63
64 fn make_writer(&'a self) -> Self::Writer {
65 match self {
66 Self::Stdout => Either::Left(io::stdout()),
67 Self::Stderr => Either::Right(io::stderr()),
68 }
69 }
70}
71
72#[must_use]
74pub fn init_logger(dst: LogDestination) -> impl Sized {
75 #[cfg(not(feature = "tracing"))]
76 {
77 if std::env::var_os("RUST_LOG").is_some() {
78 let msg = "`RUST_LOG` is set, but \"tracing\" support was not enabled at compile time";
79 DiagCtxt::new_early().warn(msg).emit();
80 }
81 if std::env::var_os("SOLAR_PROFILE").is_some() {
82 let msg =
83 "`SOLAR_PROFILE` is set, but \"tracing\" support was not enabled at compile time";
84 DiagCtxt::new_early().warn(msg).emit();
85 }
86 }
87
88 #[cfg(feature = "tracing")]
89 match try_init_logger(dst) {
90 Ok(guard) => guard,
91 Err(e) => DiagCtxt::new_early().fatal(e).emit(),
92 }
93}
94
95#[cfg(feature = "tracing")]
96fn try_init_logger(dst: LogDestination) -> Result<impl Sized, String> {
97 use tracing_subscriber::prelude::*;
98
99 let (profile_layer, guard) = match std::env::var("SOLAR_PROFILE").as_deref() {
100 Ok("chrome") => {
101 if !cfg!(feature = "tracing-chrome") {
102 return Err("chrome profiler support is not compiled in".to_string());
103 }
104 let (layer, guard) = chrome_layer();
105 (Some(layer.boxed()), Some(guard))
106 }
107 Ok("tracy") => {
108 if !cfg!(feature = "tracy") {
109 return Err("tracy profiler support is not compiled in".to_string());
110 }
111 (Some(tracy_layer().boxed()), Default::default())
112 }
113 Ok(s) => return Err(format!("unknown profiler '{s}'; valid values: 'chrome', 'tracy'")),
114 Err(_) => Default::default(),
115 };
116 tracing_subscriber::Registry::default()
117 .with(tracing_subscriber::EnvFilter::from_default_env())
118 .with(profile_layer)
119 .with(tracing_subscriber::fmt::layer().with_writer(dst))
120 .try_init()
121 .map(|()| guard)
122 .map_err(|e| e.to_string())
123}
124
125#[cfg(feature = "tracing")]
126#[cfg(feature = "tracy")]
127fn tracy_layer() -> tracing_tracy::TracyLayer<impl tracing_tracy::Config> {
128 struct Config(tracing_subscriber::fmt::format::DefaultFields);
129 impl tracing_tracy::Config for Config {
130 type Formatter = tracing_subscriber::fmt::format::DefaultFields;
131 fn formatter(&self) -> &Self::Formatter {
132 &self.0
133 }
134 fn format_fields_in_zone_name(&self) -> bool {
135 false
136 }
137 }
138
139 tracing_tracy::client::register_demangler!();
140
141 tracing_tracy::TracyLayer::new(Config(Default::default()))
142}
143
144#[cfg(feature = "tracing")]
145#[cfg(not(feature = "tracy"))]
146fn tracy_layer() -> tracing_subscriber::layer::Identity {
147 tracing_subscriber::layer::Identity::new()
148}
149
150#[cfg(feature = "tracing")]
151#[cfg(feature = "tracing-chrome")]
152fn chrome_layer<S>() -> (tracing_chrome::ChromeLayer<S>, tracing_chrome::FlushGuard)
153where
154 S: tracing::Subscriber
155 + for<'span> tracing_subscriber::registry::LookupSpan<'span>
156 + Send
157 + Sync,
158{
159 tracing_chrome::ChromeLayerBuilder::new().include_args(true).build()
160}
161
162#[cfg(feature = "tracing")]
163#[cfg(not(feature = "tracing-chrome"))]
164fn chrome_layer() -> (tracing_subscriber::layer::Identity, ()) {
165 (tracing_subscriber::layer::Identity::new(), ())
166}
167
168