use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use axum::{
Router,
extract::State,
http::StatusCode,
response::{Html, IntoResponse, Response},
routing::get,
};
use tokio::net::TcpListener;
pub const VIEWER_HTML: &str = include_str!("viewer.html");
#[derive(Clone)]
pub enum TraceSource {
File(PathBuf),
Memory(Arc<Mutex<Vec<serde_json::Value>>>),
}
#[derive(Clone)]
pub struct TraceHandle {
events: Arc<Mutex<Vec<serde_json::Value>>>,
}
impl TraceHandle {
pub fn push(&self, event: serde_json::Value) {
self.events.lock().unwrap().push(event);
}
pub fn len(&self) -> usize {
self.events.lock().unwrap().len()
}
pub fn is_empty(&self) -> bool {
self.events.lock().unwrap().is_empty()
}
}
struct AppState {
source: TraceSource,
}
pub struct TraceViewerConfig {
pub port: u16,
pub open_browser: bool,
}
impl Default for TraceViewerConfig {
fn default() -> Self {
Self {
port: 0,
open_browser: true,
}
}
}
pub async fn serve_memory(
config: TraceViewerConfig,
) -> anyhow::Result<(
TraceHandle,
impl std::future::Future<Output = anyhow::Result<()>>,
)> {
let events = Arc::new(Mutex::new(Vec::new()));
let handle = TraceHandle {
events: events.clone(),
};
let source = TraceSource::Memory(events);
let server = serve_impl(source, config);
Ok((handle, server))
}
pub async fn serve_file(path: PathBuf, config: TraceViewerConfig) -> anyhow::Result<()> {
serve_impl(TraceSource::File(path), config).await
}
async fn serve_impl(source: TraceSource, config: TraceViewerConfig) -> anyhow::Result<()> {
let state = Arc::new(AppState { source });
let app = Router::new()
.route("/", get(serve_viewer))
.route("/events", get(serve_events))
.with_state(state);
let listener = TcpListener::bind(format!("127.0.0.1:{}", config.port)).await?;
let addr = listener.local_addr()?;
eprintln!("Trace viewer at http://{}", addr);
if config.open_browser {
let url = format!("http://{}", addr);
if let Err(e) = open::that(&url) {
eprintln!("Failed to open browser: {}. Open {} manually.", e, url);
}
}
axum::serve(listener, app).await?;
Ok(())
}
async fn serve_viewer() -> Html<&'static str> {
Html(VIEWER_HTML)
}
async fn serve_events(State(state): State<Arc<AppState>>) -> Response {
match &state.source {
TraceSource::File(path) => serve_events_from_file(path).await,
TraceSource::Memory(events) => serve_events_from_memory(events),
}
}
async fn serve_events_from_file(path: &PathBuf) -> Response {
match tokio::fs::read_to_string(path).await {
Ok(content) => {
let events: Vec<serde_json::Value> = content
.lines()
.filter(|line| !line.trim().is_empty())
.filter_map(|line| serde_json::from_str(line).ok())
.collect();
match serde_json::to_string(&events) {
Ok(json) => {
(StatusCode::OK, [("content-type", "application/json")], json).into_response()
}
Err(e) => (
StatusCode::INTERNAL_SERVER_ERROR,
format!("Failed to serialize events: {}", e),
)
.into_response(),
}
}
Err(e) => (
StatusCode::INTERNAL_SERVER_ERROR,
format!("Failed to read trace file: {}", e),
)
.into_response(),
}
}
fn serve_events_from_memory(events: &Arc<Mutex<Vec<serde_json::Value>>>) -> Response {
let events = events.lock().unwrap();
match serde_json::to_string(&*events) {
Ok(json) => (StatusCode::OK, [("content-type", "application/json")], json).into_response(),
Err(e) => (
StatusCode::INTERNAL_SERVER_ERROR,
format!("Failed to serialize events: {}", e),
)
.into_response(),
}
}