nodata/crates/nodata-storage/src/lib.rs

// # Nodata Storage:
// storage is an abstraction for a log based partitiniong system on top of a given file system.
// it allows treating topics with a stream of events as logs, that are persisted for fast lookups
//
// A topic has a log corresponding log location (hash of the topic name)
// Each message is sequentially added to a log file and offsets are stored in a separate index file for fast lookup

use std::{collections::BTreeMap, sync::Arc, time::SystemTime};

use anyhow::Context;
use backend::StorageBackend;
use proto::ProtoStorage;
use sha2::{Digest, Sha256};
use tokio::sync::Mutex;

type TopicHashKey = String;

pub mod backend;

#[derive(Clone)]
pub struct Storage {
    segment_size_bytes: usize,
    buffer: Arc<Mutex<BTreeMap<TopicHashKey, Vec<Vec<u8>>>>>,
    backend: Arc<StorageBackend>,

    codec: ProtoStorage,
}

impl Storage {
    pub fn new(backend: StorageBackend) -> Self {
        Self {
            segment_size_bytes: 4096 * 1000, // 4MB
            buffer: Arc::default(),

            backend: Arc::new(backend),
            codec: ProtoStorage::default(),
        }
    }

    pub async fn push_message(&self, topic: &str, value: &[u8]) -> anyhow::Result<()> {
        // TODO: Consider caching the topic hash, such that we don't have to hash it every time we need to publish a message
        let topic_key = self.hash_topic(topic);

        if value.len() > self.segment_size_bytes {
            anyhow::bail!("value to large for segment");
        }

        {
            // FIXME: the lock is very expensive here, especially as it is over multiple wait points. Until it becomes a problem, just leave it be
            let mut buffer = self.buffer.lock().await;
            if buffer.get_mut(&topic_key).is_none() {
                buffer.insert(
                    topic_key.clone(),
                    Vec::with_capacity(self.segment_size_bytes),
                );
            }

            let log_buffer = buffer.get_mut(&topic_key).unwrap();
            if log_buffer
                .iter()
                .map(|b| b.len())
                .reduce(|acc, i| acc + i)
                .unwrap_or_default()
                >= self.segment_size_bytes
            {
                let log_buffer = buffer.remove(&topic_key).unwrap();
                let time = SystemTime::now();
                let segment_id = self
                    .backend
                    .flush_segment(&topic_key, &self.codec.format_log_message(log_buffer))
                    .await
                    .context("failed to write segment")?;
                self.backend
                    .append_index(&topic_key, &segment_id, time)
                    .await?;

                let mut log_buf = Vec::with_capacity(self.segment_size_bytes);
                log_buf.push(value.to_vec());
                buffer.insert(topic_key.clone(), log_buf);
            } else {
                log_buffer.push(value.to_vec());
            }
        }

        Ok(())
    }

    pub async fn commit_segment_offset(
        &self,
        topic: &str,
        segment_id: &str,
        time: SystemTime,
    ) -> anyhow::Result<()> {
        Ok(())
    }

    fn hash_topic(&self, input: &str) -> String {
        hex::encode(Sha256::digest(input.as_bytes()))
    }
}

mod proto {
    use prost::Message;

    mod gen {
        include!("gen/nodata.v1.rs");
    }

    #[derive(Default, Clone)]
    pub struct ProtoStorage {}

    impl ProtoStorage {
        pub fn format_log_message(&self, messages: Vec<Vec<u8>>) -> Vec<u8> {
            gen::Log { messages }.encode_to_vec()
        }
    }
}