Large files on an aindrive drive open several times faster (#3)

src/lib.rs

@@ -19,6 +19,7 @@ use serde::de::DeserializeOwned;
 use serde_json::json;
 use std::cell::{Cell, RefCell};
 use std::ops::Range;
+use std::rc::Rc;
 
 #[cfg(not(test))]
 use sdk::host;
@@ -77,8 +78,20 @@ impl Session {
   }
 }
 
+/// The last file read whole, for the slices that follow it: the host reads
+/// a file one bounded slice per call, and each slice would otherwise fetch
+/// the whole file again. A slice still stats the file first, and only a file
+/// at the same revision reuses the reply.
+struct Fetched {
+  path: String,
+  revision: Option<String>,
+  reply: Rc<Vec<Vec<u8>>>,
+  encoded: Range<usize>,
+}
+
 thread_local! {
   static SESSION: RefCell<Option<Session>> = const { RefCell::new(None) };
+  static LAST_READ: RefCell<Option<Fetched>> = const { RefCell::new(None) };
 }
 
 struct Aindrive;
@@ -204,7 +217,7 @@ fn info_of(entry: &mcp::Entry) -> FileInfo {
 /// One file read whole from the drive and handed over in pieces, decoded
 /// from the reply as each piece is asked for.
 struct Download {
-  reply: Vec<Vec<u8>>,
+  reply: Rc<Vec<Vec<u8>>>,
   /// Where the base64 contents sit in the reply, read as one run.
   encoded: Range<usize>,
   size: u64,
@@ -292,16 +305,32 @@ fn download(target: &Locator, offset: u64) -> Result<Download, ErrorCode> {
   if entry.size > MAX_READ {
     return Err(ErrorCode::Unsupported);
   }
+  let revision = entry.revision();
   let (reply, encoded) = if offset >= entry.size {
-    (Vec::new(), 0..0)
+    (Rc::default(), 0..0)
+  } else if let Some(fetched) = LAST_READ.with_borrow(|last| {
+    last
+      .as_ref()
+      .filter(|last| last.path == path && last.revision == revision)
+      .map(|last| (last.reply.clone(), last.encoded.clone()))
+  }) {
+    fetched
   } else {
     let length = usize::try_from(entry.size.div_ceil(3) * 4)
       .map_err(|_| ErrorCode::Unsupported)?;
-    let reply = post(
+    // The previous file is let go first, so two never sit in memory at once.
+    LAST_READ.set(None);
+    let reply = Rc::new(post(
       &mcp::call("read_file", &json!({ "path": path, "encoding": "base64" })),
       Effect::Read,
-    )?;
+    )?);
     let encoded = mcp::contents(&reply, length)?;
+    LAST_READ.set(Some(Fetched {
+      path,
+      revision: revision.clone(),
+      reply: reply.clone(),
+      encoded: encoded.clone(),
+    }));
     (reply, encoded)
   };
   Ok(Download {
@@ -309,7 +338,7 @@ fn download(target: &Locator, offset: u64) -> Result<Download, ErrorCode> {
     encoded,
     size: entry.size,
     position: Cell::new(offset.min(entry.size)),
-    revision: entry.revision(),
+    revision,
   })
 }
 
@@ -399,6 +428,7 @@ impl Guest for Aindrive {
     config: Vec<(String, Setting)>,
   ) -> Result<Capabilities, ErrorCode> {
     SESSION.set(None);
+    LAST_READ.set(None);
     SESSION.set(Some(Session::from_config(&config)?));
     let probe = post(mcp::TOOLS, Effect::Read)
       .and_then(|pieces| mcp::tools(&joined(pieces)));
@@ -433,6 +463,7 @@ impl Guest for Aindrive {
 
   fn close() {
     SESSION.set(None);
+    LAST_READ.set(None);
   }
 
   /// The stat tool refuses the drive root, which always exists, so the root

src/session_tests.rs

@@ -524,3 +524,39 @@ async fn only_a_server_error_leaves_a_save_unknown() -> Result<(), ErrorCode> {
   }
   Ok(())
 }
+
+#[tokio::test]
+async fn slices_of_one_revision_fetch_the_file_once() -> Result<(), ErrorCode> {
+  let bytes: Vec<u8> = (0..=255).cycle().take(30_000).collect();
+  let mut drive = sample();
+  drive
+    .files
+    .insert("data.bin".to_owned(), test_host::file(&bytes, "5"));
+  let drive = test_host::install(drive);
+  Aindrive::open("aindrive".to_owned(), config())?;
+  let mut read = Vec::new();
+  while read.len() < bytes.len() {
+    let reader = download(&at("/data.bin"), read.len() as u64)?;
+    read.extend(reader.pull(7_000)?);
+  }
+  assert_eq!(read, bytes);
+  let count = |tool: &str| {
+    drive
+      .borrow()
+      .seen
+      .iter()
+      .filter(|seen| seen.tool.as_deref() == Some(tool))
+      .count()
+  };
+  assert_eq!(count("read_file"), 1);
+  assert!(count("stat") >= 5);
+
+  // Another revision is another file, fetched afresh.
+  drive
+    .borrow_mut()
+    .files
+    .insert("data.bin".to_owned(), test_host::file(b"replaced", "6"));
+  assert_eq!(read_from("/data.bin", 0, 100)?.0, b"replaced");
+  assert_eq!(count("read_file"), 2);
+  Ok(())
+}