ifc_lite_server/routes/parse/
stream_batch.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4
5//! One streamed batch's CPU work on `POST /api/v1/parse/parquet-stream`: the
6//! cache-writer append and the client blob, in either stream mode.
7//!
8//! Split out of `parquet_stream.rs` (module-size budget) along the seam #5407
9//! added: this is the one place that decides whether a batch is planned
10//! against the whole stream (cross-batch) or serialized on its own.
11
12use crate::services::parquet_stream_shapes::StreamShapePlanner;
13use crate::services::{serialize_batch_with_layout, ParquetError, ParquetLayout, StreamingParquetCacheWriter};
14use crate::types::MeshData;
15use bytes::Bytes;
16use std::sync::Mutex;
17
18/// A batch's client blob, and `(vertex_base, index_base)` when it was planned
19/// against the whole stream.
20pub(super) type EncodedBatch = (Bytes, Option<(u32, u32)>);
21
22/// Append `meshes` to the cache fill and encode the client's blob.
23///
24/// Cross-batch (`planner` is `Some`): ONE plan feeds both the client blob and
25/// the cache row group, so the two cannot disagree on a range. Batch-local:
26/// the cache writer plans whole-stream offsets and the client blob stays
27/// batch-local, as before. A cache-writer error drops the cache fill (the slot
28/// goes `None`) and never the client's batch.
29pub(super) fn encode_batch(
30    meshes: &[MeshData],
31    layout: ParquetLayout,
32    baked_basis: Option<&[f64; 16]>,
33    planner: Option<&mut StreamShapePlanner>,
34    cache_writer: &Mutex<Option<StreamingParquetCacheWriter>>,
35) -> Result<EncodedBatch, ParquetError> {
36    let planned = planner.map(|p| p.plan(meshes, baked_basis)).transpose()?;
37    if let Ok(mut slot) = cache_writer.lock() {
38        if let Some(writer) = slot.as_mut() {
39            let appended = match &planned {
40                Some(batch) => writer.append_planned(batch),
41                None => writer.append(meshes),
42            };
43            if let Err(e) = appended {
44                tracing::error!(error = %e, "Streaming cache writer failed; skipping cache fill");
45                *slot = None;
46            }
47        }
48    }
49    match planned {
50        Some(batch) => Ok((batch.encode()?, Some((batch.vertex_base, batch.index_base)))),
51        None => Ok((serialize_batch_with_layout(meshes, layout)?, None)),
52    }
53}
54
55/// Run CPU-bound work from inside an async context without starving other
56/// connections' polls: `block_in_place` on the multi-thread runtime, inline on
57/// a current-thread one (where `block_in_place` panics, as under the
58/// `#[tokio::test]` harness). Shared by the live stream and its cached replay.
59pub(super) fn off_the_async_worker<T>(work: impl FnOnce() -> T) -> T {
60    if tokio::runtime::Handle::current().runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread {
61        tokio::task::block_in_place(work)
62    } else {
63        work()
64    }
65}
ifc_lite_server/routes/parse/stream_batch.rs

ifc_lite_server/routes/parse/
stream_batch.rs
