blob-storage/src/main.rs

#![allow(unused)]

use std::collections::LinkedList;
use std::sync::{Arc, Mutex, RwLock};

const ARRAY_SIZE: usize = 4096;

#[derive(PartialEq, Debug, Clone, Copy)]
struct Range {
    // todo: also consider there might be an off the shelf solution that does ranges
    start: usize,
    end: usize, // exclusive
}

#[derive(PartialEq, Debug)]
enum RangeLocation {
    LeftEnd,
    RightEnd,
    Middle,
    No,
}

impl Range {
    fn new(start: usize, end: usize) -> Self {
        Self { start, end }
    }

    fn where_in_range(&self, index: usize) -> RangeLocation {
        if index == self.start {
            return RangeLocation::LeftEnd;
        }
        if index >= self.end {
            return RangeLocation::No;
        }
        if index == self.end - 1 {
            return RangeLocation::RightEnd;
        }
        RangeLocation::Middle
    }

    fn contains(&self, range: &Range) -> bool {
        range.start >= self.start && range.end <= self.end
    }

    fn is_empty(&self) -> bool {
        self.start >= self.end
    }
}

#[derive(Debug)]
struct RangeNode {
    ranges: Vec<Range>, // valid ranges
    bytes: [u8; ARRAY_SIZE],
    start: usize, // global
    end: usize,   // global
    capacity: usize,
}

impl RangeNode {
    fn default() -> Self {
        RangeNode::new(0, 0)
    }

    fn new(start: usize, end: usize) -> Self {
        Self {
            ranges: Vec::new(),
            bytes: [0; ARRAY_SIZE],
            start,
            end,
            capacity: start + ARRAY_SIZE,
        }
    }

    fn read(&mut self, start: usize, len: usize) -> Result<Vec<u8>, BlobError> {
        let mut v = Vec::new();
        let read_range = Range::new(start, start + len);
        match self.ranges.last() {
            Some(last) if read_range.end <= last.end => {}
            _ => return Err(BlobError::InvalidRange),
        }

        // check bytes are not already read
        let mut contains = false;
        for valid_range in self.ranges.iter() {
            if valid_range.contains(&read_range) {
                contains = true;
                break;
            }
            if read_range.start >= valid_range.end {
                return Err(BlobError::BytesAlreadyRead);
            }
        }
        if !contains {
            // Invalid range already caught above
            return Err(BlobError::BytesAlreadyRead);
        }

        let slice = &mut self.bytes[start..start + len];
        v.extend_from_slice(slice);

        // we copy existing entries into a new list to avoid iterator invalidation
        // from inserting while iterating
        // todo: might be a better way
        let clear_range = Range {
            start,
            end: start + len,
        };
        let mut new_ranges = Vec::new();
        let mut i = 0;
        while i < self.ranges.len() {
            let mut extra_range: Option<Range> = None;
            let mut range = self.ranges[i];
            if range.where_in_range(clear_range.start) == RangeLocation::LeftEnd {
                range.start = clear_range.end;
            } else if range.where_in_range(clear_range.end - 1) == RangeLocation::RightEnd {
                range.end = clear_range.end
            } else if range.where_in_range(clear_range.start) == RangeLocation::Middle {
                let end = range.end;
                range.end = clear_range.start;
                extra_range = Some(Range::new(clear_range.end, end));
            }

            if !range.is_empty() {
                new_ranges.push(range);
            }

            if let Some(er) = extra_range
                && !er.is_empty()
            {
                new_ranges.push(er);
            }
            i += 1;
        }
        self.ranges = new_ranges;
        Ok(v)
    }

    fn append(&mut self, input: &[u8]) -> usize {
        let new_len = input.len().min(self.capacity - self.end);
        if new_len == 0 {
            return 0;
        }
        let new_range = Range {
            start: self.end,
            end: self.end + new_len,
        };

        // we are appending so only have to check the last range
        match self.ranges.last_mut() {
            Some(range) => {
                if range.where_in_range(new_range.start) != RangeLocation::No {
                    range.end = new_range.end
                } else {
                    self.ranges.push(new_range);
                }
            }
            _ => {
                self.ranges.push(new_range);
            }
        }
        self.bytes[new_range.start..new_range.end].copy_from_slice(&input[0..new_len]);
        self.end += new_len;
        new_len
    }

    fn available_bytes(&self) -> usize {
        self.bytes.len() - self.end
    }

    fn contains(&self, index: usize) -> bool {
        index >= self.start && index < self.end
    }

    fn overlaps(&self, range: &Range) -> bool {
        self.contains(range.start) || self.contains(range.end)
    }

    fn all_read(&self) -> bool {
        false
    }
}

#[derive(Debug)]
pub struct Blob {
    // todo: consider fallible_vec to detect OOM on append
    nodes: LinkedList<RangeNode>,
    end_index: usize,
}

impl Blob {
    fn new() -> Self {
        Self {
            nodes: LinkedList::new(),
            end_index: 0,
        }
    }

