1use std::io::IsTerminal as _;
22use std::sync::Arc;
23
24use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle};
25use tracing::{Event, Level, Subscriber};
26use tracing_subscriber::layer::Context;
27
28pub const POOL_TARGET: &str = "degenbot::ingest";
30
31pub const AAVE_TARGET: &str = "degenbot::aave";
33
34#[must_use]
39pub fn stderr_opt() -> Option<ProgressDrawTarget> {
40 std::io::stderr()
41 .is_terminal()
42 .then(ProgressDrawTarget::stderr)
43}
44
45#[derive(Debug)]
47pub struct Painter {
48 bar: Option<ProgressBar>,
49}
50
51impl Painter {
52 #[must_use]
56 pub fn from_draw_target(target: Option<ProgressDrawTarget>) -> Self {
57 let bar = target.map(|target| {
58 let bar = ProgressBar::with_draw_target(Some(100), target);
59 let style = ProgressStyle::with_template("{msg} [{bar:40}] {pos:>3}%")
60 .unwrap_or_else(|_| ProgressStyle::default_bar());
61 bar.set_style(style);
62 bar
63 });
64 Self { bar }
65 }
66
67 #[must_use]
69 pub const fn is_active(&self) -> bool {
70 self.bar.is_some()
71 }
72
73 pub fn paint(&self, percent: u64, message: &str) {
76 if let Some(bar) = &self.bar {
77 bar.set_length(100);
78 bar.set_position(percent.min(100));
79 bar.set_message(message.to_string());
80 }
81 }
82
83 #[must_use]
85 pub fn position(&self) -> Option<u64> {
86 self.bar.as_ref().map(ProgressBar::position)
87 }
88
89 #[must_use]
91 pub fn message(&self) -> Option<String> {
92 self.bar.as_ref().map(ProgressBar::message)
93 }
94
95 pub fn finish(&self) {
97 if let Some(bar) = &self.bar {
98 bar.finish_and_clear();
99 }
100 }
101}
102
103#[derive(Debug)]
106pub struct Layer {
107 painter: Arc<Painter>,
108}
109
110impl Layer {
111 #[must_use]
113 pub fn new(painter: Arc<Painter>) -> Self {
114 Self { painter }
115 }
116}
117
118impl<S: Subscriber> tracing_subscriber::Layer<S> for Layer {
119 fn on_event(&self, event: &Event<'_>, _context: Context<'_, S>) {
120 let metadata = event.metadata();
121 if metadata.target() != POOL_TARGET && metadata.target() != AAVE_TARGET {
122 return;
123 }
124 if *metadata.level() != Level::INFO {
125 return;
126 }
127 let mut fields = Fields::default();
128 event.record(&mut fields);
129 if let Some(percent) = fields.progress_pct {
130 self.painter
131 .paint(percent, fields.message.as_deref().unwrap_or("update"));
132 }
133 }
134}
135
136#[derive(Debug, Default)]
138struct Fields {
139 progress_pct: Option<u64>,
140 message: Option<String>,
141}
142
143impl tracing::field::Visit for Fields {
144 fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
145 if field.name() == "progress_pct" {
146 self.progress_pct = Some(value);
147 }
148 }
149
150 fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
151 if field.name() == "message" {
152 self.message = Some(value.to_string());
153 }
154 }
155
156 fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
157 if field.name() == "message" && self.message.is_none() {
158 self.message = Some(format!("{value:?}"));
159 }
160 }
161}
162
163#[cfg(test)]
164mod tests {
165 use std::sync::Arc;
166
167 use indicatif::ProgressDrawTarget;
168 use tracing_subscriber::layer::SubscriberExt as _;
169
170 use super::{Layer, Painter, POOL_TARGET};
171
172 #[test]
173 fn no_draw_target_means_no_painter_and_no_write() {
174 let painter = Painter::from_draw_target(None);
175 assert!(!painter.is_active());
176 painter.paint(100, "pool update: chunk committed");
177 assert_eq!(painter.position(), None);
178 assert_eq!(painter.message(), None);
179 painter.finish();
180 }
181
182 #[test]
183 fn progress_layer_paints_from_the_core_progress_event() {
184 let painter = Arc::new(Painter::from_draw_target(
185 Some(ProgressDrawTarget::hidden()),
186 ));
187 assert!(painter.is_active());
188 let subscriber = tracing_subscriber::registry().with(Layer::new(Arc::clone(&painter)));
189 tracing::subscriber::with_default(subscriber, || {
190 tracing::info!(
191 target: POOL_TARGET,
192 chain_id = 8453i64,
193 chunk_start = 1u64,
194 chunk_end = 500u64,
195 progress_pct = 42u64,
196 "pool update: chunk committed"
197 );
198 });
199 assert_eq!(painter.position(), Some(42));
200 assert_eq!(
201 painter.message().as_deref(),
202 Some("pool update: chunk committed")
203 );
204 painter.finish();
205 }
206
207 #[test]
208 fn unrelated_events_do_not_paint() {
209 let painter = Arc::new(Painter::from_draw_target(
210 Some(ProgressDrawTarget::hidden()),
211 ));
212 let subscriber = tracing_subscriber::registry().with(Layer::new(Arc::clone(&painter)));
213 tracing::subscriber::with_default(subscriber, || {
214 tracing::info!(target: "degenbot::rpc", progress_pct = 99u64, "not a chunk");
215 });
216 assert_eq!(painter.position(), Some(0));
217 painter.finish();
218 }
219}