fix: large directory upload (100K files)

packages/blob-index/src/sharded-dag-index.js

@@ -9,6 +9,12 @@ import { DigestMap } from './digest-map.js'
 
 export const version = 'index/sharded/dag@0.1'
 
+/**
+ * The size of the batch to process when archiving large datasets.
+ * This is a heuristic to avoid memory issues when archiving large datasets.
+ */
+const ARCHIVE_BATCH_SIZE = 10_000
+
 export const ShardedDAGIndexSchema = Schema.variant({
   [version]: Schema.struct({
     /** DAG root. */
@@ -166,23 +172,44 @@ export const create = (content) => new ShardedDAGIndex(content)
  */
 export const archive = async (model) => {
   const blocks = new Map()
-  const shards = [...model.shards.entries()].sort((a, b) =>
-    compare(a[0].digest, b[0].digest)
-  )
   const index = {
     content: model.content,
     shards: /** @type {API.Link[]} */ ([]),
   }
-  for (const s of shards) {
-    const slices = [...s[1].entries()]
-      .sort((a, b) => compare(a[0].digest, b[0].digest))
-      .map((e) => [e[0].bytes, e[1]])
-    const bytes = dagCBOR.encode([s[0].bytes, slices])
-    const digest = await sha256.digest(bytes)
-    const cid = Link.create(dagCBOR.code, digest)
-    blocks.set(cid.toString(), { cid, bytes })
-    index.shards.push(cid)
+
+  // Convert all shards to an array first
+  const allShards = [...model.shards.entries()]
+
+  // Process shards in batches
+  for (let i = 0; i < allShards.length; i += ARCHIVE_BATCH_SIZE) {
+    const batch = allShards.slice(i, i + ARCHIVE_BATCH_SIZE)
+    const sortedBatch = batch.sort((a, b) => compare(a[0].digest, b[0].digest))
+
+    for (const s of sortedBatch) {
+      // Process slices in batches
+      const allSlices = [...s[1].entries()]
+      const sortedSlices = []
+
+      // Sort slices in batches
+      for (let j = 0; j < allSlices.length; j += ARCHIVE_BATCH_SIZE) {
+        const sliceBatch = allSlices.slice(j, j + ARCHIVE_BATCH_SIZE)
+        const sortedSliceBatch = sliceBatch.sort((a, b) =>
+          compare(a[0].digest, b[0].digest)
+        )
+        sortedSlices.push(...sortedSliceBatch)
+      }
+
+      // Map the sorted slices
+      const mappedSlices = sortedSlices.map((e) => [e[0].bytes, e[1]])
+
+      const bytes = dagCBOR.encode([s[0].bytes, mappedSlices])
+      const digest = await sha256.digest(bytes)
+      const cid = Link.create(dagCBOR.code, digest)
+      blocks.set(cid.toString(), { cid, bytes })
+      index.shards.push(cid)
+    }
   }
+
   const bytes = dagCBOR.encode({ [version]: index })
   const digest = await sha256.digest(bytes)
   const cid = Link.create(dagCBOR.code, digest)

packages/blob-index/test/large-dataset.spec.js

@@ -0,0 +1,195 @@
+import { ShardedDAGIndex } from '../src/index.js'
+import { randomCAR } from './helpers/random.js'
+import * as Result from './helpers/result.js'
+import { fromShardArchives } from '../src/util.js'
+import { sha256 } from 'multiformats/hashes/sha2'
+import { base58btc } from 'multiformats/bases/base58'
+import * as Link from 'multiformats/link'
+
+/**
+ * @typedef {import('entail').Test} Test
+ * @typedef {import('entail').assert} Assert
+ */
+
+export const test = {
+  'handles normal dataset': async (/** @type {Assert} */ assert) => {
+    // Create a normal-sized dataset (under threshold)
+    const contentCAR = await randomCAR(32)
+    const contentCARBytes = new Uint8Array(await contentCAR.arrayBuffer())
+    const index = await fromShardArchives(contentCAR.roots[0], [
+      contentCARBytes,
+    ])
+
+    // Archive should use the fast path
+    const indexCAR = Result.unwrap(await index.archive())
+    const newIndex = Result.unwrap(ShardedDAGIndex.extract(indexCAR))
+
+    assert.notStrictEqual(newIndex.shards.size, 0)
+    assert.strictEqual(index.shards.size, newIndex.shards.size)
+  },
+
+  'handles large dataset': async (/** @type {Assert} */ assert) => {
+    // Create a dummy content link
+    const contentBytes = new Uint8Array(32)
+    const contentDigest = await sha256.digest(contentBytes)
+    const contentLink = Link.create(0x71, contentDigest)
+
+    const model = ShardedDAGIndex.create(contentLink)
+    const TOTAL_ENTRIES = 100001
+    /** @type {Set<string>} */ const usedBase58 = new Set()
+
+    for (let i = 0; i < TOTAL_ENTRIES; i++) {
+      // Create unique byte arrays by using the index as a counter
+      const shardBytes = new Uint8Array(32)
+      const sliceBytes = new Uint8Array(32)
+
+      // Fill with a pattern that ensures uniqueness
+      for (let j = 0; j < 32; j++) {
+        if (j < 4) {
+          // Use the index i to ensure uniqueness (4 bytes = 4.2 billion possible values)
+          shardBytes[j] = (i >> (j * 8)) & 0xff
+          sliceBytes[j] = ((i * 2) >> (j * 8)) & 0xff
+        } else {
+          // Fill remaining bytes with random values
+          shardBytes[j] = Math.floor(Math.random() * 256)
+          sliceBytes[j] = Math.floor(Math.random() * 256)
+        }
+      }
+
+      // Verify uniqueness of base58btc encoding
+      const shardBase58 = base58btc.encode(shardBytes)
+      const sliceBase58 = base58btc.encode(sliceBytes)
+
+      if (usedBase58.has(shardBase58)) {
+        throw new Error(`Duplicate shard base58: ${shardBase58}`)
+      }
+      if (usedBase58.has(sliceBase58)) {
+        throw new Error(`Duplicate slice base58: ${sliceBase58}`)
+      }
+
+      usedBase58.add(shardBase58)
+      usedBase58.add(sliceBase58)
+
+      const shard = await sha256.digest(shardBytes)
+      const slice = await sha256.digest(sliceBytes)
+      model.setSlice(shard, slice, [0, 32])
+    }
+
+    assert.strictEqual(model.shards.size, TOTAL_ENTRIES)
+  },
+
+  'maintains sorting order in large dataset': async (
+    /** @type {Assert} */ assert
+  ) => {
+    // Create a dummy content link
+    const contentBytes = new Uint8Array(32)
+    const contentDigest = await sha256.digest(contentBytes)
+    const contentLink = Link.create(0x71, contentDigest)
+
+    const model = ShardedDAGIndex.create(contentLink)
+    const TOTAL_ENTRIES = 100001
+    /** @type {Set<string>} */ const usedBase58 = new Set()
+
+    for (let i = 0; i < TOTAL_ENTRIES; i++) {
+      // Create unique byte arrays by using the index as a counter
+      const shardBytes = new Uint8Array(32)
+      const sliceBytes = new Uint8Array(32)
+
+      // Fill with a pattern that ensures uniqueness
+      for (let j = 0; j < 32; j++) {
+        if (j < 4) {
+          // Use the index i to ensure uniqueness (4 bytes = 4.2 billion possible values)
+          shardBytes[j] = (i >> (j * 8)) & 0xff
+          sliceBytes[j] = ((i * 2) >> (j * 8)) & 0xff
+        } else {
+          // Fill remaining bytes with random values
+          shardBytes[j] = Math.floor(Math.random() * 256)
+          sliceBytes[j] = Math.floor(Math.random() * 256)
+        }
+      }
+
+      // Verify uniqueness of base58btc encoding
+      const shardBase58 = base58btc.encode(shardBytes)
+      const sliceBase58 = base58btc.encode(sliceBytes)
+
+      if (usedBase58.has(shardBase58)) {
+        throw new Error(`Duplicate shard base58: ${shardBase58}`)
+      }
+      if (usedBase58.has(sliceBase58)) {
+        throw new Error(`Duplicate slice base58: ${sliceBase58}`)
+      }
+
+      usedBase58.add(shardBase58)
+      usedBase58.add(sliceBase58)
+
+      const shard = await sha256.digest(shardBytes)
+      const slice = await sha256.digest(sliceBytes)
+      model.setSlice(shard, slice, [0, 32])
+    }
+
+    assert.strictEqual(model.shards.size, TOTAL_ENTRIES)
+  },
+
+  'can archive large dataset': async (/** @type {Assert} */ assert) => {
+    // Create a dummy content link
+    const contentBytes = new Uint8Array(32)
+    const contentDigest = await sha256.digest(contentBytes)
+    const contentLink = Link.create(0x71, contentDigest)
+
+    const model = ShardedDAGIndex.create(contentLink)
+    const TOTAL_ENTRIES = 100001
+    /** @type {Set<string>} */ const usedBase58 = new Set()
+
+    for (let i = 0; i < TOTAL_ENTRIES; i++) {
+      // Create unique byte arrays by using the index as a counter
+      const shardBytes = new Uint8Array(32)
+      const sliceBytes = new Uint8Array(32)
+
+      // Fill with a pattern that ensures uniqueness
+      for (let j = 0; j < 32; j++) {
+        if (j < 4) {
+          // Use the index i to ensure uniqueness (4 bytes = 4.2 billion possible values)
+          shardBytes[j] = (i >> (j * 8)) & 0xff
+          sliceBytes[j] = ((i * 2) >> (j * 8)) & 0xff
+        } else {
+          // Fill remaining bytes with random values
+          shardBytes[j] = Math.floor(Math.random() * 256)
+          sliceBytes[j] = Math.floor(Math.random() * 256)
+        }
+      }
+
+      // Verify uniqueness of base58btc encoding
+      const shardBase58 = base58btc.encode(shardBytes)
+      const sliceBase58 = base58btc.encode(sliceBytes)
+
+      if (usedBase58.has(shardBase58)) {
+        throw new Error(`Duplicate shard base58: ${shardBase58}`)
+      }
+      if (usedBase58.has(sliceBase58)) {
+        throw new Error(`Duplicate slice base58: ${sliceBase58}`)
+      }
+
+      usedBase58.add(shardBase58)
+      usedBase58.add(sliceBase58)
+
+      const shard = await sha256.digest(shardBytes)
+      const slice = await sha256.digest(sliceBytes)
+      model.setSlice(shard, slice, [0, 32])
+    }
+
+    // Test that we can archive the large dataset
+    const result = await model.archive()
+    assert.ok(result.ok, 'Archive should succeed')
+
+    // Test that we can extract the archive
+    const extracted = ShardedDAGIndex.extract(result.ok)
+    assert.ok(extracted.ok, 'Extract should succeed')
+    assert.strictEqual(
+      extracted.ok.shards.size,
+      TOTAL_ENTRIES,
+      'Should have all shards'
+    )
+  },
+}
+
+export default test

packages/upload-api/src/index/add.js

@@ -6,6 +6,12 @@ import { Assert } from '@web3-storage/content-claims/capability'
 import { concat } from 'uint8arrays'
 import * as API from '../types.js'
 
+/**
+ * The size of the batch to process when checking shard allocations.
+ * This is a heuristic to avoid memory issues when processing large indexes.
+ */
+const ALLOCATION_BATCH_SIZE = 10_000
+
 /**
  * @param {API.IndexServiceContext} context
  * @returns {API.ServiceMethod<API.IndexAdd, API.IndexAddSuccess, API.IndexAddFailure>}
@@ -55,23 +61,43 @@ const add = async ({ capability }, context) => {
 
   const idxRes = ShardedDAGIndex.extract(concat(chunks))
   if (!idxRes.ok) return idxRes
+  return await batchProcessIndexChunks(idxRes.ok, space, context, idxLink)
+}
 
-  // ensure indexed shards are allocated in the agent's space
-  const shardDigests = [...idxRes.ok.shards.keys()]
-  const shardAllocRes = await Promise.all(
-    shardDigests.map((s) => assertAllocated(context, space, s, 'ShardNotFound'))
-  )
-  for (const res of shardAllocRes) {
-    if (res.error) return res
+/**
+ * Batch process all chunks of the index.
+ *
+ * @param {import('@web3-storage/blob-index/types').ShardedDAGIndexView} index
+ * @param {API.SpaceDID} space
+ * @param {API.IndexServiceContext} context
+ * @param {API.CARLink} idxLink
+ */
+async function batchProcessIndexChunks(index, space, context, idxLink) {
+  const shardDigests = [...index.shards.keys()]
+
+  // Process shard allocations in batches
+  for (let i = 0; i < shardDigests.length; i += ALLOCATION_BATCH_SIZE) {
+    const batch = shardDigests.slice(i, i + ALLOCATION_BATCH_SIZE)
+
+    // Each batch can be processed concurrently
+    const batchResults = await Promise.all(
+      batch.map((shard) =>
+        assertAllocated(context, space, shard, 'ShardNotFound')
+      )
+    )
+
+    for (const res of batchResults) {
+      if (res.error) return res
+    }
   }
 
   // TODO: randomly validate slices in the index correspond to slices in the blob
 
   const publishRes = await Promise.all([
     // publish the index data to IPNI
-    context.ipniService.publish(idxRes.ok),
+    context.ipniService.publish(index),
     // publish a content claim for the index
-    publishIndexClaim(context, { content: idxRes.ok.content, index: idxLink }),
+    publishIndexClaim(context, { content: index.content, index: idxLink }),
   ])
   for (const res of publishRes) {
     if (res.error) return res