Skip to main content

micromegas_ingestion/
web_ingestion_service.rs

1use crate::data_lake_config::DataLakeConfig;
2use crate::data_lake_connection::{DataLakeConnection, connect_to_data_lake};
3use crate::remote_data_lake::migrate_db;
4use anyhow::Context;
5use bytes::Buf;
6use micromegas_telemetry::block_wire_format;
7use micromegas_telemetry::property::Property;
8use micromegas_telemetry::property::make_properties;
9use micromegas_telemetry::stream_info::StreamInfo;
10use micromegas_telemetry::wire_format::encode_cbor;
11use micromegas_tracing::prelude::*;
12use micromegas_tracing::property_set;
13use std::sync::{Arc, LazyLock, Mutex};
14use std::time::Instant;
15use thiserror::Error;
16use uuid::Uuid;
17
18static EMPTY_TRANSIT_METADATA_CBOR_BYTES: LazyLock<Vec<u8>> = LazyLock::new(|| {
19    let mut buf = Vec::new();
20    ciborium::ser::into_writer(&Vec::<()>::new(), &mut buf)
21        .expect("encoding an empty Vec to CBOR is infallible");
22    buf
23});
24
25/// Sentinel for `dependencies_metadata` / `objects_metadata` on streams that
26/// don't use the transit/POD wire format (e.g. OTLP). Existing readers decode
27/// these BYTEA columns as `Vec<UserDefinedType>` and iterate; an empty Vec
28/// makes those loops no-ops without touching consumer code.
29pub fn empty_transit_metadata_cbor() -> &'static [u8] {
30    &EMPTY_TRANSIT_METADATA_CBOR_BYTES
31}
32
33/// Format string for native streams (transit-encoded payload, CBOR envelope).
34pub const FORMAT_TRANSIT: &str = "micromegas-transit";
35
36/// Stream `format` value for OTel logs (one `ResourceLogs` proto per block payload).
37pub const FORMAT_OTLP_LOGS: &str = "otlp/v1/logs";
38
39/// Stream `format` value for OTel metrics (one `ResourceMetrics` proto per block payload).
40pub const FORMAT_OTLP_METRICS: &str = "otlp/v1/metrics";
41
42/// Stream `format` value for OTel traces (one `ResourceSpans` proto per block payload).
43pub const FORMAT_OTLP_TRACES: &str = "otlp/v1/traces";
44
45/// Error type for ingestion service operations.
46/// Categorizes errors to enable proper HTTP status code mapping.
47#[derive(Error, Debug)]
48pub enum IngestionServiceError {
49    /// Client-side errors (malformed input) - maps to 400 Bad Request
50    #[error("Parse error: {0}")]
51    ParseError(String),
52
53    /// Database errors - maps to 500 Internal Server Error
54    #[error("Database error: {0}")]
55    DatabaseError(String),
56
57    /// Object storage errors - maps to 500 Internal Server Error
58    #[error("Storage error: {0}")]
59    StorageError(String),
60}
61
62#[derive(Clone)]
63pub struct WebIngestionService {
64    lake: DataLakeConnection,
65    ready_ok_until: Arc<Mutex<Option<Instant>>>,
66}
67
68impl WebIngestionService {
69    pub fn new(lake: DataLakeConnection) -> Self {
70        Self {
71            lake,
72            ready_ok_until: Arc::new(Mutex::new(None)),
73        }
74    }
75
76    pub async fn check_ready(&self) -> bool {
77        let now = Instant::now();
78        {
79            let guard = self.ready_ok_until.lock().expect("readiness cache lock");
80            if let Some(ok_until) = *guard
81                && ok_until > now
82            {
83                return true;
84            }
85        }
86
87        let probe_db = instrument_named!(
88            sqlx::query("SELECT 1").execute(&self.lake.db_pool),
89            "sql_readiness_probe"
90        );
91        let probe_blob = self.lake.blob_storage.probe();
92
93        let result = tokio::time::timeout(std::time::Duration::from_secs(2), async {
94            tokio::join!(probe_db, probe_blob)
95        })
96        .await;
97
98        match result {
99            Ok((Ok(_), Ok(()))) => {
100                let mut guard = self.ready_ok_until.lock().expect("readiness cache lock");
101                *guard = Some(Instant::now() + std::time::Duration::from_secs(1));
102                true
103            }
104            _ => {
105                let mut guard = self.ready_ok_until.lock().expect("readiness cache lock");
106                *guard = None;
107                false
108            }
109        }
110    }
111
112    /// Pre-seeds the readiness cache to `until`. Intended for testing only.
113    #[doc(hidden)]
114    pub fn set_ready_until(&self, until: Instant) {
115        let mut guard = self.ready_ok_until.lock().expect("readiness cache lock");
116        *guard = Some(until);
117    }
118
119    /// Reads MICROMEGAS_SQL_CONNECTION_STRING and MICROMEGAS_OBJECT_STORE_URI,
120    /// connects to the data lake, runs ingestion migrations, and returns
121    /// a ready-to-use service.
122    pub async fn from_env() -> anyhow::Result<Arc<Self>> {
123        let cfg = DataLakeConfig::from_env()?;
124        let lake = connect_to_data_lake(&cfg.sql_connection_string, &cfg.object_store_uri).await?;
125        migrate_db(lake.db_pool.clone())
126            .await
127            .with_context(|| "migrate_db")?;
128        Ok(Arc::new(Self::new(lake)))
129    }
130
131    #[span_fn]
132    pub async fn insert_block(&self, body: bytes::Bytes) -> Result<(), IngestionServiceError> {
133        let block: block_wire_format::Block = ciborium::from_reader(body.reader())
134            .map_err(|e| IngestionServiceError::ParseError(format!("parsing block: {e}")))?;
135        self.insert_block_typed(block).await
136    }
137
138    /// Inserts a block whose payload is already typed (no envelope round-trip on the caller side).
139    ///
140    /// The caller hands us a fully-built `Block`; we CBOR-encode the payload envelope once,
141    /// write it to object storage, and INSERT the row. Used by the OTLP adapter where
142    /// constructing the CBOR `Block` envelope just so `insert_block` could decode it
143    /// would be wasted work.
144    #[span_fn]
145    pub async fn insert_block_typed(
146        &self,
147        block: block_wire_format::Block,
148    ) -> Result<(), IngestionServiceError> {
149        let encoded_payload = encode_cbor(&block.payload)
150            .map_err(|e| IngestionServiceError::ParseError(format!("encoding payload: {e}")))?;
151        let payload_size = encoded_payload.len();
152
153        let process_id = &block.process_id;
154        let stream_id = &block.stream_id;
155        let block_id = &block.block_id;
156        let obj_path = format!("blobs/{process_id}/{stream_id}/{block_id}");
157        debug!("writing {obj_path}");
158
159        use sqlx::types::chrono::{DateTime, FixedOffset};
160        let begin_time = DateTime::<FixedOffset>::parse_from_rfc3339(&block.begin_time)
161            .map_err(|e| IngestionServiceError::ParseError(format!("parsing begin_time: {e}")))?;
162        let end_time = DateTime::<FixedOffset>::parse_from_rfc3339(&block.end_time)
163            .map_err(|e| IngestionServiceError::ParseError(format!("parsing end_time: {e}")))?;
164        {
165            let begin_put = now();
166            self.lake
167                .blob_storage
168                .put(&obj_path, encoded_payload.into())
169                .await
170                .map_err(|e| {
171                    IngestionServiceError::StorageError(format!(
172                        "writing block to blob storage: {e}"
173                    ))
174                })?;
175            imetric!("put_duration", "ticks", (now() - begin_put) as u64);
176        }
177
178        debug!("recording block_id={block_id} stream_id={stream_id} process_id={process_id}");
179        let begin_insert = now();
180        let insert_time = sqlx::types::chrono::Utc::now();
181        let result = instrument_named!(
182            sqlx::query(
183                "INSERT INTO blocks VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11) ON CONFLICT (block_id) DO NOTHING;",
184            )
185            .bind(block_id)
186            .bind(stream_id)
187            .bind(process_id)
188            .bind(begin_time)
189            .bind(block.begin_ticks)
190            .bind(end_time)
191            .bind(block.end_ticks)
192            .bind(block.nb_objects)
193            .bind(block.object_offset)
194            .bind(payload_size as i64)
195            .bind(insert_time)
196            .execute(&self.lake.db_pool),
197            "sql_insert_block"
198        )
199        .await
200        .map_err(|e| IngestionServiceError::DatabaseError(format!("inserting into blocks: {e}")))?;
201        imetric!("insert_duration", "ticks", (now() - begin_insert) as u64);
202
203        if result.rows_affected() == 0 {
204            debug!("duplicate block_id={block_id} skipped (already exists)");
205        }
206        // this measure does not benefit from a dynamic property - I just want to make sure the feature works well
207        // the cost in this context should be reasonnable
208        imetric!(
209            "payload_size_inserted",
210            "bytes",
211            property_set::PropertySet::find_or_create(vec![property_set::Property::new(
212                "target",
213                "micromegas::ingestion"
214            ),]),
215            payload_size as u64
216        );
217        debug!("recorded block_id={block_id} stream_id={stream_id} process_id={process_id}");
218
219        Ok(())
220    }
221
222    /// Registers a stream whose blocks will be ingested in the transit format.
223    #[span_fn]
224    pub async fn insert_stream(&self, body: bytes::Bytes) -> Result<(), IngestionServiceError> {
225        let stream_info: StreamInfo = ciborium::from_reader(body.reader())
226            .map_err(|e| IngestionServiceError::ParseError(format!("parsing StreamInfo: {e}")))?;
227        info!(
228            "new stream {} {:?} {:?}",
229            stream_info.stream_id, &stream_info.tags, &stream_info.properties
230        );
231        let dependencies_metadata =
232            encode_cbor(&stream_info.dependencies_metadata).map_err(|e| {
233                IngestionServiceError::ParseError(format!("encoding dependencies_metadata: {e}"))
234            })?;
235        let objects_metadata = encode_cbor(&stream_info.objects_metadata).map_err(|e| {
236            IngestionServiceError::ParseError(format!("encoding objects_metadata: {e}"))
237        })?;
238        let result = instrument_named!(
239            sqlx::query(
240                "INSERT INTO streams (stream_id, process_id, dependencies_metadata, objects_metadata, tags, properties, insert_time, format)
241             VALUES ($1,$2,$3,$4,$5,$6,$7,$8)
242             ON CONFLICT (stream_id) DO NOTHING;",
243            )
244            .bind(stream_info.stream_id)
245            .bind(stream_info.process_id)
246            .bind(dependencies_metadata)
247            .bind(objects_metadata)
248            .bind(&stream_info.tags)
249            .bind(make_properties(&stream_info.properties))
250            .bind(sqlx::types::chrono::Utc::now())
251            .bind(FORMAT_TRANSIT)
252            .execute(&self.lake.db_pool),
253            "sql_insert_stream"
254        )
255        .await
256        .map_err(|e| {
257            IngestionServiceError::DatabaseError(format!("inserting into streams: {e}"))
258        })?;
259
260        if result.rows_affected() == 0 {
261            debug!(
262                "duplicate stream_id={} skipped (already exists)",
263                stream_info.stream_id
264            );
265        }
266        Ok(())
267    }
268
269    /// Registers a stream produced by an OTLP ingestion path.
270    ///
271    /// `dependencies_metadata` and `objects_metadata` are filled with the CBOR sentinel
272    /// for an empty `Vec<UserDefinedType>` so legacy decode sites continue to work.
273    /// `format` distinguishes per-block dispatch downstream (e.g. `"otlp/v1/logs"`).
274    /// Stream `properties` are always empty for OTel — scope and per-event attrs
275    /// live on individual rows during materialization, not on the stream.
276    ///
277    /// Hack: piggybacking OTLP onto the transit-shaped `streams` row (with empty
278    /// metadata sentinels) is expedient for two formats but won't scale. To support
279    /// more formats cleanly, `dependencies_metadata`, `objects_metadata`, and `format`
280    /// should be merged into a single per-format payload column.
281    #[span_fn]
282    pub async fn register_otel_stream(
283        &self,
284        stream_id: Uuid,
285        process_id: Uuid,
286        tags: Vec<String>,
287        format: &str,
288    ) -> Result<(), IngestionServiceError> {
289        let result = instrument_named!(
290            sqlx::query(
291                "INSERT INTO streams (stream_id, process_id, dependencies_metadata, objects_metadata, tags, properties, insert_time, format)
292             VALUES ($1,$2,$3,$4,$5,$6,$7,$8)
293             ON CONFLICT (stream_id) DO NOTHING;",
294            )
295            .bind(stream_id)
296            .bind(process_id)
297            .bind(empty_transit_metadata_cbor())
298            .bind(empty_transit_metadata_cbor())
299            .bind(tags)
300            .bind(Vec::<Property>::new())
301            .bind(sqlx::types::chrono::Utc::now())
302            .bind(format)
303            .execute(&self.lake.db_pool),
304            "sql_insert_stream"
305        )
306        .await
307        .map_err(|e| {
308            IngestionServiceError::DatabaseError(format!("inserting otel stream: {e}"))
309        })?;
310
311        if result.rows_affected() == 0 {
312            debug!("duplicate otel stream_id={stream_id} skipped (already exists)");
313        }
314        Ok(())
315    }
316
317    #[span_fn]
318    pub async fn insert_process(&self, body: bytes::Bytes) -> Result<(), IngestionServiceError> {
319        let process_info: ProcessInfo = ciborium::from_reader(body.reader())
320            .map_err(|e| IngestionServiceError::ParseError(format!("parsing ProcessInfo: {e}")))?;
321
322        let insert_time = sqlx::types::chrono::Utc::now();
323        let result = instrument_named!(
324            sqlx::query(
325                "INSERT INTO processes VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13) ON CONFLICT (process_id) DO NOTHING;",
326            )
327            .bind(process_info.process_id)
328            .bind(process_info.exe)
329            .bind(process_info.username)
330            .bind(process_info.realname)
331            .bind(process_info.computer)
332            .bind(process_info.distro)
333            .bind(process_info.cpu_brand)
334            .bind(process_info.tsc_frequency)
335            .bind(process_info.start_time)
336            .bind(process_info.start_ticks)
337            .bind(insert_time)
338            .bind(process_info.parent_process_id)
339            .bind(make_properties(&process_info.properties))
340            .execute(&self.lake.db_pool),
341            "sql_insert_process"
342        )
343        .await
344        .map_err(|e| {
345            IngestionServiceError::DatabaseError(format!("inserting into processes: {e}"))
346        })?;
347
348        if result.rows_affected() == 0 {
349            debug!(
350                "duplicate process_id={} skipped (already exists)",
351                process_info.process_id
352            );
353        }
354        Ok(())
355    }
356
357    /// Registers a process originating from OTLP. Idempotent via `ON CONFLICT DO NOTHING`.
358    ///
359    /// `realname` is set equal to `username` (OTel has no separate "real name" concept).
360    /// `parent_process_id` is always NULL — OTel has no parent-process model.
361    /// `insert_time` is the server wall clock, matching the existing `insert_process` path.
362    #[span_fn]
363    #[expect(clippy::too_many_arguments, reason = "OTel process identity fields")]
364    pub async fn register_otel_process(
365        &self,
366        process_id: Uuid,
367        exe: String,
368        username: String,
369        computer: String,
370        distro: String,
371        cpu_brand: String,
372        tsc_frequency: i64,
373        start_time: sqlx::types::chrono::DateTime<sqlx::types::chrono::Utc>,
374        start_ticks: i64,
375        properties: Vec<Property>,
376    ) -> Result<(), IngestionServiceError> {
377        let insert_time = sqlx::types::chrono::Utc::now();
378        let result = instrument_named!(
379            sqlx::query(
380                "INSERT INTO processes
381             (process_id, exe, username, realname, computer, distro, cpu_brand,
382              tsc_frequency, start_time, start_ticks, insert_time, parent_process_id, properties)
383             VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,NULL,$12)
384             ON CONFLICT (process_id) DO NOTHING;",
385            )
386            .bind(process_id)
387            .bind(exe)
388            .bind(&username)
389            .bind(&username)
390            .bind(computer)
391            .bind(distro)
392            .bind(cpu_brand)
393            .bind(tsc_frequency)
394            .bind(start_time)
395            .bind(start_ticks)
396            .bind(insert_time)
397            .bind(properties)
398            .execute(&self.lake.db_pool),
399            "sql_insert_process"
400        )
401        .await
402        .map_err(|e| {
403            IngestionServiceError::DatabaseError(format!("inserting otel process: {e}"))
404        })?;
405
406        if result.rows_affected() == 0 {
407            debug!("duplicate otel process_id={process_id} skipped (already exists)");
408        }
409        Ok(())
410    }
411}