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
25pub fn empty_transit_metadata_cbor() -> &'static [u8] {
30 &EMPTY_TRANSIT_METADATA_CBOR_BYTES
31}
32
33pub const FORMAT_TRANSIT: &str = "micromegas-transit";
35
36pub const FORMAT_OTLP_LOGS: &str = "otlp/v1/logs";
38
39pub const FORMAT_OTLP_METRICS: &str = "otlp/v1/metrics";
41
42pub const FORMAT_OTLP_TRACES: &str = "otlp/v1/traces";
44
45#[derive(Error, Debug)]
48pub enum IngestionServiceError {
49 #[error("Parse error: {0}")]
51 ParseError(String),
52
53 #[error("Database error: {0}")]
55 DatabaseError(String),
56
57 #[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 #[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 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 #[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 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 #[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 #[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 #[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}