Skip to main content

dial9_viewer/
lib.rs

1pub mod cli;
2pub mod server;
3pub mod storage;
4
5use std::path::PathBuf;
6
7async fn detect_bucket_region(bucket: &str) -> Option<String> {
8    let config = aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await;
9    let client = aws_sdk_s3::Client::new(&config);
10    match client.head_bucket().bucket(bucket).send().await {
11        Ok(resp) => resp.bucket_region().map(|r| r.to_string()),
12        Err(err) => {
13            let raw = err.raw_response();
14            raw.and_then(|r| {
15                r.headers()
16                    .get("x-amz-bucket-region")
17                    .map(|v| v.to_string())
18            })
19        }
20    }
21}
22
23pub(crate) async fn serve(
24    port: u16,
25    bucket: Option<String>,
26    prefix: Option<String>,
27    local_dir: Option<PathBuf>,
28    dev: bool,
29) -> anyhow::Result<()> {
30    tracing_subscriber::fmt()
31        .with_env_filter(
32            tracing_subscriber::EnvFilter::try_from_default_env()
33                .unwrap_or_else(|_| "dial9_viewer=info".parse().unwrap()),
34        )
35        .init();
36
37    let dev_ui_dir = if dev {
38        let candidates = [PathBuf::from("ui"), PathBuf::from("dial9-viewer/ui")];
39        let dir = candidates.into_iter().find(|p| p.exists());
40        match dir {
41            Some(d) => {
42                tracing::info!(path = %d.display(), "dev mode: serving UI from disk");
43                Some(d)
44            }
45            None => {
46                anyhow::bail!(
47                    "--dev: could not find ui/ directory. Run from the dial9-viewer/ or repo root directory."
48                );
49            }
50        }
51    } else {
52        None
53    };
54
55    let app_state = if let Some(dir) = &local_dir {
56        let dir = std::fs::canonicalize(dir)?;
57        tracing::info!(path = %dir.display(), "serving traces from local directory");
58        let backend = storage::LocalBackend::new(&dir);
59        let mut state = server::AppState::new(
60            std::sync::Arc::new(backend),
61            Some("local".into()),
62            prefix.clone(),
63        );
64        if let Some(d) = dev_ui_dir {
65            state = state.with_dev_ui_dir(d);
66        }
67        state
68    } else if let Some(bucket_name) = &bucket {
69        if let Some(region) = detect_bucket_region(bucket_name).await {
70            tracing::info!(%region, bucket = %bucket_name, "detected bucket region");
71            let config = aws_config::defaults(aws_config::BehaviorVersion::latest())
72                .region(aws_sdk_s3::config::Region::new(region))
73                .load()
74                .await;
75            let client = aws_sdk_s3::Client::new(&config);
76            let backend = storage::S3Backend::from_client(client);
77            let mut state =
78                server::AppState::new(std::sync::Arc::new(backend), bucket.clone(), prefix.clone());
79            if let Some(d) = dev_ui_dir {
80                state = state.with_dev_ui_dir(d);
81            }
82            state
83        } else {
84            tracing::warn!(bucket = %bucket_name, "could not detect bucket region, using default");
85            let backend = storage::S3Backend::from_env().await;
86            let mut state =
87                server::AppState::new(std::sync::Arc::new(backend), bucket.clone(), prefix.clone());
88            if let Some(d) = dev_ui_dir {
89                state = state.with_dev_ui_dir(d);
90            }
91            state
92        }
93    } else {
94        let backend = storage::S3Backend::from_env().await;
95        let mut state =
96            server::AppState::new(std::sync::Arc::new(backend), bucket.clone(), prefix.clone());
97        if let Some(d) = dev_ui_dir {
98            state = state.with_dev_ui_dir(d);
99        }
100        state
101    };
102
103    let app = server::router(app_state);
104
105    let listener = tokio::net::TcpListener::bind(("0.0.0.0", port)).await?;
106    tracing::info!(port, dev, "dial9-viewer listening");
107    println!("\n  → http://localhost:{}\n", port);
108    if let Some(dir) = &local_dir {
109        tracing::info!(path = %dir.display(), "local directory mode");
110    } else if let Some(bucket) = &bucket {
111        tracing::info!(%bucket, "default bucket");
112    }
113
114    axum::serve(listener, app)
115        .with_graceful_shutdown(shutdown_signal())
116        .await?;
117
118    Ok(())
119}
120
121async fn shutdown_signal() {
122    tokio::signal::ctrl_c()
123        .await
124        .expect("failed to install CTRL+C handler");
125    tracing::info!("shutting down");
126}