Fetch consecutive missing blocks of streamed content in one request

fsspec's "blockcache" fetches each block of a read with its own request,
so a large read (e.g., of a dataset's data, by NWB Inspector) took one
round trip per 256 KiB, i.e., a couple of MB/s however fast the
connection.  Use instead a block cache that fetches each run of
consecutive missing blocks with a single request, and keeps up to 512
blocks (128 MiB) per open file rather than 32, which h5py's scattered
reads of metadata kept evicting and fetching again.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Qn5WBSiQgoZoL4fytF6nEr

Author: claude

dandi/support/datalad_fuse.py

@@ -191,6 +191,79 @@ def annex_fingerprint(source: Any) -> tuple[str, str] | None:
 #: Size of the blocks in which streamed content is read and cached
 STREAM_BLOCK_SIZE = 256 * 1024
 
+#: Maximum number of blocks cached for each open streamed file
+STREAM_MAX_BLOCKS = 512
+
+#: Name of the fsspec cache defined by `_coalescing_block_cache()`
+COALESCING_BLOCK_CACHE = "dandi-coalescing-blockcache"
+
+
+@cache
+def _coalescing_block_cache() -> type:
+    """
+    Define (once, so as not to import fsspec when this module is imported) and
+    register with fsspec a cache that, like fsspec's ``"blockcache"``, keeps
+    the least recently used blocks of a file, but fetches the consecutive
+    blocks missing for a read with a single request rather than with one
+    request per block, which would make reading a large dataset take one
+    round trip per `STREAM_BLOCK_SIZE` bytes
+    """
+    from collections import OrderedDict
+
+    from fsspec.caching import BaseCache, register_cache
+
+    class CoalescingBlockCache(BaseCache):  # type: ignore[misc]
+        name = COALESCING_BLOCK_CACHE
+
+        def __init__(
+            self,
+            blocksize: int,
+            fetcher: Any,
+            size: int,
+            maxblocks: int = STREAM_MAX_BLOCKS,
+        ) -> None:
+            super().__init__(blocksize, fetcher, size)
+            self.maxblocks = maxblocks
+            self.blocks: OrderedDict[int, bytes] = OrderedDict()
+
+        def _fetch(self, start: int | None, end: int | None) -> bytes:
+            if start is None:
+                start = 0
+            if end is None or end > self.size:
+                end = self.size
+            if start >= end:
+                return b""
+            first = start // self.blocksize
+            last = (end - 1) // self.blocksize
+            numbers = range(first, last + 1)
+            missing = [n for n in numbers if n not in self.blocks]
+            # Fetch each run of consecutive missing blocks in one request
+            i = 0
+            while i < len(missing):
+                j = i
+                while j + 1 < len(missing) and missing[j + 1] == missing[j] + 1:
+                    j += 1
+                run_start = missing[i] * self.blocksize
+                run_end = min((missing[j] + 1) * self.blocksize, self.size)
+                data = self.fetcher(run_start, run_end)
+                for n in range(missing[i], missing[j] + 1):
+                    offset = (n - missing[i]) * self.blocksize
+                    self.blocks[n] = data[offset : offset + self.blocksize]
+                i = j + 1
+            out = []
+            for n in numbers:
+                self.blocks.move_to_end(n)
+                out.append(self.blocks[n])
+            # Evict the least recently used blocks, but not those just read
+            while len(self.blocks) > max(self.maxblocks, len(numbers)):
+                self.blocks.popitem(last=False)
+            data = b"".join(out)
+            offset = first * self.blocksize
+            return data[start - offset : end - offset]
+
+    register_cache(CoalescingBlockCache, clobber=True)
+    return CoalescingBlockCache
+
 
 def _read_in_small_blocks(dsap: Any) -> None:
     """
@@ -206,7 +279,8 @@ def _read_in_small_blocks(dsap: Any) -> None:
     fs = getattr(dsap, "fs", None)
     if isinstance(fs, HTTPFileSystem):
         fs.block_size = STREAM_BLOCK_SIZE
-        fs.cache_type = "blockcache"
+        _coalescing_block_cache()
+        fs.cache_type = COALESCING_BLOCK_CACHE
 
 
 def get_annexed_readable(path: str | Path) -> AnnexedReadableFile | None:

dandi/support/tests/test_datalad_fuse.py

@@ -12,7 +12,12 @@
 from pynwb import NWBHDF5IO, NWBFile, TimeSeries
 import pytest
 
-from ..datalad_fuse import AnnexedReadableFile, annex_fingerprint, get_annexed_readable
+from ..datalad_fuse import (
+    AnnexedReadableFile,
+    _coalescing_block_cache,
+    annex_fingerprint,
+    get_annexed_readable,
+)
 from ...pynwb_utils import get_neurodata_types
 from ...tests.fixtures import RangeHTTPServer, make_git_annex_dandiset
 from ...tests.skip import mark
@@ -114,6 +119,62 @@ def test_annexed_readable_reads_little(
     assert 0 < server.bytes_served < size / 10
 
 
+@pytest.mark.ai_generated
+def test_coalescing_block_cache() -> None:
+    pytest.importorskip("fsspec")
+    content = bytes(range(256)) * 40
+    requests: list[tuple[int, int]] = []
+
+    def fetcher(start: int, end: int) -> bytes:
+        requests.append((start, end))
+        return content[start:end]
+
+    c = _coalescing_block_cache()(100, fetcher, len(content), maxblocks=5)
+    # The blocks missing for a read are fetched with one request
+    assert c._fetch(150, 420) == content[150:420]
+    assert requests == [(100, 500)]
+    # Cached blocks are not fetched again, and runs of missing blocks around
+    # them are fetched with one request each
+    assert c._fetch(30, 720) == content[30:720]
+    assert requests[1:] == [(0, 100), (500, 800)]
+    # Reads at the end of the content, and beyond
+    assert c._fetch(10_190, 20_000) == content[10_190:]
+    assert requests[3:] == [(10_100, 10_240)]
+    assert c._fetch(20_000, 20_100) == b""
+    # Only the 5 most recently used blocks are kept (but at least the blocks
+    # of the last read)
+    assert sorted(c.blocks) == [5, 6, 7, 101, 102]
+    requests.clear()
+    assert c._fetch(0, 10_240) == content
+    assert requests == [(0, 500), (800, 10_100)]
+    assert len(c.blocks) == 103
+
+
+@pytest.mark.ai_generated
+def test_annexed_readable_large_read(
+    range_http_server: tuple[RangeHTTPServer, str], tmp_path: Path
+) -> None:
+    """A large read of streamed content is fetched with one request"""
+    pytest.importorskip("datalad_fuse.fsspec")
+    server, base_url = range_http_server
+    local = tmp_path / "local.bin"
+    content = np.random.default_rng(0).bytes(10_000_000)
+    local.write_bytes(content)
+    shutil.copy(local, tmp_path / "served" / "content.bin")
+    make_git_annex_dandiset(
+        tmp_path / "ds", {"sub-01/data.bin": (local, f"{base_url}/content.bin")}
+    )
+    r = get_annexed_readable(tmp_path / "ds" / "sub-01" / "data.bin")
+    assert r is not None
+    with r.open() as fp:
+        fp.seek(123)
+        assert fp.read(9_000_000) == content[123:9_000_123]
+        assert server.ranged_requests == 1
+        fp.seek(5_000_000)
+        assert fp.read(100) == content[5_000_000:5_000_100]
+        assert server.ranged_requests == 1
+
+
 @pytest.mark.ai_generated
 def test_annexed_readable_none(
     streamed_nwb: tuple[Path, RangeHTTPServer], simple2_nwb: Path, tmp_path: Path