blob-storage/src/main.rs

#![allow(unused)]

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

#[derive(PartialEq, Debug, Clone, Copy)]
struct Range {
    // todo: consider just using a map instead
    // 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
    }
}

pub struct Blob {
    // todo: consider fallible_vec to detect OOM on append
    bytes: Vec<u8>,
    valid_ranges: Vec<Range>,
}

impl Blob {
    fn new() -> Self {
        Self {
            bytes: Vec::new(),
            valid_ranges: Vec::new(),
        }
    }

    fn append(&mut self, input: &[u8]) {
        let new_range = Range {
            start: self.bytes.len(),
            end: self.bytes.len() + input.len(),
        };

        // we are appending so only have to check the last range
        match self.valid_ranges.last_mut() {
            Some(range) if range.where_in_range(new_range.start) != RangeLocation::No => {
                range.end = new_range.end
            }
            _ => self.valid_ranges.push(new_range),
        }
        self.bytes.extend_from_slice(input);
    }

    // todo: better way to do the return type?
    fn read(&mut self, start: usize, len: usize, output: &mut [u8]) -> Result<(), BlobError> {
        assert!(len == output.len()); // todo: add error handling

        let read_range = Range::new(start, start + len);

        // check read within bytes len
        match self.valid_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.valid_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];
        slice.swap_with_slice(output);

        // 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.valid_ranges.len() {
            let mut extra_range: Option<Range> = None;
            let mut range = self.valid_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.valid_ranges = new_ranges;
        Ok(())
    }
}

// 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()),
        }
    }

    // todo: is usize appropriate datatype for ID?
    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 mut read = vec![0; len];
        blob.lock().unwrap().read(start, len, &mut read)?; // 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() {
        let mut b = Blob::new();
        b.append(&[b'a', b'b', b'c']);
        assert_eq!(b.valid_ranges, &[Range { start: 0, end: 3 }]);

        let mut read = vec![0; 1];
        b.read(0, 1, &mut read); // left read
        assert_eq!(b.valid_ranges, &[Range { start: 1, end: 3 }]);

        let mut read = vec![0; 1];
        b.read(0, 1, &mut read); // dupe read // todo: should error or fail
        assert_eq!(b.valid_ranges, &[Range { start: 1, end: 3 }]);

        let mut read = vec![0; 2];
        b.read(1, 2, &mut read); // right read
        assert_eq!(b.valid_ranges, &[]);

        b.append(&[b'a', b'b', b'c']);
        assert_eq!(b.valid_ranges, &[Range { start: 3, end: 6 }]);
        let mut read = vec![0; 1];
        b.read(4, 1, &mut read); // middle read
        assert_eq!(b.valid_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 = 10_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 even = {
            let bm = Arc::clone(&bm);
            thread::spawn(move || {
                for i in (0..20000usize).step_by(2) {
                    bm.create_blob(i).unwrap();
                }
            })
        };

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

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

        let blobs = bm.blobs.read().unwrap();
        assert_eq!(blobs.len(), 20000);
        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);
    }
}