spiral-sdk 0.1.7

//! Bounded Arrow IPC decoding for HTTP responses.
//!
//! The decoder requires a Tokio runtime for asynchronous input. It validates
//! framing before allocations and rejects compressed, oversized, or unfinished
//! messages. Dropping a reader cancels the body and releases its worker inputs.

use std::collections::HashMap;
use std::io::{self, Read};
use std::thread;

use arrow_array::RecordBatch;
use arrow_ipc::reader::StreamReader;
use arrow_ipc::root_as_message;
use arrow_schema::SchemaRef;
use bytes::Bytes;
use futures::StreamExt as _;
use futures::stream::BoxStream;
use tokio::sync::{mpsc, oneshot};
use tokio::task::AbortHandle;

/// Error produced while reading a streaming body.
pub(crate) type BodyError = Box<dyn std::error::Error + Send + Sync>;
/// An owned stream of body chunks. Chunk boundaries have no protocol meaning.
pub(crate) type BodyStream = BoxStream<'static, Result<Bytes, BodyError>>;

const EOS: [u8; 8] = [255, 255, 255, 255, 0, 0, 0, 0];
const MAX_CHUNK_BYTES: usize = 256 * 1024;
const MAX_METADATA_BYTES: usize = 8 * 1024 * 1024;
const MAX_BODY_BYTES: usize = 64 * 1024 * 1024;
const MAX_DICTIONARY_IDS: usize = 1024;

/// One owned, bounded Arrow IPC stream decoded on a worker thread.
pub(crate) struct ArrowIpcReader {
    schema: SchemaRef,
    batches: mpsc::Receiver<Decoded>,
    finished: bool,
    _body_task: BodyTask,
}
struct BodyTask(AbortHandle);
impl Drop for BodyTask {
    fn drop(&mut self) {
        self.0.abort();
    }
}
impl ArrowIpcReader {
    /// Read the schema without collecting the body or materializing all batches.
    ///
    /// Requires an active Tokio runtime. Cancelling this future drops the source.
    pub async fn new(body: BodyStream) -> io::Result<Self> {
        let (byte_sender, byte_receiver) = mpsc::channel(2);
        let (batch_sender, batches) = mpsc::channel(2);
        let (schema_sender, schema_receiver) = oneshot::channel();
        thread::Builder::new()
            .name("spiral-ipc-decode".to_owned())
            .spawn(move || {
                decode(byte_receiver, batch_sender, schema_sender);
            })
            .map_err(failure)?;
        let task = tokio::spawn(async move {
            let mut body = body;
            loop {
                let next = tokio::select! {
                    _ = byte_sender.closed() => return,
                    next = body.next() => next,
                };
                match next {
                    Some(Ok(bytes)) => {
                        for offset in (0..bytes.len()).step_by(MAX_CHUNK_BYTES) {
                            let end = (offset + MAX_CHUNK_BYTES).min(bytes.len());
                            if byte_sender
                                .send(Input::Bytes(bytes.slice(offset..end)))
                                .await
                                .is_err()
                            {
                                return;
                            }
                        }
                    }
                    Some(Err(error)) => {
                        drop(byte_sender.send(Input::Error(error.to_string())).await);
                        return;
                    }
                    None => {
                        drop(byte_sender.send(Input::End).await);
                        return;
                    }
                }
            }
        });
        let body_task = BodyTask(task.abort_handle());
        let schema = schema_receiver.await.map_err(failure)??;
        Ok(Self {
            schema,
            batches,
            finished: false,
            _body_task: body_task,
        })
    }
    /// Schema declared by the stream's first message.
    pub fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    /// Read a batch; a framing or transport failure is emitted before completion.
    pub async fn next_batch(&mut self) -> Option<io::Result<RecordBatch>> {
        if self.finished {
            return None;
        }
        let decoded = self.batches.recv().await;
        self.take_decoded(decoded)
    }
    fn take_decoded(&mut self, decoded: Option<Decoded>) -> Option<io::Result<RecordBatch>> {
        match decoded {
            Some(Decoded::Batch(batch)) => {
                self.finished = batch.is_err();
                Some(batch)
            }
            Some(Decoded::Complete) => {
                self.finished = true;
                None
            }
            None => {
                self.finished = true;
                Some(Err(failure(
                    "Arrow decoder ended without validated completion",
                )))
            }
        }
    }
}

