Skip to main content

tael_server/
lib.rs

1//! tael-server: OTLP ingest, tiered storage, and the REST/gRPC query API.
2//!
3//! Shipped as a library so the `tael` binary can embed it behind `tael serve`
4//! (a single `cargo install`), while still being usable as a standalone crate.
5//! [`run`] is the entrypoint; [`ServerConfig`] configures it.
6
7mod api;
8mod cluster;
9mod config;
10mod ingest;
11mod log_bus;
12mod promql;
13mod span_bus;
14mod storage;
15
16use std::sync::Arc;
17
18use anyhow::{Context, Result};
19use tokio::net::TcpListener;
20use tonic::transport::Server as TonicServer;
21use tracing_subscriber::EnvFilter;
22
23pub use config::{ServerConfig, StorageBackend};
24pub use storage::models::{
25    LogRecord, LogSeverity, MetricPoint, MetricType, Span, SpanEvent, SpanKind, SpanStatus,
26    TraceQuery,
27};
28pub use storage::{
29    BlobStore, DuckDbStore, FanoutStore, RemoteStore, RemoteWalSink, Store, TaelBackend, WalSink,
30};
31
32use log_bus::LogBus;
33use span_bus::SpanBus;
34
35/// Periodically roll spans older than the hot-tier window into the cold tier.
36/// Runs the (blocking) compaction off the async executor. The window
37/// (`retention.traces.hot_tier`, default 24h) and interval are env-tunable
38/// (`TAEL_HOT_TIER_HOURS`, `TAEL_COMPACT_INTERVAL_SECS`) until retention config
39/// lands (Phase 7); a 0-hour window compacts everything (used in tests).
40fn spawn_span_compactor(backend: Arc<TaelBackend>, blobs: Arc<BlobStore>) {
41    let window_hours: i64 = std::env::var("TAEL_HOT_TIER_HOURS")
42        .ok()
43        .and_then(|s| s.parse().ok())
44        .unwrap_or(24);
45    let interval_secs: u64 = std::env::var("TAEL_COMPACT_INTERVAL_SECS")
46        .ok()
47        .and_then(|s| s.parse().ok())
48        .unwrap_or(3600);
49    // Span metadata retention (`retention.traces.metadata`, default 365d).
50    let retention_days: i64 = std::env::var("TAEL_TRACE_RETENTION_DAYS")
51        .ok()
52        .and_then(|s| s.parse().ok())
53        .unwrap_or(365);
54    tokio::spawn(async move {
55        let mut tick = tokio::time::interval(std::time::Duration::from_secs(interval_secs));
56        loop {
57            tick.tick().await;
58            let backend = Arc::clone(&backend);
59            let blobs = Arc::clone(&blobs);
60            let result = tokio::task::spawn_blocking(move || {
61                let now = chrono::Utc::now();
62                let hot_cutoff = now - chrono::Duration::hours(window_hours);
63                let mut compacted = backend.compact_spans(hot_cutoff)?;
64                compacted += backend.compact_logs_metrics(hot_cutoff)?;
65                let dropped =
66                    backend.enforce_span_retention(now - chrono::Duration::days(retention_days))?;
67                // Payload blob GC: drop blobs no live row references (e.g. rows
68                // just removed by retention). Runs after partition drops.
69                let live = backend.collect_live_blob_hashes()?;
70                let blobs_gcd = blobs.gc(&live)?;
71                anyhow::Ok((compacted, dropped, blobs_gcd))
72            })
73            .await;
74            match result {
75                Ok(Ok((c, d, g))) if c > 0 || d > 0 || g > 0 => tracing::info!(
76                    compacted = c,
77                    partitions_dropped = d,
78                    blobs_gcd = g,
79                    "tael-backend maintenance"
80                ),
81                Ok(Ok(_)) => {}
82                Ok(Err(e)) => tracing::warn!(error = %e, "maintenance failed"),
83                Err(e) => tracing::warn!(error = %e, "maintenance task panicked"),
84            }
85        }
86    });
87}
88
89/// Start the server: OTLP gRPC + REST listeners over the configured storage
90/// backend, plus the background maintenance task when running on tael-backend.
91/// Runs until one of the listeners exits.
92pub async fn run(config: ServerConfig) -> Result<()> {
93    // Initialize tracing for the server process. `try_init` so embedding this in
94    // a binary that already set a subscriber is a no-op rather than a panic.
95    let _ = tracing_subscriber::fmt()
96        .with_env_filter(EnvFilter::from_default_env())
97        .try_init();
98
99    let blobs = Arc::new(BlobStore::new(&config.data_dir)?);
100
101    // Cluster coordination (chitchat): automatic leader election + epoch fencing
102    // of WAL replication (§5.1). On when TAEL_CLUSTER_LISTEN is set.
103    let coordinator = match &config.cluster {
104        Some(cs) => {
105            let coord = cluster::ClusterCoordinator::start(cluster::ClusterConfig {
106                node_id: cs.node_id.clone(),
107                listen_addr: cs
108                    .listen_addr
109                    .parse()
110                    .context("parsing TAEL_CLUSTER_LISTEN")?,
111                advertise_addr: cs
112                    .advertise_addr
113                    .parse()
114                    .context("parsing TAEL_CLUSTER_ADVERTISE")?,
115                seeds: cs.seeds.clone(),
116                cluster_id: cs.cluster_id.clone(),
117            })
118            .await?;
119            Some(coord)
120        }
121        None => None,
122    };
123
124    // The payload search index is shared between the ingest path (writes) and
125    // the tael-backend query path (reads); present only when that engine runs.
126    let mut search: Option<Arc<storage::SearchIndex>> = None;
127    let store: Arc<dyn Store> = if !config.query_shards.is_empty() {
128        // Stateless query-tier mode: serve reads by scatter-gather over remote
129        // shards, no local engine (`docs/tael-server-scaling-ha.md` §3, Phase 2).
130        let shards = config
131            .query_shards
132            .iter()
133            .map(|url| RemoteStore::new(url).map(|s| Arc::new(s) as Arc<dyn Store>))
134            .collect::<Result<Vec<_>>>()?;
135        tracing::info!(
136            shards = shards.len(),
137            "query fan-out mode: reads scatter-gather across remote shards (no local engine)"
138        );
139        Arc::new(FanoutStore::new(shards)?)
140    } else {
141        match config.storage {
142            StorageBackend::Duckdb => Arc::new(DuckDbStore::new(&config.data_dir)?),
143            StorageBackend::TaelBackend => {
144                // WAL replication: when standbys are configured, this node is a
145                // leader that ships every appended record to them before acking
146                // (§5.1). With no standbys the write path is unchanged.
147                let sinks: Vec<Arc<dyn WalSink>> = config
148                    .wal_standbys
149                    .iter()
150                    .map(|url| {
151                        // Stamp the leader epoch (for standby fencing) when a
152                        // coordinator is running; otherwise ship unfenced.
153                        let sink = match &coordinator {
154                            Some(c) => RemoteWalSink::with_epoch(url, c.leader_epoch_handle()),
155                            None => RemoteWalSink::new(url),
156                        };
157                        sink.map(|s| Arc::new(s) as Arc<dyn WalSink>)
158                    })
159                    .collect::<Result<Vec<_>>>()?;
160                let backend = Arc::new(if sinks.is_empty() {
161                    TaelBackend::new(&config.data_dir)?
162                } else {
163                    tracing::info!(
164                        standbys = sinks.len(),
165                        required_acks = ?config.wal_required_acks,
166                        "WAL replication enabled: shipping to standbys (leader)"
167                    );
168                    TaelBackend::with_wal_key_and_sinks(
169                        &config.data_dir,
170                        "tael-backend",
171                        sinks,
172                        config.wal_required_acks,
173                    )?
174                });
175                search = Some(backend.search_index());
176                spawn_span_compactor(Arc::clone(&backend), Arc::clone(&blobs));
177                backend as Arc<dyn Store>
178            }
179        }
180    };
181    let bus = Arc::new(SpanBus::new()?);
182    let log_bus = Arc::new(LogBus::new()?);
183
184    tracing::info!(
185        otlp_grpc = %config.otlp_grpc_addr,
186        rest_api = %config.rest_api_addr,
187        data_dir = %config.data_dir,
188        storage = ?config.storage,
189        "starting tael server"
190    );
191
192    let grpc_handle = tokio::spawn({
193        let store = Arc::clone(&store);
194        let blobs = Arc::clone(&blobs);
195        let bus = Arc::clone(&bus);
196        let log_bus = Arc::clone(&log_bus);
197        let addr = config.otlp_grpc_addr.parse()?;
198        async move {
199            let trace_service = ingest::otlp::OtlpTraceService::new(
200                Arc::clone(&store),
201                Arc::clone(&blobs),
202                search.clone(),
203                bus,
204            );
205            let logs_service = ingest::otlp_logs::OtlpLogsService::new(
206                Arc::clone(&store),
207                Arc::clone(&blobs),
208                log_bus,
209            );
210            let metrics_service = ingest::otlp_metrics::OtlpMetricsService::new(store);
211            TonicServer::builder()
212                .add_service(
213                    opentelemetry_proto::tonic::collector::trace::v1::trace_service_server::TraceServiceServer::new(trace_service),
214                )
215                .add_service(
216                    opentelemetry_proto::tonic::collector::logs::v1::logs_service_server::LogsServiceServer::new(logs_service),
217                )
218                .add_service(
219                    opentelemetry_proto::tonic::collector::metrics::v1::metrics_service_server::MetricsServiceServer::new(metrics_service),
220                )
221                .serve_with_shutdown(addr, shutdown_signal())
222                .await
223                .expect("gRPC server failed");
224        }
225    });
226
227    let rest_handle = tokio::spawn({
228        let store = Arc::clone(&store);
229        let blobs = Arc::clone(&blobs);
230        let bus = Arc::clone(&bus);
231        let log_bus = Arc::clone(&log_bus);
232        let cluster = coordinator.clone();
233        let addr = config.rest_api_addr.clone();
234        async move {
235            let app = api::rest::router(store, blobs, bus, log_bus, cluster);
236            let listener = TcpListener::bind(&addr)
237                .await
238                .expect("failed to bind REST addr");
239            tracing::info!(%addr, "REST API listening");
240            axum::serve(listener, app)
241                .with_graceful_shutdown(shutdown_signal())
242                .await
243                .expect("REST server failed");
244        }
245    });
246
247    print_startup_banner(&config);
248
249    // Both listeners drain on SIGTERM/Ctrl-C; await both so in-flight requests
250    // finish before we flush and exit (`docs/tael-server-scaling-ha.md` §5.4).
251    let (grpc_res, rest_res) = tokio::join!(grpc_handle, rest_handle);
252    grpc_res?;
253    rest_res?;
254
255    // Best-effort flush so a restart/standby replays less WAL. Durability is
256    // already guaranteed by the per-write WAL fsync.
257    if let Err(e) = store.flush() {
258        tracing::warn!(error = %e, "flush on shutdown failed");
259    }
260    tracing::info!("tael server stopped");
261
262    Ok(())
263}
264
265/// Friendly stdout banner shown on startup so a user running `tael serve`
266/// (with or without `--port`) immediately sees where to connect a CLI and
267/// where to point an OTLP exporter. Goes through `println!` so it's visible
268/// regardless of `RUST_LOG`.
269fn print_startup_banner(config: &ServerConfig) {
270    let rest = &config.rest_api_addr;
271    let otlp = &config.otlp_grpc_addr;
272    let connect_flag = cli_connect_flag(rest);
273
274    println!("tael server starting");
275    println!("  REST API     http://{rest}");
276    println!("  OTLP gRPC    {otlp}");
277    println!("  data dir     {}", config.data_dir);
278    println!("  storage      {:?}", config.storage);
279    println!();
280    println!("Connect a CLI from this machine:");
281    println!("  tael{connect_flag} services");
282    println!("  tael{connect_flag} live");
283    println!();
284    println!("Point a service at this server (OTLP):");
285    println!("  export OTEL_EXPORTER_OTLP_ENDPOINT=http://{otlp}");
286    println!("  export OTEL_EXPORTER_OTLP_PROTOCOL=grpc");
287    println!("  export OTEL_SERVICE_NAME=<your-service>");
288    println!();
289}
290
291/// Pick the CLI flag (if any) needed to reach this REST listener. Empty when
292/// REST is on the CLI default `127.0.0.1:7701`; `--port-rest N` when only the
293/// port differs; full `--server …` otherwise.
294fn cli_connect_flag(rest_addr: &str) -> String {
295    let (host, port) = match rest_addr.rsplit_once(':') {
296        Some((h, p)) => (h, p),
297        None => return String::new(),
298    };
299    let local = matches!(
300        host,
301        "127.0.0.1" | "localhost" | "0.0.0.0" | "::1" | "[::1]"
302    );
303    match (local, port) {
304        (true, "7701") => String::new(),
305        (true, p) => format!(" --port-rest {p}"),
306        (false, _) => format!(" --server http://{rest_addr}"),
307    }
308}
309
310/// Resolve when the process is asked to stop: Ctrl-C, or SIGTERM on Unix
311/// (the orchestrator's graceful-stop signal). Both listeners await their own
312/// copy; the OS delivers the signal to every registered handler.
313async fn shutdown_signal() {
314    let ctrl_c = async {
315        let _ = tokio::signal::ctrl_c().await;
316    };
317
318    #[cfg(unix)]
319    let terminate = async {
320        match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
321            Ok(mut s) => {
322                s.recv().await;
323            }
324            Err(e) => {
325                tracing::warn!(error = %e, "failed to install SIGTERM handler");
326                std::future::pending::<()>().await;
327            }
328        }
329    };
330
331    #[cfg(not(unix))]
332    let terminate = std::future::pending::<()>();
333
334    tokio::select! {
335        _ = ctrl_c => {}
336        _ = terminate => {}
337    }
338    tracing::info!("shutdown signal received; draining listeners");
339}