@s3/store: add backpressure + download incomplete part first (#561)

Co-authored-by: Merlijn Vos <[email protected]>

packages/s3-store/README.md

@@ -91,6 +91,22 @@ you need to provide a cache implementation that is shared between all instances
 
 See the exported [KV stores][kvstores] from `@tus/server` for more information.
 
+#### `options.maxConcurrentPartUploads`
+
+This setting determines the maximum number of simultaneous part uploads to an S3 storage service.
+The default value is 60. This default is chosen in conjunction with the typical partSize of 8MiB, aiming for an effective transfer rate of 3.84Gbit/s.
+
+**Considerations:**
+The ideal value for `maxConcurrentPartUploads` varies based on your `partSize` and the upload bandwidth to your S3 bucket. A larger partSize means less overall upload bandwidth available for other concurrent uploads.
+
+- **Lowering the Value**: Reducing `maxConcurrentPartUploads` decreases the number of simultaneous upload requests to S3. This can be beneficial for conserving memory, CPU, and disk I/O resources, especially in environments with limited system resources or where the upload speed it low or the part size is large.
+
+
+- **Increasing the Value**: A higher value potentially enhances the data transfer rate to the server, but at the cost of increased resource usage (memory, CPU, and disk I/O). This can be advantageous when the goal is to maximize throughput, and sufficient system resources are available.
+
+
+- **Bandwidth Considerations**: It's important to note that if your upload bandwidth to S3 is a limiting factor, increasing `maxConcurrentPartUploads` won’t lead to higher throughput. Instead, it will result in additional resource consumption without proportional gains in transfer speed.
+
 ## Extensions
 
 The tus protocol supports optional [extensions][]. Below is a table of the supported extensions in `@tus/s3-store`.

packages/s3-store/index.ts

@@ -2,7 +2,6 @@ import os from 'node:os'
 import fs, {promises as fsProm} from 'node:fs'
 import stream, {promises as streamProm} from 'node:stream'
 import type {Readable} from 'node:stream'
-import http from 'node:http'
 
 import AWS, {NoSuchKey, NotFound, S3, S3ClientConfig} from '@aws-sdk/client-s3'
 import debug from 'debug'
@@ -17,6 +16,11 @@ import {
   MemoryKvStore,
 } from '@tus/utils'
 
+import {Semaphore, Permit} from '@shopify/semaphore'
+import MultiStream from 'multistream'
+import crypto from 'node:crypto'
+import path from 'node:path'
+
 const log = debug('tus-node-server:stores:s3store')
 
 type Options = {
@@ -25,6 +29,7 @@ type Options = {
   // but may increase it to not exceed the S3 10K parts limit.
   partSize?: number
   useTags?: boolean
+  maxConcurrentPartUploads?: number
   cache?: KvStore<MetadataValue>
   expirationPeriodInMilliseconds?: number
   // Options to pass to the AWS S3 SDK.
@@ -82,6 +87,7 @@ export class S3Store extends DataStore {
   private preferredPartSize: number
   private expirationPeriodInMilliseconds = 0
   private useTags = true
+  private partUploadSemaphore: Semaphore
   public maxMultipartParts = 10_000 as const
   public minPartSize = 5_242_880 as const // 5MiB
   public maxUploadSize = 5_497_558_138_880 as const // 5TiB
@@ -101,8 +107,9 @@ export class S3Store extends DataStore {
     this.preferredPartSize = partSize || 8 * 1024 * 1024
     this.expirationPeriodInMilliseconds = options.expirationPeriodInMilliseconds ?? 0
     this.useTags = options.useTags ?? true
-    this.client = new S3(restS3ClientConfig)
     this.cache = options.cache ?? new MemoryKvStore<MetadataValue>()
+    this.client = new S3(restS3ClientConfig)
+    this.partUploadSemaphore = new Semaphore(options.maxConcurrentPartUploads ?? 60)
   }
 
   protected shouldUseExpirationTags() {
@@ -233,6 +240,61 @@ export class S3Store extends DataStore {
     return data.ETag as string
   }
 
+  private async downloadIncompletePart(id: string) {
+    const incompletePart = await this.getIncompletePart(id)
+
+    if (!incompletePart) {
+      return
+    }
+    const filePath = await this.uniqueTmpFileName('tus-s3-incomplete-part-')
+
+    try {
+      let incompletePartSize = 0
+
+      const byteCounterTransform = new stream.Transform({
+        transform(chunk, _, callback) {
+          incompletePartSize += chunk.length
+          callback(null, chunk)
+        },
+      })
+
+      // write to temporary file
+      await streamProm.pipeline(
+        incompletePart,
+        byteCounterTransform,
+        fs.createWriteStream(filePath)
+      )
+
+      const createReadStream = (options: {cleanUpOnEnd: boolean}) => {
+        const fileReader = fs.createReadStream(filePath)
+
+        if (options.cleanUpOnEnd) {
+          fileReader.on('end', () => {
+            fs.unlink(filePath, () => {
+              // ignore
+            })
+          })
+
+          fileReader.on('error', (err) => {
+            fileReader.destroy(err)
+            fs.unlink(filePath, () => {
+              // ignore
+            })
+          })
+        }
+
+        return fileReader
+      }
+
+      return {size: incompletePartSize, path: filePath, createReader: createReadStream}
+    } catch (err) {
+      fsProm.rm(filePath).catch(() => {
+        /* ignore */
+      })
+      throw err
+    }
+  }
+
   private async getIncompletePart(id: string): Promise<Readable | undefined> {
     try {
       const data = await this.client.getObject({
@@ -271,102 +333,50 @@ export class S3Store extends DataStore {
     })
   }
 
-  private async prependIncompletePart(
-    newChunkPath: string,
-    previousIncompletePart: Readable
-  ): Promise<number> {
-    const tempPath = `${newChunkPath}-prepend`
-    try {
-      let incompletePartSize = 0
-
-      const byteCounterTransform = new stream.Transform({
-        transform(chunk, _, callback) {
-          incompletePartSize += chunk.length
-          callback(null, chunk)
-        },
-      })
-
-      // write to temporary file, truncating if needed
-      await streamProm.pipeline(
-        previousIncompletePart,
-        byteCounterTransform,
-        fs.createWriteStream(tempPath)
-      )
-      // append to temporary file
-      await streamProm.pipeline(
-        fs.createReadStream(newChunkPath),
-        fs.createWriteStream(tempPath, {flags: 'a'})
-      )
-      // overwrite existing file
-      await fsProm.rename(tempPath, newChunkPath)
-
-      return incompletePartSize
-    } catch (err) {
-      fsProm.rm(tempPath).catch(() => {
-        /* ignore */
-      })
-      throw err
-    }
-  }
-
   /**
    * Uploads a stream to s3 using multiple parts
    */
-  private async processUpload(
+  private async uploadParts(
     metadata: MetadataValue,
-    readStream: http.IncomingMessage | fs.ReadStream,
+    readStream: stream.Readable,
     currentPartNumber: number,
     offset: number
   ): Promise<number> {
     const size = metadata.file.size
     const promises: Promise<void>[] = []
     let pendingChunkFilepath: string | null = null
     let bytesUploaded = 0
-    let currentChunkNumber = 0
+    let permit: Permit | undefined = undefined
 
     const splitterStream = new StreamSplitter({
       chunkSize: this.calcOptimalPartSize(size),
       directory: os.tmpdir(),
     })
+      .on('beforeChunkStarted', async () => {
+        permit = await this.partUploadSemaphore.acquire()
+      })
       .on('chunkStarted', (filepath) => {
         pendingChunkFilepath = filepath
       })
       .on('chunkFinished', ({path, size: partSize}) => {
         pendingChunkFilepath = null
 
         const partNumber = currentPartNumber++
-        const chunkNumber = currentChunkNumber++
+        const acquiredPermit = permit
 
         offset += partSize
 
-        const isFirstChunk = chunkNumber === 0
         const isFinalPart = size === offset
 
         // eslint-disable-next-line no-async-promise-executor
         const deferred = new Promise<void>(async (resolve, reject) => {
           try {
-            let incompletePartSize = 0
             // Only the first chunk of each PATCH request can prepend
             // an incomplete part (last chunk) from the previous request.
-            if (isFirstChunk) {
-              // If we received a chunk under the minimum part size in a previous iteration,
-              // we used a regular S3 upload to save it in the bucket. We try to get the incomplete part here.
-
-              const incompletePart = await this.getIncompletePart(metadata.file.id)
-              if (incompletePart) {
-                // We found an incomplete part, prepend it to the chunk on disk we were about to upload,
-                // and delete the incomplete part from the bucket. This can be done in parallel.
-                incompletePartSize = await this.prependIncompletePart(
-                  path,
-                  incompletePart
-                )
-                await this.deleteIncompletePart(metadata.file.id)
-              }
-            }
-
             const readable = fs.createReadStream(path)
             readable.on('error', reject)
-            if (partSize + incompletePartSize >= this.minPartSize || isFinalPart) {
+
+            if (partSize >= this.minPartSize || isFinalPart) {
               await this.uploadPart(metadata, readable, partNumber)
             } else {
               await this.uploadIncompletePart(metadata.file.id, readable)
@@ -380,11 +390,15 @@ export class S3Store extends DataStore {
             fsProm.rm(path).catch(() => {
               /* ignore */
             })
+            acquiredPermit?.release()
           }
         })
 
         promises.push(deferred)
       })
+      .on('chunkError', () => {
+        permit?.release()
+      })
 
     try {
       await streamProm.pipeline(readStream, splitterStream)
@@ -533,26 +547,30 @@ export class S3Store extends DataStore {
   /**
    * Write to the file, starting at the provided offset
    */
-  public async write(
-    readable: http.IncomingMessage | fs.ReadStream,
-    id: string,
-    offset: number
-  ): Promise<number> {
+  public async write(src: stream.Readable, id: string, offset: number): Promise<number> {
     // Metadata request needs to happen first
     const metadata = await this.getMetadata(id)
     const parts = await this.retrieveParts(id)
     // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
     const partNumber: number = parts.length > 0 ? parts[parts.length - 1].PartNumber! : 0
     const nextPartNumber = partNumber + 1
 
-    const bytesUploaded = await this.processUpload(
-      metadata,
-      readable,
-      nextPartNumber,
-      offset
-    )
+    const incompletePart = await this.downloadIncompletePart(id)
+    const requestedOffset = offset
+
+    if (incompletePart) {
+      // once the file is on disk, we delete the incomplete part
+      await this.deleteIncompletePart(id)
+
+      offset = requestedOffset - incompletePart.size
+      src = new MultiStream([incompletePart.createReader({cleanUpOnEnd: true}), src])
+    }
+
+    const bytesUploaded = await this.uploadParts(metadata, src, nextPartNumber, offset)
 
-    const newOffset = offset + bytesUploaded
+    // The size of the incomplete part should not be counted, because the
+    // process of the incomplete part should be fully transparent to the user.
+    const newOffset = requestedOffset + bytesUploaded - (incompletePart?.size ?? 0)
 
     if (metadata.file.size === newOffset) {
       try {
@@ -741,4 +759,29 @@ export class S3Store extends DataStore {
 
     return deleted
   }
+
+  private async uniqueTmpFileName(template: string): Promise<string> {
+    let tries = 0
+    const maxTries = 10
+
+    while (tries < maxTries) {
+      const fileName =
+        template + crypto.randomBytes(10).toString('base64url').slice(0, 10)
+      const filePath = path.join(os.tmpdir(), fileName)
+
+      try {
+        await fsProm.lstat(filePath)
+        // If no error, file exists, so try again
+        tries++
+      } catch (e) {
+        if (e.code === 'ENOENT') {
+          // File does not exist, return the path
+          return filePath
+        }
+        throw e // For other errors, rethrow
+      }
+    }
+
+    throw new Error(`Could not find a unique file name after ${maxTries} tries`)
+  }
 }

packages/s3-store/package.json

@@ -22,12 +22,15 @@
   },
   "dependencies": {
     "@aws-sdk/client-s3": "^3.490.0",
+    "@shopify/semaphore": "^3.0.2",
     "@tus/utils": "workspace:*",
-    "debug": "^4.3.4"
+    "debug": "^4.3.4",
+    "multistream": "^4.1.0"
   },
   "devDependencies": {
     "@types/debug": "^4.1.12",
     "@types/mocha": "^10.0.6",
+    "@types/multistream": "^4.1.3",
     "@types/node": "^20.11.5",
     "eslint": "^8.56.0",
     "eslint-config-custom": "workspace:*",

packages/s3-store/test.ts

@@ -1,5 +1,4 @@
 import path from 'node:path'
-import fs from 'node:fs/promises'
 import assert from 'node:assert/strict'
 import {Readable} from 'node:stream'
 
@@ -39,17 +38,6 @@ describe('S3DataStore', function () {
     assert.strictEqual(Number.isFinite(store.calcOptimalPartSize(undefined)), true)
   })
 
-  it('should correctly prepend a buffer to a file', async function () {
-    const p = path.resolve(fixturesPath, 'foo.txt')
-    await fs.writeFile(p, 'world!')
-    await this.datastore.prependIncompletePart(
-      p,
-      Readable.from([new TextEncoder().encode('Hello, ')])
-    )
-    assert.strictEqual(await fs.readFile(p, 'utf8'), 'Hello, world!')
-    await fs.unlink(p)
-  })
-
   it('should store in between chunks under the minimum part size and prepend it to the next call', async function () {
     const store = this.datastore
     const size = 1024

packages/server/src/handlers/BaseHandler.ts

@@ -149,9 +149,11 @@ export class BaseHandler extends EventEmitter {
         reject(err.name === 'AbortError' ? ERRORS.ABORTED : err)
       })
 
-      req.on('error', (err) => {
+      req.on('error', () => {
         if (!proxy.closed) {
-          proxy.destroy(err)
+          // we end the stream gracefully here so that we can upload the remaining bytes to the store
+          // as an incompletePart
+          proxy.end()
         }
       })