perf: skip maxDepth filtering if natively supported (#560)

unjs · Jan 3, 2025 · aec28a8 · aec28a8
1 parent cc0b0ca
commit aec28a8
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 2 deletions.
diff --git a/src/drivers/fs-lite.ts b/src/drivers/fs-lite.ts
@@ -40,6 +40,9 @@ export default defineDriver((opts: FSStorageOptions = {}) => {
   return {
     name: DRIVER_NAME,
     options: opts,
+    flags: {
+      maxDepth: true,
+    },
     hasItem(key) {
       return existsSync(r(key));
     },

diff --git a/src/drivers/fs.ts b/src/drivers/fs.ts
@@ -59,6 +59,9 @@ export default defineDriver((opts: FSStorageOptions = {}) => {
   return {
     name: DRIVER_NAME,
     options: opts,
+    flags: {
+      maxDepth: true,
+    },
     hasItem(key) {
       return existsSync(r(key));
     },

diff --git a/src/storage.ts b/src/storage.ts
@@ -348,7 +348,11 @@ export function createStorage<T extends StorageValue>(
       const mounts = getMounts(base, true);
       let maskedMounts: string[] = [];
       const allKeys = [];
+      let allMountsSupportMaxDepth = true;
       for (const mount of mounts) {
+        if (!mount.driver.flags?.maxDepth) {
+          allMountsSupportMaxDepth = false;
+        }
         const rawKeys = await asyncCall(
           mount.driver.getKeys,
           mount.relativeBase,
@@ -368,10 +372,11 @@ export function createStorage<T extends StorageValue>(
           ...maskedMounts.filter((p) => !p.startsWith(mount.mountpoint)),
         ];
       }
+      const shouldFilterByDepth =
+        opts.maxDepth !== undefined && !allMountsSupportMaxDepth;
       return allKeys.filter(
         (key) =>
-          (opts.maxDepth === undefined ||
-            filterKeyByDepth(key, opts.maxDepth)) &&
+          (!shouldFilterByDepth || filterKeyByDepth(key, opts.maxDepth)) &&
           filterKeyByBase(key, base)
       );
     },

diff --git a/test/storage.test.ts b/test/storage.test.ts
@@ -1,11 +1,13 @@
 import { describe, it, expect, vi } from "vitest";
+import { resolve } from "node:path";
 import {
   createStorage,
   snapshot,
   restoreSnapshot,
   prefixStorage,
 } from "../src";
 import memory from "../src/drivers/memory";
+import fs from "../src/drivers/fs";
 
 const data = {
   "etc:conf": "test",
@@ -222,4 +224,29 @@ describe("Regression", () => {
     await pStorage.remove("y");
     expect(await pStorage.has("y")).toBe(false);
   });
+
+  it("getKeys supports maxDepth with mixed native support", async () => {
+    const base = resolve(__dirname, "tmp/fs");
+    const mainStorage = memory();
+    const secondaryStorage = fs({ base });
+    const storage = createStorage({ driver: mainStorage });
+
+    storage.mount("/storage_b", secondaryStorage);
+
+    try {
+      await storage.setItem("/storage_a/file_depth1", "contents");
+      await storage.setItem("/storage_a/depth1/file_depth2", "contents");
+      await storage.setItem("/storage_b/file_depth1", "contents");
+      await storage.setItem("/storage_b/depth1/file_depth2", "contents");
+
+      const keys = await storage.getKeys(undefined, { maxDepth: 1 });
+
+      expect(keys.sort()).toMatchObject([
+        "storage_a:file_depth1",
+        "storage_b:file_depth1",
+      ]);
+    } finally {
+      await storage.clear();
+    }
+  });
 });