fix: handling consistent cancellation across stream and locks

Author: fenos

packages/server/src/handlers/BaseHandler.ts

@@ -1,6 +1,6 @@
 import EventEmitter from 'node:events'
 import stream from 'node:stream/promises'
-import {addAbortSignal, PassThrough} from 'node:stream'
+import {PassThrough, Readable} from 'node:stream'
 import type http from 'node:http'
 
 import type {ServerOptions} from '../types'
@@ -119,7 +119,7 @@ export class BaseHandler extends EventEmitter {
   ) {
     const locker = await this.getLocker(req)
 
-    const lock = locker.newLock(id)
+    const lock = locker.newLock(id, context.signal)
 
     await lock.lock(() => {
       context.cancel()
@@ -129,7 +129,7 @@ export class BaseHandler extends EventEmitter {
   }
 
   protected writeToStore(
-    req: http.IncomingMessage,
+    data: Readable,
     upload: Upload,
     maxFileSize: number,
     context: CancellationContext
@@ -145,10 +145,19 @@ export class BaseHandler extends EventEmitter {
       // Create a PassThrough stream as a proxy to manage the request stream.
       // This allows for aborting the write process without affecting the incoming request stream.
       const proxy = new PassThrough()
-      addAbortSignal(context.signal, proxy)
+
+      // gracefully terminate the proxy stream when the request is aborted
+      const onAbort = () => {
+        data.unpipe(proxy)
+
+        if (!proxy.closed) {
+          proxy.end()
+        }
+      }
+      context.signal.addEventListener('abort', onAbort, { once: true })
 
       proxy.on('error', (err) => {
-        req.unpipe(proxy)
+        data.unpipe(proxy)
         reject(err.name === 'AbortError' ? ERRORS.ABORTED : err)
       })
 
@@ -166,23 +175,18 @@ export class BaseHandler extends EventEmitter {
         postReceive(tempOffset)
       })
 
-      req.on('error', () => {
-        if (!proxy.closed) {
-          // we end the stream gracefully here so that we can upload the remaining bytes to the store
-          // as an incompletePart
-          proxy.end()
-        }
-      })
-
       // Pipe the request stream through the proxy. We use the proxy instead of the request stream directly
       // to ensure that errors in the pipeline do not cause the request stream to be destroyed,
       // which would result in a socket hangup error for the client.
       stream
-        .pipeline(req.pipe(proxy), new StreamLimiter(maxFileSize), async (stream) => {
+        .pipeline(data.pipe(proxy), new StreamLimiter(maxFileSize), async (stream) => {
           return this.store.write(stream as StreamLimiter, upload.id, upload.offset)
         })
         .then(resolve)
         .catch(reject)
+        .finally(() => {
+          context.signal.removeEventListener('abort', onAbort)
+        })
     })
   }
 

packages/server/src/lockers/MemoryLocker.ts

@@ -49,11 +49,14 @@ class MemoryLock implements Lock {
     private timeout: number = 1000 * 30
   ) {}
 
-  async lock(requestRelease: RequestRelease): Promise<void> {
+  async lock(requestRelease: RequestRelease, signal: AbortSignal): Promise<void> {
     const abortController = new AbortController()
+
+    const abortSignal = AbortSignal.any([signal, abortController.signal])
+
     const lock = await Promise.race([
-      this.waitTimeout(abortController.signal),
-      this.acquireLock(this.id, requestRelease, abortController.signal),
+      this.waitTimeout(abortSignal),
+      this.acquireLock(this.id, requestRelease, abortSignal),
     ])
 
     abortController.abort()
@@ -104,6 +107,12 @@ class MemoryLock implements Lock {
     this.locker.locks.delete(this.id)
   }
 
+  protected waitAbortSignal(signal: AbortSignal) {
+    return new Promise((resolve, reject) => {
+
+    })
+  }
+
   protected waitTimeout(signal: AbortSignal) {
     return new Promise<boolean>((resolve) => {
       const timeout = setTimeout(() => {

packages/server/src/server.ts

@@ -151,6 +151,11 @@ export class Server extends EventEmitter {
   ): Promise<http.ServerResponse | stream.Writable | void> {
     const context = this.createContext(req)
 
+    // Once the request is closed we abort the context to clean up underline resources
+    req.on('close', () => {
+      context.abort()
+    })
+
     log(`[TusServer] handle: ${req.method} ${req.url}`)
     // Allow overriding the HTTP method. The reason for this is
     // that some libraries/environments to not support PATCH and
@@ -289,6 +294,17 @@ export class Server extends EventEmitter {
 
     res.writeHead(status, headers)
     res.write(body)
+
+    // Abort the context once the response is sent.
+    // Useful for clean-up when the server uses keep-alive
+    if (!isAborted) {
+      res.on('finish', () => {
+        if (!req.closed) {
+          context.abort()
+        }
+      })
+    }
+
     return res.end()
   }
 

packages/utils/src/models/Locker.ts

@@ -26,6 +26,6 @@ export interface Locker {
  *
  */
 export interface Lock {
-  lock(cancelReq: RequestRelease): Promise<void>
+  lock(cancelReq: RequestRelease, signal: AbortSignal): Promise<void>
   unlock(): Promise<void>
 }