use std::sync::Mutex;
use structfs_core_store::Value;
use crate::gate::{CancelToken, Cancelled, Gate};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ByteChunk {
pub bytes: Vec<u8>,
pub offset: u64,
pub eof: bool,
}
impl ByteChunk {
pub fn into_value(self) -> Value {
Value::Bytes(self.bytes)
}
pub fn next_offset(&self) -> u64 {
self.offset + self.bytes.len() as u64
}
}
struct StreamState {
data: Vec<u8>,
closed: bool,
}
pub struct ByteStream {
state: Mutex<StreamState>,
gate: Gate,
}
impl Default for ByteStream {
fn default() -> Self {
Self::new()
}
}
impl ByteStream {
pub fn new() -> Self {
Self {
state: Mutex::new(StreamState {
data: Vec::new(),
closed: false,
}),
gate: Gate::new(),
}
}
fn lock(&self) -> std::sync::MutexGuard<'_, StreamState> {
self.state.lock().unwrap_or_else(|e| e.into_inner())
}
pub fn push(&self, bytes: &[u8]) -> bool {
{
let mut state = self.lock();
if state.closed {
return false;
}
state.data.extend_from_slice(bytes);
}
self.gate.notify();
true
}
pub fn close(&self) {
self.lock().closed = true;
self.gate.notify();
}
pub fn is_closed(&self) -> bool {
self.lock().closed
}
pub fn len(&self) -> u64 {
self.lock().data.len() as u64
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
fn chunk_at(state: &StreamState, offset: u64, max: usize) -> ByteChunk {
let start = (offset as usize).min(state.data.len());
let end = start.saturating_add(max).min(state.data.len());
ByteChunk {
bytes: state.data[start..end].to_vec(),
offset: start as u64,
eof: state.closed && end == state.data.len(),
}
}
pub fn snapshot_at(&self, offset: u64, max: usize) -> ByteChunk {
Self::chunk_at(&self.lock(), offset, max)
}
pub async fn read_at(&self, offset: u64, max: usize) -> ByteChunk {
self.gate
.wait_until(|| {
let state = self.lock();
if (state.data.len() as u64) > offset || state.closed {
Some(Self::chunk_at(&state, offset, max))
} else {
None
}
})
.await
}
pub async fn read_at_cancellable(
&self,
offset: u64,
max: usize,
token: &CancelToken,
) -> Result<ByteChunk, Cancelled> {
self.gate
.wait_until_cancellable(token, || {
let state = self.lock();
if (state.data.len() as u64) > offset || state.closed {
Some(Self::chunk_at(&state, offset, max))
} else {
None
}
})
.await
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
#[tokio::test]
async fn ranged_reads_return_available_bytes() {
let stream = ByteStream::new();
stream.push(b"hello world");
let chunk = stream.read_at(0, 5).await;
assert_eq!(chunk.bytes, b"hello");
assert_eq!(chunk.next_offset(), 5);
assert!(!chunk.eof);
let chunk = stream.read_at(chunk.next_offset(), 100).await;
assert_eq!(chunk.bytes, b" world");
assert!(!chunk.eof);
}
#[tokio::test]
async fn read_parks_until_push() {
let stream = Arc::new(ByteStream::new());
let reader = {
let stream = stream.clone();
tokio::spawn(async move { stream.read_at(0, 16).await })
};
tokio::task::yield_now().await;
stream.push(b"data");
let chunk = reader.await.unwrap();
assert_eq!(chunk.bytes, b"data");
}
#[tokio::test]
async fn empty_read_means_eof() {
let stream = Arc::new(ByteStream::new());
stream.push(b"tail");
let reader = {
let stream = stream.clone();
tokio::spawn(async move {
let first = stream.read_at(0, 100).await;
let second = stream.read_at(first.next_offset(), 100).await;
(first, second)
})
};
tokio::task::yield_now().await;
stream.close();
let (first, second) = reader.await.unwrap();
assert_eq!(first.bytes, b"tail");
assert!(second.bytes.is_empty());
assert!(second.eof);
}
#[tokio::test]
async fn data_racing_close_arrives_before_eof() {
let stream = Arc::new(ByteStream::new());
let reader = {
let stream = stream.clone();
tokio::spawn(async move { stream.read_at(0, 100).await })
};
tokio::task::yield_now().await;
stream.push(b"last words");
stream.close();
let chunk = reader.await.unwrap();
assert_eq!(chunk.bytes, b"last words");
assert!(chunk.eof);
}
#[tokio::test]
async fn stale_offset_clamps() {
let stream = ByteStream::new();
stream.push(b"abc");
stream.close();
let chunk = stream.read_at(999, 10).await;
assert!(chunk.bytes.is_empty());
assert_eq!(chunk.offset, 3);
assert!(chunk.eof);
}
#[test]
fn push_after_close_is_dropped() {
let stream = ByteStream::new();
assert!(stream.push(b"a"));
stream.close();
assert!(!stream.push(b"b"));
assert_eq!(stream.len(), 1);
}
#[test]
fn snapshot_is_nonblocking() {
let stream = ByteStream::new();
let chunk = stream.snapshot_at(0, 10);
assert!(chunk.bytes.is_empty());
assert!(!chunk.eof); }
#[tokio::test]
async fn cancellation_fails_parked_read() {
let stream = Arc::new(ByteStream::new());
let token = CancelToken::new();
let reader = {
let stream = stream.clone();
let token = token.clone();
tokio::spawn(async move { stream.read_at_cancellable(0, 10, &token).await })
};
tokio::task::yield_now().await;
token.cancel();
assert!(reader.await.unwrap().is_err());
}
#[test]
fn value_encoding_is_bytes() {
let chunk = ByteChunk {
bytes: vec![1, 2, 3],
offset: 0,
eof: false,
};
assert_eq!(chunk.into_value(), Value::Bytes(vec![1, 2, 3]));
}
}