feat: add memory budgeting

Author: Alan Shaw

package-lock.json

@@ -28,7 +28,7 @@
     "@ipld/dag-json": "^8.0.11",
     "@ipld/dag-pb": "^2.1.18",
     "@web3-storage/fast-unixfs-exporter": "^0.2.0",
-    "@web3-storage/gateway-lib": "^1.3.0",
+    "@web3-storage/gateway-lib": "^2.0.0",
     "cardex": "^1.0.0",
     "chardet": "^1.5.0",
     "dagula": "^4.1.0",

package.json

@@ -1,5 +1,6 @@
 import type { CID } from 'multiformats/cid'
 import type { Context } from '@web3-storage/gateway-lib'
+import type { MemoryBudget } from './lib/mem-budget'
 
 export {}
 
@@ -14,6 +15,10 @@ export interface CarCidsContext extends Context {
   carCids: CID[]
 }
 
+export interface MemoryBudgetContext extends Context {
+  memoryBudget: MemoryBudget
+}
+
 export interface R2GetOptions {
   range?: {
     offset: number

src/bindings.d.ts

@@ -6,14 +6,21 @@ import {
   withHttpGet,
   withCdnCache,
   withParsedIpfsUrl,
+  withFixedLengthStream,
   composeMiddleware
 } from '@web3-storage/gateway-lib/middleware'
 import {
   handleUnixfs,
   handleBlock,
   handleCar
 } from '@web3-storage/gateway-lib/handlers'
-import { withDagula, withCarCids, withUnsupportedFeaturesHandler } from './middleware.js'
+import {
+  withDagula,
+  withCarCids,
+  withUnsupportedFeaturesHandler,
+  withMemoryBudget,
+  withResponseMemoryRelease
+} from './middleware.js'
 
 /**
  * @typedef {import('./bindings').Environment} Environment
@@ -35,7 +42,10 @@ export default {
       withHttpGet,
       withParsedIpfsUrl,
       withCarCids,
-      withDagula
+      withMemoryBudget,
+      withDagula,
+      withResponseMemoryRelease,
+      withFixedLengthStream
     )
     return middleware(handler)(request, env, ctx)
   }

src/index.js

@@ -25,9 +25,11 @@ export class R2Blockstore {
    * @param {R2Bucket} dataBucket
    * @param {R2Bucket} indexBucket
    * @param {CID[]} carCids
+   * @param {import('./mem-budget.js').MemoryBudget} memoryBudget
    */
-  constructor (dataBucket, indexBucket, carCids) {
+  constructor (dataBucket, indexBucket, carCids, memoryBudget) {
     this._dataBucket = dataBucket
+    this._memoryBudget = memoryBudget
     this._idx = new MultiCarIndex()
     for (const carCid of carCids) {
       this._idx.addIndex(carCid, new StreamingCarIndex((async function * () {
@@ -54,14 +56,10 @@ export class R2Blockstore {
 
     const reader = res.body.getReader()
     const bytesReader = asyncIterableReader((async function * () {
-      try {
-        while (true) {
-          const { done, value } = await reader.read()
-          if (done) return
-          yield value
-        }
-      } finally {
-        reader.releaseLock()
+      while (true) {
+        const { done, value } = await reader.read()
+        if (done) return
+        yield value
       }
     })())
 
@@ -110,6 +108,7 @@ export class BatchingR2Blockstore extends R2Blockstore {
   }
 
   async #processBatch () {
+    console.log('processing batch')
     const batches = this.#batches
     const batchBlocks = this.#batchBlocks
     this.#batches = new Map()
@@ -121,7 +120,9 @@ export class BatchingR2Blockstore extends R2Blockstore {
         if (!batch) break
         const carPath = `${carCid}/${carCid}.car`
         const range = { offset: batch[0], length: batch[batch.length - 1] - batch[0] + MAX_BLOCK_LENGTH }
-        console.log(`requesting ${batch.length} blocks from ${carCid} (${range.length} bytes @ ${range.offset})`)
+        await this._memoryBudget.request(range.length)
+
+        console.log(`fetching ${batch.length} blocks from ${carCid} (${range.length} bytes @ ${range.offset})`)
         const res = await this._dataBucket.get(carPath, { range })
         if (!res) {
           for (const blocks of batchBlocks.values()) {
@@ -139,6 +140,7 @@ export class BatchingR2Blockstore extends R2Blockstore {
           }
         })())
 
+        let bytesResolved = 0
         while (true) {
           try {
             const blockHeader = await readBlockHead(bytesReader)
@@ -151,15 +153,19 @@ export class BatchingR2Blockstore extends R2Blockstore {
               // console.log(`got wanted block for ${blockHeader.cid}`)
               blocks.forEach(b => b.resolve({ cid: blockHeader.cid, bytes }))
               batchBlocks.delete(key)
+              bytesResolved += bytes.length
             }
           } catch {
             break
           }
         }
 
+        // release the bytes we didn't send on (they are released later)
+        this._memoryBudget.release(range.length - bytesResolved)
         reader.cancel()
       }
     }
+    console.log('finished processing batch')
   }
 
   /** @param {CID} cid */

src/lib/blockstore.js

@@ -0,0 +1,55 @@
+import defer from 'p-defer'
+
+export class MemoryBudget {
+  /** @type {Array<{ size: number, deferred: import('p-defer').DeferredPromise<void> }>} */
+  #requests = []
+
+  /**
+   * Allocated bytes.
+   * @type {number}
+   */
+  #allocated = 0
+
+  /**
+   * Maximum number of bytes that can be allocated.
+   * @type {number}
+   */
+  #max
+
+  /**
+   * @param {number} max
+   */
+  constructor (max) {
+    this.#max = max
+  }
+
+  /**
+   * @param {number} size
+   */
+  async request (size) {
+    if (!this.#requests.length && this.#allocated + size < this.#max) {
+      this.#allocated += size
+      console.log(`allocating ${size} bytes, ${Math.floor(this.#allocated / this.#max * 100)}% allocated`)
+      return
+    }
+    const deferred = defer()
+    this.#requests.push({ size, deferred })
+    return deferred.promise
+  }
+
+  /**
+   * @param {number} size
+   */
+  release (size) {
+    this.#allocated = Math.max(0, this.#allocated - size)
+    console.log(`released ${size} bytes, ${Math.floor(this.#allocated / this.#max * 100)}% allocated`)
+    while (this.#requests.length) {
+      const req = this.#requests[0]
+      if (this.#allocated + req.size >= this.#max) break
+      this.#allocated += req.size
+      console.log(`allocating ${req.size} bytes, ${Math.floor(this.#allocated / this.#max * 100)}% allocated`)
+      req.deferred.resolve()
+      this.#requests.shift()
+    }
+  }
+}

src/lib/mem-budget.js

@@ -3,15 +3,18 @@ import { Dagula } from 'dagula'
 import { CarReader } from '@ipld/car'
 import { parseCid, HttpError, toIterable } from '@web3-storage/gateway-lib/util'
 import { BatchingR2Blockstore } from './lib/blockstore.js'
+import { MemoryBudget } from './lib/mem-budget.js'
 
 const MAX_CAR_BYTES_IN_MEMORY = 1024 * 1024 * 5
 const CAR_CODE = 0x0202
+const MAX_MEMORY_BUDGET = 1024 * 1024 * 50
 
 /**
  * @typedef {import('./bindings').Environment} Environment
  * @typedef {import('@web3-storage/gateway-lib').IpfsUrlContext} IpfsUrlContext
  * @typedef {import('./bindings').CarCidsContext} CarCidsContext
  * @typedef {import('@web3-storage/gateway-lib').DagulaContext} DagulaContext
+ * @typedef {import('./bindings').MemoryBudgetContext} MemoryBudgetContext
  */
 
 /**
@@ -75,15 +78,50 @@ export function withCarCids (handler) {
   }
 }
 
+/**
+ * @type {import('@web3-storage/gateway-lib').Middleware<MemoryBudgetContext>}
+ */
+export function withMemoryBudget (handler) {
+  return async (request, env, ctx) => {
+    const memoryBudget = new MemoryBudget(MAX_MEMORY_BUDGET)
+    return handler(request, env, { ...ctx, memoryBudget })
+  }
+}
+
+/**
+ * @type {import('@web3-storage/gateway-lib').Middleware<MemoryBudgetContext, MemoryBudgetContext>}
+ */
+export function withResponseMemoryRelease (handler) {
+  return async (request, env, ctx) => {
+    const response = await handler(request, env, ctx)
+
+    const body = response.body
+    if (!body) return response
+
+    console.log('adding response memory release transform...')
+
+    return new Response(
+      body.pipeThrough(new TransformStream({
+        transform (chunk, controller) {
+          ctx.memoryBudget.release(chunk.length)
+          controller.enqueue(chunk)
+        }
+      })),
+      response
+    )
+  }
+}
+
 /**
  * Creates a dagula instance backed by the R2 blockstore.
- * @type {import('@web3-storage/gateway-lib').Middleware<DagulaContext & CarCidsContext & IpfsUrlContext, CarCidsContext & IpfsUrlContext, Environment>}
+ * @type {import('@web3-storage/gateway-lib').Middleware<DagulaContext & MemoryBudgetContext & CarCidsContext & IpfsUrlContext, MemoryBudgetContext & CarCidsContext & IpfsUrlContext, Environment>}
  */
 export function withDagula (handler) {
   return async (request, env, ctx) => {
-    const { carCids, searchParams } = ctx
+    const { carCids, searchParams, memoryBudget } = ctx
     if (!carCids) throw new Error('missing CAR CIDs in context')
     if (!searchParams) throw new Error('missing URL search params in context')
+    if (!memoryBudget) throw new Error('missing memory budget instance')
 
     /** @type {import('dagula').Blockstore?} */
     let blockstore = null
@@ -99,7 +137,7 @@ export function withDagula (handler) {
     }
 
     if (!blockstore) {
-      blockstore = new BatchingR2Blockstore(env.CARPARK, env.SATNAV, carCids)
+      blockstore = new BatchingR2Blockstore(env.CARPARK, env.SATNAV, carCids, memoryBudget)
     }
 
     const dagula = new Dagula(blockstore)

src/middleware.js

@@ -2,7 +2,8 @@ name = "freeway"
 main = "./dist/worker.mjs"
 compatibility_date = "2022-10-03"
 compatibility_flags = [
-  "streams_enable_constructors"
+  "streams_enable_constructors",
+  "transformstream_enable_standard_constructor"
 ]
 r2_buckets = [
   { binding = "CARPARK", bucket_name = "carpark-dev-0" },