ifc_lite_server/routes/parse/

parquet_stream.rs

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

//! SSE Parquet-batch streaming parse endpoint.

use super::cache_keys::{
    data_model_cache_key, parquet_geometry_key, parquet_metadata_key,
    request_cache_key,
};
use super::parquet::ParquetMetadataHeader;
use super::stream_event::ParquetStreamEvent;
use super::stream_progress::{cache_stream_progress, StreamProgressRecorder};
use super::{cache_symbolic_data_off_runtime, extract_file, ParseQuery};
use crate::error::ApiError;
use crate::services::{extract_data_model, process_streaming, serialize_data_model_to_parquet};
use crate::types::StreamEvent;
use crate::AppState;
use axum::{
    extract::{Multipart, Query, State},
    response::sse::{Event, KeepAlive, Sse},
};
use std::convert::Infallible;

/// POST /api/v1/parse/parquet-stream - Streaming parse with Parquet batches.
///
/// Returns SSE events with Parquet-encoded geometry batches for progressive rendering.
/// Each batch can be decoded and rendered immediately without waiting for the full response.
///
/// Events:
/// - `start`: Initial event with `total_estimate` and `cache_key`
/// - `progress`: Progress updates with `processed` and `total` counts
/// - `batch`: Geometry batch with base64-encoded Parquet `data`, `mesh_count`, `batch_number`
/// - `complete`: Final event with `stats` and `metadata`
/// - `error`: Error event with `message`
///
/// After `complete`, client should fetch data model via `/api/v1/data-model/{cache_key}`.
///
/// ## Two ways to ask
///
/// With a multipart `file` body: the normal path. The cache key is the SHA-256
/// of the RECEIVED bytes, so the upload always completes first, and a `sha256`
/// query parameter sent alongside a body is IGNORED - the bytes on the wire
/// decide which entry is read and written, never the client's claim about them.
///
/// With `?sha256={hex}` and no body: a probe (issue #3901). It replays the
/// cached stream if, and only if, every entry the replay needs already exists
/// under that key, and otherwise answers `404` meaning "upload it". See
/// [`cached_replay::replay_by_client_hash`]. This is what lets a 40 MB cache
/// hit cost no upload.
pub async fn parse_parquet_stream(
    State(state): State<AppState>,
    Query(query): Query<ParseQuery>,
    multipart: Option<Multipart>,
) -> Result<axum::response::Response, ApiError> {
    use crate::services::parquet_stream_shapes::StreamShapePlanner;
    use crate::services::{StreamShapes, StreamingParquetCacheWriter};
    use axum::response::IntoResponse;
    use base64::{engine::general_purpose::STANDARD, Engine};
    use futures::StreamExt;
    use std::sync::{Arc, Mutex};

    let tessellation_quality = query.resolved_tessellation_quality()?;
    let stream_shapes = query.resolved_stream_shapes()?;

    // Hash-only probe: no body was sent, so there is nothing to extract and
    // nothing to parse. It is checked before the gate below only because that
    // gate reserves an upload that does not exist. The probe takes admission
    // itself, on the hit path where it has real work to bound -- see
    // `replay_by_client_hash`.
    let Some(mut multipart) = multipart else {
        let Some(sha256) = query.sha256.as_deref() else {
            // No body and no hash: there is nothing to identify a file with.
            // `MissingFile` (400) is what a body with no `file` field already
            // answers, and it says the same thing here.
            return Err(ApiError::MissingFile);
        };
        return super::cached_replay::replay_by_client_hash(
            &state,
            &query,
            tessellation_quality,
            stream_shapes,
            sha256,
        )
        .await;
    };

    // Extract file
    // Admission gate (bounded concurrency + byte budget): acquired BEFORE the
    // upload is buffered, reserving the max upload size since multipart rarely
    // declares a length up front. Held for the request's whole lifetime so a
    // disconnected-but-still-running job keeps its memory slot.
    let admission_guard = state
        .admission
        .acquire(state.config.max_file_size_mb as u64 * 1024 * 1024)
        .await?;
    let data = extract_file(&mut multipart, state.config.max_file_size_mb).await?;

    // Generate cache key before processing (include opening filter + quality).
    // From the RECEIVED BYTES, always: a `sha256` parameter that arrived
    // alongside a body has no say here, so a client whose claimed hash does not
    // describe what it uploaded still reads and writes the entry its bytes name.
    let cache_key = request_cache_key(&data, &query, tessellation_quality);
    let cache_key_clone = cache_key.clone();
    // Without `stream_shapes=cross-batch` this route SHARES nothing -- each
    // batch must decode on its own -- so the layout only decides whether the
    // mesh table carries identity `rot0..rot8`, and which cache namespace the
    // result lands in. A default request therefore still produces
    // byte-identical v5 output. With it (#5407), `planner` below carries the
    // shapes already emitted from one batch to the next.
    let layout = query.parquet_layout;

    // OPTIMIZATION: Check cache first and fast-path return if available
    // This avoids re-processing files that are already cached (see
    // `cached_replay.rs`; a short/corrupt blob falls through as a miss).
    if let Some(response) =
        super::cached_replay::try_cached_replay(&state, &cache_key, layout, stream_shapes, query.data_model_entities).await?
    {
        // Cached replay: no parse work runs, so holding the admission
        // guard (and its CPU slot) while a slow client drains the SSE
        // would starve real parses for nothing. The replay blob is
        // already materialized and bounded by cache content, far below
        // a parse working set.
        drop(admission_guard);
        return Ok(response);
    }

    tracing::info!(
        cache_key = %cache_key,
        size = data.len(),
        "Streaming cache MISS - processing file"
    );

    let content = data;
    let initial_batch_size = state.config.initial_batch_size;
    let max_batch_size = state.config.max_batch_size;
    let cache = state.cache.clone();

    // Incremental cache writer: each batch's columns are appended as Parquet
    // row groups (GLOBAL offsets) and the meshes dropped, replacing the old
    // Arc<Mutex<Vec<MeshData>>> accumulator that held a FULL second copy of
    // the model's geometry until Complete. `None` after a writer error (the
    // cache fill is skipped; the client stream is unaffected).
    let cache_writer: Arc<Mutex<Option<StreamingParquetCacheWriter>>> =
        Arc::new(Mutex::new(match StreamingParquetCacheWriter::new(layout) {
            Ok(w) => Some(w),
            Err(e) => {
                tracing::error!(error = %e, "Failed to create streaming cache writer");
                None
            }
        }));
    let cache_writer_for_stream = cache_writer.clone();
    // Cross-batch streams only: the whole-stream offsets and the shapes
    // emitted so far (#5407). Owned by the stream, NOT by the cache writer,
    // because the client's batches depend on it and must keep flowing when a
    // cache-writer error has dropped the cache fill.
    let mut planner = (stream_shapes == StreamShapes::CrossBatch)
        .then(|| StreamShapePlanner::new(layout, stream_shapes));
    // Job-unit progress checkpoints for the cache-hit replay: `progress`
    // events report the pipeline's `processed_jobs` / `total_jobs`, which the
    // geometry blob does not record (issue #3897).
    let mut progress_recorder = StreamProgressRecorder::default();
    let cache_for_geometry = cache.clone();
    let cache_key_for_geometry = cache_key.clone();

    // Create streaming response that yields Parquet batches
    let stream = process_streaming(
        content.clone(),
        initial_batch_size,
        max_batch_size,
        query.opening_filter,
        tessellation_quality,
        Some(admission_guard),
    )
    .map(move |event: StreamEvent| {
        let sse_event = match event {
            StreamEvent::Start { total_estimate } => {
                ParquetStreamEvent::Start {
                    total_estimate,
                    cache_key: cache_key_clone.clone(),
                }
            }
            StreamEvent::Progress { processed, total, .. } => {
                progress_recorder.on_progress(processed, total);
                ParquetStreamEvent::Progress { processed, total }
            }
            StreamEvent::Batch { meshes, batch_number, baked_basis } => {
                progress_recorder.on_batch();
                // Per-batch CPU work (client-blob serialization + cache-writer
                // append) runs inside this stream map, i.e. on an async worker,
                // so it steps off the async pool for it.
                let serialized = super::stream_batch::off_the_async_worker(|| {
                    super::stream_batch::encode_batch(
                        &meshes,
                        layout,
                        baked_basis.as_ref(),
                        planner.as_mut(),
                        &cache_writer_for_stream,
                    )
                });

                match serialized {
                    Ok((parquet_bytes, bases)) => {
                        let base64_data = STANDARD.encode(&parquet_bytes);
                        ParquetStreamEvent::Batch {
                            data: base64_data,
                            mesh_count: meshes.len(),
                            batch_number,
                            vertex_base: bases.map(|(v, _)| v),
                            index_base: bases.map(|(_, i)| i),
                        }
                    }
                    Err(e) => {
                        ParquetStreamEvent::Error {
                            message: format!("Failed to serialize batch: {}", e),
                        }
                    }
                }
            }
            StreamEvent::Complete { stats, metadata, mesh_coordinate_space, site_transform, building_transform, symbolic_data, .. } => {
                // Cache the symbolic stream so the cached-geometry fast-path and
                // `GET /api/v1/parse/symbolic/{cache_key}` reach parity (issue #900).
                // Reuses the value already computed inside `process_streaming` —
                // no re-extraction.
                tokio::spawn(cache_symbolic_data_off_runtime(
                    cache_for_geometry.clone(),
                    cache_key_for_geometry.clone(),
                    symbolic_data.clone(),
                ));

                // Finish the incremental writer: the cache blob was built row
                // group by row group as batches streamed, so nothing is
                // re-serialized and no second copy of the geometry exists.
                let cache = cache_for_geometry.clone();
                let key = cache_key_for_geometry.clone();
                let stats_clone = stats.clone();
                let metadata_clone = metadata.clone();
                let writer_for_cache = cache_writer.clone();
                let recorded_progress = progress_recorder.take();
                let coord_space = mesh_coordinate_space;
                let site_tf = site_transform.clone();
                let building_tf = building_transform.clone();

                tokio::spawn(async move {
                    // Take the writer (moves out of Arc<Mutex>).
                    let writer = {
                        match writer_for_cache.lock() {
                            Ok(mut guard) => guard.take(),
                            Err(_) => {
                                tracing::error!("Failed to lock streaming cache writer");
                                return;
                            }
                        }
                    };
                    let Some(writer) = writer else {
                        tracing::warn!("Streaming cache writer unavailable; skipping cache fill");
                        return;
                    };
                    if writer.mesh_count() == 0 {
                        tracing::warn!("No meshes streamed; skipping cache fill");
                        return;
                    }

                    tracing::info!(
                        mesh_count = writer.mesh_count(),
                        "Caching streamed geometry (incremental writer, no re-serialization)"
                    );

                    // `finish_combined` writes the outer `[geo_len][geo_bytes]
                    // [dm_len=0]` framing (same as non-streaming endpoint,
                    // format: [geometry_len: u32][geometry_data][data_model_len: u32])
                    // directly, instead of framing the inner geometry blob and
                    // then copying it a second time into an outer buffer.
                    let finish_result =
                        tokio::task::spawn_blocking(move || writer.finish_combined()).await;

                    if let Ok(Ok(combined_parquet)) = finish_result {
                        // Cache geometry (same format as non-streaming)
                        let parquet_cache_key = parquet_geometry_key(&key, layout);
                        if let Err(e) = cache.set_bytes(&parquet_cache_key, &combined_parquet).await {
                            tracing::error!(error = %e, "Failed to cache geometry from stream");
                        } else {
                            tracing::info!(
                                cache_key = %parquet_cache_key,
                                size = combined_parquet.len(),
                                "Geometry cached from stream (optimized - no re-processing)"
                            );
                        }

                        // Cache metadata
                        let metadata_header = ParquetMetadataHeader {
                            cache_key: key.clone(),
                            metadata: metadata_clone,
                            stats: stats_clone,
                            mesh_coordinate_space: coord_space,
                            site_transform: site_tf,
                            building_transform: building_tf,
                            data_model_stats: None, // Data model cached separately via data model endpoint
                        };
                        if let Ok(metadata_json) = serde_json::to_vec(&metadata_header) {
                            let metadata_cache_key = parquet_metadata_key(&key);
                            if let Err(e) = cache.set_bytes(&metadata_cache_key, &metadata_json).await {
                                tracing::error!(error = %e, "Failed to cache metadata from stream");
                            } else {
                                tracing::debug!(cache_key = %metadata_cache_key, "Metadata cached from stream");
                            }
                        }

                        // The job-unit progress a replay must reproduce.
                        cache_stream_progress(&cache, &key, &recorded_progress).await;
                    } else {
                        tracing::error!("Failed to serialize accumulated meshes for caching");
                    }
                });

                ParquetStreamEvent::Complete { stats, metadata, symbolic_data }
            }
            StreamEvent::Error { message } => {
                ParquetStreamEvent::Error { message }
            }
        };

        let json = serde_json::to_string(&sse_event).unwrap_or_else(|e| {
            serde_json::to_string(&ParquetStreamEvent::Error {
                message: e.to_string(),
            })
            .unwrap()
        });
        Ok(Event::default().data(json))
    });

    // Spawn background task to extract and cache data model. Marked
    // in-flight for the task's whole lifetime (#5129): `get_data_model`
    // answers 202 only for a key in this set, so the marker must cover every
    // way the task can end -- completion, a `set_bytes` failure, and the
    // admission-saturated early return alike -- or a client polling during
    // exactly that window sees a 404 for a fill that is, in fact, still
    // possible on retry. `InFlightGuard::drop` handles all three uniformly.
    let content_for_cache = content.clone();
    // Keyed by the data-model ENTRY this fill writes, as `get_data_model`
    // asks (#6034): a fill of one variant says nothing about the other.
    let data_model_entities = query.data_model_entities;
    let dm_key = data_model_cache_key(&cache_key, data_model_entities);
    let cache_for_dm = cache.clone();
    let admission_for_dm = state.admission.clone();
    let in_flight = state.data_model_in_flight.begin(dm_key.clone());
    tokio::spawn(async move {
        let _in_flight = in_flight;
        // The data-model extraction re-parses the whole upload, so it must
        // pass admission like any parse job. It is a cache-fill optimization:
        // when the server is saturated, skipping it (the next request rebuilds
        // it inline) beats bypassing the gate.
        let _dm_admission = match admission_for_dm
            .acquire(content_for_cache.len() as u64)
            .await
        {
            Ok(guard) => guard,
            Err(_) => {
                tracing::debug!("Skipping data-model cache fill: admission saturated");
                return;
            }
        };
        // Run data model extraction in blocking task
        let dm_result =
            tokio::task::spawn_blocking(move || extract_data_model(&content_for_cache)).await;

        if let Ok(mut data_model) = dm_result {
            // Serialize and cache
            let serialize_result = tokio::task::spawn_blocking(move || {
                data_model_entities.apply(&mut data_model);
                serialize_data_model_to_parquet(&data_model)
            }).await;

            if let Ok(Ok(parquet_data)) = serialize_result {
                if let Err(e) = cache_for_dm.set_bytes(&dm_key, &parquet_data).await {
                    tracing::error!(error = %e, "Failed to cache data model from stream");
                } else {
                    tracing::info!(cache_key = %dm_key, size = parquet_data.len(), "Data model cached from stream");
                }
            }
        }
        // `_in_flight` drops here, clearing the marker only after the cache
        // write (success or failure) is fully resolved.
    });

    let boxed_stream: std::pin::Pin<
        Box<dyn futures::Stream<Item = Result<Event, Infallible>> + Send>,
    > = Box::pin(stream);
    Ok(Sse::new(boxed_stream)
        .keep_alive(KeepAlive::default())
        .into_response())
}

#[cfg(test)]
#[path = "parquet_stream_tests.rs"]
mod parquet_stream_tests;