feat: allow skip full block download in first few page read in block

Author: darkskygit

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "sqlite-vfs-http"
-version = "0.1.1"
+version = "0.1.2"
 edition = "2021"
 authors = ["DarkSky <[email protected]>"]
 license = "AGPL-3.0-only"

src/buffer.rs

@@ -35,6 +35,7 @@ type Fetch = Box<dyn Fn(usize, usize) -> Result<Vec<u8>> + Send + Sync>;
 
 pub struct LazyBuffer {
     blocks: HashMap<usize, LazyBlock>,
+    blocks_read: HashMap<usize, usize>,
     block_size: usize,
     download_threshold: usize,
     length: usize,
@@ -45,6 +46,7 @@ impl LazyBuffer {
     pub fn new(length: usize, block_size: usize, download_threshold: usize, fetch: Fetch) -> Self {
         Self {
             blocks: HashMap::new(),
+            blocks_read: HashMap::new(),
             block_size,
             download_threshold,
             length,
@@ -76,13 +78,25 @@ impl LazyBuffer {
         }
         let block_index = offset / self.block_size;
         let block_offset = offset - block_index * self.block_size;
-        if !self.blocks.contains_key(&block_index) && buf.len() <= self.download_threshold {
+        let block_read = self
+            .blocks_read
+            .get(&block_index)
+            .copied()
+            .unwrap_or_default();
+        let buffer_size = buf.len();
+        if !self.blocks.contains_key(&block_index)
+            && block_read + buffer_size <= self.download_threshold
+        {
             // skip full block fetch if buf is smaller than download_threshold
-            let data = (self.fetch)(offset, buf.len())?;
-            if data.len() != buf.len() {
+            let data = (self.fetch)(offset, buffer_size)?;
+            if data.len() != buffer_size {
                 return Err(Error::new(ErrorKind::UnexpectedEof, "read out of bounds"));
             }
             buf.copy_from_slice(&data);
+            self.blocks_read
+                .entry(block_index)
+                .and_modify(|size| *size += buffer_size)
+                .or_insert(buffer_size);
         } else {
             self.get_block(block_index)?.read(buf, block_offset)?;
         }

src/lib.rs

@@ -10,23 +10,44 @@ use std::sync::{Arc, Once, RwLock};
 use utils::AtomicRuntime;
 use vfs::HttpVfs;
 
-pub use vfs::HTTP_VFS;
+const SQLITE_PAGE_SIZE: usize = 1024 * 4;
+pub const HTTP_VFS: &str = "http";
 
 pub struct HttpVfsRegister {
+    /// how many pages in block, default is 8MB, 2048 pages
     block_size: usize,
+    /// read the first few pages of each block without downloading the entire block
     download_threshold: usize,
+    /// sqlite's page size is 4KB by default
+    page_size: usize,
 }
 
 impl HttpVfsRegister {
     pub fn new() -> Self {
         Self {
-            block_size: 1024 * 1024 * 4,
-            download_threshold: 1024,
+            block_size: SQLITE_PAGE_SIZE * 1024 * 2,
+            download_threshold: 0,
+            page_size: SQLITE_PAGE_SIZE,
         }
     }
 
-    pub fn with_block_size(self, block_size: usize) -> Self {
-        Self { block_size, ..self }
+    pub fn with_block_size(self, page_num: usize) -> Self {
+        Self {
+            block_size: page_num * self.page_size,
+            ..self
+        }
+    }
+
+    /// Set how many page read don't download full block
+    pub fn with_download_threshold(self, page_num: usize) -> Self {
+        Self {
+            download_threshold: page_num * self.page_size,
+            ..self
+        }
+    }
+
+    pub fn with_page_size(self, page_size: usize) -> Self {
+        Self { page_size, ..self }
     }
 
     pub fn register(self) {

src/vfs.rs

@@ -6,8 +6,6 @@ use std::{
     time::Duration,
 };
 
-pub const HTTP_VFS: &str = "http";
-
 pub struct HttpVfs {
     pub(crate) block_size: usize,
     pub(crate) download_threshold: usize,