enum Decoded {
    Batch(io::Result<RecordBatch>),
    Complete,
}
enum Input {
    Bytes(Bytes),
    Error(String),
    End,
}
struct ByteReader {
    input: mpsc::Receiver<Input>,
    current: Bytes,
    framing: Framing,
    ended: bool,
}
impl ByteReader {
    fn finish(&mut self) -> io::Result<()> {
        if !self.current.is_empty() {
            return Err(invalid("bytes followed the Arrow end marker"));
        }
        if !self.ended {
            loop {
                match self.input.blocking_recv() {
                    Some(Input::Bytes(bytes)) if bytes.is_empty() => continue,
                    Some(Input::Bytes(_)) => {
                        return Err(invalid("bytes followed the Arrow end marker"));
                    }
                    Some(Input::Error(error)) => return Err(invalid(error)),
                    Some(Input::End) => {
                        self.ended = true;
                        break;
                    }
                    None => return Err(invalid("HTTP body ended without completion")),
                }
            }
        }
        if !matches!(self.framing.phase, Phase::Complete) {
            return Err(invalid("Arrow stream is missing its end marker"));
        }
        Ok(())
    }
}
impl Read for ByteReader {
    fn read(&mut self, output: &mut [u8]) -> io::Result<usize> {
        if output.is_empty() {
            return Ok(0);
        }
        while self.current.is_empty() && !self.ended {
            match self.input.blocking_recv() {
                Some(Input::Bytes(bytes)) => self.current = bytes,
                Some(Input::Error(error)) => return Err(invalid(error)),
                Some(Input::End) => {
                    self.ended = true;
                    self.finish()?;
                }
                None => return Err(invalid("HTTP body ended without completion")),
            }
        }
        let count = self.current.len().min(output.len());
        output[..count].copy_from_slice(&self.current[..count]);
        self.current = self.current.slice(count..);
        self.framing.accept(&output[..count])?;
        Ok(count)
    }
}
#[derive(Default)]
struct Framing {
    phase: Phase,
    dictionaries: DictionaryBudget,
}

#[derive(Default)]
struct DictionaryBudget {
    sizes: HashMap<i64, usize>,
    total: usize,
}
impl DictionaryBudget {
    // Encoded metadata is charged conservatively alongside body buffers, so
    // delta dictionaries cannot accumulate unbounded retained values.
    fn record(&mut self, id: i64, delta: bool, bytes: usize) -> io::Result<()> {
        let previous = self.sizes.get(&id).copied().unwrap_or_default();
        if !self.sizes.contains_key(&id) && self.sizes.len() >= MAX_DICTIONARY_IDS {
            return Err(invalid("Arrow IPC exceeds 1024 dictionary IDs"));
        }
        let replacement = if delta { previous + bytes } else { bytes };
        let total = self.total - previous + replacement;
        if total > MAX_BODY_BYTES {
            return Err(invalid("Arrow IPC retained dictionaries exceed 64 MiB"));
        }
        self.sizes.insert(id, replacement);
        self.total = total;
        Ok(())
    }
}
enum Phase {
    Prefix(Vec<u8>),
    Metadata { length: usize, bytes: Vec<u8> },
    Body(usize),
    Complete,
}
impl Default for Phase {
    fn default() -> Self {
        Self::Prefix(Vec::with_capacity(8))
    }
}
impl Framing {
    fn accept(&mut self, mut bytes: &[u8]) -> io::Result<()> {
        while !bytes.is_empty() {
            match &mut self.phase {
                Phase::Prefix(prefix) => {
                    let count = (8 - prefix.len()).min(bytes.len());
                    prefix.extend_from_slice(&bytes[..count]);
                    bytes = &bytes[count..];
                    if prefix.len() == 8 {
                        if prefix[..4] != EOS[..4] {
                            return Err(invalid("Arrow IPC continuation marker is required"));
                        }
                        let length =
                            u32::from_le_bytes(prefix[4..8].try_into().expect("four prefix bytes"))
                                as usize;
                        if length > MAX_METADATA_BYTES {
                            return Err(invalid("Arrow IPC metadata exceeds 8 MiB"));
                        }
                        self.phase = if length == 0 {
                            Phase::Complete
                        } else {
                            Phase::Metadata {
                                length,
                                bytes: Vec::with_capacity(length),
                            }
                        };
                    }
                }
                Phase::Metadata {
                    length,
                    bytes: metadata,
                } => {
                    let count = (*length - metadata.len()).min(bytes.len());
                    metadata.extend_from_slice(&bytes[..count]);
                    bytes = &bytes[count..];
                    if metadata.len() == *length {
                        let message = root_as_message(metadata).map_err(invalid)?;
                        let body = usize::try_from(message.bodyLength()).map_err(invalid)?;
                        if body > MAX_BODY_BYTES {
                            return Err(invalid("Arrow IPC batch body exceeds 64 MiB"));
                        }
                        let compression = message
                            .header_as_record_batch()
                            .and_then(|batch| batch.compression())
                            .or_else(|| {
                                message
                                    .header_as_dictionary_batch()
                                    .and_then(|batch| batch.data())
                                    .and_then(|batch| batch.compression())
                            });
                        if compression.is_some() {
                            return Err(invalid("compressed Arrow IPC is unsupported"));
                        }
                        if let Some(dictionary) = message.header_as_dictionary_batch() {
                            self.dictionaries.record(
                                dictionary.id(),
                                dictionary.isDelta(),
                                body + *length,
                            )?;
                        }
                        self.phase = if body == 0 {
                            Phase::default()
                        } else {
                            Phase::Body(body)
                        };
                    }
                }
                Phase::Body(remaining) => {
                    let count = (*remaining).min(bytes.len());
                    bytes = &bytes[count..];
                    *remaining -= count;
                    if *remaining == 0 {
                        self.phase = Phase::default();
                    }
                }
                Phase::Complete => return Err(invalid("bytes followed the Arrow end marker")),
            }
        }
        Ok(())
    }
}

