Sync memory between emulator thread and UI

src/memory.rs

@@ -1,15 +1,18 @@
 use std::{
     collections::HashMap,
-    sync::{Arc, Mutex, MutexGuard},
+    fmt::Debug,
+    sync::{atomic::AtomicU64, Arc, RwLock, RwLockReadGuard, TryLockError, Weak},
 };
 
 use bytemuck::BoxBytes;
+use itertools::Itertools;
+use tracing::warn;
 
 use crate::emulator::{EmulatorClient, EmulatorCommand, SimId};
 
 pub struct MemoryMonitor {
     client: EmulatorClient,
-    regions: HashMap<MemoryRegion, Arc<Mutex<BoxBytes>>>,
+    regions: HashMap<MemoryRange, Weak<MemoryRegion>>,
 }
 
 impl MemoryMonitor {
@@ -21,20 +24,19 @@ impl MemoryMonitor {
     }
 
     pub fn view(&mut self, sim: SimId, start: u32, length: usize) -> MemoryView {
-        let region = MemoryRegion { sim, start, length };
-        let memory = self.regions.entry(region).or_insert_with(|| {
-            let mut buf = aligned_memory(start, length);
-            let (tx, rx) = oneshot::channel();
-            self.client
-                .send_command(EmulatorCommand::ReadMemory(sim, start, length, vec![], tx));
-            let bytes = pollster::block_on(rx).unwrap();
-            buf.copy_from_slice(&bytes);
-            #[expect(clippy::arc_with_non_send_sync)] // TODO: remove after bytemuck upgrade
-            Arc::new(Mutex::new(buf))
-        });
-        MemoryView {
-            memory: memory.clone(),
-        }
+        let range = MemoryRange { sim, start, length };
+        let region = self
+            .regions
+            .get(&range)
+            .and_then(|r| r.upgrade())
+            .unwrap_or_else(|| {
+                let region = Arc::new(MemoryRegion::new(start, length));
+                self.regions.insert(range, Arc::downgrade(&region));
+                self.client
+                    .send_command(EmulatorCommand::WatchMemory(range, Arc::downgrade(&region)));
+                region
+            });
+        MemoryView { region }
     }
 }
 
@@ -52,21 +54,17 @@ fn aligned_memory(start: u32, length: usize) -> BoxBytes {
 }
 
 pub struct MemoryView {
-    memory: Arc<Mutex<BoxBytes>>,
+    region: Arc<MemoryRegion>,
 }
-// SAFETY: BoxBytes is supposed to be Send+Sync, will be in a new version
-unsafe impl Send for MemoryView {}
 
 impl MemoryView {
     pub fn borrow(&self) -> MemoryRef<'_> {
-        MemoryRef {
-            inner: self.memory.lock().unwrap(),
-        }
+        self.region.borrow()
     }
 }
 
 pub struct MemoryRef<'a> {
-    inner: MutexGuard<'a, BoxBytes>,
+    inner: RwLockReadGuard<'a, BoxBytes>,
 }
 
 impl MemoryRef<'_> {
@@ -83,9 +81,102 @@ impl MemoryRef<'_> {
     }
 }
 
-#[derive(Clone, Copy, PartialEq, Eq, Hash)]
-struct MemoryRegion {
-    sim: SimId,
-    start: u32,
-    length: usize,
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct MemoryRange {
+    pub sim: SimId,
+    pub start: u32,
+    pub length: usize,
+}
+
+const BUFFERS: usize = 4;
+pub struct MemoryRegion {
+    gens: [AtomicU64; BUFFERS],
+    bufs: [RwLock<BoxBytes>; BUFFERS],
+}
+impl Debug for MemoryRegion {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("MemoryRegion")
+            .field("gens", &self.gens)
+            .finish_non_exhaustive()
+    }
+}
+// SAFETY: BoxBytes is meant to be Send+Sync, will be in a future version
+unsafe impl Send for MemoryRegion {}
+// SAFETY: BoxBytes is meant to be Send+Sync, will be in a future version
+unsafe impl Sync for MemoryRegion {}
+
+impl MemoryRegion {
+    fn new(start: u32, length: usize) -> Self {
+        Self {
+            gens: std::array::from_fn(|i| AtomicU64::new(i as u64)),
+            bufs: std::array::from_fn(|_| RwLock::new(aligned_memory(start, length))),
+        }
+    }
+
+    pub fn borrow(&self) -> MemoryRef<'_> {
+        /*
+         * When reading memory, a thread will grab the newest buffer (with the highest gen)
+         * It will only fail to grab the lock if the writer already has it,
+         * but the writer prioritizes older buffers (with lower gens).
+         * So this method will only block if the writer produces three full buffers
+         * in the time it takes the reader to do four atomic reads and grab a lock.
+         * In the unlikely event this happens... just try again.
+         */
+        loop {
+            let newest_index = self
+                .gens
+                .iter()
+                .map(|i| i.load(std::sync::atomic::Ordering::Acquire))
+                .enumerate()
+                .max_by_key(|(_, gen)| *gen)
+                .map(|(i, _)| i)
+                .unwrap();
+            let inner = match self.bufs[newest_index].try_read() {
+                Ok(inner) => inner,
+                Err(TryLockError::Poisoned(e)) => e.into_inner(),
+                Err(TryLockError::WouldBlock) => {
+                    continue;
+                }
+            };
+            break MemoryRef { inner };
+        }
+    }
+
+    pub fn update(&self, data: &[u8]) {
+        let gens: Vec<u64> = self
+            .gens
+            .iter()
+            .map(|i| i.load(std::sync::atomic::Ordering::Acquire))
+            .collect();
+        let next_gen = gens.iter().max().unwrap() + 1;
+        let indices = gens
+            .into_iter()
+            .enumerate()
+            .sorted_by_key(|(_, val)| *val)
+            .map(|(i, _)| i);
+        for index in indices {
+            let mut lock = match self.bufs[index].try_write() {
+                Ok(inner) => inner,
+                Err(TryLockError::Poisoned(e)) => e.into_inner(),
+                Err(TryLockError::WouldBlock) => {
+                    continue;
+                }
+            };
+            lock.copy_from_slice(data);
+            self.gens[index].store(next_gen, std::sync::atomic::Ordering::Release);
+            return;
+        }
+        /*
+         * We have four buffers, and (at time of writing) only three threads interacting with memory:
+         * - The UI thread, reading small regions of memory
+         * - The "vram renderer" thread, reading large regions of memory
+         * - The emulation thread, writing memory every so often
+         * So it should be impossible for all four buffers to have a read lock at the same time,
+         *  and (because readers always read the newest buffer) at least one of the oldest three
+         *  buffers will be free the entire time we're in this method.
+         * TL;DR this should never happen.
+         * But if it does, do nothing. This isn't medical software, better to show stale data than crash.
+         */
+        warn!("all buffers were locked by a reader at the same time")
+    }
 }