perf(isrc): bound retained indexes and stream FLAC comments

This commit is contained in:
zarzet committed 2026-10-04 01:40:58 +07:00
1 parent 3bc12d4483
commit 644d035661
4 files changed
+517 -65

No files matched your search

+336 -36
View File
@@ -1,10 +1,12 @@
use crate::matching::uppercase;
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, MutexGuard, TryLockError, Weak};
use std::time::{Duration, Instant};
type Check<'a> = &'a (dyn Fn() -> Result<(), String> + Sync);
const MAX_DIRECTORIES: usize = 32;
const MAX_INDEX_BYTES: usize = 16 << 20;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct FileStamp {
@@ -24,7 +26,6 @@ pub trait IndexFiles: Sync {
fn stat(&self, path: &str) -> Result<Option<FileStamp>, String>;
}
#[derive(Clone)]
struct FileEntry {
stamp: FileStamp,
isrc: String,
@@ -32,13 +33,113 @@ struct FileEntry {
struct Index {
entries: BTreeMap<String, String>,
files: BTreeMap<String, FileEntry>,
files: Arc<BTreeMap<String, Arc<FileEntry>>>,
bytes: usize,
built_at: Instant,
}
#[derive(Default)]
struct Builder {
lock: Mutex<()>,
generation: AtomicU64,
}
struct Cached {
index: Arc<Mutex<Index>>,
bytes: usize,
used: u64,
}
#[derive(Default)]
struct State {
indexes: BTreeMap<String, Cached>,
builders: BTreeMap<String, Weak<Builder>>,
bytes: usize,
sequence: u64,
}
impl State {
fn remove(&mut self, directory: &str) {
if let Some(entry) = self.indexes.remove(directory) {
self.bytes -= entry.bytes;
}
}
fn store(&mut self, directory: &str, index: Arc<Mutex<Index>>, bytes: usize) {
self.remove(directory);
// Include a conservative allowance for tree nodes/control blocks. This
// bounds retained cache memory, not snapshots owned by active callers.
let bytes = bytes.saturating_add(directory.len() + std::mem::size_of::<Cached>() + 64);
if bytes > MAX_INDEX_BYTES {
return;
}
while self.indexes.len() >= MAX_DIRECTORIES || self.bytes + bytes > MAX_INDEX_BYTES {
let oldest = self
.indexes
.iter()
.min_by_key(|(_, entry)| entry.used)
.map(|(directory, _)| directory.clone())
.expect("nonempty ISRC cache");
self.remove(&oldest);
}
self.sequence += 1;
self.bytes += bytes;
self.indexes.insert(
directory.into(),
Cached {
index,
bytes,
used: self.sequence,
},
);
}
}
impl Index {
fn insert_match(&mut self, isrc: String, path: String) {
// Replacing a value retains the map's existing key allocation.
match self.entries.entry(isrc) {
std::collections::btree_map::Entry::Occupied(mut entry) => {
self.bytes -= entry.get().capacity();
self.bytes += path.capacity();
entry.insert(path);
}
std::collections::btree_map::Entry::Vacant(entry) => {
self.bytes += std::mem::size_of::<(String, String)>()
+ 64
+ entry.key().capacity()
+ path.capacity();
entry.insert(path);
}
}
}
fn insert_file(&mut self, path: String, file: Arc<FileEntry>) {
let cost = |file: &FileEntry| {
std::mem::size_of::<FileEntry>()
+ 32
+ file.stamp.path.capacity()
+ file.isrc.capacity()
};
match Arc::make_mut(&mut self.files).entry(path) {
std::collections::btree_map::Entry::Occupied(mut entry) => {
self.bytes -= cost(entry.get());
self.bytes += cost(&file);
entry.insert(file);
}
std::collections::btree_map::Entry::Vacant(entry) => {
self.bytes += std::mem::size_of::<(String, Arc<FileEntry>)>()
+ 64
+ entry.key().capacity()
+ cost(&file);
entry.insert(file);
}
}
}
}
pub struct IndexCache {
indexes: Mutex<BTreeMap<String, Arc<Mutex<Index>>>>,
builders: Mutex<BTreeMap<String, Weak<Mutex<()>>>>,
state: Mutex<State>,
ttl: Duration,
}
@@ -51,18 +152,19 @@ impl Default for IndexCache {
impl IndexCache {
pub fn with_ttl(ttl: Duration) -> Self {
Self {
indexes: Mutex::default(),
builders: Mutex::default(),
state: Mutex::default(),
ttl,
}
}
fn cached(&self, directory: &str) -> Option<Arc<Mutex<Index>>> {
self.indexes
.lock()
.expect("ISRC cache lock")
.get(directory)
.cloned()
let mut state = self.state.lock().expect("ISRC cache lock");
state.sequence += 1;
let sequence = state.sequence;
state.indexes.get_mut(directory).map(|entry| {
entry.used = sequence;
Arc::clone(&entry.index)
})
}
fn fresh(&self, directory: &str) -> Option<Arc<Mutex<Index>>> {
@@ -82,25 +184,31 @@ impl IndexCache {
return Ok(index);
}
let builder = {
let mut builders = self.builders.lock().expect("ISRC builders lock");
let entry = builders.entry(directory.into()).or_default();
let mut state = self.state.lock().expect("ISRC cache lock");
state
.builders
.retain(|_, builder| builder.strong_count() > 0);
let entry = state.builders.entry(directory.into()).or_default();
entry.upgrade().unwrap_or_else(|| {
let lock = Arc::new(Mutex::new(()));
let lock = Arc::new(Builder::default());
*entry = Arc::downgrade(&lock);
lock
})
};
let _building = wait_for_builder(&builder, check)?;
let _building = wait_for_builder(&builder.lock, check)?;
if !force && let Some(index) = self.fresh(directory) {
return Ok(index);
}
let generation = builder.generation.load(Ordering::Acquire);
let previous = self
.cached(directory)
.map(|index| index.lock().expect("ISRC index lock").files.clone())
.unwrap_or_default();
let mut index = Index {
entries: BTreeMap::new(),
files: BTreeMap::new(),
files: Arc::default(),
// Account for the two maps' sparsely populated root nodes too.
bytes: std::mem::size_of::<Index>() + 2048,
built_at: Instant::now(),
};
let mut changed = Vec::new();
@@ -115,9 +223,9 @@ impl IndexCache {
&& entry.stamp == stamp
{
if !entry.isrc.is_empty() {
index.entries.insert(entry.isrc.clone(), stamp.path.clone());
index.insert_match(entry.isrc.clone(), stamp.path.clone());
}
index.files.insert(stamp.path.clone(), entry.clone());
index.insert_file(stamp.path, Arc::clone(entry));
} else {
changed.push(stamp);
}
@@ -157,19 +265,18 @@ impl IndexCache {
{
check()?;
if !isrc.is_empty() {
index.entries.insert(isrc.clone(), stamp.path.clone());
index.insert_match(isrc.clone(), stamp.path.clone());
}
index
.files
.insert(stamp.path.clone(), FileEntry { stamp, isrc });
index.insert_file(stamp.path.clone(), Arc::new(FileEntry { stamp, isrc }));
}
check()?;
let bytes = index.bytes;
let index = Arc::new(Mutex::new(index));
if !directory.is_empty() {
self.indexes
.lock()
.expect("ISRC cache lock")
.insert(directory.into(), Arc::clone(&index));
let mut state = self.state.lock().expect("ISRC cache lock");
if builder.generation.load(Ordering::Acquire) == generation {
state.store(directory, Arc::clone(&index), bytes);
}
}
Ok(index)
}
@@ -221,29 +328,51 @@ impl IndexCache {
if directory.is_empty() || isrc.is_empty() || path.is_empty() {
return Ok(());
}
if let Some(index) = self.cached(directory) {
if let Some(index_owner) = self.cached(directory) {
let stamp = files.stat(path)?;
check()?;
let mut index = index.lock().expect("ISRC index lock");
let mut index = index_owner.lock().expect("ISRC index lock");
let isrc = uppercase(isrc);
index.entries.insert(isrc.clone(), path.into());
index.insert_match(isrc.clone(), path.into());
if let Some(stamp) = stamp {
index.files.insert(path.into(), FileEntry { stamp, isrc });
index.insert_file(path.into(), Arc::new(FileEntry { stamp, isrc }));
}
index.built_at = Instant::now();
let bytes = index.bytes;
drop(index);
let mut state = self.state.lock().expect("ISRC cache lock");
if let Some(current) = state.indexes.get(directory)
&& Arc::ptr_eq(&current.index, &index_owner)
{
state.store(directory, index_owner, bytes);
}
}
Ok(())
}
pub fn invalidate(&self, directory: &str) {
self.indexes
.lock()
.expect("ISRC cache lock")
.remove(directory);
let mut state = self.state.lock().expect("ISRC cache lock");
state.remove(directory);
if let Some(builder) = state.builders.get(directory).and_then(Weak::upgrade) {
builder.generation.fetch_add(1, Ordering::AcqRel);
}
state
.builders
.retain(|_, builder| builder.strong_count() > 0);
}
pub fn clear(&self) {
self.indexes.lock().expect("ISRC cache lock").clear();
let mut state = self.state.lock().expect("ISRC cache lock");
state.indexes.clear();
state.bytes = 0;
state.builders.retain(|_, builder| {
if let Some(builder) = builder.upgrade() {
builder.generation.fetch_add(1, Ordering::AcqRel);
true
} else {
false
}
});
}
}
@@ -260,3 +389,174 @@ fn wait_for_builder<'a>(
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Barrier;
struct Files {
listed: Mutex<Vec<FileStamp>>,
reads: AtomicUsize,
lists: AtomicUsize,
block_first: Option<(Arc<Barrier>, Arc<Barrier>)>,
}
impl Files {
fn new() -> Self {
Self {
listed: Mutex::new(vec![FileStamp {
path: "example.flac".into(),
size: 1,
modified_ns: 1,
directory: false,
}]),
reads: AtomicUsize::new(0),
lists: AtomicUsize::new(0),
block_first: None,
}
}
}
impl IndexFiles for Files {
fn list(&self, _: &str, check: Check<'_>) -> Result<Vec<FileStamp>, String> {
if self.lists.fetch_add(1, Ordering::SeqCst) == 0
&& let Some((started, release)) = &self.block_first
{
started.wait();
release.wait();
}
check()?;
Ok(self.listed.lock().unwrap().clone())
}
fn read(&self, _: &str, check: Check<'_>) -> Result<String, String> {
check()?;
self.reads.fetch_add(1, Ordering::SeqCst);
Ok("EXAMPLE".into())
}
fn stat(&self, path: &str) -> Result<Option<FileStamp>, String> {
Ok(self
.listed
.lock()
.unwrap()
.iter()
.find(|stamp| stamp.path == path)
.cloned())
}
}
#[test]
fn rebuild_shares_unchanged_entries_and_keeps_changed_duplicate_walk_order() {
let files = Files::new();
files.listed.lock().unwrap().push(FileStamp {
path: "second.flac".into(),
size: 1,
modified_ns: 1,
directory: false,
});
let cache = IndexCache::default();
let first = cache.index("directory", &files, false, &|| Ok(())).unwrap();
let previous = Arc::clone(&first.lock().unwrap().files);
let second = cache.index("directory", &files, true, &|| Ok(())).unwrap();
let next = Arc::clone(&second.lock().unwrap().files);
assert!(Arc::ptr_eq(
&previous["example.flac"],
&next["example.flac"]
));
assert_eq!(files.reads.load(Ordering::SeqCst), 2);
assert_eq!(
cache
.check("directory", "example", &files, &|| Ok(()))
.unwrap(),
"second.flac"
);
files.listed.lock().unwrap()[0].size = 2;
cache.index("directory", &files, true, &|| Ok(())).unwrap();
assert_eq!(
cache
.check("directory", "example", &files, &|| Ok(()))
.unwrap(),
"example.flac"
);
assert_eq!(files.reads.load(Ordering::SeqCst), 3);
}
#[test]
fn cache_bounds_directories_bytes_and_dead_builder_keys() {
let files = Files::new();
let cache = IndexCache::default();
for directory in 0..MAX_DIRECTORIES * 2 {
cache
.index(&directory.to_string(), &files, false, &|| Ok(()))
.unwrap();
}
{
let state = cache.state.lock().unwrap();
assert_eq!(state.indexes.len(), MAX_DIRECTORIES);
assert!(state.bytes <= MAX_INDEX_BYTES);
assert!(state.builders.len() <= 1);
}
let last = cache
.cached(&(MAX_DIRECTORIES * 2 - 1).to_string())
.unwrap();
let mut state = cache.state.lock().unwrap();
state.store("oversized", last, MAX_INDEX_BYTES);
assert!(!state.indexes.contains_key("oversized"));
drop(state);
cache.clear();
let state = cache.state.lock().unwrap();
assert!(state.indexes.is_empty());
assert!(state.builders.is_empty());
assert_eq!(state.bytes, 0);
}
#[test]
fn invalidation_and_clear_preserve_active_builder_and_reject_its_stale_publish() {
for clear in [false, true] {
let started = Arc::new(Barrier::new(2));
let release = Arc::new(Barrier::new(2));
let mut files = Files::new();
files.block_first = Some((started.clone(), release.clone()));
let cache = IndexCache::default();
std::thread::scope(|scope| {
let first = scope.spawn(|| cache.index("directory", &files, false, &|| Ok(())));
started.wait();
let original = cache.state.lock().unwrap().builders["directory"]
.upgrade()
.unwrap();
if clear {
cache.clear();
} else {
cache.invalidate("directory");
}
let same = cache.state.lock().unwrap().builders["directory"]
.upgrade()
.unwrap();
assert!(Arc::ptr_eq(&original, &same));
release.wait();
first.join().unwrap().unwrap();
assert!(cache.cached("directory").is_none());
cache.index("directory", &files, false, &|| Ok(())).unwrap();
assert!(cache.cached("directory").is_some());
assert_eq!(files.lists.load(Ordering::SeqCst), 2);
});
}
}
#[test]
fn cancelled_rebuild_never_replaces_previous_snapshot() {
let files = Files::new();
let cache = IndexCache::default();
let previous = cache.index("directory", &files, false, &|| Ok(())).unwrap();
let checks = AtomicUsize::new(0);
let result = cache.index("directory", &files, true, &|| {
if checks.fetch_add(1, Ordering::SeqCst) >= 3 {
Err("cancelled".into())
} else {
Ok(())
}
});
assert!(matches!(result, Err(error) if error == "cancelled"));
assert!(Arc::ptr_eq(&previous, &cache.cached("directory").unwrap()));
}
}
+171 -29
View File
@@ -64,12 +64,7 @@ fn read_flac(
reader.read_exact(&mut header).map_err(|e| e.to_string())?;
let length = u32::from_be_bytes([0, header[1], header[2], header[3]]);
if header[0] & 0x7f == 4 {
let mut payload = vec![0; length as usize];
for chunk in payload.chunks_mut(64 * 1024) {
check()?;
reader.read_exact(chunk).map_err(|e| e.to_string())?;
}
return Ok(comment_isrc(&payload, check));
return comment_isrc(reader, length, check);
}
if header[0] & 0x80 != 0 {
return Ok(String::new());
@@ -80,32 +75,71 @@ fn read_flac(
}
}
fn comment_isrc(mut payload: &[u8], check: &dyn Fn() -> Result<(), String>) -> String {
fn integer(payload: &mut &[u8]) -> Option<usize> {
let (value, rest) = payload.split_at_checked(4)?;
*payload = rest;
Some(u32::from_le_bytes(value.try_into().ok()?) as usize)
fn comment_isrc(
reader: &mut (impl Read + Seek),
mut remaining: u32,
check: &dyn Fn() -> Result<(), String>,
) -> Result<String, String> {
fn integer(reader: &mut impl Read, remaining: &mut u32) -> Result<u32, String> {
*remaining = remaining.checked_sub(4).ok_or("truncated FLAC comment")?;
let mut bytes = [0; 4];
reader.read_exact(&mut bytes).map_err(|e| e.to_string())?;
Ok(u32::from_le_bytes(bytes))
}
let mut parse = || {
let vendor = integer(&mut payload)?;
payload = payload.get(vendor..)?;
let count = integer(&mut payload)?;
for _ in 0..count {
check().ok()?;
let length = integer(&mut payload)?;
let (comment, rest) = payload.split_at_checked(length)?;
payload = rest;
let Some(equal) = comment.iter().position(|byte| *byte == b'=') else {
continue;
};
let key = std::str::from_utf8(&comment[..equal]).ok();
if key.is_some_and(isrc_key) {
return Some(crate::text::utf8(&comment[equal + 1..]).trim().to_owned());
// The old reader rejected an incomplete block even when its first field
// contained ISRC. Validate its extent before skipping irrelevant bytes.
check()?;
let start = reader.stream_position().map_err(|e| e.to_string())?;
let size = reader.seek(SeekFrom::End(0)).map_err(|e| e.to_string())?;
if size.saturating_sub(start) < u64::from(remaining) {
return Err("truncated FLAC comment block".into());
}
reader
.seek(SeekFrom::Start(start))
.map_err(|e| e.to_string())?;
let vendor = integer(reader, &mut remaining)?;
remaining = remaining
.checked_sub(vendor)
.ok_or("truncated FLAC vendor")?;
reader
.seek(SeekFrom::Current(i64::from(vendor)))
.map_err(|e| e.to_string())?;
let count = integer(reader, &mut remaining)?;
for _ in 0..count {
check()?;
let length = integer(reader, &mut remaining)?;
remaining = remaining
.checked_sub(length)
.ok_or("truncated FLAC field")?;
// Accepted ISRC keys are four ASCII characters, except that long s
// may replace S. Six bytes cover the longest accepted key and '='.
let mut prefix = [0; 6];
let prefix_len = (length as usize).min(prefix.len());
reader
.read_exact(&mut prefix[..prefix_len])
.map_err(|e| e.to_string())?;
if let Some(equal) = prefix[..prefix_len].iter().position(|byte| *byte == b'=')
&& std::str::from_utf8(&prefix[..equal])
.ok()
.is_some_and(isrc_key)
{
let mut value = Vec::with_capacity(length as usize - equal - 1);
value.extend_from_slice(&prefix[equal + 1..prefix_len]);
let tail = length as usize - prefix_len;
value.resize(value.len() + tail, 0);
let offset = prefix_len - equal - 1;
for chunk in value[offset..].chunks_mut(64 * 1024) {
check()?;
reader.read_exact(chunk).map_err(|e| e.to_string())?;
}
return Ok(crate::text::utf8(&value).trim().to_owned());
}
None
};
parse().unwrap_or_default()
reader
.seek(SeekFrom::Current(i64::from(length) - prefix_len as i64))
.map_err(|e| e.to_string())?;
}
Ok(String::new())
}
fn isrc_key(key: &str) -> bool {
@@ -115,3 +149,111 @@ fn isrc_key(key: &str) -> bool {
actual.eq_ignore_ascii_case(&expected) || (actual == 'ſ' && expected == 'S')
})
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Cursor;
use std::sync::atomic::{AtomicUsize, Ordering};
fn flac(comments: &[&[u8]]) -> Vec<u8> {
let mut payload = 7_u32.to_le_bytes().to_vec();
payload.extend_from_slice(b"Example");
payload.extend_from_slice(&(comments.len() as u32).to_le_bytes());
for comment in comments {
payload.extend_from_slice(&(comment.len() as u32).to_le_bytes());
payload.extend_from_slice(comment);
}
let mut bytes = b"fLaC\x84".to_vec();
bytes.extend_from_slice(&(payload.len() as u32).to_be_bytes()[1..]);
bytes.extend(payload);
bytes
}
#[test]
fn streaming_comments_preserve_first_isrc_alias_and_utf8_contracts() {
for (comments, expected) in [
(
vec![
b"TITLE=Example".as_slice(),
b"ISRC= example ",
b"ISRC=second",
],
"example",
),
(vec![b"ISRC=".as_slice(), b"ISRC=second"], ""),
(vec!["IſRC=long-s".as_bytes()], "long-s"),
(vec!["ıSRC=wrong".as_bytes(), b"isrc=correct"], "correct"),
(vec![b"ISRC=\xffx".as_slice()], "\u{fffd}x"),
(
vec![b"ISRCX=wrong".as_slice(), b"LONGKEYISRC=wrong", b"ISRC"],
"",
),
] {
assert_eq!(
read_isrc(&mut Cursor::new(flac(&comments)), "flac", &|| Ok(())).unwrap(),
expected
);
}
}
#[test]
fn incomplete_block_and_malformed_field_lengths_remain_empty() {
let bytes = flac(&[b"ISRC=example", b"TITLE=Example"]);
for end in 0..bytes.len() {
assert_eq!(
read_isrc(&mut Cursor::new(&bytes[..end]), "flac", &|| Ok(())).unwrap(),
""
);
}
for offset in [8, 23] {
let mut invalid = bytes.clone();
invalid[offset..offset + 4].copy_from_slice(&u32::MAX.to_le_bytes());
assert_eq!(
read_isrc(&mut Cursor::new(invalid), "flac", &|| Ok(())).unwrap(),
""
);
}
}
#[test]
fn irrelevant_large_fields_are_skipped_and_reads_observe_cancellation() {
struct Counted {
source: Cursor<Vec<u8>>,
bytes: usize,
}
impl Read for Counted {
fn read(&mut self, bytes: &mut [u8]) -> std::io::Result<usize> {
let count = self.source.read(bytes)?;
self.bytes += count;
Ok(count)
}
}
impl Seek for Counted {
fn seek(&mut self, from: SeekFrom) -> std::io::Result<u64> {
self.source.seek(from)
}
}
let unrelated = vec![b'x'; 2 << 20];
let mut source = Counted {
source: Cursor::new(flac(&[&unrelated, b"ISRC=example"])),
bytes: 0,
};
assert_eq!(
read_isrc(&mut source, "flac", &|| Ok(())).unwrap(),
"example"
);
assert!(source.bytes < 128, "read {} irrelevant bytes", source.bytes);
let checks = AtomicUsize::new(0);
assert_eq!(
read_isrc(&mut source, "flac", &|| {
if checks.fetch_add(1, Ordering::SeqCst) >= 3 {
Err("cancelled".into())
} else {
Ok(())
}
}),
Err("cancelled".into())
);
}
}
@@ -225,6 +225,10 @@ impl Backend {
.cleanup_connections()
.map_err(|error| error.to_string())?;
if under_pressure {
self.manager
.environment()
.clear_isrc_cache()
.map_err(|error| error.to_string())?;
self.cover.clear();
self.lyrics
.drop_memory()
@@ -44,4 +44,10 @@ impl ExtensionEnvironment {
self.isrc.invalidate(directory);
Ok(())
}
pub(crate) fn clear_isrc_cache(&self) -> Result<(), EnvironmentError> {
let _operation = self.enter()?;
self.isrc.clear();
Ok(())
}
}