fn decode(
    input: mpsc::Receiver<Input>,
    output: mpsc::Sender<Decoded>,
    schema: oneshot::Sender<io::Result<SchemaRef>>,
) {
    let reader = ByteReader {
        input,
        current: Bytes::new(),
        framing: Framing::default(),
        ended: false,
    };
    let mut reader = match StreamReader::try_new(reader, None) {
        Ok(reader) => reader,
        Err(error) => {
            drop(schema.send(Err(failure(error))));
            return;
        }
    };
    if schema.send(Ok(reader.schema())).is_err() {
        return;
    }
    for batch in reader.by_ref() {
        let failed = batch.is_err();
        if output
            .blocking_send(Decoded::Batch(batch.map_err(failure)))
            .is_err()
            || failed
        {
            return;
        }
    }
    match reader.get_mut().finish() {
        Ok(()) => drop(output.blocking_send(Decoded::Complete)),
        Err(error) => drop(output.blocking_send(Decoded::Batch(Err(failure(error))))),
    }
}
fn invalid(error: impl ToString) -> io::Error {
    io::Error::new(io::ErrorKind::InvalidData, error.to_string())
}
fn failure(error: impl ToString) -> io::Error {
    io::Error::other(error.to_string())
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use arrow_array::builder::StringDictionaryBuilder;
    use arrow_array::types::Int32Type;
    use arrow_array::{Int64Array, RecordBatch};
    use arrow_ipc::writer::{DictionaryHandling, IpcWriteOptions, StreamWriter};
    use arrow_schema::{DataType, Field, Schema};
    use futures::channel::mpsc::unbounded;
    use futures::stream;

    use super::*;

    fn body<E: Into<BodyError> + 'static>(
        source: impl futures::Stream<Item = Result<Bytes, E>> + Send + 'static,
    ) -> BodyStream {
        source.map(|chunk| chunk.map_err(Into::into)).boxed()
    }

    fn encoded() -> Vec<u8> {
        let schema = Arc::new(Schema::new(vec![Field::new(
            "value",
            DataType::Int64,
            false,
        )]));
        let batch =
            RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(vec![1, 2]))])
                .unwrap();
        let mut bytes = Vec::new();
        let mut writer = StreamWriter::try_new(&mut bytes, &schema).unwrap();
        writer.write(&batch).unwrap();
        writer.finish().unwrap();
        bytes
    }

    #[tokio::test]
    async fn remote_reader_preserves_rows_and_rejects_missing_completion() {
        let bytes = encoded();
        let mut with_trailing_bytes = bytes.clone();
        with_trailing_bytes.push(1);
        let truncated = bytes[..bytes.len() - EOS.len()].to_vec();
        for (source, completed) in [
            (bytes, true),
            (truncated, false),
            (with_trailing_bytes, false),
        ] {
            let chunks = source
                .chunks(7)
                .map(Bytes::copy_from_slice)
                .map(Ok::<_, io::Error>)
                .collect::<Vec<_>>();
            let mut scan = ArrowIpcReader::new(body(stream::iter(chunks)))
                .await
                .unwrap();
            assert_eq!(scan.next_batch().await.unwrap().unwrap().num_rows(), 2);
            if completed {
                assert!(scan.next_batch().await.is_none());
            } else {
                assert!(scan.next_batch().await.unwrap().is_err());
            }
        }
        let mut oversized_prefix = EOS;
        oversized_prefix[4..].copy_from_slice(&((MAX_METADATA_BYTES + 1) as u32).to_le_bytes());
        let source = stream::iter([Ok::<_, io::Error>(Bytes::copy_from_slice(
            &oversized_prefix,
        ))]);
        let result = ArrowIpcReader::new(body(source)).await;
        assert!(
            result.is_err(),
            "metadata limits must fail before schema allocation"
        );
    }

    #[tokio::test]
    async fn dictionary_updates_remain_bounded_before_body_allocation() {
        let schema = Arc::new(Schema::new(vec![Field::new(
            "value",
            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
            false,
        )]));
        let mut builder = StringDictionaryBuilder::<Int32Type>::new();
        let mut bytes = Vec::new();
        let options =
            IpcWriteOptions::default().with_dictionary_handling(DictionaryHandling::Delta);
        let mut writer = StreamWriter::try_new_with_options(&mut bytes, &schema, options).unwrap();
        for value in ["first", "second"] {
            builder.append(value).unwrap();
            let batch = RecordBatch::try_new(
                schema.clone(),
                vec![Arc::new(builder.finish_preserve_values())],
            )
            .unwrap();
            writer.write(&batch).unwrap();
        }
        writer.finish().unwrap();
        let mut reader = ArrowIpcReader::new(body(stream::iter([Ok::<_, io::Error>(
            Bytes::copy_from_slice(&bytes),
        )])))
        .await
        .unwrap();
        assert_eq!(reader.next_batch().await.unwrap().unwrap().num_rows(), 1);
        assert_eq!(reader.next_batch().await.unwrap().unwrap().num_rows(), 1);
        assert!(reader.next_batch().await.is_none());

        let mut offset = 0;
        let mut checked = 0;
        loop {
            let length =
                u32::from_le_bytes(bytes[offset + 4..offset + 8].try_into().unwrap()) as usize;
            if length == 0 {
                break;
            }
            let end = offset + 8 + length;
            let message = root_as_message(&bytes[offset + 8..end]).unwrap();
            if let Some(dictionary) = message.header_as_dictionary_batch() {
                let mut framing = Framing::default();
                // Seed a full retained dictionary; parsing the next header must
                // reject the update before Arrow requests any body allocation.
                let existing_id = if dictionary.isDelta() {
                    dictionary.id()
                } else {
                    -1
                };
                framing
                    .dictionaries
                    .record(existing_id, false, MAX_BODY_BYTES)
                    .unwrap();
                let error = framing.accept(&bytes[offset..end]).unwrap_err();
                assert!(error.to_string().contains("retained dictionaries"));
                checked += 1;
            }
            offset = end + message.bodyLength() as usize;
        }
        assert_eq!(
            checked, 2,
            "both new and delta dictionaries must be checked"
        );
        let mut budget = DictionaryBudget::default();
        budget.record(0, false, MAX_BODY_BYTES).unwrap();
        budget.record(0, false, 1).unwrap();
        budget.record(1, false, MAX_BODY_BYTES - 1).unwrap();
        assert_eq!(
            budget.total, MAX_BODY_BYTES,
            "replacement dictionaries release their budget"
        );
        let mut budget = DictionaryBudget::default();
        for id in 0..MAX_DICTIONARY_IDS as i64 {
            budget.record(id, false, 0).unwrap();
        }
        assert!(budget.record(MAX_DICTIONARY_IDS as i64, false, 0).is_err());
    }

    #[tokio::test]
    async fn dropping_reader_cancels_a_stalled_http_body() {
        let (sender, receiver) = unbounded::<Result<Bytes, io::Error>>();
        let mut bytes = encoded();
        bytes.truncate(bytes.len() - EOS.len());
        sender.unbounded_send(Ok(Bytes::from(bytes))).unwrap();
        let mut scan = ArrowIpcReader::new(body(receiver)).await.unwrap();
        assert_eq!(scan.next_batch().await.unwrap().unwrap().num_rows(), 2);
        drop(scan);
        tokio::time::timeout(std::time::Duration::from_secs(1), async {
            while !sender.is_closed() {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("dropping the reader must drop its pending HTTP source");
    }

    #[tokio::test]
    async fn remote_reader_reports_transport_failure_after_a_complete_batch() {
        let mut bytes = encoded();
        bytes.truncate(bytes.len() - EOS.len());
        let chunks = vec![
            Ok(Bytes::from(bytes)),
            Err(io::Error::other("upstream failed")),
        ];
        let mut scan = ArrowIpcReader::new(body(stream::iter(chunks)))
            .await
            .unwrap();
        assert_eq!(scan.next_batch().await.unwrap().unwrap().num_rows(), 2);
        assert!(
            scan.next_batch()
                .await
                .unwrap()
                .unwrap_err()
                .to_string()
                .contains("upstream failed")
        );
    }
}