nordisk/edda

crates/nornir-flight/src/store.rs

//! A default in-memory [`FlightSink`] + [`FlightSource`].
//!
//! `MemFlightStore` is the reference implementation of the two seams: `put`
//! stashes the batches under the key, `get` hands them back. It is what the
//! round-trip test drives, and a perfectly good real sink for a single-box
//! "produce here, replay here" deployment (the design's in-process default).

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use arrow_array::RecordBatch;
use async_trait::async_trait;

use crate::error::{Error, Result};
use crate::traits::{FlightSink, FlightSource};

/// An in-memory, key→batches store implementing both transport seams. Cheap to
/// clone (shares one `Arc<Mutex<..>>`), so the same handle can be a sink for a
/// producer and a source for a consumer.
#[derive(Clone, Default)]
pub struct MemFlightStore {
    inner: Arc<Mutex<HashMap<String, Vec<RecordBatch>>>>,
}

impl MemFlightStore {
    /// A fresh empty store.
    pub fn new() -> Self {
        Self::default()
    }

    /// Number of keys currently held (test/introspection helper).
    pub fn len(&self) -> usize {
        self.inner.lock().expect("store lock").len()
    }

    /// Whether the store holds nothing.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

#[async_trait]
impl FlightSink for MemFlightStore {
    async fn put(&self, key: String, batches: Vec<RecordBatch>) -> Result<()> {
        self.inner
            .lock()
            .map_err(|_| Error::Route("store lock poisoned".into()))?
            .entry(key)
            .or_default()
            .extend(batches);
        Ok(())
    }
}

#[async_trait]
impl FlightSource for MemFlightStore {
    async fn get(&self, key: String) -> Result<Vec<RecordBatch>> {
        self.inner
            .lock()
            .map_err(|_| Error::Route("store lock poisoned".into()))?
            .get(&key)
            .cloned()
            .ok_or_else(|| Error::Route(format!("no stream for key {key:?}")))
    }
}