nvrl/proton-stream
1
0 Fork 0
mirror of https://github.com/narrrl/proton-stream.git synced 2026-11-03 16:41:26 +01:00
Code Issues Packages Projects Releases Wiki Activity

init

Browse Source
This commit is contained in:
nvrl committed 2026-08-02 19:37:03 +02:00
commit a1b345a029
80 files changed
+28033
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/pstr-stream/src/block.rs
+281
View File
@@ -0,0 +1,281 @@
//! What the stream layer needs from whatever is holding the bytes.
//!
//! A Proton revision is a list of independently-encrypted content blocks, and
//! every layer above this one — the memory ring, the disk cache, the read-ahead
//! — thinks purely in whole blocks. So the seam is a *block* source, not a byte
//! source: [`BlockSource`] hands over one decrypted block at a time, and
//! [`BlockMap`] turns a byte range into the blocks that overlap it.
//!
//! The trait exists for a second reason, and it is the one that pays: with it,
//! everything in this crate is testable against a deterministic fake, with no
//! account, no network and no share. Only [`crate::reader`] touches the SDK.
use std::sync::Arc;
use async_trait::async_trait;
use crate::error::Result;
/// A revision's decrypted blocks, addressed by index.
///
/// Implementors must be cheap to keep alive — one is held per open stream.
#[async_trait]
pub trait BlockSource: Send + Sync + 'static {
/// The revision these blocks belong to.
///
/// This is the cache validity key. Unlike a modification time it advances
/// **iff** a new revision was sealed, so bytes cached under one revision id
/// can never be served for another one's content.
fn revision_id(&self) -> &str;
/// Plaintext size of each block, in block order.
///
/// Never assume 4 MiB. Proton's block size is a property of the *uploader*,
/// the last block is short, and a padded size vector would silently serve
/// bytes from the wrong offset.
fn block_sizes(&self) -> &[u64];
/// Fetch and decrypt one block.
async fn read_block(&self, index: usize) -> Result<Vec<u8>>;
}
/// Where each block starts in the plaintext, so a byte range can be resolved to
/// blocks without a scan.
#[derive(Debug, Clone)]
pub struct BlockMap {
/// Plaintext offset of each block. `starts[i] + sizes[i] == starts[i + 1]`.
starts: Vec<u64>,
sizes: Vec<u64>,
size: u64,
}
impl BlockMap {
pub fn new(sizes: &[u64]) -> Self {
let mut starts = Vec::with_capacity(sizes.len());
let mut offset = 0_u64;
for &size in sizes {
starts.push(offset);
offset = offset.saturating_add(size);
}
Self {
starts,
sizes: sizes.to_vec(),
size: offset,
}
}
/// Total plaintext size of the revision.
pub fn size(&self) -> u64 {
self.size
}
/// How many blocks the revision has.
pub fn len(&self) -> usize {
self.sizes.len()
}
pub fn is_empty(&self) -> bool {
self.sizes.is_empty()
}
/// Plaintext offset where block `index` starts.
pub fn start_of(&self, index: usize) -> Option<u64> {
self.starts.get(index).copied()
}
/// Plaintext size of block `index`.
pub fn size_of(&self, index: usize) -> Option<u64> {
self.sizes.get(index).copied()
}
/// The block holding byte `offset`, or `None` at or past EOF.
///
/// Binary search rather than a divide: block sizes are uniform in practice
/// but nothing guarantees it, and a wrong division here is the exact bug the
/// SDK warns about — bytes served from the wrong offset, silently.
pub fn block_at(&self, offset: u64) -> Option<usize> {
if offset >= self.size {
return None;
}
match self.starts.binary_search(&offset) {
Ok(index) => Some(index),
// `starts` is sorted and starts at 0, so a miss is never at 0.
Err(index) => Some(index - 1),
}
}
/// Every block index overlapping `[offset, end)`, clamped to the revision.
pub fn blocks_in(&self, offset: u64, end: u64) -> std::ops::Range<usize> {
let end = end.min(self.size);
if offset >= end {
return 0..0;
}
let first = self.block_at(offset).unwrap_or(0);
// `end` is exclusive, so the last byte read is `end - 1`.
let last = self.block_at(end - 1).unwrap_or(first);
first..last + 1
}
}
/// Copy the part of `block` that falls inside `[offset, end)` into `out`.
///
/// `block_start` is where this block begins in the plaintext. Returns how many
/// bytes were written.
pub(crate) fn splice(
out: &mut [u8],
out_base: u64,
block: &[u8],
block_start: u64,
offset: u64,
end: u64,
) -> usize {
let block_end = block_start.saturating_add(block.len() as u64);
let from = offset.max(block_start);
let to = end.min(block_end);
if from >= to {
return 0;
}
let src = (from - block_start) as usize..(to - block_start) as usize;
let dst_start = (from - out_base) as usize;
let len = src.end - src.start;
out[dst_start..dst_start + len].copy_from_slice(&block[src]);
len
}
/// A `BlockSource` over bytes already in memory. Used by the crate's own tests
/// and by the benchmark's warm-up path.
#[derive(Debug)]
pub struct MemoryBlocks {
revision: String,
blocks: Vec<Vec<u8>>,
sizes: Vec<u64>,
}
impl MemoryBlocks {
pub fn new(revision: impl Into<String>, blocks: Vec<Vec<u8>>) -> Self {
let sizes = blocks.iter().map(|block| block.len() as u64).collect();
Self {
revision: revision.into(),
blocks,
sizes,
}
}
}
#[async_trait]
impl BlockSource for MemoryBlocks {
fn revision_id(&self) -> &str {
&self.revision
}
fn block_sizes(&self) -> &[u64] {
&self.sizes
}
async fn read_block(&self, index: usize) -> Result<Vec<u8>> {
self.blocks
.get(index)
.cloned()
.ok_or_else(|| crate::Error::NotFound(format!("block {index} is past the end")))
}
}
/// Boxed for storage in the stream.
pub type SharedBlocks = Arc<dyn BlockSource>;
#[cfg(test)]
mod tests {
use super::*;
/// The realistic shape: uniform blocks with a short tail.
fn map() -> BlockMap {
BlockMap::new(&[4096, 4096, 1000])
}
#[test]
fn a_block_map_sums_to_the_revision_size() {
assert_eq!(map().size(), 9192);
assert_eq!(map().len(), 3);
}
#[test]
fn an_offset_resolves_to_the_block_holding_it() {
let map = map();
assert_eq!(map.block_at(0), Some(0));
assert_eq!(map.block_at(4095), Some(0));
assert_eq!(map.block_at(4096), Some(1));
assert_eq!(map.block_at(8191), Some(1));
assert_eq!(map.block_at(8192), Some(2));
assert_eq!(map.block_at(9191), Some(2));
}
#[test]
fn an_offset_at_or_past_the_end_resolves_to_nothing() {
assert_eq!(map().block_at(9192), None);
assert_eq!(map().block_at(u64::MAX), None);
}
#[test]
fn a_range_covers_every_block_it_touches() {
let map = map();
assert_eq!(map.blocks_in(0, 1), 0..1);
assert_eq!(map.blocks_in(4095, 4097), 0..2);
assert_eq!(map.blocks_in(0, 9192), 0..3);
assert_eq!(map.blocks_in(8192, 9192), 2..3);
}
/// A read ending exactly on a boundary must not pull the next block — that
/// would be a 4 MiB fetch for zero bytes.
#[test]
fn a_range_ending_on_a_boundary_stops_short_of_the_next_block() {
assert_eq!(map().blocks_in(0, 4096), 0..1);
assert_eq!(map().blocks_in(4096, 8192), 1..2);
}
#[test]
fn an_empty_or_past_the_end_range_covers_nothing() {
let map = map();
assert!(map.blocks_in(100, 100).is_empty());
assert!(map.blocks_in(9192, 10000).is_empty());
assert!(map.blocks_in(500, 100).is_empty());
}
/// Non-uniform block sizes are the case a divide-by-block-size would get
/// wrong, so they are pinned explicitly.
#[test]
fn non_uniform_blocks_still_resolve_correctly() {
let map = BlockMap::new(&[10, 1000, 5]);
assert_eq!(map.block_at(9), Some(0));
assert_eq!(map.block_at(10), Some(1));
assert_eq!(map.block_at(1009), Some(1));
assert_eq!(map.block_at(1010), Some(2));
assert_eq!(map.size(), 1015);
}
#[test]
fn a_spliced_block_lands_at_the_right_place_in_the_output() {
// Reading bytes 4090..4106 of the map above: 6 bytes from block 0, 10
// from block 1.
let (offset, end) = (4090_u64, 4106_u64);
let mut out = vec![0_u8; (end - offset) as usize];
let block0 = vec![b'a'; 4096];
let block1 = vec![b'b'; 4096];
assert_eq!(splice(&mut out, offset, &block0, 0, offset, end), 6);
assert_eq!(splice(&mut out, offset, &block1, 4096, offset, end), 10);
assert_eq!(&out[..6], b"aaaaaa");
assert_eq!(&out[6..], b"bbbbbbbbbb");
}
/// A block wholly outside the range writes nothing rather than corrupting
/// the buffer.
#[test]
fn a_block_outside_the_range_writes_nothing() {
let mut out = vec![0_u8; 8];
let block = vec![b'x'; 4096];
assert_eq!(splice(&mut out, 0, &block, 8192, 0, 8), 0);
assert_eq!(out, vec![0_u8; 8]);
}
}
crates/pstr-stream/src/disk.rs
+571
View File
@@ -0,0 +1,571 @@
//! The on-disk block cache: re-watching an episode should not re-download it.
//!
//! ## Why a sidecar, and why it is written last
//!
//! An entry is two files — `b{index}` holding the decrypted block, and
//! `b{index}.meta` holding its length. The block is written and synced **first**,
//! the sidecar **last**. That ordering is the whole validation scheme: a crash
//! or a full disk mid-write leaves a block with no sidecar, or a sidecar whose
//! length disagrees with the file, and both read as a miss. There is no state in
//! which a truncated block is served as if it were whole — which for a video
//! stream would not error, it would just decode garbage.
//!
//! Lifted from `proton-drive-linux`'s `pdfs-core/src/cache.rs`, which learned
//! the ordering the hard way.
//!
//! ## Why the path is hashed
//!
//! `{sha256(uid + revision)}` keyed by revision, not by name or mtime. A file
//! that gains a new revision lands in an entirely different directory, so stale
//! bytes are unreachable rather than merely unlikely. Names never touch the
//! filesystem: they are the one part of a share that is genuinely private, and
//! writing them into cache paths would leak the library's contents to anyone who
//! can list a directory.
use std::path::{Path, PathBuf};
use std::sync::{Mutex, PoisonError};
use lru::LruCache;
use sha2::{Digest, Sha256};
use crate::error::Result;
use crate::ring::BlockKey;
/// How much disk the cache may use, and where.
#[derive(Debug, Clone)]
pub struct DiskCacheConfig {
pub root: PathBuf,
pub budget_bytes: u64,
}
impl DiskCacheConfig {
/// 4 GiB — roughly two 1080p episodes' worth of blocks, small enough to be
/// unremarkable on any machine that plays video.
pub const DEFAULT_BUDGET_BYTES: u64 = 4 * 1024 * 1024 * 1024;
pub fn new(root: impl Into<PathBuf>) -> Self {
Self {
root: root.into(),
budget_bytes: Self::DEFAULT_BUDGET_BYTES,
}
}
pub fn with_budget(mut self, bytes: u64) -> Self {
self.budget_bytes = bytes;
self
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct DiskStats {
pub hits: u64,
pub misses: u64,
/// Entries dropped to stay inside the budget.
pub evictions: u64,
/// Entries found unreadable or inconsistent, and therefore ignored.
pub rejected: u64,
pub stored_bytes: u64,
pub entries: usize,
}
pub(crate) struct DiskCache {
root: PathBuf,
budget: u64,
index: Mutex<Index>,
}
struct Index {
/// Keyed by the block file's path, because that is all a startup scan can
/// recover — the key it was written under is hashed away by design.
lru: LruCache<PathBuf, u64>,
bytes: u64,
stats: DiskStats,
}
impl DiskCache {
/// Open the cache, adopting whatever a previous run left behind.
///
/// The scan is ordered by modification time so the recovered LRU is at least
/// approximately right; within a session, real access order takes over.
pub(crate) async fn open(config: DiskCacheConfig) -> Result<Self> {
let root = config.root.clone();
let scanned = tokio::task::spawn_blocking(move || scan(&root)).await;
let entries = match scanned {
Ok(Ok(entries)) => entries,
// A cache that cannot be scanned is a cache that starts empty. It is
// disposable by definition; refusing to play a film over it would be
// absurd.
Ok(Err(e)) => {
tracing::warn!(error = %e, "block cache could not be scanned; starting empty");
Vec::new()
}
Err(e) => {
tracing::warn!(error = %e, "block cache scan failed to run; starting empty");
Vec::new()
}
};
let mut lru = LruCache::unbounded();
let mut bytes = 0_u64;
for (path, len) in entries {
bytes += len;
lru.put(path, len);
}
let cache = Self {
root: config.root,
budget: config.budget_bytes,
index: Mutex::new(Index {
lru,
bytes,
stats: DiskStats {
stored_bytes: bytes,
..DiskStats::default()
},
}),
};
cache.evict_to_budget().await;
Ok(cache)
}
/// Read a block back, or `None` when it is absent or fails validation.
pub(crate) async fn get(&self, key: &BlockKey) -> Option<Vec<u8>> {
let path = self.block_path(key);
match read_entry(&path).await {
Ok(Some(block)) => {
let mut index = self.lock();
index.stats.hits += 1;
// Promotes it, and adopts an entry a concurrent writer added.
let len = block.len() as u64;
if index.lru.get(&path).is_none() {
index.bytes += len;
index.lru.put(path, len);
}
Some(block)
}
Ok(None) => {
self.lock().stats.misses += 1;
None
}
Err(e) => {
tracing::debug!(path = %path.display(), error = %e, "cached block rejected");
// Scoped, not `drop`ped: a guard merely *moved* before an await
// is still captured by the generated future, which would make
// every read non-`Send` and unspawnable.
{
let mut index = self.lock();
index.stats.misses += 1;
index.stats.rejected += 1;
}
// A half-written entry is never going to become valid.
self.discard(&path).await;
None
}
}
}
/// Store a block. Failures are logged and swallowed: a cache that cannot be
/// written is a slower app, not a broken one.
pub(crate) async fn put(&self, key: &BlockKey, block: &[u8]) {
let len = block.len() as u64;
if len > self.budget {
return;
}
let path = self.block_path(key);
if let Err(e) = write_entry(&path, block).await {
tracing::warn!(path = %path.display(), error = %e, "could not cache block to disk");
self.discard(&path).await;
return;
}
{
let mut index = self.lock();
if let Some(previous) = index.lru.put(path, len) {
index.bytes -= previous;
}
index.bytes += len;
}
self.evict_to_budget().await;
}
pub(crate) fn stats(&self) -> DiskStats {
let index = self.lock();
DiskStats {
stored_bytes: index.bytes,
entries: index.lru.len(),
..index.stats
}
}
/// Drop entries, least-recently-used first, until inside the budget.
async fn evict_to_budget(&self) {
loop {
let doomed = {
let mut index = self.lock();
if index.bytes <= self.budget {
return;
}
match index.lru.pop_lru() {
Some((path, len)) => {
index.bytes -= len;
index.stats.evictions += 1;
path
}
None => {
index.bytes = 0;
return;
}
}
};
self.discard(&doomed).await;
}
}
/// Remove both halves of an entry. The sidecar goes first, so an
/// interrupted deletion leaves an entry that fails validation rather than
/// one that claims to be valid.
async fn discard(&self, path: &Path) {
let _ = tokio::fs::remove_file(meta_path(path)).await;
let _ = tokio::fs::remove_file(path).await;
}
/// `{root}/{aa}/{hash}/b{index}` — two levels of fan-out so a large library
/// does not put tens of thousands of entries in one directory.
fn block_path(&self, key: &BlockKey) -> PathBuf {
let mut hasher = Sha256::new();
hasher.update(key.uid.to_string().as_bytes());
hasher.update(b"\0");
hasher.update(key.revision.as_bytes());
let digest = hasher.finalize();
let hash = hex(&digest);
self.root
.join(&hash[..2])
.join(&hash)
.join(format!("b{}", key.index))
}
fn lock(&self) -> std::sync::MutexGuard<'_, Index> {
self.index.lock().unwrap_or_else(PoisonError::into_inner)
}
}
fn meta_path(block: &Path) -> PathBuf {
let mut name = block.as_os_str().to_os_string();
name.push(".meta");
PathBuf::from(name)
}
fn hex(bytes: &[u8]) -> String {
use std::fmt::Write;
bytes.iter().fold(String::new(), |mut out, byte| {
let _ = write!(out, "{byte:02x}");
out
})
}
/// Read an entry, validating it against its sidecar.
///
/// `Ok(None)` is "not cached". `Err` is "cached, but wrong" — the caller deletes
/// those.
async fn read_entry(path: &Path) -> Result<Option<Vec<u8>>> {
let claimed = match tokio::fs::read_to_string(meta_path(path)).await {
Ok(text) => text,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
Err(e) => return Err(e.into()),
};
let claimed_len: u64 = claimed.trim().parse().map_err(|_| {
crate::Error::NotFound(format!(
"{} holds no usable length",
meta_path(path).display()
))
})?;
let block = match tokio::fs::read(path).await {
Ok(block) => block,
// Sidecar without a block: the block was evicted or half-deleted.
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
return Err(crate::Error::NotFound("cached block is missing".into()));
}
Err(e) => return Err(e.into()),
};
if block.len() as u64 != claimed_len {
return Err(crate::Error::NotFound(format!(
"cached block is {} bytes, sidecar claims {claimed_len}",
block.len()
)));
}
Ok(Some(block))
}
/// Write an entry: block, synced, then sidecar. See the module note.
async fn write_entry(path: &Path, block: &[u8]) -> Result<()> {
use tokio::io::AsyncWriteExt;
let parent = path
.parent()
.ok_or_else(|| crate::Error::NotFound(format!("{} has no parent", path.display())))?;
tokio::fs::create_dir_all(parent).await?;
// Through a temp file so a concurrent reader never sees a partial block even
// momentarily.
let temp = meta_path(path).with_extension("part");
{
let mut file = tokio::fs::File::create(&temp).await?;
file.write_all(block).await?;
file.sync_all().await?;
}
tokio::fs::rename(&temp, path).await?;
let mut sidecar = tokio::fs::File::create(meta_path(path)).await?;
sidecar
.write_all(block.len().to_string().as_bytes())
.await?;
sidecar.sync_all().await?;
Ok(())
}
/// Walk the cache tree, returning `(block path, size)` oldest-first.
///
/// Blocking, so it runs on the blocking pool. Anything that is not a valid pair
/// is skipped rather than deleted — a cache directory the user pointed at by
/// mistake should not be eaten.
fn scan(root: &Path) -> std::io::Result<Vec<(PathBuf, u64)>> {
let mut found: Vec<(PathBuf, u64, std::time::SystemTime)> = Vec::new();
let shards = match std::fs::read_dir(root) {
Ok(shards) => shards,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
Err(e) => return Err(e),
};
for shard in shards.flatten() {
let Ok(revisions) = std::fs::read_dir(shard.path()) else {
continue;
};
for revision in revisions.flatten() {
let Ok(blocks) = std::fs::read_dir(revision.path()) else {
continue;
};
for block in blocks.flatten() {
let path = block.path();
// Sidecars and interrupted writes are accounted through their
// block, not on their own.
if path.extension().is_some() {
continue;
}
let Ok(metadata) = block.metadata() else {
continue;
};
if !metadata.is_file() || !meta_path(&path).exists() {
continue;
}
let modified = metadata.modified().unwrap_or(std::time::UNIX_EPOCH);
found.push((path, metadata.len(), modified));
}
}
}
found.sort_by_key(|(_, _, modified)| *modified);
Ok(found
.into_iter()
.map(|(path, len, _)| (path, len))
.collect())
}
#[cfg(test)]
mod tests {
use super::*;
use proton_sdk::ids::{LinkId, NodeUid, VolumeId};
fn temp_root(tag: &str) -> PathBuf {
let dir = std::env::temp_dir().join(format!(
"pstr-disk-{tag}-{}-{:?}",
std::process::id(),
std::thread::current().id()
));
let _ = std::fs::remove_dir_all(&dir);
dir
}
fn key(link: &str, revision: &str, index: usize) -> BlockKey {
BlockKey::new(
&NodeUid::new(VolumeId::new("vol"), LinkId::new(link)),
revision,
index,
)
}
#[tokio::test]
async fn a_stored_block_reads_back() {
let root = temp_root("roundtrip");
let cache = DiskCache::open(DiskCacheConfig::new(&root)).await.unwrap();
let block = vec![9_u8; 4096];
cache.put(&key("a", "rev1", 3), &block).await;
assert_eq!(cache.get(&key("a", "rev1", 3)).await, Some(block));
assert_eq!(cache.stats().hits, 1);
std::fs::remove_dir_all(&root).ok();
}
#[tokio::test]
async fn an_absent_block_is_a_miss_and_not_an_error() {
let root = temp_root("absent");
let cache = DiskCache::open(DiskCacheConfig::new(&root)).await.unwrap();
assert!(cache.get(&key("a", "rev1", 0)).await.is_none());
assert_eq!(cache.stats().misses, 1);
assert_eq!(cache.stats().rejected, 0);
std::fs::remove_dir_all(&root).ok();
}
/// The reason the sidecar is written last. A truncated block must never be
/// served — for video it would not fail, it would decode as corruption.
#[tokio::test]
async fn a_block_that_disagrees_with_its_sidecar_is_rejected_and_removed() {
let root = temp_root("torn");
let cache = DiskCache::open(DiskCacheConfig::new(&root)).await.unwrap();
let k = key("a", "rev1", 0);
cache.put(&k, &[1_u8; 100]).await;
// Simulate the torn write the ordering is designed to catch.
let path = cache.block_path(&k);
std::fs::write(&path, vec![1_u8; 40]).unwrap();
assert!(cache.get(&k).await.is_none());
assert_eq!(cache.stats().rejected, 1);
assert!(!path.exists(), "invalid entry should be cleaned up");
std::fs::remove_dir_all(&root).ok();
}
/// A block with no sidecar at all — the crash-between-writes case.
#[tokio::test]
async fn a_block_without_a_sidecar_is_a_miss() {
let root = temp_root("nosidecar");
let cache = DiskCache::open(DiskCacheConfig::new(&root)).await.unwrap();
let k = key("a", "rev1", 0);
cache.put(&k, &[1_u8; 100]).await;
std::fs::remove_file(meta_path(&cache.block_path(&k))).unwrap();
assert!(cache.get(&k).await.is_none());
std::fs::remove_dir_all(&root).ok();
}
#[tokio::test]
async fn the_cache_stays_inside_its_budget() {
let root = temp_root("budget");
let cache = DiskCache::open(DiskCacheConfig::new(&root).with_budget(300))
.await
.unwrap();
for index in 0..10 {
cache.put(&key("a", "rev1", index), &[1_u8; 100]).await;
assert!(cache.stats().stored_bytes <= 300, "over budget at {index}");
}
assert_eq!(cache.stats().entries, 3);
assert!(cache.stats().evictions > 0);
std::fs::remove_dir_all(&root).ok();
}
/// Eviction must delete the files, not merely forget them, or the budget is
/// a fiction and the disk fills up anyway.
#[tokio::test]
async fn eviction_removes_the_files_from_disk() {
let root = temp_root("evictfiles");
let cache = DiskCache::open(DiskCacheConfig::new(&root).with_budget(150))
.await
.unwrap();
let first = key("a", "rev1", 0);
cache.put(&first, &[1_u8; 100]).await;
cache.put(&key("a", "rev1", 1), &[2_u8; 100]).await;
let path = cache.block_path(&first);
assert!(!path.exists(), "evicted block still on disk");
assert!(!meta_path(&path).exists(), "evicted sidecar still on disk");
std::fs::remove_dir_all(&root).ok();
}
/// The point of the cache: a second run reuses the first run's bytes.
#[tokio::test]
async fn a_reopened_cache_adopts_what_the_previous_run_wrote() {
let root = temp_root("reopen");
{
let cache = DiskCache::open(DiskCacheConfig::new(&root)).await.unwrap();
cache.put(&key("a", "rev1", 0), &[7_u8; 256]).await;
}
let cache = DiskCache::open(DiskCacheConfig::new(&root)).await.unwrap();
assert_eq!(cache.stats().entries, 1);
assert_eq!(cache.stats().stored_bytes, 256);
assert_eq!(cache.get(&key("a", "rev1", 0)).await, Some(vec![7_u8; 256]));
std::fs::remove_dir_all(&root).ok();
}
/// Reopening with a smaller budget must shrink the cache, not sit over it
/// until enough new blocks arrive to trigger eviction.
#[tokio::test]
async fn reopening_with_a_smaller_budget_evicts_immediately() {
let root = temp_root("shrink");
{
let cache = DiskCache::open(DiskCacheConfig::new(&root)).await.unwrap();
for index in 0..5 {
cache.put(&key("a", "rev1", index), &[1_u8; 100]).await;
}
}
let cache = DiskCache::open(DiskCacheConfig::new(&root).with_budget(250))
.await
.unwrap();
assert!(cache.stats().stored_bytes <= 250);
std::fs::remove_dir_all(&root).ok();
}
/// Two revisions of the same file are different entries, so stale bytes are
/// unreachable rather than merely unlikely.
#[tokio::test]
async fn a_new_revision_does_not_read_the_old_ones_blocks() {
let root = temp_root("revision");
let cache = DiskCache::open(DiskCacheConfig::new(&root)).await.unwrap();
cache.put(&key("a", "rev1", 0), &[1_u8; 64]).await;
assert!(cache.get(&key("a", "rev2", 0)).await.is_none());
std::fs::remove_dir_all(&root).ok();
}
/// Nothing about the library should be legible from the cache directory.
#[tokio::test]
async fn cache_paths_carry_no_names() {
let root = temp_root("opaque");
let cache = DiskCache::open(DiskCacheConfig::new(&root)).await.unwrap();
let path = cache.block_path(&key("Frieren-S01E01", "rev1", 0));
let rendered = path.display().to_string();
assert!(
!rendered.contains("Frieren"),
"path leaks content: {rendered}"
);
std::fs::remove_dir_all(&root).ok();
}
/// A cache root that does not exist yet is the first-run case.
#[tokio::test]
async fn a_missing_cache_root_opens_empty() {
let root = temp_root("firstrun");
let cache = DiskCache::open(DiskCacheConfig::new(&root)).await.unwrap();
assert_eq!(cache.stats().entries, 0);
}
}
crates/pstr-stream/src/error.rs
+32
View File
@@ -0,0 +1,32 @@
//! One error type for the stream layer.
/// What can go wrong opening a stream or reading bytes out of it.
#[derive(Debug, thiserror::Error)]
pub enum Error {
/// A share could not be opened, or its catalog could not be read.
#[error(transparent)]
Core(#[from] pstr_core::Error),
/// The Drive SDK failed — network, API refusal, or decrypt.
#[error(transparent)]
Drive(#[from] proton_sdk::error::ProtonError),
/// The block cache's filesystem failed.
#[error("block cache: {0}")]
Io(#[from] std::io::Error),
/// The caller asked for something that does not exist.
#[error("{0}")]
NotFound(String),
/// A failure that happened in *another* task.
///
/// Block fetches are deduplicated: when several readers want the same block
/// only one of them fetches it, and the rest wait. The waiters cannot be
/// handed the leader's typed error — errors are not `Clone` — so they get
/// its message. The leader itself always sees the real error.
#[error("{0}")]
Shared(String),
}
pub type Result<T> = std::result::Result<T, Error>;
crates/pstr-stream/src/lib.rs
+45
View File
@@ -0,0 +1,45 @@
//! Seekable, cached byte streams over Proton Drive revisions.
//!
//! Proton stores a file as a list of independently-encrypted 4 MiB blocks, each
//! a self-contained PGP packet under one session key. Nothing about that is
//! sequential — which is the fact this whole crate rests on. Seeking is not an
//! approximation of streaming here; block 400 costs exactly what block 0 costs.
//!
//! What sits between the SDK and the player:
//!
//! ```text
//! mpv demuxer VideoStream BlockRing → DiskCache → Proton
//! (small, forward, → byte range → → memory disk network
//! seeking reads) blocks (128 MiB) (4 GiB)
//! ```
//!
//! * [`StreamSource`] opens files and keeps the open ones in an LRU.
//! * [`VideoStream`] answers byte ranges, reading ahead of the player.
//! * The ring and the disk cache are byte-budgeted, shared across streams, and
//! keyed by revision id so a resealed file can never serve stale bytes.
//!
//! Only [`reader`] touches the SDK. Everything else runs against a
//! [`BlockSource`], which is what makes the caching and read-ahead testable
//! without an account, a network or a share.
#![forbid(unsafe_code)]
pub mod block;
pub mod disk;
pub mod error;
pub mod reader;
pub mod ring;
mod single_flight;
pub mod source;
pub mod stream;
pub use block::{BlockMap, BlockSource, MemoryBlocks, SharedBlocks};
pub use disk::{DiskCacheConfig, DiskStats};
pub use error::{Error, Result};
/// Re-exported because every caller of this crate names one — `VideoStream::uid`
/// returns it and `StreamSource::open` takes it — and taking a dependency on the
/// whole SDK for a newtype is not a reasonable price for that.
pub use proton_sdk::ids::NodeUid;
pub use reader::{LibraryOpener, RevisionBlocks};
pub use ring::{DEFAULT_RING_BYTES, RingStats};
pub use source::{RevisionOpener, StreamConfig, StreamSource};
pub use stream::{DEFAULT_READAHEAD_BLOCKS, StreamStats, VideoStream};
crates/pstr-stream/src/reader.rs
+107
View File
@@ -0,0 +1,107 @@
//! The bridge to the SDK: a real Proton revision as a [`BlockSource`].
//!
//! Everything else in this crate is toolkit-free and account-free. This is the
//! one module that knows a `RevisionReader` exists.
//!
//! Reads are issued **block-aligned**. `RevisionReader::read_at` fetches every
//! block a range overlaps, so an unaligned read would pull two blocks to serve
//! one, and the caching layer above would then hold neither of them whole.
//! Aligning here means each block is fetched, decrypted and cached exactly once.
use std::sync::Arc;
use async_trait::async_trait;
use proton_drive_rs::ProtonDrivePublicLinkClient;
use proton_sdk::ids::NodeUid;
use crate::block::{BlockMap, BlockSource, SharedBlocks};
use crate::error::{Error, Result};
use crate::source::RevisionOpener;
/// One open revision on a public link.
pub struct RevisionBlocks {
reader: proton_drive_rs::RevisionReader,
map: BlockMap,
}
impl RevisionBlocks {
pub fn new(reader: proton_drive_rs::RevisionReader) -> Self {
let map = BlockMap::new(reader.block_sizes());
Self { reader, map }
}
}
#[async_trait]
impl BlockSource for RevisionBlocks {
fn revision_id(&self) -> &str {
self.reader.revision_id()
}
fn block_sizes(&self) -> &[u64] {
self.reader.block_sizes()
}
async fn read_block(&self, index: usize) -> Result<Vec<u8>> {
let (Some(start), Some(size)) = (self.map.start_of(index), self.map.size_of(index)) else {
return Err(Error::NotFound(format!(
"block {index} is past the end of revision {}",
self.reader.revision_id()
)));
};
let block = self.reader.read_at(start, size).await?;
// A short block here is not a tail — the map says how long it is. It
// means the revision changed underneath us or the block table lied, and
// serving it would silently shift every later byte.
if block.len() as u64 != size {
return Err(Error::NotFound(format!(
"block {index} came back {} bytes, expected {size}",
block.len()
)));
}
Ok(block)
}
}
/// Opens revisions out of the configured shares.
pub struct LibraryOpener {
library: Arc<pstr_core::SharedLibrary>,
}
impl LibraryOpener {
pub fn new(library: Arc<pstr_core::SharedLibrary>) -> Self {
Self { library }
}
fn client(&self, share_id: &str) -> Result<&ProtonDrivePublicLinkClient> {
self.library
.client(share_id)
.ok_or_else(|| Error::NotFound(format!("share {share_id} is not open")))
}
}
#[async_trait]
impl RevisionOpener for LibraryOpener {
async fn open(&self, share_id: &str, uid: &NodeUid) -> Result<SharedBlocks> {
let reader = self.client(share_id)?.open_revision(uid).await?;
Ok(Arc::new(RevisionBlocks::new(reader)))
}
}
#[cfg(test)]
mod tests {
use super::*;
/// The player hands its stream to mpv's demuxer thread, so everything from
/// the opener down has to cross threads. A compile-time guard, because the
/// failure mode is an unhelpful `FnOnce is not general enough` error deep in
/// a caller rather than here.
#[test]
fn the_library_opener_is_send_and_sync() {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<LibraryOpener>();
assert_send_sync::<RevisionBlocks>();
assert_send_sync::<crate::StreamSource>();
assert_send_sync::<crate::VideoStream>();
}
}
crates/pstr-stream/src/ring.rs
+338
View File
@@ -0,0 +1,338 @@
//! The in-memory block cache: a byte-budgeted LRU shared by every open stream.
//!
//! Sized in **bytes**, not entries. A block is 4 MiB in practice but nothing
//! guarantees it, and a count-based bound would turn into a wildly different
//! memory ceiling on a library uploaded by a different client. The budget is the
//! app's actual resident-memory promise, so it is what gets counted.
//!
//! Shared across streams rather than one per stream: a viewer who abandons an
//! episode after ninety seconds should not keep 128 MiB of it resident while the
//! next one starves.
use std::num::NonZeroUsize;
use std::sync::{Arc, Mutex, PoisonError};
use lru::LruCache;
use proton_sdk::ids::NodeUid;
/// Identifies one decrypted block, globally.
///
/// The revision id is part of the key on purpose: a file that gains a new
/// revision must not serve a byte of the old one out of cache. It is a stable
/// identity that advances iff a new revision was sealed.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(crate) struct BlockKey {
pub uid: NodeUid,
pub revision: String,
pub index: usize,
}
impl BlockKey {
pub(crate) fn new(uid: &NodeUid, revision: &str, index: usize) -> Self {
Self {
uid: uid.clone(),
revision: revision.to_string(),
index,
}
}
}
/// A decrypted block, shared by reference so a cache hit costs no copy.
pub(crate) type CachedBlock = Arc<[u8]>;
/// Hit/miss counters, for the benchmark and the settings screen.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct RingStats {
pub hits: u64,
pub misses: u64,
/// Blocks dropped to stay inside the budget.
pub evictions: u64,
/// Bytes currently resident.
pub resident_bytes: u64,
pub blocks: usize,
}
pub(crate) struct BlockRing {
inner: Mutex<Inner>,
budget: u64,
}
struct Inner {
/// Unbounded by entry count — eviction is driven by `bytes` against the
/// budget, below.
lru: LruCache<BlockKey, CachedBlock>,
bytes: u64,
hits: u64,
misses: u64,
evictions: u64,
}
impl BlockRing {
pub(crate) fn new(budget_bytes: u64) -> Self {
Self {
inner: Mutex::new(Inner {
lru: LruCache::unbounded(),
bytes: 0,
hits: 0,
misses: 0,
evictions: 0,
}),
budget: budget_bytes,
}
}
/// Look a block up, marking it most-recently-used on a hit.
pub(crate) fn get(&self, key: &BlockKey) -> Option<CachedBlock> {
let mut inner = self.lock();
match inner.lru.get(key) {
Some(block) => {
let block = Arc::clone(block);
inner.hits += 1;
Some(block)
}
None => {
inner.misses += 1;
None
}
}
}
/// Insert a block, evicting until the budget is met.
///
/// A block larger than the whole budget is *not* stored — storing it would
/// evict everything else and then still be over. It is still returned to the
/// caller that fetched it; it simply is not cached.
pub(crate) fn insert(&self, key: BlockKey, block: CachedBlock) {
let len = block.len() as u64;
if len > self.budget {
return;
}
let mut inner = self.lock();
if let Some(previous) = inner.lru.put(key, block) {
inner.bytes -= previous.len() as u64;
}
inner.bytes += len;
while inner.bytes > self.budget {
match inner.lru.pop_lru() {
Some((_, evicted)) => {
inner.bytes -= evicted.len() as u64;
inner.evictions += 1;
}
// Unreachable while `bytes` and `lru` agree; not worth a panic.
None => {
inner.bytes = 0;
break;
}
}
}
}
/// Whether a block is resident, *without* counting a hit or touching its
/// recency. Read-ahead asks this to decide what to prefetch, and it must not
/// look like demand traffic in the stats or in the LRU order.
pub(crate) fn contains(&self, key: &BlockKey) -> bool {
self.lock().lru.contains(key)
}
/// The byte budget this ring was built with.
pub(crate) fn budget(&self) -> u64 {
self.budget
}
pub(crate) fn stats(&self) -> RingStats {
let inner = self.lock();
RingStats {
hits: inner.hits,
misses: inner.misses,
evictions: inner.evictions,
resident_bytes: inner.bytes,
blocks: inner.lru.len(),
}
}
/// Drop every block of one revision. Used when a stream is closed and its
/// bytes are known to be dead.
pub(crate) fn forget(&self, uid: &NodeUid, revision: &str) {
let mut inner = self.lock();
let doomed: Vec<BlockKey> = inner
.lru
.iter()
.filter(|(key, _)| &key.uid == uid && key.revision == revision)
.map(|(key, _)| key.clone())
.collect();
for key in doomed {
if let Some(block) = inner.lru.pop(&key) {
inner.bytes -= block.len() as u64;
}
}
}
fn lock(&self) -> std::sync::MutexGuard<'_, Inner> {
// Non-poisoning for the same reason the rest of the stack is: a panic in
// one read must not make every later read fail. The cache holds no
// invariant a panic could leave half-updated.
self.inner.lock().unwrap_or_else(PoisonError::into_inner)
}
}
/// The default in-memory budget: 128 MiB, matching the FUSE client's stream
/// ring. At 4 MiB a block that is 32 blocks — about two minutes of a 1080p
/// episode, enough to absorb a scrub backwards without refetching.
pub const DEFAULT_RING_BYTES: u64 = 128 * 1024 * 1024;
/// `lru` wants a non-zero capacity for bounded caches, and a configured zero
/// must clamp rather than panic.
pub(crate) fn non_zero(value: usize) -> NonZeroUsize {
NonZeroUsize::new(value.max(1)).expect("max(1) is non-zero")
}
#[cfg(test)]
mod tests {
use super::*;
use proton_sdk::ids::{LinkId, VolumeId};
fn uid(link: &str) -> NodeUid {
NodeUid::new(VolumeId::new("vol"), LinkId::new(link))
}
fn block(byte: u8, len: usize) -> CachedBlock {
vec![byte; len].into()
}
#[test]
fn a_stored_block_reads_back() {
let ring = BlockRing::new(1024);
let key = BlockKey::new(&uid("a"), "rev1", 0);
ring.insert(key.clone(), block(7, 16));
assert_eq!(ring.get(&key).expect("hit").as_ref(), &[7_u8; 16]);
assert_eq!(ring.stats().hits, 1);
}
/// The budget is the app's memory promise. Exceeding it is the one thing
/// this type must not do.
#[test]
fn the_ring_never_exceeds_its_byte_budget() {
let ring = BlockRing::new(100);
for index in 0..10 {
ring.insert(BlockKey::new(&uid("a"), "rev1", index), block(1, 30));
assert!(
ring.stats().resident_bytes <= 100,
"over budget at block {index}"
);
}
assert_eq!(ring.stats().blocks, 3);
assert!(ring.stats().evictions > 0);
}
/// Eviction order must be least-recently-*used*, not insertion order, or a
/// viewer scrubbing back and forth loses the block they keep returning to.
#[test]
fn the_least_recently_used_block_is_evicted_first() {
let ring = BlockRing::new(60);
let (a, b) = (
BlockKey::new(&uid("x"), "rev1", 0),
BlockKey::new(&uid("x"), "rev1", 1),
);
ring.insert(a.clone(), block(1, 30));
ring.insert(b.clone(), block(2, 30));
// Touch `a`, making `b` the eviction candidate.
assert!(ring.get(&a).is_some());
ring.insert(BlockKey::new(&uid("x"), "rev1", 2), block(3, 30));
assert!(ring.get(&a).is_some(), "recently used block was evicted");
assert!(ring.get(&b).is_none(), "stale block survived");
}
/// The reason the revision id is in the key: a resealed file must never
/// serve a byte of its previous content.
#[test]
fn a_new_revision_does_not_hit_the_old_revisions_blocks() {
let ring = BlockRing::new(1024);
ring.insert(BlockKey::new(&uid("a"), "rev1", 0), block(1, 8));
assert!(ring.get(&BlockKey::new(&uid("a"), "rev2", 0)).is_none());
}
/// And neither must two different files at the same block index.
#[test]
fn two_files_do_not_collide_at_the_same_block_index() {
let ring = BlockRing::new(1024);
ring.insert(BlockKey::new(&uid("a"), "rev1", 0), block(1, 8));
ring.insert(BlockKey::new(&uid("b"), "rev1", 0), block(2, 8));
assert_eq!(
ring.get(&BlockKey::new(&uid("a"), "rev1", 0)).unwrap()[0],
1
);
assert_eq!(
ring.get(&BlockKey::new(&uid("b"), "rev1", 0)).unwrap()[0],
2
);
}
/// Read-ahead polls this constantly; it must not distort demand statistics
/// or promote a block it only *considered* fetching.
#[test]
fn a_containment_check_records_no_traffic_and_no_recency() {
let ring = BlockRing::new(60);
let (a, b) = (
BlockKey::new(&uid("x"), "rev1", 0),
BlockKey::new(&uid("x"), "rev1", 1),
);
ring.insert(a.clone(), block(1, 30));
ring.insert(b.clone(), block(2, 30));
assert!(ring.contains(&a));
assert_eq!(ring.stats().hits, 0);
assert_eq!(ring.stats().misses, 0);
// `a` is still the LRU victim despite having been asked about.
ring.insert(BlockKey::new(&uid("x"), "rev1", 2), block(3, 30));
assert!(!ring.contains(&a));
}
/// A block bigger than the whole budget must not flush the cache to make
/// room for something that cannot fit anyway.
#[test]
fn a_block_larger_than_the_budget_is_not_stored() {
let ring = BlockRing::new(50);
let keeper = BlockKey::new(&uid("a"), "rev1", 0);
ring.insert(keeper.clone(), block(1, 40));
ring.insert(BlockKey::new(&uid("a"), "rev1", 1), block(2, 500));
assert!(ring.get(&keeper).is_some(), "existing block was flushed");
assert_eq!(ring.stats().blocks, 1);
}
#[test]
fn forgetting_a_revision_reclaims_its_bytes() {
let ring = BlockRing::new(1024);
for index in 0..4 {
ring.insert(BlockKey::new(&uid("a"), "rev1", index), block(1, 32));
}
ring.insert(BlockKey::new(&uid("b"), "rev1", 0), block(2, 32));
ring.forget(&uid("a"), "rev1");
let stats = ring.stats();
assert_eq!(stats.blocks, 1);
assert_eq!(stats.resident_bytes, 32);
}
/// Re-inserting the same key must not double-count its bytes; a refetched
/// block would otherwise inflate the accounting until the ring evicted
/// everything.
#[test]
fn reinserting_a_block_does_not_double_count_its_bytes() {
let ring = BlockRing::new(1024);
let key = BlockKey::new(&uid("a"), "rev1", 0);
ring.insert(key.clone(), block(1, 64));
ring.insert(key.clone(), block(1, 64));
assert_eq!(ring.stats().resident_bytes, 64);
assert_eq!(ring.stats().blocks, 1);
}
}
crates/pstr-stream/src/single_flight.rs
+263
View File
@@ -0,0 +1,263 @@
//! One in-flight call per key, however many callers ask for it.
//!
//! Two places need this and both are hot:
//!
//! * **Opening a stream.** A cold open costs a link-details fetch, an S2K
//! ancestor-key unlock and a revision listing. The UI's "play" click and the
//! player's first read arrive within milliseconds of each other; without
//! deduplication they pay for that twice and race to insert into the LRU.
//! * **Fetching a block.** Read-ahead runs *ahead* of the demand read by
//! definition, so the two collide constantly on the block the player is about
//! to want. Fetching it twice doubles the bytes off the wire for no gain.
//!
//! The entry is removed by a guard rather than by the leader's happy path, so a
//! cancelled call (a dropped future — the player abandoning a seek, say) cannot
//! leave a key permanently marked as in flight.
use std::collections::HashMap;
use std::future::Future;
use std::hash::Hash;
use std::sync::{Arc, Mutex, PoisonError};
use tokio::sync::watch;
use crate::error::{Error, Result};
/// The result as a waiter sees it: the value, or the leader's error message.
type Shared<V> = Option<std::result::Result<V, String>>;
/// Deduplicates concurrent calls that share a key.
pub(crate) struct SingleFlight<K, V> {
calls: Mutex<HashMap<K, watch::Receiver<Shared<V>>>>,
}
impl<K, V> SingleFlight<K, V>
where
K: Eq + Hash + Clone,
V: Clone,
{
pub(crate) fn new() -> Self {
Self {
calls: Mutex::new(HashMap::new()),
}
}
/// Run `work` for `key`, or wait for the call already running it.
///
/// Whoever arrives first is the leader and actually runs the future; the
/// rest wait on its result. The leader gets its own typed error, waiters get
/// [`Error::Shared`].
pub(crate) async fn run<F, Fut>(self: &Arc<Self>, key: K, work: F) -> Result<V>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<V>>,
{
let role = {
let mut calls = self.lock();
match calls.get(&key) {
Some(receiver) => Role::Waiter(receiver.clone()),
None => {
let (sender, receiver) = watch::channel(None);
calls.insert(key.clone(), receiver);
Role::Leader(sender)
}
}
};
match role {
Role::Leader(sender) => {
// Removes the map entry on *every* exit path, including the one
// where this future is dropped mid-await.
let _guard = LeaderGuard {
flight: Arc::clone(self),
key: Some(key),
};
let result = work().await;
let shared = match &result {
Ok(value) => Ok(value.clone()),
Err(e) => Err(e.to_string()),
};
// Waiters that already hold a receiver are woken here; ones that
// arrive after the guard drops start a fresh call.
let _ = sender.send(Some(shared));
result
}
Role::Waiter(mut receiver) => loop {
// Cloned out before the await: holding the borrow across it
// would deadlock the sender.
let seen = receiver.borrow_and_update().clone();
if let Some(result) = seen {
return result.map_err(Error::Shared);
}
if receiver.changed().await.is_err() {
return Err(Error::Shared(
"the task loading this was cancelled before it finished".into(),
));
}
},
}
}
fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<K, watch::Receiver<Shared<V>>>> {
// A panicking caller must not wedge every future call on this key. The
// map is a plain index with no invariant a panic could break.
self.calls.lock().unwrap_or_else(PoisonError::into_inner)
}
}
enum Role<V> {
Leader(watch::Sender<Shared<V>>),
Waiter(watch::Receiver<Shared<V>>),
}
struct LeaderGuard<K: Eq + Hash + Clone, V: Clone> {
flight: Arc<SingleFlight<K, V>>,
key: Option<K>,
}
impl<K: Eq + Hash + Clone, V: Clone> Drop for LeaderGuard<K, V> {
fn drop(&mut self) {
if let Some(key) = self.key.take() {
self.flight.lock().remove(&key);
}
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use super::*;
/// The point of the whole file: N callers, one execution.
#[tokio::test]
async fn concurrent_callers_of_one_key_run_the_work_once() {
let flight: Arc<SingleFlight<u32, u32>> = Arc::new(SingleFlight::new());
let runs = Arc::new(AtomicUsize::new(0));
let gate = Arc::new(tokio::sync::Notify::new());
let mut handles = Vec::new();
for _ in 0..8 {
let flight = Arc::clone(&flight);
let runs = Arc::clone(&runs);
let gate = Arc::clone(&gate);
handles.push(tokio::spawn(async move {
flight
.run(7, || async {
runs.fetch_add(1, Ordering::SeqCst);
// Held open so every caller is definitely waiting.
gate.notified().await;
Ok(42)
})
.await
}));
}
// Let all eight reach the map before the leader completes.
tokio::task::yield_now().await;
gate.notify_waiters();
for handle in handles {
assert_eq!(handle.await.expect("join").expect("value"), 42);
}
assert_eq!(runs.load(Ordering::SeqCst), 1, "work ran more than once");
}
/// Different keys must not serialize against each other.
#[tokio::test]
async fn different_keys_run_independently() {
let flight: Arc<SingleFlight<u32, u32>> = Arc::new(SingleFlight::new());
let runs = Arc::new(AtomicUsize::new(0));
for key in 0..4 {
let runs = Arc::clone(&runs);
let value = flight
.run(key, || async {
runs.fetch_add(1, Ordering::SeqCst);
Ok(key * 2)
})
.await
.expect("value");
assert_eq!(value, key * 2);
}
assert_eq!(runs.load(Ordering::SeqCst), 4);
}
/// A waiter gets the leader's message, not a success and not a hang.
#[tokio::test]
async fn a_waiter_sees_the_leaders_failure() {
let flight: Arc<SingleFlight<u32, u32>> = Arc::new(SingleFlight::new());
let gate = Arc::new(tokio::sync::Notify::new());
let leader = {
let flight = Arc::clone(&flight);
let gate = Arc::clone(&gate);
tokio::spawn(async move {
flight
.run(1, || async {
gate.notified().await;
Err(Error::NotFound("no such block".into()))
})
.await
})
};
tokio::task::yield_now().await;
let waiter = {
let flight = Arc::clone(&flight);
tokio::spawn(async move { flight.run(1, || async { Ok(0) }).await })
};
tokio::task::yield_now().await;
gate.notify_waiters();
assert!(leader.await.expect("join").is_err());
let error = waiter
.await
.expect("join")
.expect_err("waiter must fail too");
assert!(
error.to_string().contains("no such block"),
"waiter should carry the leader's reason: {error}"
);
}
/// A failed call must not poison the key — the next caller retries.
#[tokio::test]
async fn a_failed_call_leaves_the_key_retryable() {
let flight: Arc<SingleFlight<u32, u32>> = Arc::new(SingleFlight::new());
let first = flight
.run(1, || async { Err(Error::NotFound("transient".into())) })
.await;
assert!(first.is_err());
let second = flight.run(1, || async { Ok(9) }).await.expect("retry");
assert_eq!(second, 9);
}
/// A dropped leader must not leave the key marked in-flight forever — the
/// player cancels reads whenever the viewer seeks away.
#[tokio::test]
async fn a_cancelled_leader_does_not_wedge_the_key() {
let flight: Arc<SingleFlight<u32, u32>> = Arc::new(SingleFlight::new());
{
let pending = flight.run(1, || async {
std::future::pending::<()>().await;
Ok(0)
});
// Poll once so it becomes the leader, then drop it unfinished.
futures::pin_mut!(pending);
let mut context = std::task::Context::from_waker(futures::task::noop_waker_ref());
assert!(std::future::Future::poll(pending.as_mut(), &mut context).is_pending());
}
assert!(
flight.lock().is_empty(),
"entry outlived the cancelled call"
);
assert_eq!(flight.run(1, || async { Ok(5) }).await.expect("retry"), 5);
}
}
crates/pstr-stream/src/source.rs
+334
View File
@@ -0,0 +1,334 @@
//! Opening streams, and keeping the open ones around.
//!
//! A cold open is expensive: link details, an ancestor-key unlock (an S2K
//! derivation, tens of milliseconds), then a revision listing — all before the
//! first byte. Two things make that cost once instead of repeatedly:
//!
//! * an **LRU of open streams**, so going back to the episode you paused, or
//! scrubbing after the player reopened its demuxer, is instant;
//! * **single-flight on the open itself**, because the UI's "play" click and the
//! player's first read arrive milliseconds apart and would otherwise both pay
//! for the same handshake, then race to insert into the LRU.
use std::sync::{Arc, Mutex, PoisonError};
use async_trait::async_trait;
use lru::LruCache;
use proton_sdk::ids::NodeUid;
use crate::block::SharedBlocks;
use crate::disk::{DiskCache, DiskCacheConfig};
use crate::error::Result;
use crate::ring::{BlockRing, DEFAULT_RING_BYTES, non_zero};
use crate::single_flight::SingleFlight;
use crate::stream::{DEFAULT_READAHEAD_BLOCKS, VideoStream};
/// How the source turns a node into blocks.
///
/// A trait so the cache, ring and read-ahead can be exercised without an
/// account: the only implementation that talks to Proton is
/// [`crate::reader::LibraryOpener`].
#[async_trait]
pub trait RevisionOpener: Send + Sync + 'static {
/// Open the active revision of `uid` inside `share_id`.
async fn open(&self, share_id: &str, uid: &NodeUid) -> Result<SharedBlocks>;
}
/// How much memory, how much disk, and how far ahead to read.
#[derive(Debug, Clone)]
pub struct StreamConfig {
/// Resident block budget, shared by every open stream.
pub ring_bytes: u64,
/// Blocks to pull ahead of the reader. Zero disables read-ahead.
pub readahead_blocks: usize,
/// How many opened streams to keep. Cheap to hold — the bytes live in the
/// ring, not here — so this is generous.
pub max_open_streams: usize,
/// The on-disk block cache, when enabled.
pub disk_cache: Option<DiskCacheConfig>,
}
impl Default for StreamConfig {
fn default() -> Self {
Self {
ring_bytes: DEFAULT_RING_BYTES,
readahead_blocks: DEFAULT_READAHEAD_BLOCKS,
max_open_streams: 64,
disk_cache: None,
}
}
}
impl StreamConfig {
pub fn with_disk_cache(mut self, cache: DiskCacheConfig) -> Self {
self.disk_cache = Some(cache);
self
}
pub fn with_readahead(mut self, blocks: usize) -> Self {
self.readahead_blocks = blocks;
self
}
pub fn with_ring_bytes(mut self, bytes: u64) -> Self {
self.ring_bytes = bytes;
self
}
}
/// Identifies an open stream. The share is part of the key: the same file can
/// legitimately be reachable through two different links.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct StreamKey {
share: String,
uid: NodeUid,
}
/// Opens and caches [`VideoStream`]s.
#[derive(Clone)]
pub struct StreamSource {
inner: Arc<Inner>,
}
struct Inner {
opener: Arc<dyn RevisionOpener>,
ring: Arc<BlockRing>,
disk: Option<Arc<DiskCache>>,
readahead: usize,
open: Mutex<LruCache<StreamKey, VideoStream>>,
flight: Arc<SingleFlight<StreamKey, VideoStream>>,
}
impl StreamSource {
/// Build a source, opening the disk cache if one is configured.
pub async fn new(opener: Arc<dyn RevisionOpener>, config: StreamConfig) -> Result<Self> {
let disk = match config.disk_cache {
Some(cache) => Some(Arc::new(DiskCache::open(cache).await?)),
None => None,
};
Ok(Self {
inner: Arc::new(Inner {
opener,
ring: Arc::new(BlockRing::new(config.ring_bytes)),
disk,
readahead: config.readahead_blocks,
open: Mutex::new(LruCache::new(non_zero(config.max_open_streams))),
flight: Arc::new(SingleFlight::new()),
}),
})
}
/// Open a file for playback, or hand back the stream already open on it.
pub async fn open(&self, share_id: &str, uid: &NodeUid) -> Result<VideoStream> {
let key = StreamKey {
share: share_id.to_string(),
uid: uid.clone(),
};
if let Some(stream) = self.inner.lock().get(&key) {
return Ok(stream.clone());
}
let inner = Arc::clone(&self.inner);
let flight = Arc::clone(&self.inner.flight);
let share = share_id.to_string();
let uid = uid.clone();
let cached_key = key.clone();
flight
.run(key, move || async move {
let blocks = inner.opener.open(&share, &uid).await?;
let stream = VideoStream::new(
uid,
blocks,
Arc::clone(&inner.ring),
inner.disk.clone(),
inner.readahead,
);
// Evicting a stream here only drops the handle; its blocks stay
// in the ring under the ring's own budget, so a viewer who
// reopens still gets them.
inner.lock().put(cached_key, stream.clone());
Ok(stream)
})
.await
}
/// Drop a stream and everything cached for it in memory.
///
/// The disk cache is deliberately untouched: it exists precisely so the next
/// viewing is cheap.
pub fn close(&self, share_id: &str, uid: &NodeUid) {
let key = StreamKey {
share: share_id.to_string(),
uid: uid.clone(),
};
if let Some(stream) = self.inner.lock().pop(&key) {
self.inner.ring.forget(uid, stream.revision_id());
}
}
/// Ring statistics, across every stream.
pub fn ring_stats(&self) -> crate::ring::RingStats {
self.inner.ring.stats()
}
/// Disk-cache statistics, when one is configured.
pub fn disk_stats(&self) -> Option<crate::disk::DiskStats> {
self.inner.disk.as_ref().map(|disk| disk.stats())
}
}
impl Inner {
fn lock(&self) -> std::sync::MutexGuard<'_, LruCache<StreamKey, VideoStream>> {
self.open.lock().unwrap_or_else(PoisonError::into_inner)
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use proton_sdk::ids::{LinkId, VolumeId};
use super::*;
use crate::block::MemoryBlocks;
fn uid(link: &str) -> NodeUid {
NodeUid::new(VolumeId::new("vol"), LinkId::new(link))
}
struct CountingOpener {
opens: AtomicUsize,
delay: std::time::Duration,
}
impl CountingOpener {
fn new(delay: std::time::Duration) -> Arc<Self> {
Arc::new(Self {
opens: AtomicUsize::new(0),
delay,
})
}
fn opens(&self) -> usize {
self.opens.load(Ordering::SeqCst)
}
}
#[async_trait]
impl RevisionOpener for CountingOpener {
async fn open(&self, _share: &str, _uid: &NodeUid) -> Result<SharedBlocks> {
self.opens.fetch_add(1, Ordering::SeqCst);
if !self.delay.is_zero() {
tokio::time::sleep(self.delay).await;
}
Ok(Arc::new(MemoryBlocks::new(
"rev1",
vec![vec![1_u8; 64], vec![2_u8; 64]],
)))
}
}
async fn source(opener: Arc<CountingOpener>) -> StreamSource {
StreamSource::new(opener, StreamConfig::default().with_readahead(0))
.await
.expect("source")
}
#[tokio::test]
async fn an_opened_stream_reads_its_blocks() {
let source = source(CountingOpener::new(std::time::Duration::ZERO)).await;
let stream = source.open("share", &uid("a")).await.unwrap();
assert_eq!(stream.size(), 128);
assert_eq!(
stream.read_range(60, 8).await.unwrap(),
vec![1, 1, 1, 1, 2, 2, 2, 2]
);
}
/// Reopening the episode you paused must not re-run the handshake.
#[tokio::test]
async fn reopening_the_same_file_reuses_the_open_stream() {
let opener = CountingOpener::new(std::time::Duration::ZERO);
let source = source(Arc::clone(&opener)).await;
source.open("share", &uid("a")).await.unwrap();
source.open("share", &uid("a")).await.unwrap();
source.open("share", &uid("a")).await.unwrap();
assert_eq!(opener.opens(), 1);
}
/// The race the single-flight exists for: the UI click and the player's
/// first read, milliseconds apart, against a slow open.
#[tokio::test]
async fn concurrent_opens_of_one_file_handshake_once() {
let opener = CountingOpener::new(std::time::Duration::from_millis(40));
let source = source(Arc::clone(&opener)).await;
let mut handles = Vec::new();
for _ in 0..6 {
let source = source.clone();
handles.push(tokio::spawn(async move {
source.open("share", &uid("a")).await
}));
}
for handle in handles {
handle.await.unwrap().unwrap();
}
assert_eq!(opener.opens(), 1);
}
/// The same file behind two links is two streams — the shares may have
/// different lifetimes and different credentials.
#[tokio::test]
async fn the_same_file_in_two_shares_is_two_streams() {
let opener = CountingOpener::new(std::time::Duration::ZERO);
let source = source(Arc::clone(&opener)).await;
source.open("share-one", &uid("a")).await.unwrap();
source.open("share-two", &uid("a")).await.unwrap();
assert_eq!(opener.opens(), 2);
}
/// Closing must reclaim the memory, or a long session leaks a ring's worth
/// of every episode watched.
#[tokio::test]
async fn closing_a_stream_releases_its_blocks() {
let source = source(CountingOpener::new(std::time::Duration::ZERO)).await;
let stream = source.open("share", &uid("a")).await.unwrap();
stream.read_range(0, 128).await.unwrap();
assert!(source.ring_stats().resident_bytes > 0);
source.close("share", &uid("a"));
assert_eq!(source.ring_stats().resident_bytes, 0);
}
/// Beyond the LRU's capacity the oldest handle is dropped; the next open
/// re-runs the handshake rather than failing.
#[tokio::test]
async fn the_open_stream_lru_is_bounded() {
let opener = CountingOpener::new(std::time::Duration::ZERO);
let source = StreamSource::new(
Arc::clone(&opener) as Arc<dyn RevisionOpener>,
StreamConfig {
max_open_streams: 2,
readahead_blocks: 0,
..StreamConfig::default()
},
)
.await
.unwrap();
for link in ["a", "b", "c"] {
source.open("share", &uid(link)).await.unwrap();
}
// "a" was evicted, so this is a fourth open.
source.open("share", &uid("a")).await.unwrap();
assert_eq!(opener.opens(), 4);
}
}
crates/pstr-stream/src/stream.rs
+835
View File
@@ -0,0 +1,835 @@
//! A seekable byte stream over one revision, with caching and read-ahead.
//!
//! This is what the player reads from. mpv's demuxer asks for small, mostly
//! forward, occasionally wildly-seeking ranges; Proton serves 4 MiB encrypted
//! blocks. Everything between those two facts lives here:
//!
//! * a **byte range → blocks** mapping that never assumes a block size,
//! * a **memory ring** and an optional **disk cache**, checked in that order,
//! * **deduplication**, so read-ahead and the demand read never fetch the same
//! block twice,
//! * **forward read-ahead**, so the demuxer's next read is usually already in
//! memory rather than a 4 MiB round-trip away.
//!
//! ## Read-ahead has to yield to seeks
//!
//! Measured against a real 761 MiB episode over a public link, naive read-ahead
//! made everything *worse*: worst-case cold seek 2895 ms with six blocks of
//! prefetch versus 572 ms with none, and sustained throughput 5.0 MiB/s versus
//! 7.2. The cause is not subtle — prefetches queued from the position the viewer
//! just left keep running, and the block they are fetching is exactly the
//! bandwidth the seek needs.
//!
//! So a read that is not a continuation of the previous one cancels every
//! outstanding prefetch before it fetches anything. Read-ahead resumes from
//! wherever the seek settles. What is thrown away is speculative by definition;
//! what is protected is the only latency the viewer can feel.
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, PoisonError};
use futures::stream::{self, StreamExt, TryStreamExt};
use proton_sdk::ids::NodeUid;
use tokio::sync::Semaphore;
use crate::block::{BlockMap, SharedBlocks, splice};
use crate::disk::{DiskCache, DiskStats};
use crate::error::Result;
use crate::ring::{BlockKey, BlockRing, CachedBlock, RingStats};
use crate::single_flight::SingleFlight;
/// Blocks fetched concurrently to satisfy one read.
///
/// A read spanning more than a couple of blocks only happens on a cold seek;
/// steady-state playback reads well inside one. Kept modest because the SDK
/// enforces its own client-wide in-flight ceiling underneath this.
const MAX_CONCURRENT_BLOCKS_PER_READ: usize = 4;
/// How much of the file to pull ahead of the reader, in blocks.
///
/// Measured on a 761 MiB episode over a real public link, reading in 64 KiB
/// chunks against the default 128 MiB ring:
///
/// | depth | sustained | note |
/// |------:|----------:|------|
/// | 0 | 7.2 MiB/s | strictly one block at a time — latency-bound |
/// | 6 | 5.8 MiB/s | *worse* than serial: too shallow to overlap, still competing |
/// | 12 | 8.7 MiB/s | no evictions |
/// | 16 | 8.1–9.6 MiB/s | occasional evictions |
/// | 32 | 7.8 MiB/s | regresses — see [`clamp_readahead`] |
///
/// Run-to-run spread is wide enough that 12 and 16 are not meaningfully
/// different; 12 is chosen because it never evicted and wasted a third less
/// bandwidth on cancelled prefetches. Well past a 1080p bitrate either way.
pub const DEFAULT_READAHEAD_BLOCKS: usize = 12;
/// How far a read may land from where the previous one ended and still count as
/// a continuation rather than a seek.
///
/// One 4 MiB block's worth. Demuxers routinely re-read a little behind
/// themselves for headers and index entries, and a small forward hop is a
/// container skipping a chunk, not a viewer moving the seek bar. Treating either
/// as a seek would cancel read-ahead constantly and give back everything it
/// buys.
const SEQUENTIAL_SLACK: u64 = 4 * 1024 * 1024;
/// What a stream did, for the benchmark and the settings screen.
#[derive(Debug, Clone, Copy)]
pub struct StreamStats {
pub ring: RingStats,
pub disk: Option<DiskStats>,
/// Blocks that actually came off the network.
pub blocks_fetched: u64,
pub bytes_fetched: u64,
/// Read-ahead fetches that were started.
pub readahead_blocks: u64,
/// Read-ahead fetches cancelled by a seek.
pub readahead_cancelled: u64,
/// Reads that were not a continuation of the previous one.
pub seeks: u64,
}
/// One open revision, readable at any offset.
///
/// Cheap to clone — clones share the same reader, cache state and read-ahead.
#[derive(Clone)]
pub struct VideoStream {
inner: Arc<Inner>,
}
struct Inner {
uid: NodeUid,
blocks: SharedBlocks,
map: BlockMap,
ring: Arc<BlockRing>,
disk: Option<Arc<DiskCache>>,
/// One fetch per block, however many readers want it.
flight: Arc<SingleFlight<BlockKey, CachedBlock>>,
readahead: usize,
/// Bounds detached read-ahead work, so a seek-happy viewer cannot queue an
/// unbounded pile of prefetches.
readahead_slots: Arc<Semaphore>,
/// Live prefetches, so a seek can cancel them. See the module note.
readahead_tasks: Mutex<Vec<tokio::task::JoinHandle<()>>>,
/// Plaintext offset just past the previous read, for seek detection.
last_end: AtomicU64,
fetched_blocks: AtomicU64,
fetched_bytes: AtomicU64,
readahead_started: AtomicU64,
readahead_cancelled: AtomicU64,
seeks: AtomicU64,
}
impl VideoStream {
pub(crate) fn new(
uid: NodeUid,
blocks: SharedBlocks,
ring: Arc<BlockRing>,
disk: Option<Arc<DiskCache>>,
readahead: usize,
) -> Self {
let map = BlockMap::new(blocks.block_sizes());
let readahead = clamp_readahead(readahead, ring.budget(), blocks.block_sizes());
Self {
inner: Arc::new(Inner {
uid,
blocks,
map,
ring,
disk,
flight: Arc::new(SingleFlight::new()),
readahead,
readahead_slots: Arc::new(Semaphore::new(readahead.max(1))),
readahead_tasks: Mutex::new(Vec::new()),
last_end: AtomicU64::new(0),
fetched_blocks: AtomicU64::new(0),
fetched_bytes: AtomicU64::new(0),
readahead_started: AtomicU64::new(0),
readahead_cancelled: AtomicU64::new(0),
seeks: AtomicU64::new(0),
}),
}
}
/// Total plaintext size of the revision.
pub fn size(&self) -> u64 {
self.inner.map.size()
}
/// The node this stream reads.
pub fn uid(&self) -> &NodeUid {
&self.inner.uid
}
/// The revision this stream is pinned to. It does not follow a reseal —
/// reopen for that.
pub fn revision_id(&self) -> &str {
self.inner.blocks.revision_id()
}
/// Plaintext size of each block, in block order.
pub fn block_sizes(&self) -> &[u64] {
self.inner.blocks.block_sizes()
}
/// Fill `buf` from `offset`, returning how many bytes were written.
///
/// A read at or past EOF returns 0 rather than failing — that is how a
/// demuxer discovers the end of a file.
pub async fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<usize> {
let inner = &self.inner;
if buf.is_empty() || offset >= inner.map.size() {
return Ok(0);
}
let end = offset
.saturating_add(buf.len() as u64)
.min(inner.map.size());
// Before anything is fetched: a seek must not queue behind the
// speculation it just invalidated. See the module note — this is worth
// roughly 5x on worst-case seek latency.
let previous_end = inner.last_end.swap(end, Ordering::Relaxed);
if !is_continuation(offset, previous_end) {
inner.seeks.fetch_add(1, Ordering::Relaxed);
self.cancel_read_ahead();
}
let wanted = inner.map.blocks_in(offset, end);
// Taken by value: a closure borrowing from the range would give each
// fetch future a higher-ranked lifetime, which `tokio::spawn` rejects in
// callers with "implementation of `FnOnce` is not general enough".
let indices: Vec<usize> = wanted.clone().collect();
let last = wanted.end;
let mut fetches = stream::iter(
indices
.into_iter()
.map(|index| block_at(Arc::clone(inner), index)),
)
.buffered(MAX_CONCURRENT_BLOCKS_PER_READ);
let mut written = 0_usize;
let mut index = wanted.start;
// `buffered` yields in input order, so blocks arrive in file order and
// each one splices into its own part of `buf`.
while let Some(block) = fetches.try_next().await? {
let start = inner.map.start_of(index).unwrap_or(0);
written += splice(buf, offset, &block, start, offset, end);
index += 1;
}
self.read_ahead_from(last);
Ok(written)
}
/// [`Self::read_at`] into a fresh buffer. Convenient for tests and for the
/// benchmark; the player uses the buffer form.
pub async fn read_range(&self, offset: u64, length: u64) -> Result<Vec<u8>> {
let end = offset.saturating_add(length).min(self.size());
if length == 0 || offset >= end {
return Ok(Vec::new());
}
let mut buf = vec![0_u8; (end - offset) as usize];
let read = self.read_at(offset, &mut buf).await?;
buf.truncate(read);
Ok(buf)
}
pub fn stats(&self) -> StreamStats {
StreamStats {
ring: self.inner.ring.stats(),
disk: self.inner.disk.as_ref().map(|disk| disk.stats()),
blocks_fetched: self.inner.fetched_blocks.load(Ordering::Relaxed),
bytes_fetched: self.inner.fetched_bytes.load(Ordering::Relaxed),
readahead_blocks: self.inner.readahead_started.load(Ordering::Relaxed),
readahead_cancelled: self.inner.readahead_cancelled.load(Ordering::Relaxed),
seeks: self.inner.seeks.load(Ordering::Relaxed),
}
}
/// Abandon every outstanding prefetch.
///
/// Aborting mid-fetch throws away whatever that block had downloaded. That
/// is the intended trade: those bytes were speculative, and the request slot
/// they occupy is what the seek is waiting for. A demand read that later
/// wants an aborted block simply refetches it — the single-flight entry is
/// released by its guard when the task is dropped, so nothing waits on a
/// call that will never finish.
fn cancel_read_ahead(&self) {
let tasks = std::mem::take(&mut *self.inner.tasks());
let mut cancelled = 0_u64;
for task in tasks {
if !task.is_finished() {
task.abort();
cancelled += 1;
}
}
if cancelled > 0 {
self.inner
.readahead_cancelled
.fetch_add(cancelled, Ordering::Relaxed);
}
}
/// Queue the next few blocks after `from`.
///
/// Detached and best-effort by design: a prefetch that fails, or that never
/// gets a slot, costs the next demand read a round-trip and nothing else.
/// Blocks already resident are skipped without touching their recency, so
/// speculation cannot reorder the LRU against real playback.
fn read_ahead_from(&self, from: usize) {
let inner = &self.inner;
if inner.readahead == 0 {
return;
}
let revision = inner.blocks.revision_id();
let upto = (from + inner.readahead).min(inner.map.len());
for index in from..upto {
let key = BlockKey::new(&inner.uid, revision, index);
if inner.ring.contains(&key) {
continue;
}
let Ok(slot) = Arc::clone(&inner.readahead_slots).try_acquire_owned() else {
// Already as far ahead as this stream is allowed to run.
return;
};
inner.readahead_started.fetch_add(1, Ordering::Relaxed);
let spawned = Arc::clone(inner);
let task = tokio::spawn(async move {
let _slot = slot;
if let Err(e) = block_at(spawned, index).await {
tracing::debug!(index, error = %e, "read-ahead block failed");
}
});
let mut tasks = inner.tasks();
// Finished handles are dead weight; drop them while we are here so
// the list stays the size of the read-ahead window.
tasks.retain(|task| !task.is_finished());
tasks.push(task);
}
}
}
/// Cap read-ahead at half the ring's capacity in blocks.
///
/// Reading further ahead than the ring can hold is self-defeating: the blocks at
/// the front of the window are evicted before the player reaches them, so they
/// are fetched twice and the bandwidth spent on them is lost. Measured at depth
/// 32 against a 128 MiB ring (32 blocks), sustained throughput fell from
/// 8.7 MiB/s to 7.8 with six evictions — worse than depth 12.
///
/// Half, not all, because the ring also has to hold what the player has just
/// *passed*: a scrub a few seconds backwards is common and should not refetch.
fn clamp_readahead(requested: usize, ring_budget: u64, block_sizes: &[u64]) -> usize {
let block = block_sizes.iter().copied().max().unwrap_or(0);
if block == 0 {
return requested;
}
let capacity = (ring_budget / block) as usize;
requested.min((capacity / 2).max(1))
}
/// Whether a read starting at `offset` continues from where the last one ended.
///
/// Deliberately generous in both directions — see [`SEQUENTIAL_SLACK`]. The
/// first read of a stream (`previous_end == 0`, `offset == 0`) is a
/// continuation, so opening a file does not count as a seek.
fn is_continuation(offset: u64, previous_end: u64) -> bool {
offset >= previous_end.saturating_sub(SEQUENTIAL_SLACK)
&& offset <= previous_end.saturating_add(SEQUENTIAL_SLACK)
}
impl Inner {
fn tasks(&self) -> std::sync::MutexGuard<'_, Vec<tokio::task::JoinHandle<()>>> {
self.readahead_tasks
.lock()
.unwrap_or_else(PoisonError::into_inner)
}
}
/// Resolve one block: memory, then disk, then the network.
///
/// Takes an owned `Arc` so the future is `'static` and can be both spawned for
/// read-ahead and buffered inside a read.
async fn block_at(inner: Arc<Inner>, index: usize) -> Result<CachedBlock> {
let key = BlockKey::new(&inner.uid, inner.blocks.revision_id(), index);
if let Some(block) = inner.ring.get(&key) {
return Ok(block);
}
let flight = Arc::clone(&inner.flight);
let fetch_key = key.clone();
flight
.run(key, move || async move {
if let Some(disk) = inner.disk.as_ref()
&& let Some(bytes) = disk.get(&fetch_key).await
{
let block: CachedBlock = bytes.into();
inner.ring.insert(fetch_key, Arc::clone(&block));
return Ok(block);
}
let bytes = inner.blocks.read_block(index).await?;
inner.fetched_blocks.fetch_add(1, Ordering::Relaxed);
inner
.fetched_bytes
.fetch_add(bytes.len() as u64, Ordering::Relaxed);
if let Some(disk) = inner.disk.as_ref() {
disk.put(&fetch_key, &bytes).await;
}
let block: CachedBlock = bytes.into();
inner.ring.insert(fetch_key, Arc::clone(&block));
Ok(block)
})
.await
}
#[cfg(test)]
mod tests {
use std::sync::atomic::AtomicUsize;
use async_trait::async_trait;
use proton_sdk::ids::{LinkId, VolumeId};
use super::*;
use crate::block::BlockSource;
use crate::ring::DEFAULT_RING_BYTES;
fn uid() -> NodeUid {
NodeUid::new(VolumeId::new("vol"), LinkId::new("link"))
}
/// Deterministic content: byte `i` of the file is `(i % 251) as u8`, so any
/// range can be checked without holding the whole file.
fn expected(offset: u64, len: usize) -> Vec<u8> {
(0..len)
.map(|i| ((offset + i as u64) % 251) as u8)
.collect()
}
/// A block source that counts fetches per block and can be made slow or
/// failing.
///
/// Per-block counters rather than one total: read-ahead runs detached, so a
/// total is inherently racy, whereas "block 1 was fetched exactly once" is
/// the property the tests actually care about and is stable regardless of
/// when the prefetch tasks happen to be scheduled.
struct Fake {
sizes: Vec<u64>,
starts: Vec<u64>,
per_block: Vec<AtomicUsize>,
delay: std::time::Duration,
fail_block: Option<usize>,
}
impl Fake {
fn new(sizes: &[u64]) -> Arc<Self> {
let mut starts = Vec::new();
let mut offset = 0;
for &size in sizes {
starts.push(offset);
offset += size;
}
Arc::new(Self {
sizes: sizes.to_vec(),
starts,
per_block: sizes.iter().map(|_| AtomicUsize::new(0)).collect(),
delay: std::time::Duration::ZERO,
fail_block: None,
})
}
fn slow(sizes: &[u64], delay: std::time::Duration) -> Arc<Self> {
let mut fake = Fake::new(sizes);
Arc::get_mut(&mut fake).unwrap().delay = delay;
fake
}
fn failing(sizes: &[u64], block: usize) -> Arc<Self> {
let mut fake = Fake::new(sizes);
Arc::get_mut(&mut fake).unwrap().fail_block = Some(block);
fake
}
fn fetches_of(&self, index: usize) -> usize {
self.per_block[index].load(Ordering::SeqCst)
}
fn fetches(&self) -> usize {
self.per_block
.iter()
.map(|count| count.load(Ordering::SeqCst))
.sum()
}
}
#[async_trait]
impl BlockSource for Fake {
fn revision_id(&self) -> &str {
"rev1"
}
fn block_sizes(&self) -> &[u64] {
&self.sizes
}
async fn read_block(&self, index: usize) -> Result<Vec<u8>> {
if Some(index) == self.fail_block {
return Err(crate::Error::NotFound(format!("block {index} is broken")));
}
self.per_block[index].fetch_add(1, Ordering::SeqCst);
if !self.delay.is_zero() {
tokio::time::sleep(self.delay).await;
}
let start = self.starts[index];
Ok(expected(start, self.sizes[index] as usize))
}
}
fn stream_over(fake: Arc<Fake>, readahead: usize) -> VideoStream {
VideoStream::new(
uid(),
fake,
Arc::new(BlockRing::new(DEFAULT_RING_BYTES)),
None,
readahead,
)
}
#[tokio::test]
async fn a_stream_reports_the_revisions_size() {
let stream = stream_over(Fake::new(&[1000, 1000, 500]), 0);
assert_eq!(stream.size(), 2500);
assert_eq!(stream.block_sizes(), &[1000, 1000, 500]);
}
#[tokio::test]
async fn a_read_inside_one_block_returns_the_right_bytes() {
let stream = stream_over(Fake::new(&[1000, 1000, 500]), 0);
assert_eq!(stream.read_range(10, 20).await.unwrap(), expected(10, 20));
}
/// The case that matters for playback: a read straddling a block boundary
/// must stitch two blocks without a seam.
#[tokio::test]
async fn a_read_across_a_boundary_stitches_blocks_correctly() {
let stream = stream_over(Fake::new(&[1000, 1000, 500]), 0);
assert_eq!(stream.read_range(990, 30).await.unwrap(), expected(990, 30));
}
#[tokio::test]
async fn a_read_spanning_every_block_returns_the_whole_file() {
let stream = stream_over(Fake::new(&[1000, 1000, 500]), 0);
assert_eq!(stream.read_range(0, 2500).await.unwrap(), expected(0, 2500));
}
/// A cold seek to the middle must fetch only the block it lands in — that is
/// the entire reason this app can exist over a 4 MiB block store.
#[tokio::test]
async fn a_seek_fetches_only_the_blocks_it_touches() {
let fake = Fake::new(&[1000; 100]);
let stream = stream_over(Arc::clone(&fake), 0);
stream.read_range(75_000, 100).await.unwrap();
assert_eq!(fake.fetches(), 1, "seek pulled more than the landing block");
}
/// The short tail. Reading past it yields what exists, not an error.
#[tokio::test]
async fn a_read_past_the_end_is_clamped_rather_than_failing() {
let stream = stream_over(Fake::new(&[1000, 500]), 0);
let bytes = stream.read_range(1400, 1000).await.unwrap();
assert_eq!(bytes, expected(1400, 100));
}
#[tokio::test]
async fn a_read_at_or_past_the_end_returns_nothing() {
let stream = stream_over(Fake::new(&[100]), 0);
assert!(stream.read_range(100, 10).await.unwrap().is_empty());
assert!(stream.read_range(5_000, 10).await.unwrap().is_empty());
}
/// Rewatching a scene must come out of memory, not off the wire.
#[tokio::test]
async fn a_repeated_read_is_served_from_the_ring() {
let fake = Fake::new(&[1000, 1000]);
let stream = stream_over(Arc::clone(&fake), 0);
stream.read_range(0, 100).await.unwrap();
stream.read_range(200, 100).await.unwrap();
stream.read_range(500, 100).await.unwrap();
assert_eq!(fake.fetches(), 1);
assert!(stream.stats().ring.hits >= 2);
}
/// Read-ahead exists so the demuxer's *next* read is already resident.
#[tokio::test]
async fn read_ahead_pulls_the_following_blocks() {
let fake = Fake::new(&[1000; 10]);
let stream = stream_over(Arc::clone(&fake), 3);
stream.read_range(0, 100).await.unwrap();
// Detached tasks; give them a chance to run.
for _ in 0..50 {
if fake.fetches() >= 4 {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
}
for index in 0..4 {
assert_eq!(fake.fetches_of(index), 1, "block {index} not read ahead");
}
assert_eq!(fake.fetches_of(4), 0, "read-ahead ran past its window");
// And the next read is served from what read-ahead already pulled.
assert_eq!(
stream.read_range(1000, 50).await.unwrap(),
expected(1000, 50)
);
assert_eq!(fake.fetches_of(1), 1, "read-ahead's block was refetched");
}
/// Read-ahead must stop at EOF rather than fetching blocks that do not
/// exist.
#[tokio::test]
async fn read_ahead_stops_at_the_end_of_the_file() {
let fake = Fake::new(&[1000, 1000]);
let stream = stream_over(Arc::clone(&fake), 8);
stream.read_range(0, 10).await.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
assert_eq!(fake.fetches(), 2, "fetched past the last block");
}
/// The collision read-ahead makes inevitable: the demand read wants the
/// block a prefetch is already fetching. It must wait for it, not start a
/// second fetch.
#[tokio::test]
async fn a_demand_read_joins_an_in_flight_read_ahead_instead_of_refetching() {
let fake = Fake::slow(&[1000; 10], std::time::Duration::from_millis(80));
let stream = stream_over(Arc::clone(&fake), 4);
stream.read_range(0, 10).await.unwrap();
// Read-ahead for blocks 1..5 is now in flight against an 80 ms source.
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
assert_eq!(
stream.read_range(1000, 10).await.unwrap(),
expected(1000, 10)
);
assert_eq!(
fake.fetches_of(1),
1,
"the demand read started a second fetch of an in-flight block"
);
}
/// Read-ahead is speculative; its failures must never surface to the player.
#[tokio::test]
async fn a_failing_read_ahead_does_not_fail_the_read() {
let fake = Fake::failing(&[1000; 5], 2);
let stream = stream_over(Arc::clone(&fake), 4);
assert_eq!(stream.read_range(0, 10).await.unwrap(), expected(0, 10));
tokio::time::sleep(std::time::Duration::from_millis(30)).await;
// And the demand read for that block still reports the real failure.
assert!(stream.read_range(2000, 10).await.is_err());
}
/// A source failure on the demand path is an error, not silence — silence
/// would look like EOF and the player would stop mid-episode.
#[tokio::test]
async fn a_failing_block_surfaces_as_an_error() {
let fake = Fake::failing(&[1000; 3], 0);
let stream = stream_over(fake, 0);
assert!(stream.read_range(0, 10).await.is_err());
}
/// Concurrent readers of one stream — mpv's demuxer and its seek probe do
/// exactly this — must not multiply the fetches.
#[tokio::test]
async fn concurrent_reads_of_the_same_block_fetch_it_once() {
let fake = Fake::slow(&[4096], std::time::Duration::from_millis(50));
let stream = stream_over(Arc::clone(&fake), 0);
let mut handles = Vec::new();
for i in 0..8_u64 {
let stream = stream.clone();
handles.push(tokio::spawn(
async move { stream.read_range(i * 10, 10).await },
));
}
for (i, handle) in handles.into_iter().enumerate() {
let bytes = handle.await.unwrap().unwrap();
assert_eq!(bytes, expected(i as u64 * 10, 10));
}
assert_eq!(fake.fetches(), 1);
}
/// Non-uniform block sizes must read correctly — the SDK warns that
/// assuming 4 MiB serves bytes from the wrong offset, silently.
#[tokio::test]
async fn non_uniform_block_sizes_read_correctly() {
let fake = Fake::new(&[13, 1021, 7, 4096, 1]);
let stream = stream_over(fake, 0);
assert_eq!(stream.size(), 5138);
assert_eq!(stream.read_range(0, 5138).await.unwrap(), expected(0, 5138));
assert_eq!(
stream.read_range(1030, 40).await.unwrap(),
expected(1030, 40)
);
}
/// The memory promise holds even for a file far larger than the ring.
#[tokio::test]
async fn streaming_a_file_larger_than_the_ring_stays_inside_the_budget() {
let fake = Fake::new(&[1024; 200]);
let stream = VideoStream::new(uid(), fake, Arc::new(BlockRing::new(8 * 1024)), None, 0);
let mut offset = 0;
while offset < stream.size() {
stream.read_range(offset, 512).await.unwrap();
assert!(stream.stats().ring.resident_bytes <= 8 * 1024);
offset += 512;
}
}
#[tokio::test]
async fn read_at_fills_the_callers_buffer() {
let stream = stream_over(Fake::new(&[1000, 1000]), 0);
let mut buf = vec![0_u8; 64];
assert_eq!(stream.read_at(990, &mut buf).await.unwrap(), 64);
assert_eq!(buf, expected(990, 64));
}
/// Reading further ahead than the ring can hold evicts the window's own
/// blocks before the player reaches them.
#[test]
fn read_ahead_is_capped_at_half_the_rings_capacity() {
// 128 MiB of 4 MiB blocks is 32 blocks, so at most 16 ahead.
let four_mib = 4 * 1024 * 1024;
assert_eq!(clamp_readahead(32, 128 * 1024 * 1024, &[four_mib; 200]), 16);
assert_eq!(clamp_readahead(12, 128 * 1024 * 1024, &[four_mib; 200]), 12);
}
/// A ring too small for even one block must still read ahead by one, or
/// playback degenerates to a stall per block with nothing overlapping.
#[test]
fn a_tiny_ring_still_reads_one_block_ahead() {
assert_eq!(clamp_readahead(12, 1024, &[4 * 1024 * 1024]), 1);
}
/// A stream configured with no read-ahead must keep none — the clamp raises
/// a floor for what to *fetch*, not for whether to fetch at all.
#[tokio::test]
async fn read_ahead_stays_off_when_it_is_configured_off() {
let fake = Fake::new(&[1024; 32]);
let stream = VideoStream::new(
uid(),
Arc::clone(&fake) as SharedBlocks,
Arc::new(BlockRing::new(1024)),
None,
0,
);
stream.read_range(0, 100).await.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
assert_eq!(fake.fetches(), 1);
}
#[test]
fn a_forward_read_within_one_block_continues_the_previous_one() {
assert!(is_continuation(0, 0), "the first read is not a seek");
assert!(is_continuation(1_000, 1_000));
assert!(is_continuation(1_000 + SEQUENTIAL_SLACK, 1_000));
}
/// Demuxers re-read a little behind themselves for headers and index
/// entries. Calling that a seek would cancel read-ahead constantly.
#[test]
fn a_small_backward_read_still_continues_the_previous_one() {
assert!(is_continuation(100_000_000 - SEQUENTIAL_SLACK, 100_000_000));
assert!(is_continuation(0, SEQUENTIAL_SLACK));
}
#[test]
fn a_jump_past_the_slack_is_a_seek() {
assert!(!is_continuation(100_000_000, 1_000));
assert!(!is_continuation(0, 100_000_000));
}
/// The measured bug: prefetches from the position the viewer left were
/// holding the bandwidth the seek needed, costing 5x on worst-case latency.
#[tokio::test]
async fn a_seek_cancels_the_read_ahead_it_invalidates() {
let fake = Fake::slow(&[1024 * 1024; 200], std::time::Duration::from_millis(200));
let stream = stream_over(Arc::clone(&fake), 4);
stream.read_range(0, 1_000).await.unwrap();
// Prefetches for blocks 1..5 are now in flight against a slow source.
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
assert!(stream.stats().readahead_blocks > 0, "nothing to cancel");
// Far past the slack: a seek-bar drag.
stream.read_range(150 * 1024 * 1024, 1_000).await.unwrap();
let stats = stream.stats();
assert_eq!(stats.seeks, 1);
assert!(
stats.readahead_cancelled > 0,
"seek did not cancel the stale prefetches"
);
}
/// And the mirror image: steady playback must not keep throwing its own
/// read-ahead away.
#[tokio::test]
async fn sequential_playback_never_cancels_its_own_read_ahead() {
let fake = Fake::new(&[1024; 64]);
let stream = stream_over(fake, 4);
let mut offset = 0;
while offset < stream.size() {
stream.read_range(offset, 256).await.unwrap();
offset += 256;
}
let stats = stream.stats();
assert_eq!(stats.seeks, 0, "sequential reads read as seeks");
assert_eq!(stats.readahead_cancelled, 0);
}
/// A cancelled prefetch must not leave the block unfetchable — the
/// single-flight entry has to be released when the task is dropped.
#[tokio::test]
async fn a_block_whose_prefetch_was_cancelled_can_still_be_read() {
let fake = Fake::slow(&[1024 * 1024; 200], std::time::Duration::from_millis(100));
let stream = stream_over(Arc::clone(&fake), 4);
stream.read_range(0, 1_000).await.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
stream.read_range(150 * 1024 * 1024, 1_000).await.unwrap();
// Block 1 was almost certainly aborted mid-fetch; reading it must work.
assert_eq!(
stream.read_range(1024 * 1024, 16).await.unwrap(),
expected(1024 * 1024, 16)
);
}
/// The player passes a fixed-size buffer at every offset, including the last
/// one; a short read there is normal and must be reported honestly.
#[tokio::test]
async fn a_short_read_at_the_tail_reports_what_it_wrote() {
let stream = stream_over(Fake::new(&[100]), 0);
let mut buf = vec![0_u8; 64];
assert_eq!(stream.read_at(80, &mut buf).await.unwrap(), 20);
assert_eq!(&buf[..20], &expected(80, 20)[..]);
}
}