use scope::metrics::{Metric, Sample};
use std::collections::VecDeque;
use std::fmt;
use std::io::{self, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Instant;
type Sink = Arc<Mutex<dyn Write + Send>>;
#[derive(Debug, Clone, Default)]
pub struct RunProgress {
pub run: u64,
pub total: Option<u64>,
pub next_at: Option<Instant>,
}
pub type ProgressHook = Arc<Mutex<RunProgress>>;
#[derive(Debug, Clone)]
pub struct RunRecord {
pub run: u64,
pub at: chrono::DateTime<chrono::Utc>,
pub sample: Sample,
pub snapshot: String,
pub view: Option<Arc<SummaryView>>,
}
#[derive(Debug, Clone)]
pub struct SummaryView {
pub summary: scope::market::MarketSummary,
pub venue: String,
pub tokens: Vec<scope::market::assets::TokenAddress>,
}
pub type ViewHook = Arc<Mutex<Option<SummaryView>>>;
#[derive(Debug, Default)]
pub struct Series {
pub metrics: Vec<Metric>,
pub runs: VecDeque<RunRecord>,
pub dropped: u64,
}
pub type SeriesHook = Arc<Mutex<Series>>;
pub struct Ask {
pub prompt: String,
pub options: Vec<String>,
pub reply: tokio::sync::oneshot::Sender<Option<usize>>,
}
pub type Asker = Arc<Mutex<Option<Ask>>>;
pub const MAX_SERIES_RUNS: usize = 500;
#[derive(Clone)]
pub struct Output {
out: Sink,
err: Sink,
progress: bool,
hook: Option<ProgressHook>,
series: Option<SeriesHook>,
tui: bool,
prompts: bool,
asker: Option<Asker>,
view: Option<ViewHook>,
}
impl Output {
pub fn stdio() -> Self {
Self {
out: Arc::new(Mutex::new(io::stdout())),
err: Arc::new(Mutex::new(io::stderr())),
progress: true,
hook: None,
series: None,
tui: false,
prompts: true,
asker: None,
view: None,
}
}
pub fn capture() -> (Self, Captured) {
let out = Arc::new(Mutex::new(Vec::new()));
let err = Arc::new(Mutex::new(Vec::new()));
let captured = Captured {
data: out.clone(),
out: out.clone(),
err: err.clone(),
trimmed: Arc::default(),
};
let output = Self {
out,
err,
progress: false,
hook: None,
series: None,
tui: false,
prompts: false,
asker: None,
view: None,
};
(output, captured)
}
pub fn capture_merged() -> (Self, Captured) {
let buf = Arc::new(Mutex::new(Vec::new()));
let data = Arc::new(Mutex::new(Vec::new()));
let captured = Captured {
out: buf.clone(),
err: buf.clone(),
data: data.clone(),
trimmed: Arc::default(),
};
let output = Self {
out: Arc::new(Mutex::new(Tee {
merged: buf.clone(),
data,
})),
err: buf,
progress: false,
hook: Some(ProgressHook::default()),
series: Some(SeriesHook::default()),
tui: true,
prompts: false,
asker: None,
view: Some(ViewHook::default()),
};
(output, captured)
}
pub fn teed(&self, limit: usize) -> (Self, Captured) {
let buf = Arc::new(Mutex::new(Vec::new()));
let data = Arc::new(Mutex::new(Vec::new()));
let trimmed = Arc::new(AtomicBool::new(false));
let captured = Captured {
out: buf.clone(),
err: buf.clone(),
data: data.clone(),
trimmed: trimmed.clone(),
};
let copy = |real: &Sink, data: Option<Arc<Mutex<Vec<u8>>>>| -> Sink {
Arc::new(Mutex::new(TeeSink {
real: real.clone(),
merged: buf.clone(),
data,
limit,
trimmed: trimmed.clone(),
}))
};
let output = Self {
out: copy(&self.out, Some(data.clone())),
err: copy(&self.err, None),
progress: self.progress,
hook: None,
series: None,
tui: true,
prompts: self.prompts,
asker: None,
view: None,
};
(output, captured)
}
#[cfg(test)]
pub(crate) fn from_writers(
out: impl Write + Send + 'static,
err: impl Write + Send + 'static,
) -> Self {
Self {
out: Arc::new(Mutex::new(out)),
err: Arc::new(Mutex::new(err)),
progress: false,
hook: None,
series: None,
tui: false,
prompts: true,
asker: None,
view: None,
}
}
pub fn progress_enabled(&self) -> bool {
self.progress
}
pub fn without_prompts(mut self) -> Self {
self.prompts = false;
self
}
pub fn with_asker(mut self, asker: Asker) -> Self {
self.asker = Some(asker);
self
}
pub fn asker(&self) -> Option<Asker> {
self.asker.clone()
}
pub fn prompts_allowed(&self) -> bool {
self.prompts
}
pub fn is_tui(&self) -> bool {
self.tui
}
pub fn progress_hook(&self) -> Option<ProgressHook> {
self.hook.clone()
}
pub fn series_hook(&self) -> Option<SeriesHook> {
self.series.clone()
}
pub fn view_hook(&self) -> Option<ViewHook> {
self.view.clone()
}
pub fn record_view(&self, view: SummaryView) -> bool {
let Some(slot) = &self.view else {
return false;
};
*slot.lock().unwrap_or_else(|e| e.into_inner()) = Some(view);
true
}
pub fn record_run(&self, metrics: &[Metric], record: RunRecord) -> bool {
let Some(series) = &self.series else {
return false;
};
let mut series = lock(series);
if series.metrics.is_empty() {
series.metrics = metrics.to_vec();
}
series.runs.push_back(record);
while series.runs.len() > MAX_SERIES_RUNS {
series.runs.pop_front();
series.dropped += 1;
}
true
}
pub fn update_progress(&self, change: impl FnOnce(&mut RunProgress)) {
if let Some(hook) = &self.hook {
change(&mut lock(hook));
}
}
pub fn write_out(&self, args: fmt::Arguments<'_>) -> io::Result<()> {
emit(&self.out, args)
}
pub fn write_err(&self, args: fmt::Arguments<'_>) -> io::Result<()> {
emit(&self.err, args)
}
pub fn note(&self, args: fmt::Arguments<'_>) {
let _ = emit(&self.err, args);
}
pub fn note_saved(&self, what: &str, path: impl fmt::Display) {
crate::note!(self, "{what} saved to: {path}");
}
}
#[derive(Clone)]
pub struct Captured {
data: Arc<Mutex<Vec<u8>>>,
out: Arc<Mutex<Vec<u8>>>,
err: Arc<Mutex<Vec<u8>>>,
trimmed: Arc<AtomicBool>,
}
impl Captured {
pub fn trimmed(&self) -> bool {
self.trimmed.load(Ordering::Relaxed)
}
pub fn out(&self) -> String {
String::from_utf8_lossy(&lock(&self.out)).into_owned()
}
pub fn data(&self) -> String {
String::from_utf8_lossy(&lock(&self.data)).into_owned()
}
pub fn err(&self) -> String {
String::from_utf8_lossy(&lock(&self.err)).into_owned()
}
pub fn out_sink(&self) -> scope::diag::Sink {
self.out.clone()
}
}
struct Tee {
merged: Arc<Mutex<Vec<u8>>>,
data: Arc<Mutex<Vec<u8>>>,
}
impl Write for Tee {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
lock(&self.merged).extend_from_slice(buf);
lock(&self.data).extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
struct TeeSink {
real: Sink,
merged: Arc<Mutex<Vec<u8>>>,
data: Option<Arc<Mutex<Vec<u8>>>>,
limit: usize,
trimmed: Arc<AtomicBool>,
}
impl Write for TeeSink {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
for copy in std::iter::once(&self.merged).chain(self.data.as_ref()) {
let mut copy = lock(copy);
copy.extend_from_slice(buf);
if keep_tail(&mut copy, self.limit) {
self.trimmed.store(true, Ordering::Relaxed);
}
}
lock(&self.real).write_all(buf)?;
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
lock(&self.real).flush()
}
}
fn keep_tail(buf: &mut Vec<u8>, limit: usize) -> bool {
if buf.len() <= limit {
return false;
}
let cut = buf.len() - limit / 2;
let start = buf[cut..]
.iter()
.position(|&b| b == b'\n')
.map(|i| cut + i + 1)
.filter(|&start| start < buf.len())
.unwrap_or_else(|| whole_cut(buf, cut));
buf.drain(..start);
true
}
fn whole_cut(buf: &[u8], cut: usize) -> usize {
let mut i = 0;
while i < cut {
i = crate::cli::ansi::escape_end(buf, i).unwrap_or(i + 1);
}
let continuation = |b: u8| b & 0b1100_0000 == 0b1000_0000;
buf[i.min(buf.len())..]
.iter()
.position(|&b| !continuation(b))
.map_or(buf.len(), |n| i + n)
}
fn lock<T: ?Sized>(m: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
m.lock().unwrap_or_else(|poisoned| poisoned.into_inner())
}
fn emit(sink: &Mutex<dyn Write + Send>, args: fmt::Arguments<'_>) -> io::Result<()> {
lock(sink).write_fmt(args)
}
#[macro_export]
macro_rules! outln {
($o:expr) => {
$o.write_out(format_args!("\n"))
};
($o:expr, $($arg:tt)*) => {
$o.write_out(format_args!("{}\n", format_args!($($arg)*)))
};
}
#[macro_export]
macro_rules! out {
($o:expr, $($arg:tt)*) => {
$o.write_out(format_args!($($arg)*))
};
}
#[macro_export]
macro_rules! errln {
($o:expr) => {
$o.write_err(format_args!("\n"))
};
($o:expr, $($arg:tt)*) => {
$o.write_err(format_args!("{}\n", format_args!($($arg)*)))
};
}
#[macro_export]
macro_rules! note {
($o:expr) => {
$o.note(format_args!("\n"))
};
($o:expr, $($arg:tt)*) => {
$o.note(format_args!("{}\n", format_args!($($arg)*)))
};
}
#[macro_export]
macro_rules! err {
($o:expr, $($arg:tt)*) => {
$o.write_err(format_args!($($arg)*))
};
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
#[test]
fn test_capture_separates_data_from_diagnostics() {
let (o, cap) = Output::capture();
outln!(o, "{{\"ok\":{}}}", true).unwrap();
errln!(o, "Warning: {}", "partial data").unwrap();
assert_eq!(cap.out(), "{\"ok\":true}\n");
assert_eq!(cap.err(), "Warning: partial data\n");
}
#[test]
fn test_a_note_goes_to_the_diagnostic_channel_and_never_fails() {
let (o, cap) = Output::capture();
note!(o, "Fetching {}...", "USDC");
assert_eq!(cap.err(), "Fetching USDC...\n");
assert_eq!(cap.out(), "");
let closed = Output::from_writers(Vec::new(), ClosedPipe);
note!(closed, "dropped");
}
#[test]
fn test_note_saved_is_one_line_on_the_diagnostic_channel() {
let (o, cap) = Output::capture();
o.note_saved("Report", "r.md");
assert_eq!(cap.err(), "Report saved to: r.md\n");
assert_eq!(cap.out(), "");
}
#[test]
fn test_clones_share_writers() {
let (o, cap) = Output::capture();
let o2 = o.clone();
out!(o, "a").unwrap();
out!(o2, "b").unwrap();
outln!(o).unwrap();
assert_eq!(cap.out(), "ab\n");
}
pub(crate) struct ClosedPipe;
impl Write for ClosedPipe {
fn write(&mut self, _: &[u8]) -> io::Result<usize> {
Err(io::Error::new(io::ErrorKind::BrokenPipe, "closed"))
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
#[test]
fn test_merged_capture_keeps_data_apart_from_diagnostics() {
let (o, cap) = Output::capture_merged();
outln!(o, "data").unwrap();
errln!(o, "warn").unwrap();
assert_eq!(cap.data(), "data\n");
assert_eq!(cap.out(), "data\nwarn\n");
}
#[test]
fn test_failed_write_is_returned_not_dropped() {
let o = Output::from_writers(ClosedPipe, ClosedPipe);
let e = outln!(o, "data").unwrap_err();
assert_eq!(e.kind(), io::ErrorKind::BrokenPipe);
assert_eq!(
errln!(o, "warn").unwrap_err().kind(),
io::ErrorKind::BrokenPipe
);
}
#[test]
fn test_only_the_tui_capture_is_tui() {
assert!(Output::capture_merged().0.is_tui());
assert!(!Output::capture().0.is_tui());
assert!(!Output::stdio().is_tui());
assert!(Output::stdio().teed(1).0.is_tui());
}
#[test]
fn test_capture_merged_keeps_write_order() {
let (o, cap) = Output::capture_merged();
outln!(o, "row 1").unwrap();
errln!(o, "warning").unwrap();
outln!(o, "row 2").unwrap();
assert_eq!(cap.out(), "row 1\nwarning\nrow 2\n");
}
#[test]
fn test_teed_output_reaches_the_terminal_and_keeps_a_copy() {
let (terminal, screen) = Output::capture();
let (o, copy) = terminal.teed(1024);
outln!(o, "row").unwrap();
errln!(o, "warn").unwrap();
assert_eq!(screen.out(), "row\n");
assert_eq!(screen.err(), "warn\n");
assert_eq!(copy.out(), "row\nwarn\n");
assert_eq!(copy.data(), "row\n");
assert!(!copy.trimmed());
assert!(o.is_tui());
assert_eq!(o.progress_enabled(), terminal.progress_enabled());
assert!(Output::stdio().teed(1).0.progress_enabled());
assert!(o.progress_hook().is_none());
}
#[test]
fn test_teed_copy_keeps_whole_lines_from_the_end_within_its_limit() {
let (terminal, screen) = Output::capture();
let (o, copy) = terminal.teed(100);
for i in 0..1000 {
outln!(o, "line {i:04}").unwrap();
}
assert!(copy.trimmed());
let kept = copy.out();
assert!(kept.len() <= 100, "{}", kept.len());
assert!(kept.starts_with("line "), "whole lines: {kept:?}");
assert!(kept.ends_with("line 0999\n"), "{kept:?}");
assert!(copy.data().ends_with("line 0999\n"));
assert_eq!(screen.out().lines().count(), 1000, "the terminal got all");
}
#[test]
fn test_teed_copy_keeps_the_tail_of_one_long_line() {
let (terminal, _screen) = Output::capture();
let (o, copy) = terminal.teed(100);
let blob: String = (0..300)
.map(|i| char::from(b'a' + (i % 26) as u8))
.collect();
out!(o, "{blob}").unwrap();
assert!(copy.trimmed());
let kept = copy.out();
assert!(!kept.is_empty() && kept.len() <= 100, "{}", kept.len());
assert!(blob.ends_with(&kept), "the end of the line: {kept:?}");
}
#[test]
fn test_trim_of_one_line_splits_no_character_and_no_escape() {
let mut chars = "é".repeat(60).into_bytes(); assert!(keep_tail(&mut chars, 102)); assert!(String::from_utf8(chars.clone()).is_ok(), "{chars:?}");
assert_eq!(chars.len(), 50);
let mut colored = b"x".repeat(60);
colored.extend_from_slice(b"\x1b[38;5;196mred");
assert!(keep_tail(&mut colored, 20));
assert_eq!(colored, b"red");
let mut link = b"y".repeat(60);
link.extend_from_slice(b"\x1b]8;;https://example.com/a/long/path\x1b\\text");
let cut_at = link.len() - 40 / 2;
assert!(cut_at > 60 && cut_at < 98, "{cut_at}");
assert!(keep_tail(&mut link, 40));
assert_eq!(link, b"text");
}
#[test]
fn test_teed_copy_keeps_text_the_terminal_refused() {
let (o, copy) = Output::from_writers(ClosedPipe, ClosedPipe).teed(64);
assert!(outln!(o, "data").is_err());
assert_eq!(copy.out(), "data\n");
}
#[test]
fn test_capture_disables_progress() {
let (o, _) = Output::capture();
assert!(!o.progress_enabled());
assert!(Output::stdio().progress_enabled());
}
}