Merge pull request #16628 from Budibase/app-export-memory-use

Use bounded concurrency for retrieveDirectory

Author: samwho

packages/backend-core/src/objectStore/objectStore.ts

@@ -25,6 +25,7 @@ import { ReadableStream } from "stream/web"
 import { NodeJsClient } from "@smithy/types"
 import tracer from "dd-trace"
 import { pipeline } from "stream/promises"
+import { utils } from "@budibase/shared-core"
 
 // use this as a temporary store of buckets that are being created
 const STATE = {
@@ -440,22 +441,30 @@ export async function retrieveDirectory(bucketName: string, path: string) {
     await fsp.mkdir(writePath, { recursive: true })
 
     let numObjects = 0
-    for await (const object of listAllObjects(bucketName, path)) {
-      numObjects++
-
-      const filename = object.Key!
-      const stream = await getReadStream(bucketName, filename)
-      const possiblePath = filename.split("/")
-      const dirs = possiblePath.slice(0, possiblePath.length - 1)
-      const possibleDir = join(writePath, ...dirs)
-      if (possiblePath.length > 1 && !fs.existsSync(possibleDir)) {
-        await fsp.mkdir(possibleDir, { recursive: true })
-      }
-      await pipeline(
-        stream,
-        fs.createWriteStream(join(writePath, ...possiblePath), { mode: 0o644 })
-      )
-    }
+    await utils.parallelForeach(
+      listAllObjects(bucketName, path),
+      async object => {
+        numObjects++
+        await tracer.trace("retrieveDirectory.object", async span => {
+          const filename = object.Key!
+          span.addTags({ filename })
+          const stream = await getReadStream(bucketName, filename)
+          const possiblePath = filename.split("/")
+          const dirs = possiblePath.slice(0, possiblePath.length - 1)
+          const possibleDir = join(writePath, ...dirs)
+          if (possiblePath.length > 1 && !fs.existsSync(possibleDir)) {
+            await fsp.mkdir(possibleDir, { recursive: true })
+          }
+          await pipeline(
+            stream,
+            fs.createWriteStream(join(writePath, ...possiblePath), {
+              mode: 0o644,
+            })
+          )
+        })
+      },
+      3 /* max concurrency */
+    )
 
     span.addTags({ numObjects })
     return writePath
@@ -592,18 +601,25 @@ export async function getReadStream(
   bucketName: string,
   path: string
 ): Promise<Readable> {
-  bucketName = sanitizeBucket(bucketName)
-  path = sanitizeKey(path)
-  const client = ObjectStore()
-  const params = {
-    Bucket: bucketName,
-    Key: path,
-  }
-  const response = await client.getObject(params)
-  if (!response.Body || !(response.Body instanceof stream.Readable)) {
-    throw new Error("Unable to retrieve stream - invalid response")
-  }
-  return response.Body
+  return await tracer.trace("getReadStream", async span => {
+    bucketName = sanitizeBucket(bucketName)
+    path = sanitizeKey(path)
+    span.addTags({ bucketName, path })
+    const client = ObjectStore()
+    const params = {
+      Bucket: bucketName,
+      Key: path,
+    }
+    const response = await client.getObject(params)
+    if (!response.Body || !(response.Body instanceof stream.Readable)) {
+      throw new Error("Unable to retrieve stream - invalid response")
+    }
+    span.addTags({
+      contentLength: response.ContentLength,
+      contentType: response.ContentType,
+    })
+    return response.Body
+  })
 }
 
 export async function getObjectMetadata(

packages/shared-core/src/utils.ts

@@ -39,24 +39,32 @@ export function unreachable(
 }
 
 export async function parallelForeach<T>(
-  items: T[],
+  items: Iterable<T> | AsyncIterable<T>,
   task: (item: T) => Promise<void>,
   maxConcurrency: number
 ): Promise<void> {
   const results: Promise<void>[] = []
-  let index = 0
 
-  const executeNext = async (): Promise<void> => {
-    while (index < items.length) {
-      const currentIndex = index++
-      const item = items[currentIndex]
+  // Check if it's an async iterable
+  const isAsyncIterable = Symbol.asyncIterator in items
+  const iterator = isAsyncIterable
+    ? (items as AsyncIterable<T>)[Symbol.asyncIterator]()
+    : (items as Iterable<T>)[Symbol.iterator]()
 
-      await task(item)
+  const executeNext = async (): Promise<void> => {
+    let result = await (isAsyncIterable
+      ? (iterator as AsyncIterator<T>).next()
+      : Promise.resolve((iterator as Iterator<T>).next()))
+
+    while (!result.done) {
+      await task(result.value)
+      result = await (isAsyncIterable
+        ? (iterator as AsyncIterator<T>).next()
+        : Promise.resolve((iterator as Iterator<T>).next()))
     }
   }
 
-  // Start up to maxConcurrency workers
-  for (let i = 0; i < Math.min(maxConcurrency, items.length); i++) {
+  for (let i = 0; i < maxConcurrency; i++) {
     results.push(executeNext())
   }