Skip to main content

micromegas_analytics/
payload.rs

1use anyhow::{Context, Result};
2use bumpalo::Bump;
3use micromegas_telemetry::{blob_storage::BlobStorage, compression::decompress};
4use micromegas_tracing::{parsing::make_custom_readers, prelude::*};
5use micromegas_transit::{CustomReaderMap, parse_object_buffer, read_dependencies, value::Value};
6use std::sync::Arc;
7
8use crate::metadata::StreamMetadata;
9
10thread_local! {
11    /// The custom-reader map is identical for every block, so build it once per
12    /// worker thread instead of rebuilding a `HashMap` of `Arc<dyn Fn>` on every
13    /// `parse_block` call.
14    static CUSTOM_READERS: CustomReaderMap = make_custom_readers();
15}
16
17/// Fetches the payload of a block from blob storage.
18#[span_fn]
19pub async fn fetch_block_payload(
20    blob_storage: Arc<BlobStorage>,
21    process_id: sqlx::types::Uuid,
22    stream_id: sqlx::types::Uuid,
23    block_id: sqlx::types::Uuid,
24) -> Result<micromegas_telemetry::block_wire_format::BlockPayload> {
25    let obj_path = format!("blobs/{process_id}/{stream_id}/{block_id}");
26    let buffer: Vec<u8> = blob_storage
27        .read_blob(&obj_path)
28        .await
29        .with_context(|| "reading block payload from blob storage")?
30        .into();
31    {
32        span_scope!("decode");
33        let payload: micromegas_telemetry::block_wire_format::BlockPayload =
34            ciborium::from_reader(&buffer[..])
35                .with_context(|| format!("reading payload {}", block_id))?;
36        Ok(payload)
37    }
38}
39
40/// Parses a block of telemetry data, calling a function for each object in the block.
41///
42/// Each parsed `Value` borrows from a per-block bump arena (and the decompressed
43/// buffers) that live only for the duration of this call. The higher-ranked
44/// `FnMut(Value<'_>)` bound forbids the callback from retaining a `Value` beyond
45/// its invocation — anything that must outlive the block (e.g. an Arrow append)
46/// must copy out inside the callback.
47// parse_block calls fun for each object in the block until fun returns `false`
48#[span_fn]
49pub fn parse_block<F>(
50    stream: &StreamMetadata,
51    payload: &micromegas_telemetry::block_wire_format::BlockPayload,
52    mut fun: F,
53) -> Result<bool>
54where
55    F: for<'a> FnMut(Value<'a>) -> Result<bool>,
56{
57    let dep_udts = &stream.dependencies_metadata;
58    let obj_udts = &stream.objects_metadata;
59    let process_id = stream.process_id;
60    let stream_id = stream.stream_id;
61    // A corrupt block is unexpected enough to be a potential attack indicator,
62    // so every occurrence is logged here regardless of what the caller does
63    // with the propagated `Err`.
64    let log_decompress_err = |e: &anyhow::Error| {
65        error!("corrupt block payload: process_id={process_id} stream_id={stream_id} error={e:?}");
66    };
67    // Bind the decompressed buffers and the arena to locals so every parsed
68    // Value borrows from storage that outlives the parse below.
69    let deps_buf = decompress(&payload.dependencies)
70        .with_context(|| "decompressing dependencies payload")
71        .inspect_err(log_decompress_err)?;
72    let objs_buf = decompress(&payload.objects)
73        .with_context(|| "decompressing objects payload")
74        .inspect_err(log_decompress_err)?;
75    let bump = Bump::new();
76    CUSTOM_READERS.with(|custom_readers| {
77        let log_parse_err = |e: &anyhow::Error| {
78            error!("corrupt block: process_id={process_id} stream_id={stream_id} error={e:?}");
79        };
80        let dependencies = read_dependencies(&bump, custom_readers, dep_udts, &deps_buf)
81            .with_context(|| "reading dependencies")
82            .inspect_err(log_parse_err)?;
83        let continue_iterating = parse_object_buffer(
84            &bump,
85            custom_readers,
86            &dependencies,
87            obj_udts,
88            &objs_buf,
89            &mut fun,
90        )
91        .with_context(|| "parsing object buffer")
92        .inspect_err(log_parse_err)?;
93        Ok(continue_iterating)
94    })
95}