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