    fn append(&mut self, input: &[u8]) -> usize {
        let mut remaining = input;
        let mut written_total = 0;
        while !remaining.is_empty() {
            // add nodes if necessary
            let mut new_node_start = 0;
            let need_new = match self.nodes.back() {
                Some(n) => {
                    new_node_start = n.end;
                    n.available_bytes() == 0
                }
                None => true,
            };
            if need_new {
                self.nodes.push_back(RangeNode::default());
            }
            // write to each node's array
            let node = self.nodes.back_mut().unwrap();
            let written = node.append(remaining);
            remaining = &remaining[written..];
            written_total += written;
        }
        self.end_index += input.len();
        written_total
    }

    fn read(&mut self, start: usize, len: usize) -> Result<Vec<u8>, BlobError> {
        // early check for out-of-range
        if start + len > self.end_index {
            return Err(BlobError::InvalidRange);
        }

        // read the bytes from each node
        let mut remaining_start = start;
        let mut remaining_len = len;
        let mut read_vec = Vec::with_capacity(len);
        let mut node_global_start = 0;
        // works, but this is SLOW
        for n in self.nodes.iter_mut() {
            if read_vec.len() == len {
                break;
            }
            // if !n.overlaps(&Range::new(remaining_start, remaining_start+remaining_len)) {
            //    continue;
            // }

            // todo: this works, but is a bit too clunky
            let node_global_end = node_global_start + n.end;
            let want = start + read_vec.len(); // next global byte we still need
            if want < node_global_end {
                let local_start = want - node_global_start;
                let take = (len - read_vec.len()).min(n.end - local_start);
                let chunk = n.read(local_start, take)?;
                read_vec.extend_from_slice(&chunk);
            }
            node_global_start = node_global_end;
        }
        Ok(read_vec)
    }
}

// combine all errors into one, since there aren't that many
#[derive(Debug, PartialEq)]
enum BlobError {
    BlobDNE,
    BlobExists,
    InvalidRange,
    BytesAlreadyRead,
}

struct BlobManager {
    // fine if we assume:
    // * IDs are sequentially generated
    // * and IDs can be re-used
    blobs: RwLock<Vec<Option<Arc<Mutex<Blob>>>>>,
    // Also note RwLock (create) starvation might be a concern if caller is
    // append/read heavy due to platform-defined fairness
    // todo: better approach?
    // RwLock benefits shrink as create traffic becomes heavy
}

// todo: blob indexing is left up to the caller maybe add ID generator
//       for cross-thread synchronization or change approach
impl BlobManager {
    pub fn new() -> Self {
        Self {
            blobs: RwLock::new(Vec::new()),
        }
    }

    pub fn create_blob(&self, id: usize) -> Result<(), BlobError> {
        let mut blobs = self.blobs.write().unwrap();
        if let Some(Some(_)) = blobs.get(id) {
            return Err(BlobError::BlobExists);
        } else if id >= blobs.len() {
            // todo: better strategy possible?
            blobs.resize_with(id + 1, || None);
        }
        blobs[id] = Some(Arc::new(Mutex::new(Blob::new())));
        Ok(())
    }

    pub fn append(&self, id: usize, input: &[u8]) -> Result<(), BlobError> {
        let blob = {
            let blobs = self.blobs.read().unwrap();
            match blobs.get(id) {
                Some(Some(b)) => Arc::clone(b),
                _ => return Err(BlobError::BlobDNE),
            }
        };
        blob.lock().unwrap().append(input);
        Ok(())
    }

    pub fn read(&self, id: usize, start: usize, len: usize) -> Result<Vec<u8>, BlobError> {
        let blob = {
            let blobs = self.blobs.read().unwrap();
            match blobs.get(id) {
                Some(Some(b)) => Arc::clone(b),
                _ => return Err(BlobError::BlobDNE),
            }
        };
        let read = blob.lock().unwrap().read(start, len)?; // propagate error to caller
        // todo: if read clears a blob entirely should the blob be removed?
        Ok(read)
    }
}

