Skip to main content

max / audiofiles

8.5 KB · 200 lines History Blame Raw
1 //! Background search worker: builds and queries the similarity / near-duplicate
2 //! VP-trees off the GUI thread.
3 //!
4 //! Building a VP-tree over the whole library is CPU-bound (the ultra-fuzz audit
5 //! flagged the lazy first-query build freezing the egui frame on large
6 //! libraries). This worker owns both indexes and its own read-only `Database`
7 //! connection (WAL lets it read alongside the main connection), so the build and
8 //! query happen entirely off the render thread. Results return via [`SearchEvent`]
9 //! and are applied in `poll_workers`.
10 //!
11 //! The indexes are cached across queries and dropped on [`SearchCommand::Invalidate`]
12 //! (sent when new analysis lands), so they rebuild from fresh data on the next query.
13
14 use std::path::PathBuf;
15
16 use audiofiles_core::db::Database;
17 use audiofiles_core::worker_runtime::{WorkerCtx, WorkerHandle, spawn_worker};
18 use audiofiles_core::{fingerprint, similarity};
19
20 /// Command sent from the GUI thread to the search worker.
21 pub enum SearchCommand {
22 /// Find samples similar to `hash` (builds the similarity index if needed).
23 FindSimilar { hash: String, limit: usize },
24 /// Find near-duplicates of `hash` by fingerprint (builds the index if needed).
25 FindNearDuplicates { hash: String, limit: usize },
26 /// Drop both cached indexes so the next query rebuilds from fresh data.
27 Invalidate,
28 }
29
30 /// Event sent from the search worker back to the GUI thread.
31 pub enum SearchEvent {
32 /// Ranked similar-sample hashes for `source`.
33 SimilarResults { source: String, hashes: Vec<String> },
34 /// Ranked near-duplicate hashes for `source`.
35 NearDuplicateResults { source: String, hashes: Vec<String> },
36 /// A search failed (DB open / load / build error).
37 Error { error: String },
38 }
39
40 /// Handle for communicating with the background search worker. Search has no
41 /// Cancel command, so the runtime handle is used directly.
42 pub type SearchWorkerHandle = WorkerHandle<SearchCommand, SearchEvent>;
43
44 /// Persistent worker state: the read-only DB connection and the two cached
45 /// VP-tree indexes, reused across queries and dropped on `Invalidate`.
46 ///
47 /// Eviction policy: each index is bounded by library size (one VP-tree node per
48 /// analyzed sample, order ~100 MB at a million samples), and **at most one is
49 /// resident at a time**. A similarity query evicts the fingerprint index and vice
50 /// versa, so peak index memory is a single tree, not both. The evicted index
51 /// rebuilds on its next query from fresh data, so results are never truncated,
52 /// only the cache is. Combined with `Invalidate` (on new analysis), this is the
53 /// ceiling: ≤ one library-sized tree held, released when the other kind of search
54 /// runs or when analysis lands.
55 struct SearchState {
56 db: Database,
57 similarity_index: Option<similarity::SimilarityIndex>,
58 fingerprint_index: Option<fingerprint::FingerprintIndex>,
59 }
60
61 /// Spawn the background search worker. It opens its own read-only connection to
62 /// the database at `db_path` on the worker thread.
63 pub fn spawn_search_worker(db_path: PathBuf) -> std::io::Result<SearchWorkerHandle> {
64 spawn_worker(
65 "search-worker",
66 move || {
67 Database::open(&db_path).map(|db| SearchState {
68 db,
69 similarity_index: None,
70 fingerprint_index: None,
71 })
72 },
73 |e| SearchEvent::Error {
74 error: e.to_string(),
75 },
76 |_state| SearchEvent::Error {
77 error: "search worker panicked (internal error)".to_string(),
78 },
79 search_step,
80 )
81 }
82
83 fn search_step(state: &mut SearchState, cmd: SearchCommand, ctx: &WorkerCtx<SearchEvent>) {
84 match cmd {
85 SearchCommand::Invalidate => {
86 state.similarity_index = None;
87 state.fingerprint_index = None;
88 }
89 SearchCommand::FindSimilar { hash, limit } => {
90 // A malformed feature row that panics is caught by the runtime and
91 // reported as Error; the worker survives for the next query.
92 // Evict the other index: at most one VP-tree resident at a time.
93 state.fingerprint_index = None;
94 let result: Result<Vec<String>, audiofiles_core::error::CoreError> = (|| {
95 if state.similarity_index.is_none() {
96 let data = similarity::SimilarityIndex::load_data(&state.db)?;
97 state.similarity_index =
98 Some(similarity::SimilarityIndex::build_from_data(data));
99 }
100 let built = state.similarity_index.as_ref().expect("set above");
101 let features = similarity::load_features(&state.db, &hash)?;
102 Ok(built
103 .find_similar(&hash, &features, limit)
104 .into_iter()
105 .map(|r| r.hash)
106 .collect())
107 })();
108 ctx.emit(match result {
109 Ok(hashes) => SearchEvent::SimilarResults {
110 source: hash,
111 hashes,
112 },
113 Err(e) => SearchEvent::Error {
114 error: e.to_string(),
115 },
116 });
117 }
118 SearchCommand::FindNearDuplicates { hash, limit } => {
119 // Evict the other index: at most one VP-tree resident at a time.
120 state.similarity_index = None;
121 let result: Result<Vec<String>, audiofiles_core::error::CoreError> = (|| {
122 if state.fingerprint_index.is_none() {
123 let entries = fingerprint::FingerprintIndex::load_data(&state.db)?;
124 state.fingerprint_index =
125 Some(fingerprint::FingerprintIndex::build_from_data(entries));
126 }
127 let built = state.fingerprint_index.as_ref().expect("set above");
128 let reference = fingerprint::load_fingerprint(&state.db, &hash)?;
129 Ok(built
130 .find_near_duplicates(&hash, &reference.envelope, limit)
131 .into_iter()
132 .map(|r| r.hash)
133 .collect())
134 })();
135 ctx.emit(match result {
136 Ok(hashes) => SearchEvent::NearDuplicateResults {
137 source: hash,
138 hashes,
139 },
140 Err(e) => SearchEvent::Error {
141 error: e.to_string(),
142 },
143 });
144 }
145 }
146 }
147
148 #[cfg(test)]
149 mod tests {
150 use super::*;
151
152 /// Block (with a generous wall-clock bound) for one event so a regression
153 /// that wedges the worker fails the test instead of hanging the suite. The
154 /// bound is time-based, not iteration-based, so it stays reliable under the
155 /// load of the full parallel test run.
156 fn recv_one(handle: &SearchWorkerHandle) -> SearchEvent {
157 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
158 while std::time::Instant::now() < deadline {
159 if let Some(ev) = handle.try_recv() {
160 return ev;
161 }
162 std::thread::sleep(std::time::Duration::from_millis(1));
163 }
164 panic!("search worker produced no event");
165 }
166
167 #[test]
168 fn worker_responds_on_empty_db() {
169 let dir = tempfile::tempdir().unwrap();
170 let db_path = dir.path().join("audiofiles.db");
171 // Create + migrate the file so the worker opens a valid (empty) schema.
172 let _ = Database::open(&db_path).unwrap();
173
174 let handle = spawn_search_worker(db_path).unwrap();
175 // A query against an empty library must still produce an event (empty
176 // results or a clean error), never a hang or a panic that kills the worker.
177 assert!(handle.send(SearchCommand::FindSimilar {
178 hash: "deadbeef".to_string(),
179 limit: 10,
180 }));
181 match recv_one(&handle) {
182 SearchEvent::SimilarResults { source, hashes } => {
183 assert_eq!(source, "deadbeef");
184 assert!(hashes.is_empty());
185 }
186 SearchEvent::Error { .. } => {} // acceptable: no features for the hash
187 SearchEvent::NearDuplicateResults { .. } => panic!("wrong event kind"),
188 }
189
190 // Invalidate + a second query still respond (worker survived).
191 assert!(handle.send(SearchCommand::Invalidate));
192 assert!(handle.send(SearchCommand::FindNearDuplicates {
193 hash: "deadbeef".to_string(),
194 limit: 10,
195 }));
196 let _ = recv_one(&handle);
197 drop(handle); // Drop joins the thread; a wedged worker would hang here.
198 }
199 }
200