From 5ea214083a7fab51b0db2fd7fdb97cb97c5bc9c9 Mon Sep 17 00:00:00 2001
From: Pekka Enberg <penberg@iki.fi>
Date: Sat, 13 Jan 2024 09:03:21 +0200
Subject: [PATCH] core: Switch to using Arc<Completion>

We need to be able to leak a reference for io_uring so switch to
Arc<Completion> in the code.
---
 core/io/darwin.rs      |  6 ++++--
 core/io/linux.rs       |  8 ++++++--
 core/io/mod.rs         | 21 ++++++++++++++++++---
 core/sqlite3_ondisk.rs | 11 ++++++-----
 core/storage.rs        |  8 ++++----
 5 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/core/io/darwin.rs b/core/io/darwin.rs
index eac13e84d..f291cb929 100644
--- a/core/io/darwin.rs
+++ b/core/io/darwin.rs
@@ -27,10 +27,12 @@ pub struct File {
 }
 
 impl File {
-    pub fn pread(&self, pos: usize, c: &mut Completion) -> Result<()> {
+    pub fn pread(&self, pos: usize, c: Arc<Completion>) -> Result<()> {
         let mut file = self.file.borrow_mut();
         file.seek(std::io::SeekFrom::Start(pos as u64))?;
-        file.read_exact(c.buf.as_mut_slice())?;
+        let buf = c.buf();
+        let mut buf = buf.as_mut_slice();
+        file.read_exact(buf)?;
         Ok(())
     }
 }
\ No newline at end of file
diff --git a/core/io/linux.rs b/core/io/linux.rs
index 338f75f41..0ad77bcaa 100644
--- a/core/io/linux.rs
+++ b/core/io/linux.rs
@@ -3,6 +3,7 @@ use anyhow::Result;
 use std::cell::RefCell;
 use std::os::unix::io::AsRawFd;
 use std::rc::Rc;
+use std::sync::Arc;
 
 pub struct IO {
     ring: Rc<RefCell<io_uring::IoUring>>,
@@ -38,9 +39,12 @@ pub struct File {
 }
 
 impl File {
-    pub fn pread(&self, pos: usize, c: &mut Completion) -> Result<()> {
+    pub fn pread(&self, pos: usize, c: Arc<Completion>) -> Result<()> {
         let fd = io_uring::types::Fd(self.file.as_raw_fd());
-        let read_e = io_uring::opcode::Read::new(fd, c.buf.as_mut_ptr(), c.buf.len() as u32 )
+        let mut buf = c.buf_mut();
+        let len = buf.len();
+        let buf = buf.as_mut_ptr();
+        let read_e = io_uring::opcode::Read::new(fd, buf, len as u32 )
             .offset(pos as u64)
             .build();
         let mut ring = self.ring.borrow_mut();
diff --git a/core/io/mod.rs b/core/io/mod.rs
index 02594f426..e4ec1a739 100644
--- a/core/io/mod.rs
+++ b/core/io/mod.rs
@@ -1,20 +1,35 @@
 use cfg_block::cfg_block;
-use std::{mem::ManuallyDrop, pin::Pin, sync::Arc};
+use std::{
+    cell::{Ref, RefCell, RefMut},
+    mem::ManuallyDrop,
+    pin::Pin,
+    sync::Arc,
+};
 
 pub type Complete = dyn Fn(&Buffer) + Send + Sync;
 
 pub struct Completion {
-    pub buf: Buffer,
+    pub buf: RefCell<Buffer>,
     pub complete: Box<Complete>,
 }
 
 impl Completion {
     pub fn new(buf: Buffer, complete: Box<Complete>) -> Self {
+        let buf = RefCell::new(buf);
         Self { buf, complete }
     }
 
+    pub fn buf<'a>(&'a self) -> Ref<'a, Buffer> {
+        self.buf.borrow()
+    }
+
+    pub fn buf_mut<'a>(&'a self) -> RefMut<'a, Buffer> {
+        self.buf.borrow_mut()
+    }
+
     pub fn complete(&self) {
-        (self.complete)(&self.buf);
+        let buf = self.buf.borrow_mut();
+        (self.complete)(&buf);
     }
 }
 
diff --git a/core/sqlite3_ondisk.rs b/core/sqlite3_ondisk.rs
index c011c9282..1750c22bd 100644
--- a/core/sqlite3_ondisk.rs
+++ b/core/sqlite3_ondisk.rs
@@ -65,9 +65,10 @@ pub fn read_database_header(storage: &Storage) -> Result<DatabaseHeader> {
     let drop_fn = Arc::new(|_buf| {});
     let buf = Buffer::allocate(512, drop_fn);
     let complete = Box::new(move |_buf: &Buffer| {});
-    let mut c = Completion::new(buf, complete);
-    storage.get(1, &mut c)?;
-    let buf = c.buf.as_slice();
+    let c = Arc::new(Completion::new(buf, complete));
+    storage.get(1, c.clone())?;
+    let buf = c.buf();
+    let buf = buf.as_slice();
     let mut header = DatabaseHeader::default();
     header.magic.copy_from_slice(&buf[0..16]);
     header.page_size = u16::from_be_bytes([buf[16], buf[17]]);
@@ -152,8 +153,8 @@ pub fn begin_read_btree_page(
             page.set_error();
         }
     });
-    let mut c = Completion::new(buf, complete);
-    storage.get(page_idx, &mut c)?;
+    let c = Arc::new(Completion::new(buf, complete));
+    storage.get(page_idx, c.clone())?;
     c.complete();
     Ok(())
 }
diff --git a/core/storage.rs b/core/storage.rs
index 7f49e2201..f7e4c6f12 100644
--- a/core/storage.rs
+++ b/core/storage.rs
@@ -18,13 +18,13 @@ impl Storage {
         }
     }
 
-    pub fn get(&self, page_idx: usize, c: &mut Completion) -> Result<()> {
+    pub fn get(&self, page_idx: usize, c: Arc<Completion>) -> Result<()> {
         self.io.get(page_idx, c)
     }
 }
 
 pub trait StorageIO {
-    fn get(&self, page_idx: usize, c: &mut Completion) -> Result<()>;
+    fn get(&self, page_idx: usize, c: Arc<Completion>) -> Result<()>;
 }
 
 #[cfg(feature = "fs")]
@@ -34,8 +34,8 @@ struct FileStorage {
 
 #[cfg(feature = "fs")]
 impl StorageIO for FileStorage {
-    fn get(&self, page_idx: usize, c: &mut Completion) -> Result<()> {
-        let page_size = c.buf.len();
+    fn get(&self, page_idx: usize, c: Arc<Completion>) -> Result<()> {
+        let page_size = c.buf().len();
         assert!(page_idx > 0);
         assert!(page_size >= 512);
         assert!(page_size <= 65536);