re_viewer 0.31.4

use arrow::array::RecordBatch;

use re_chunk::{Chunk, TimeInt};
use re_entity_db::EntityDb;
use re_log_types::TimelinePoint;
use re_redap_client::{ApiResult, ConnectionClient};
use re_viewer_context::TimeControl;

pub fn prefetch_chunks_for_recording(
    egui_ctx: &egui::Context,
    chunk_budget: re_memory::MemoryLimit,
    recording: &mut EntityDb,
    time_ctrl: Option<&TimeControl>,
    connection_registry: &re_redap_client::ConnectionRegistryHandle,
) {
    re_tracing::profile_function!();

    let Some(redap_uri) = recording.redap_uri().cloned() else {
        return;
    };
    let origin = redap_uri.origin.clone();

    let time_cursor = time_ctrl.and_then(|time_ctrl| {
        let current_time = time_ctrl.time_i64()?;
        let timeline = *time_ctrl.timeline()?;
        let start_time = TimeInt::new_temporal(current_time);

        Some(TimelinePoint::from((timeline, start_time)))
    });

    if !recording.can_fetch_chunks_from_redap() {
        return;
    }

    let options = re_entity_db::ChunkPrefetchOptions {
        total_uncompressed_byte_budget: chunk_budget.as_bytes(),
        ..Default::default()
    };

    let (rrd_manifest, storage_engine) = recording.rrd_manifest_index_mut_and_storage_engine();

    let load_chunks = &|rb| {
        egui_ctx.request_repaint();
        let connection_registry = connection_registry.clone();
        let origin = origin.clone();

        let fut = async move {
            let mut client = connection_registry.client(origin).await.map_err(|err| {
                re_log::warn_once!("Failed to connect to remote: {err}");
            })?;
            load_chunks(&mut client, &rb).await.map_err(|err| {
                re_log::warn_once!("{err}");
            })
        };

        // Annoying poll_promise API difference:
        cfg_if::cfg_if! {
            if #[cfg(target_arch = "wasm32")] {
                poll_promise::Promise::spawn_local(fut)
            } else {
                poll_promise::Promise::spawn_async(fut)
            }
        }
    };

    if let Err(err) =
        rrd_manifest.prefetch_chunks(storage_engine.store(), &options, time_cursor, load_chunks)
    {
        re_log::warn_once!("prefetch_chunks failed: {err}");
    }
}

/// Takes a dataframe that looks like an [`re_log_encoding::RrdManifest`] (has a `chunk_key` column).
async fn load_chunks(client: &mut ConnectionClient, batch: &RecordBatch) -> ApiResult<Vec<Chunk>> {
    use tokio_stream::StreamExt as _;

    if batch.num_rows() == 0 {
        return Ok(vec![]);
    }

    re_log::trace!("Requesting {} chunk(s) from server…", batch.num_rows());

    let chunk_stream = client.fetch_segment_chunks_by_id(batch).await?;
    let mut chunk_stream =
        re_redap_client::fetch_chunks_response_to_chunk_and_segment_id(chunk_stream);
    let mut all_chunks = Vec::new();
    while let Some(chunks) = chunk_stream.next().await {
        for (chunk, _partition_id) in chunks? {
            all_chunks.push(chunk);
        }
    }

    re_log::trace!("Finished downloading {} chunk(s).", batch.num_rows());

    Ok(all_chunks)
}