micromegas_analytics/lakehouse/
net_spans_view.rs

1use super::{
2    blocks_view::BlocksView,
3    dataframe_time_bounds::{DataFrameTimeBounds, NamedColumnsTimeBounds},
4    jit_partitions::{
5        JitPartitionConfig, generate_process_jit_partitions, is_jit_partition_up_to_date,
6    },
7    lakehouse_context::LakehouseContext,
8    partition_cache::PartitionCache,
9    partition_source_data::{SourceDataBlocksInMemory, hash_to_object_count},
10    view::{PartitionSpec, View, ViewMetadata},
11    view_factory::{ViewFactory, ViewMaker},
12};
13use crate::{
14    lakehouse::write_partition::{PartitionRowSet, write_partition_from_rows},
15    metadata::{StreamMetadata, find_process_with_latest_timing},
16    net_span_tree::make_net_span_tree,
17    net_spans_table::{NetSpanRecordBuilder, net_spans_table_schema},
18    response_writer::ResponseWriter,
19    time::{ConvertTicks, TimeRange, datetime_to_scalar, make_time_converter_from_latest_timing},
20};
21use anyhow::{Context, Result};
22use async_trait::async_trait;
23use chrono::{DateTime, Utc};
24use datafusion::logical_expr::{BinaryExpr, Expr, Operator};
25use datafusion::{arrow::datatypes::Schema, logical_expr::expr_fn::col};
26use micromegas_ingestion::data_lake_connection::DataLakeConnection;
27use micromegas_telemetry::{blob_storage::BlobStorage, types::block::BlockMetadata};
28use micromegas_tracing::prelude::*;
29use std::sync::Arc;
30use uuid::Uuid;
31
32const VIEW_SET_NAME: &str = "net_spans";
33const SCHEMA_VERSION: u8 = 1;
34const NET_STREAM_TAG: &str = "net";
35
36lazy_static::lazy_static! {
37    static ref BEGIN_TIME_COLUMN: Arc<String> = Arc::new(String::from("begin_time"));
38    static ref END_TIME_COLUMN: Arc<String> = Arc::new(String::from("end_time"));
39}
40
41/// A `ViewMaker` for creating `NetSpansView` instances.
42#[derive(Debug)]
43pub struct NetSpansViewMaker {
44    view_factory: Arc<ViewFactory>,
45}
46
47impl NetSpansViewMaker {
48    pub fn new(view_factory: Arc<ViewFactory>) -> Self {
49        Self { view_factory }
50    }
51}
52
53impl ViewMaker for NetSpansViewMaker {
54    fn make_view(&self, view_instance_id: &str) -> Result<Arc<dyn View>> {
55        Ok(Arc::new(NetSpansView::new(
56            view_instance_id,
57            self.view_factory.clone(),
58        )?))
59    }
60
61    fn get_schema_hash(&self) -> Vec<u8> {
62        vec![SCHEMA_VERSION]
63    }
64
65    fn get_schema(&self) -> Arc<Schema> {
66        Arc::new(net_spans_table_schema())
67    }
68}
69
70/// A view of network bandwidth spans (Connection / Object / Property / RPC).
71#[derive(Debug)]
72pub struct NetSpansView {
73    view_set_name: Arc<String>,
74    view_instance_id: Arc<String>,
75    process_id: sqlx::types::Uuid,
76    view_factory: Arc<ViewFactory>,
77}
78
79impl NetSpansView {
80    pub fn new(view_instance_id: &str, view_factory: Arc<ViewFactory>) -> Result<Self> {
81        if view_instance_id == "global" {
82            anyhow::bail!("NetSpansView does not support global view access");
83        }
84        let process_id = Uuid::parse_str(view_instance_id).with_context(|| "Uuid::parse_str")?;
85        Ok(Self {
86            view_set_name: Arc::new(String::from(VIEW_SET_NAME)),
87            view_instance_id: Arc::new(view_instance_id.into()),
88            process_id,
89            view_factory,
90        })
91    }
92}
93
94#[span_fn]
95async fn append_net_span_tree(
96    record_builder: &mut NetSpanRecordBuilder,
97    convert_ticks: &ConvertTicks,
98    blocks: &[BlockMetadata],
99    blob_storage: Arc<BlobStorage>,
100    stream: &StreamMetadata,
101    process_id: Arc<String>,
102) -> Result<()> {
103    make_net_span_tree(
104        blocks,
105        record_builder,
106        blob_storage,
107        stream,
108        process_id,
109        convert_ticks.clone(),
110    )
111    .await
112    .with_context(|| "make_net_span_tree")
113}
114
115/// Writes a partition from a set of blocks.
116#[span_fn]
117async fn write_partition(
118    lake: Arc<DataLakeConnection>,
119    view_meta: ViewMetadata,
120    schema: Arc<Schema>,
121    convert_ticks: &ConvertTicks,
122    spec: &SourceDataBlocksInMemory,
123    process_id: Arc<String>,
124) -> Result<()> {
125    let nb_events = hash_to_object_count(&spec.block_ids_hash)? as usize;
126    info!("nb_events: {nb_events}");
127    if spec.blocks.is_empty() {
128        anyhow::bail!("empty partition spec");
129    }
130    let min_insert_time = spec.blocks[0].block.insert_time;
131    let max_insert_time = spec.blocks[spec.blocks.len() - 1].block.insert_time;
132
133    let (tx, rx) = tokio::sync::mpsc::channel(1);
134    let null_response_writer = Arc::new(ResponseWriter::new(None));
135    let join_handle = spawn_with_context(write_partition_from_rows(
136        lake.clone(),
137        view_meta,
138        schema,
139        TimeRange::new(min_insert_time, max_insert_time),
140        spec.block_ids_hash.clone(),
141        None,
142        rx,
143        null_response_writer,
144    ));
145
146    let build_result: Result<Option<PartitionRowSet>> = async {
147        let mut record_builder = NetSpanRecordBuilder::with_capacity(nb_events / 2);
148        let stream = spec.blocks[0].stream.clone();
149        for b in &spec.blocks {
150            anyhow::ensure!(
151                b.stream.stream_id == stream.stream_id,
152                "net_spans partition contains multiple streams ({} and {}); expected one per process",
153                stream.stream_id,
154                b.stream.stream_id,
155            );
156        }
157        let mut blocks_to_process: Vec<BlockMetadata> = vec![];
158        let mut last_end: Option<i64> = None;
159        for block in &spec.blocks {
160            let contiguous = last_end
161                .map(|e| block.block.begin_ticks == e)
162                .unwrap_or(true);
163            if !contiguous {
164                append_net_span_tree(
165                    &mut record_builder,
166                    convert_ticks,
167                    &blocks_to_process,
168                    lake.blob_storage.clone(),
169                    &stream,
170                    process_id.clone(),
171                )
172                .await?;
173                blocks_to_process = vec![];
174            }
175            blocks_to_process.push(block.block.clone());
176            last_end = Some(block.block.end_ticks);
177        }
178        if !blocks_to_process.is_empty() {
179            append_net_span_tree(
180                &mut record_builder,
181                convert_ticks,
182                &blocks_to_process,
183                lake.blob_storage.clone(),
184                &stream,
185                process_id.clone(),
186            )
187            .await?;
188        }
189        let min_time_row = convert_ticks.delta_ticks_to_time(spec.blocks[0].block.begin_ticks);
190        let max_time_row =
191            convert_ticks.delta_ticks_to_time(spec.blocks[spec.blocks.len() - 1].block.end_ticks);
192        let rows_time_range = record_builder
193            .get_time_range()
194            .unwrap_or(TimeRange::new(min_time_row, max_time_row));
195        let nb_rows = record_builder.len();
196        let rows = record_builder
197            .finish()
198            .with_context(|| "record_builder.finish()")?;
199        info!("writing {nb_rows} rows");
200        if nb_rows > 0 {
201            Ok(Some(PartitionRowSet { rows_time_range, rows }))
202        } else {
203            Ok(None)
204        }
205    }
206    .await;
207
208    match build_result {
209        Ok(Some(row_set)) => {
210            tx.send(Ok(row_set)).await?;
211            drop(tx);
212            join_handle.await??;
213            Ok(())
214        }
215        Ok(None) => {
216            drop(tx);
217            join_handle.await??;
218            Ok(())
219        }
220        Err(e) => {
221            warn!(
222                "aborting net-spans partition write for block {:?}: {e:?}",
223                spec.block_ids_hash
224            );
225            let _ = tx
226                .send(Err(anyhow::anyhow!("net-spans build aborted")))
227                .await;
228            drop(tx);
229            match join_handle.await {
230                Ok(Ok(())) => {}
231                Ok(Err(writer_err)) => {
232                    debug!("net-spans writer task error during abort: {writer_err:?}");
233                }
234                Err(join_err) => {
235                    warn!("net-spans writer task panicked during abort: {join_err:?}");
236                }
237            }
238            Err(e)
239        }
240    }
241}
242
243/// Rebuilds the partition if it's missing or out of date.
244#[span_fn]
245async fn update_partition(
246    lake: Arc<DataLakeConnection>,
247    view_meta: ViewMetadata,
248    schema: Arc<Schema>,
249    convert_ticks: &ConvertTicks,
250    spec: &SourceDataBlocksInMemory,
251    process_id: Arc<String>,
252) -> Result<()> {
253    if is_jit_partition_up_to_date(&lake.db_pool, view_meta.clone(), spec).await? {
254        return Ok(());
255    }
256    write_partition(lake, view_meta, schema, convert_ticks, spec, process_id)
257        .await
258        .with_context(|| "write_partition")?;
259    Ok(())
260}
261
262#[async_trait]
263impl View for NetSpansView {
264    fn get_view_set_name(&self) -> Arc<String> {
265        self.view_set_name.clone()
266    }
267
268    fn get_view_instance_id(&self) -> Arc<String> {
269        self.view_instance_id.clone()
270    }
271
272    async fn make_batch_partition_spec(
273        &self,
274        _lakehouse: Arc<LakehouseContext>,
275        _existing_partitions: Arc<PartitionCache>,
276        _insert_range: TimeRange,
277    ) -> Result<Arc<dyn PartitionSpec>> {
278        anyhow::bail!("NetSpansView does not support batch partition specs")
279    }
280
281    fn get_file_schema_hash(&self) -> Vec<u8> {
282        vec![SCHEMA_VERSION]
283    }
284
285    fn get_file_schema(&self) -> Arc<Schema> {
286        Arc::new(net_spans_table_schema())
287    }
288
289    #[span_fn]
290    async fn jit_update(
291        &self,
292        lakehouse: Arc<LakehouseContext>,
293        query_range: Option<TimeRange>,
294    ) -> Result<()> {
295        let (process, last_block_end_ticks, last_block_end_time) = find_process_with_latest_timing(
296            lakehouse.clone(),
297            self.view_factory.clone(),
298            &self.process_id,
299            query_range,
300        )
301        .await
302        .with_context(|| "find_process_with_latest_timing")?;
303
304        let process = Arc::new(process);
305        let query_range =
306            query_range.unwrap_or_else(|| TimeRange::new(process.start_time, last_block_end_time));
307
308        let convert_ticks = make_time_converter_from_latest_timing(
309            &process,
310            last_block_end_ticks,
311            last_block_end_time,
312        )
313        .with_context(|| "make_time_converter_from_latest_timing")?;
314
315        let blocks_view = BlocksView::new()?;
316        let all_partitions = generate_process_jit_partitions(
317            &JitPartitionConfig::default(),
318            lakehouse.clone(),
319            &blocks_view,
320            &query_range,
321            process.clone(),
322            NET_STREAM_TAG,
323        )
324        .await
325        .with_context(|| "generate_process_jit_partitions")?;
326
327        let process_id_str = Arc::new(self.process_id.to_string());
328        for part in &all_partitions {
329            update_partition(
330                lakehouse.lake().clone(),
331                ViewMetadata {
332                    view_set_name: self.get_view_set_name(),
333                    view_instance_id: self.get_view_instance_id(),
334                    file_schema_hash: self.get_file_schema_hash(),
335                },
336                self.get_file_schema(),
337                &convert_ticks,
338                part,
339                process_id_str.clone(),
340            )
341            .await
342            .with_context(|| "update_partition")?;
343        }
344        Ok(())
345    }
346
347    fn make_time_filter(&self, begin: DateTime<Utc>, end: DateTime<Utc>) -> Result<Vec<Expr>> {
348        Ok(vec![
349            Expr::BinaryExpr(BinaryExpr::new(
350                col("begin_time").into(),
351                Operator::LtEq,
352                Expr::Literal(datetime_to_scalar(end), None).into(),
353            )),
354            Expr::BinaryExpr(BinaryExpr::new(
355                col("end_time").into(),
356                Operator::GtEq,
357                Expr::Literal(datetime_to_scalar(begin), None).into(),
358            )),
359        ])
360    }
361
362    fn get_time_bounds(&self) -> Arc<dyn DataFrameTimeBounds> {
363        Arc::new(NamedColumnsTimeBounds::new(
364            BEGIN_TIME_COLUMN.clone(),
365            END_TIME_COLUMN.clone(),
366        ))
367    }
368
369    fn get_update_group(&self) -> Option<i32> {
370        None
371    }
372}
micromegas_analytics/lakehouse/net_spans_view.rs

micromegas_analytics/lakehouse/
net_spans_view.rs