fn main() {
    println!("run cargo test instead");
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_range() {
        let range = Range { start: 0, end: 3 };
        assert_eq!(range, range); // identity

        let range2 = Range { start: 0, end: 3 };
        assert_eq!(range, range2); // similar

        assert_eq!(range.where_in_range(0), RangeLocation::LeftEnd);
        assert_eq!(range.where_in_range(1), RangeLocation::Middle);
        assert_eq!(range.where_in_range(2), RangeLocation::RightEnd);
        assert_eq!(range.where_in_range(3), RangeLocation::No);

        assert!(range.contains(&range2));
        assert!(range2.contains(&range));

        let range3 = Range::new(0, 1);
        assert!(range.contains(&range3));
        assert!(!range3.contains(&range));

        let range4 = Range::new(2, 3);
        assert!(range.contains(&range4));
        assert!(!range4.contains(&range));

        let range5 = Range::new(2, 4);
        assert!(!range.contains(&range5));
        assert!(!range5.contains(&range));
    }

    #[test]
    fn test_blob_core() {
        let mut b = Blob::new();
        let written = b.append(&[b'a', b'b', b'c']);
        assert_eq!(written, 3);

        let read = b.read(0, 1); // left read
        assert!(read.is_ok());
        assert_eq!(read.unwrap(), vec![b'a']);
        assert_eq!(
            b.nodes.back().unwrap().ranges,
            &[Range { start: 1, end: 3 }]
        );

        let read = b.read(0, 1); // dupe read
        assert!(read.is_err());
        assert_eq!(
            b.nodes.back().unwrap().ranges,
            &[Range { start: 1, end: 3 }]
        );

        let read = b.read(1, 2); // right read
        assert!(read.is_ok());
        assert_eq!(b.nodes.back().unwrap().ranges, &[]);

        let written = b.append(&[b'a', b'b', b'c']);
        assert_eq!(written, 3);
        assert_eq!(
            b.nodes.back().unwrap().ranges,
            &[Range { start: 3, end: 6 }]
        );

        let read = b.read(4, 1); // middle read
        assert!(read.is_ok());
        assert_eq!(
            b.nodes.back().unwrap().ranges,
            &[Range::new(3, 4), Range::new(5, 6)]
        );
    }

    #[test]
    fn test_range_node() {
        let mut b = RangeNode::default();
        let written = b.append(&[b'a', b'b', b'c']);
        assert_eq!(written, 3);
        assert_eq!(b.ranges, &[Range { start: 0, end: 3 }]);

        let read = b.read(0, 1); // left read
        assert!(read.is_ok());
        assert_eq!(read.unwrap(), &[b'a']);
        assert_eq!(b.ranges, &[Range { start: 1, end: 3 }]);

        let read = b.read(0, 1); // dupe read
        assert!(read.is_err());
        assert_eq!(b.ranges, &[Range { start: 1, end: 3 }]);

        let read = b.read(1, 2); // right read
        assert!(read.is_ok());
        assert_eq!(read.unwrap(), &[b'b', b'c']);
        assert_eq!(b.ranges, &[]);

        let written = b.append(&[b'a', b'b', b'c']);
        assert_eq!(written, 3);
        assert_eq!(b.ranges, &[Range { start: 3, end: 6 }]);

        let read = b.read(4, 1); // middle read
        assert!(read.is_ok());
        assert_eq!(read.unwrap(), &[b'b']);
        assert_eq!(b.ranges, &[Range::new(3, 4), Range::new(5, 6)]);
    }

    #[test]
    fn test_blob_manager_threaded_read() {
        use std::thread;

        let bm = Arc::new(BlobManager::new());

        let handles: Vec<_> = (0..2usize)
            .map(|i| {
                let bm = Arc::clone(&bm);
                thread::spawn(move || {
                    let len = 1_000_000; // 100 MB test takes a while
                    let data = vec![i as u8; len];
                    bm.create_blob(i).unwrap();
                    bm.append(i, &data).unwrap();
                    for j in 0..len as usize {
                        let byte = bm.read(i, j, 1).unwrap();
                        assert_eq!(byte[0], i as u8);
                    }
                })
            })
            .collect();

        for handle in handles {
            handle.join().unwrap();
        }
    }

    #[test]
    fn test_blob_manager_threaded_create() {
        use std::thread;

        let bm = Arc::new(BlobManager::new());

        let num_creates = 20_000;

        let even = {
            let bm = Arc::clone(&bm);
            thread::spawn(move || {
                for i in (0..num_creates).step_by(2) {
                    bm.create_blob(i).unwrap();
                }
            })
        };

        let odd = {
            let bm = Arc::clone(&bm);
            thread::spawn(move || {
                for i in (1..num_creates).step_by(2) {
                    bm.create_blob(i).unwrap();
                }
            })
        };

        even.join().unwrap();
        odd.join().unwrap();

        let blobs = bm.blobs.read().unwrap();
        assert_eq!(blobs.len(), num_creates);
        assert!(blobs.iter().all(|b| b.is_some()));
    }

    #[test]
    fn test_blob_manager() {
        let bm = BlobManager::new();
        bm.create_blob(0);
        bm.append(0, &[b'a', b'b', b'c']);
        let v = Vec::from([b'a', b'b']);
        assert_eq!(bm.read(0, 0, 2).unwrap(), v);

        bm.create_blob(1);
        bm.append(1, &[b'a', b'b', b'c']);
        assert_eq!(bm.read(1, 0, 4).unwrap_err(), BlobError::InvalidRange);

        let read = bm.read(1, 0, 2).unwrap();
        assert_eq!(read, v);
        assert_eq!(bm.read(1, 0, 2).unwrap_err(), BlobError::BytesAlreadyRead);
        assert_eq!(bm.read(1, 1, 2).unwrap_err(), BlobError::BytesAlreadyRead);

        assert_eq!(bm.read(2, 0, 2).unwrap_err(), BlobError::BlobDNE);
        assert_eq!(bm.create_blob(1).unwrap_err(), BlobError::BlobExists);
    }
}