micromegas_analytics/lakehouse/
async_events_block_processor.rs

1use super::{
2    block_partition_spec::BlockProcessor, partition_source_data::PartitionSourceBlock,
3    write_partition::PartitionRowSet,
4};
5use crate::{
6    async_block_processing::{AsyncBlockProcessor, parse_async_block_payload},
7    async_events_table::{AsyncEventRecord, AsyncEventRecordBuilder},
8    payload::fetch_block_payload,
9    scope::BorrowedScopeDesc,
10    time::ConvertTicks,
11};
12use anyhow::{Context, Result};
13use async_trait::async_trait;
14use micromegas_telemetry::blob_storage::BlobStorage;
15use micromegas_tracing::prelude::*;
16use std::sync::Arc;
17
18const BEGIN_EVENT_TYPE: &str = "begin";
19const END_EVENT_TYPE: &str = "end";
20
21/// A `BlockProcessor` implementation for processing async event blocks.
22#[derive(Debug)]
23pub struct AsyncEventsBlockProcessor {
24    convert_ticks: Arc<ConvertTicks>,
25}
26
27impl AsyncEventsBlockProcessor {
28    pub fn new(convert_ticks: Arc<ConvertTicks>) -> Self {
29        Self { convert_ticks }
30    }
31}
32
33/// Helper struct to collect async events during processing.
34struct AsyncEventCollector {
35    record_builder: AsyncEventRecordBuilder,
36    stream_id: Arc<String>,
37    block_id: Arc<String>,
38    convert_ticks: Arc<ConvertTicks>,
39}
40
41impl AsyncEventCollector {
42    fn new(
43        capacity: usize,
44        stream_id: Arc<String>,
45        block_id: Arc<String>,
46        convert_ticks: Arc<ConvertTicks>,
47    ) -> Self {
48        Self {
49            record_builder: AsyncEventRecordBuilder::with_capacity(capacity),
50            stream_id,
51            block_id,
52            convert_ticks,
53        }
54    }
55}
56
57impl AsyncBlockProcessor for AsyncEventCollector {
58    fn on_begin_async_scope(
59        &mut self,
60        _block_id: &str,
61        scope: BorrowedScopeDesc<'_>,
62        ts: i64,
63        span_id: i64,
64        parent_span_id: i64,
65        depth: u32,
66    ) -> Result<bool> {
67        let time_ns = self.convert_ticks.ticks_to_nanoseconds(ts);
68        let record = AsyncEventRecord {
69            stream_id: self.stream_id.clone(),
70            block_id: self.block_id.clone(),
71            time: time_ns,
72            event_type: BEGIN_EVENT_TYPE,
73            span_id,
74            parent_span_id,
75            depth,
76            hash: scope.hash,
77            name: scope.name,
78            filename: scope.filename,
79            target: scope.target,
80            line: scope.line,
81        };
82        self.record_builder.append(&record)?;
83        Ok(true)
84    }
85
86    fn on_end_async_scope(
87        &mut self,
88        _block_id: &str,
89        scope: BorrowedScopeDesc<'_>,
90        ts: i64,
91        span_id: i64,
92        parent_span_id: i64,
93        depth: u32,
94    ) -> Result<bool> {
95        let time_ns = self.convert_ticks.ticks_to_nanoseconds(ts);
96        let record = AsyncEventRecord {
97            stream_id: self.stream_id.clone(),
98            block_id: self.block_id.clone(),
99            time: time_ns,
100            event_type: END_EVENT_TYPE,
101            span_id,
102            parent_span_id,
103            depth,
104            hash: scope.hash,
105            name: scope.name,
106            filename: scope.filename,
107            target: scope.target,
108            line: scope.line,
109        };
110        self.record_builder.append(&record)?;
111        Ok(true)
112    }
113}
114
115#[async_trait]
116impl BlockProcessor for AsyncEventsBlockProcessor {
117    #[span_fn]
118    async fn process(
119        &self,
120        blob_storage: Arc<BlobStorage>,
121        src_block: Arc<PartitionSourceBlock>,
122    ) -> Result<Option<PartitionRowSet>> {
123        // Use the shared ConvertTicks instance instead of creating a new one per block
124        let convert_ticks = self.convert_ticks.clone();
125        // Use nb_objects as initial capacity estimate (may contain non-async events)
126        let estimated_capacity = src_block.block.nb_objects;
127        let mut collector = AsyncEventCollector::new(
128            estimated_capacity as usize,
129            Arc::new(format!("{}", src_block.stream.stream_id)),
130            Arc::new(format!("{}", src_block.block.block_id)),
131            convert_ticks,
132        );
133        let payload = fetch_block_payload(
134            blob_storage,
135            src_block.process.process_id,
136            src_block.stream.stream_id,
137            src_block.block.block_id,
138        )
139        .await
140        .with_context(|| "fetch_block_payload")?;
141        let block_id_str = src_block
142            .block
143            .block_id
144            .hyphenated()
145            .encode_lower(&mut sqlx::types::uuid::Uuid::encode_buffer())
146            .to_owned();
147        parse_async_block_payload(
148            &block_id_str,
149            0,
150            &payload,
151            &src_block.stream,
152            &mut collector,
153        )
154        .with_context(|| "parse_async_block_payload")?;
155        if let Some(time_range) = collector.record_builder.get_time_range() {
156            let record_batch = collector.record_builder.finish()?;
157            Ok(Some(PartitionRowSet {
158                rows_time_range: time_range,
159                rows: record_batch,
160            }))
161        } else {
162            Ok(None)
163        }
164    }
165}
micromegas_analytics/lakehouse/async_events_block_processor.rs

micromegas_analytics/lakehouse/
async_events_block_processor.rs
