eredu_runtime/cache/
worker.rs

1//! Backend-neutral bounded physical worker for cache backing-store tasks.
2
3use std::{
4    collections::HashMap,
5    panic::{catch_unwind, AssertUnwindSafe},
6    sync::{mpsc, Arc, Condvar, Mutex},
7    thread::{self, JoinHandle},
8};
9
10use super::{
11    CacheIoAdmission, CacheIoCompletionDisposition, CacheIoExecutionState,
12    CacheIoExecutionStateError, CacheIoOperationKey, CacheIoPreparation, CacheIoStartDisposition,
13};
14
15enum CacheIoWorkerRequest<Task, Output> {
16    Operation {
17        key: CacheIoOperationKey,
18        task: Box<Task>,
19        completion: Arc<CacheIoCompletion<Output>>,
20    },
21    Stop,
22}
23
24#[derive(Debug, Clone)]
25enum CacheIoCompletionState<Output> {
26    Finished(Result<Output, String>),
27    Cancelled,
28}
29
30#[derive(Debug)]
31struct CacheIoCompletion<Output> {
32    state: Mutex<Option<CacheIoCompletionState<Output>>>,
33    ready: Condvar,
34    released: Mutex<bool>,
35    released_ready: Condvar,
36}
37
38impl<Output> Default for CacheIoCompletion<Output> {
39    fn default() -> Self {
40        Self {
41            state: Mutex::new(None),
42            ready: Condvar::new(),
43            released: Mutex::new(false),
44            released_ready: Condvar::new(),
45        }
46    }
47}
48
49impl<Output> CacheIoCompletion<Output> {
50    fn finish(&self, result: Result<Output, String>) {
51        if let Ok(mut state) = self.state.lock() {
52            if state.is_none() {
53                *state = Some(CacheIoCompletionState::Finished(result));
54                self.ready.notify_all();
55            }
56        }
57    }
58
59    fn cancel(&self) -> bool {
60        let Ok(mut state) = self.state.lock() else {
61            return false;
62        };
63        if state.is_some() {
64            return false;
65        }
66        *state = Some(CacheIoCompletionState::Cancelled);
67        self.ready.notify_all();
68        true
69    }
70
71    fn is_ready(&self) -> bool {
72        self.state.lock().map_or(true, |state| state.is_some())
73    }
74
75    fn release_task_resources(&self) {
76        if let Ok(mut released) = self.released.lock() {
77            *released = true;
78            self.released_ready.notify_all();
79        }
80    }
81
82    fn wait_for_task_resources(&self) -> Result<(), CacheIoWorkerError> {
83        let mut released = self
84            .released
85            .lock()
86            .map_err(|_| CacheIoWorkerError::Poisoned)?;
87        while !*released {
88            released = self
89                .released_ready
90                .wait(released)
91                .map_err(|_| CacheIoWorkerError::Poisoned)?;
92        }
93        Ok(())
94    }
95}
96
97impl<Output: Clone> CacheIoCompletion<Output> {
98    fn wait(&self, generation: u64) -> Result<Output, CacheIoWorkerError> {
99        let mut state = self
100            .state
101            .lock()
102            .map_err(|_| CacheIoWorkerError::Poisoned)?;
103        while state.is_none() {
104            state = self
105                .ready
106                .wait(state)
107                .map_err(|_| CacheIoWorkerError::Poisoned)?;
108        }
109        match state.as_ref().expect("completion state was awaited") {
110            CacheIoCompletionState::Finished(Ok(output)) => Ok(output.clone()),
111            CacheIoCompletionState::Finished(Err(error)) => {
112                Err(CacheIoWorkerError::OperationFailed(error.clone()))
113            }
114            CacheIoCompletionState::Cancelled => Err(CacheIoWorkerError::Cancelled { generation }),
115        }
116    }
117}
118
119#[derive(Debug)]
120struct CacheIoWorkerShared<Output> {
121    in_flight: Mutex<HashMap<CacheIoOperationKey, Arc<CacheIoCompletion<Output>>>>,
122    execution: Mutex<CacheIoExecutionState>,
123    space_available: Condvar,
124}
125
126impl<Output> CacheIoWorkerShared<Output> {
127    fn new(capacity: usize) -> Result<Self, CacheIoWorkerError> {
128        Ok(Self {
129            in_flight: Mutex::new(HashMap::new()),
130            execution: Mutex::new(CacheIoExecutionState::new(capacity)?),
131            space_available: Condvar::new(),
132        })
133    }
134}
135
136/// Exact completion ownership for one coalesced cache I/O operation.
137pub struct CacheIoTicket<Output> {
138    /// Exact logical operation identity.
139    pub key: CacheIoOperationKey,
140    completion: Arc<CacheIoCompletion<Output>>,
141    shared: Arc<CacheIoWorkerShared<Output>>,
142}
143
144impl<Output> Clone for CacheIoTicket<Output> {
145    fn clone(&self) -> Self {
146        Self {
147            key: self.key.clone(),
148            completion: Arc::clone(&self.completion),
149            shared: Arc::clone(&self.shared),
150        }
151    }
152}
153
154impl<Output> std::fmt::Debug for CacheIoTicket<Output> {
155    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
156        formatter
157            .debug_struct("CacheIoTicket")
158            .field("key", &self.key)
159            .finish_non_exhaustive()
160    }
161}
162
163impl<Output: Clone> CacheIoTicket<Output> {
164    /// Waits for the logical output or cancellation.
165    pub fn wait(&self) -> Result<Output, CacheIoWorkerError> {
166        self.completion.wait(self.key.generation)
167    }
168
169    /// Cancels prepared, queued, or in-flight work exactly once.
170    pub fn cancel(&self) -> bool {
171        let Ok(mut execution) = self.shared.execution.lock() else {
172            return false;
173        };
174        let cancelled = execution.cancel(&self.key) && self.completion.cancel();
175        self.shared.space_available.notify_all();
176        cancelled
177    }
178
179    /// Waits until all backend task inputs and retained resources are dropped.
180    pub fn wait_for_task_resources(&self) -> Result<(), CacheIoWorkerError> {
181        self.completion.wait_for_task_resources()
182    }
183
184    /// Returns whether two tickets join the same exact completion owner.
185    pub fn shares_completion_with(&self, other: &Self) -> bool {
186        Arc::ptr_eq(&self.completion, &other.completion)
187    }
188}
189
190/// Prepared cache I/O that admits physical work only when explicitly enqueued.
191pub struct CacheIoSubmission<Task, Output> {
192    /// Ticket shared by the operation owner and all exact-key joiners.
193    pub ticket: CacheIoTicket<Output>,
194    sender: mpsc::Sender<CacheIoWorkerRequest<Task, Output>>,
195    shared: Arc<CacheIoWorkerShared<Output>>,
196    unsent: Option<CacheIoWorkerRequest<Task, Output>>,
197    joined_task: Option<Task>,
198    /// Whether this submission joined an already prepared exact operation.
199    pub joined: bool,
200}
201
202/// Physical admission observations for one submission.
203#[derive(Debug, Clone, Copy, Eq, PartialEq)]
204pub struct CacheIoSubmissionOutcome {
205    /// Whether this submission joined an existing exact operation.
206    pub joined: bool,
207    /// Whether finite queue capacity delayed physical admission.
208    pub backpressure: bool,
209    /// Largest observed physical queue occupancy.
210    pub peak_occupancy: usize,
211}
212
213impl<Task, Output: Clone> CacheIoSubmission<Task, Output> {
214    /// Returns the unused backend task when this submission joined existing work.
215    ///
216    /// Backends may disarm task-local rollback guards before the unused task is
217    /// dropped; the task is never physically executed.
218    pub fn joined_task_mut(&mut self) -> Option<&mut Task> {
219        self.joined_task.as_mut()
220    }
221
222    /// Admits this prepared task, blocking only on finite queue capacity.
223    pub fn enqueue(mut self) -> Result<CacheIoSubmissionOutcome, CacheIoWorkerError> {
224        let mut backpressure = false;
225        if let Some(request) = self.unsent.take() {
226            let mut execution = match self.shared.execution.lock() {
227                Ok(execution) => execution,
228                Err(_) => {
229                    drop(request);
230                    self.ticket.completion.release_task_resources();
231                    return Err(CacheIoWorkerError::Poisoned);
232                }
233            };
234            loop {
235                match execution.admit(&self.ticket.key)? {
236                    CacheIoAdmission::Admitted => {
237                        if self.sender.send(request).is_err() {
238                            execution.rollback_admission(&self.ticket.key)?;
239                            self.ticket
240                                .completion
241                                .finish(Err("cache I/O physical worker stopped".into()));
242                            self.ticket.completion.release_task_resources();
243                        }
244                        break;
245                    }
246                    CacheIoAdmission::AtCapacity => {
247                        backpressure = true;
248                        execution = match self.shared.space_available.wait(execution) {
249                            Ok(execution) => execution,
250                            Err(_) => {
251                                drop(request);
252                                self.ticket.completion.release_task_resources();
253                                return Err(CacheIoWorkerError::Poisoned);
254                            }
255                        };
256                    }
257                    CacheIoAdmission::Cancelled => {
258                        drop(request);
259                        self.ticket.completion.release_task_resources();
260                        break;
261                    }
262                }
263            }
264            drop(execution);
265        }
266        Ok(CacheIoSubmissionOutcome {
267            joined: self.joined,
268            backpressure,
269            peak_occupancy: self
270                .shared
271                .execution
272                .lock()
273                .map_err(|_| CacheIoWorkerError::Poisoned)?
274                .peak_queued(),
275        })
276    }
277}
278
279impl<Task, Output> Drop for CacheIoSubmission<Task, Output> {
280    fn drop(&mut self) {
281        let Some(request) = self.unsent.take() else {
282            return;
283        };
284        if let Ok(mut execution) = self.shared.execution.lock() {
285            execution.cancel(&self.ticket.key);
286        }
287        drop(request);
288        self.ticket.completion.release_task_resources();
289        retire_completion(&self.shared, &self.ticket.key, &self.ticket.completion);
290    }
291}
292
293/// Generic bounded background worker over opaque backend task and output types.
294pub struct CacheIoWorker<Task, Output> {
295    sender: mpsc::Sender<CacheIoWorkerRequest<Task, Output>>,
296    handle: Mutex<Option<JoinHandle<()>>>,
297    shared: Arc<CacheIoWorkerShared<Output>>,
298}
299
300impl<Task, Output> std::fmt::Debug for CacheIoWorker<Task, Output> {
301    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
302        formatter
303            .debug_struct("CacheIoWorker")
304            .finish_non_exhaustive()
305    }
306}
307
308impl<Task, Output> CacheIoWorker<Task, Output>
309where
310    Task: Send + 'static,
311    Output: Clone + Send + 'static,
312{
313    /// Starts a bounded worker using statically dispatched task and cleanup functions.
314    pub fn new(
315        capacity: usize,
316        thread_name: impl Into<String>,
317        execute: fn(Task) -> Result<Output, String>,
318        discard: fn(Output),
319    ) -> Result<Self, CacheIoWorkerError> {
320        let thread_name = thread_name.into();
321        let (sender, receiver) = mpsc::channel::<CacheIoWorkerRequest<Task, Output>>();
322        let shared = Arc::new(CacheIoWorkerShared::new(capacity)?);
323        let worker_shared = Arc::clone(&shared);
324        let handle = thread::Builder::new()
325            .name(thread_name.clone())
326            .spawn(move || {
327                while let Ok(request) = receiver.recv() {
328                    match request {
329                        CacheIoWorkerRequest::Operation {
330                            key,
331                            task,
332                            completion,
333                        } => {
334                            let start = worker_shared
335                                .execution
336                                .lock()
337                                .map_err(|_| CacheIoWorkerError::Poisoned)
338                                .and_then(|mut execution| {
339                                    execution.begin(&key).map_err(Into::into)
340                                });
341                            worker_shared.space_available.notify_all();
342                            match start {
343                                Ok(CacheIoStartDisposition::Execute) => {}
344                                Ok(CacheIoStartDisposition::Discard) => {
345                                    drop(task);
346                                    completion.release_task_resources();
347                                    retire_completion(&worker_shared, &key, &completion);
348                                    continue;
349                                }
350                                Err(error) => {
351                                    drop(task);
352                                    completion.finish(Err(error.to_string()));
353                                    completion.release_task_resources();
354                                    retire_completion(&worker_shared, &key, &completion);
355                                    continue;
356                                }
357                            }
358                            let result = catch_unwind(AssertUnwindSafe(|| execute(*task)))
359                                .unwrap_or_else(|_| {
360                                    Err("cache I/O physical worker operation panicked".into())
361                                });
362                            let disposition = worker_shared
363                                .execution
364                                .lock()
365                                .map_err(|_| CacheIoWorkerError::Poisoned)
366                                .and_then(|mut execution| {
367                                    execution.complete(&key).map_err(Into::into)
368                                });
369                            if !matches!(disposition, Ok(CacheIoCompletionDisposition::Publish))
370                                || completion.is_ready()
371                            {
372                                if let Ok(output) = result {
373                                    discard(output);
374                                }
375                            } else {
376                                completion.finish(result);
377                            }
378                            completion.release_task_resources();
379                            retire_completion(&worker_shared, &key, &completion);
380                        }
381                        CacheIoWorkerRequest::Stop => break,
382                    }
383                }
384            })
385            .map_err(|source| CacheIoWorkerError::Spawn {
386                thread_name,
387                source,
388            })?;
389        Ok(Self {
390            sender,
391            handle: Mutex::new(Some(handle)),
392            shared,
393        })
394    }
395
396    /// Prepares new work or joins an exact operation already owned by the worker.
397    pub fn prepare(
398        &self,
399        key: CacheIoOperationKey,
400        task: Task,
401    ) -> Result<CacheIoSubmission<Task, Output>, CacheIoWorkerError> {
402        let mut execution = self
403            .shared
404            .execution
405            .lock()
406            .map_err(|_| CacheIoWorkerError::Poisoned)?;
407        let preparation = execution.prepare(key.clone());
408        let mut completions = self
409            .shared
410            .in_flight
411            .lock()
412            .map_err(|_| CacheIoWorkerError::Poisoned)?;
413        if preparation == CacheIoPreparation::Joined {
414            let completion = completions
415                .get(&key)
416                .expect("runtime joined key has an exact completion");
417            return Ok(CacheIoSubmission {
418                ticket: CacheIoTicket {
419                    key,
420                    completion: Arc::clone(completion),
421                    shared: Arc::clone(&self.shared),
422                },
423                sender: self.sender.clone(),
424                shared: Arc::clone(&self.shared),
425                unsent: None,
426                joined_task: Some(task),
427                joined: true,
428            });
429        }
430        let completion = Arc::new(CacheIoCompletion::default());
431        completions.insert(key.clone(), Arc::clone(&completion));
432        drop(completions);
433        drop(execution);
434        let request = CacheIoWorkerRequest::Operation {
435            key: key.clone(),
436            task: Box::new(task),
437            completion: Arc::clone(&completion),
438        };
439        Ok(CacheIoSubmission {
440            ticket: CacheIoTicket {
441                key,
442                completion,
443                shared: Arc::clone(&self.shared),
444            },
445            sender: self.sender.clone(),
446            shared: Arc::clone(&self.shared),
447            unsent: Some(request),
448            joined_task: None,
449            joined: false,
450        })
451    }
452
453    /// Releases exact-key ownership after task resources are safe to drop.
454    pub fn retire(&self, ticket: &CacheIoTicket<Output>) {
455        retire_completion(&self.shared, &ticket.key, &ticket.completion);
456    }
457}
458
459impl<Task, Output> Drop for CacheIoWorker<Task, Output> {
460    fn drop(&mut self) {
461        let _ = self.sender.send(CacheIoWorkerRequest::Stop);
462        if let Ok(handle) = self.handle.get_mut() {
463            if let Some(handle) = handle.take() {
464                let _ = handle.join();
465            }
466        }
467    }
468}
469
470fn retire_completion<Output>(
471    shared: &CacheIoWorkerShared<Output>,
472    key: &CacheIoOperationKey,
473    completion: &Arc<CacheIoCompletion<Output>>,
474) {
475    let retired = if let Ok(mut execution) = shared.execution.lock() {
476        execution.retire(key).unwrap_or(false)
477    } else {
478        false
479    };
480    if retired {
481        shared.space_available.notify_all();
482        if let Ok(mut in_flight) = shared.in_flight.lock() {
483            if in_flight
484                .get(key)
485                .is_some_and(|current| Arc::ptr_eq(current, completion))
486            {
487                in_flight.remove(key);
488            }
489        }
490    }
491}
492
493/// Failure in generic cache I/O worker coordination or task execution.
494#[derive(Debug, thiserror::Error)]
495pub enum CacheIoWorkerError {
496    /// The worker's synchronization state was poisoned.
497    #[error("cache I/O worker synchronization state is poisoned")]
498    Poisoned,
499    /// A task returned a backend-specific failure string.
500    #[error("cache I/O operation failed: {0}")]
501    OperationFailed(String),
502    /// Cancellation won for this generation.
503    #[error("cache I/O operation was cancelled for generation {generation}")]
504    Cancelled {
505        /// Cancelled model/cache generation.
506        generation: u64,
507    },
508    /// The physical worker thread could not be started.
509    #[error("failed to start cache I/O worker {thread_name}: {source}")]
510    Spawn {
511        /// Requested worker thread name.
512        thread_name: String,
513        /// Underlying thread creation failure.
514        #[source]
515        source: std::io::Error,
516    },
517    /// The exact admission/cancellation state transition was invalid.
518    #[error(transparent)]
519    Execution(#[from] CacheIoExecutionStateError),
520}
521
522#[cfg(test)]
523mod tests {
524    use super::*;
525    use crate::cache::CacheIoOperationKind;
526    use eredu_core::cache::{CacheBlockId, CacheRepresentation};
527    use std::{sync::mpsc, time::Duration};
528
529    enum Task {
530        Value(u64),
531        Pause(mpsc::Sender<()>, mpsc::Receiver<()>),
532        Panic,
533    }
534
535    fn execute(task: Task) -> Result<u64, String> {
536        match task {
537            Task::Value(value) => Ok(value),
538            Task::Pause(started, release) => {
539                let _ = started.send(());
540                let _ = release.recv();
541                Ok(0)
542            }
543            Task::Panic => panic!("injected worker panic"),
544        }
545    }
546
547    fn discard(_value: u64) {}
548
549    fn key(block: i64) -> CacheIoOperationKey {
550        CacheIoOperationKey {
551            generation: 7,
552            id: CacheBlockId {
553                session_id: 1,
554                global_layer: 0,
555                representation: CacheRepresentation::KeyValue,
556                start: block,
557                end: block + 1,
558                rank: None,
559            },
560            kind: CacheIoOperationKind::Read,
561        }
562    }
563
564    #[test]
565    fn worker_coalesces_and_contains_task_panics() {
566        let worker = CacheIoWorker::new(1, "cache-worker-test", execute, discard).unwrap();
567        let first = worker.prepare(key(0), Task::Value(9)).unwrap();
568        let first_ticket = first.ticket.clone();
569        let joined = worker.prepare(key(0), Task::Value(10)).unwrap();
570        let joined_ticket = joined.ticket.clone();
571        assert!(joined.joined);
572        first.enqueue().unwrap();
573        joined.enqueue().unwrap();
574        assert_eq!(first_ticket.wait().unwrap(), 9);
575        assert_eq!(joined_ticket.wait().unwrap(), 9);
576        assert!(first_ticket.shares_completion_with(&joined_ticket));
577        worker.retire(&first_ticket);
578
579        let panicking = worker.prepare(key(1), Task::Panic).unwrap();
580        let ticket = panicking.ticket.clone();
581        panicking.enqueue().unwrap();
582        assert!(matches!(
583            ticket.wait(),
584            Err(CacheIoWorkerError::OperationFailed(message))
585                if message.contains("operation panicked")
586        ));
587        worker.retire(&ticket);
588    }
589
590    #[test]
591    fn cancellation_wakes_a_backpressured_submission() {
592        let worker =
593            Arc::new(CacheIoWorker::new(1, "cache-worker-cancel-test", execute, discard).unwrap());
594        let (started_tx, started_rx) = mpsc::channel();
595        let (release_tx, release_rx) = mpsc::channel();
596        let blocker = worker
597            .prepare(key(0), Task::Pause(started_tx, release_rx))
598            .unwrap();
599        let blocker_ticket = blocker.ticket.clone();
600        blocker.enqueue().unwrap();
601        started_rx.recv().unwrap();
602
603        let queued = worker.prepare(key(1), Task::Value(1)).unwrap();
604        queued.enqueue().unwrap();
605        let blocked = worker.prepare(key(2), Task::Value(2)).unwrap();
606        let blocked_ticket = blocked.ticket.clone();
607        let (outcome_tx, outcome_rx) = mpsc::channel();
608        let enqueue = std::thread::spawn(move || outcome_tx.send(blocked.enqueue()).unwrap());
609        assert!(outcome_rx.recv_timeout(Duration::from_millis(20)).is_err());
610        assert!(blocked_ticket.cancel());
611        assert!(
612            outcome_rx
613                .recv_timeout(Duration::from_secs(1))
614                .unwrap()
615                .unwrap()
616                .backpressure
617        );
618        enqueue.join().unwrap();
619        assert!(matches!(
620            blocked_ticket.wait(),
621            Err(CacheIoWorkerError::Cancelled { generation: 7 })
622        ));
623        release_tx.send(()).unwrap();
624        assert_eq!(blocker_ticket.wait().unwrap(), 0);
625    }
626}
eredu_runtime/cache/worker.rs

eredu_runtime/cache/
worker.rs
