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,16 +145,25 @@ 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)
       })
 
       const postReceive = throttle(
         (offset: number) => {
-          this.emit(EVENTS.POST_RECEIVE_V2, req, {...upload, offset})
+          this.emit(EVENTS.POST_RECEIVE_V2, data, {...upload, offset})
         },
         this.options.postReceiveInterval,
         {leading: false}
@@ -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

@@ -37,23 +37,27 @@ export class MemoryLocker implements Locker {
     this.timeout = options?.acquireLockTimeout ?? 1000 * 30
   }
 
-  newLock(id: string) {
-    return new MemoryLock(id, this, this.timeout)
+  newLock(id: string, signal: AbortSignal) {
+    return new MemoryLock(id, this, this.timeout, signal)
   }
 }
 
 class MemoryLock implements Lock {
   constructor(
     private id: string,
     private locker: MemoryLocker,
-    private timeout: number = 1000 * 30
+    private timeout: number = 1000 * 30,
+    private abortSignal: AbortSignal
   ) {}
 
   async lock(requestRelease: RequestRelease): Promise<void> {
     const abortController = new AbortController()
+
+    const abortSignal = AbortSignal.any([this.abortSignal, 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 +108,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/server/test/Locker.test.ts

@@ -6,12 +6,13 @@ describe('MemoryLocker', () => {
   it('will acquire a lock by notifying another to release it', async () => {
     const locker = new MemoryLocker()
     const lockId = 'upload-id-1'
+    const abortController = new AbortController()
 
     const cancel = sinon.spy()
     const cancel2 = sinon.spy()
 
-    const lock1 = locker.newLock(lockId)
-    const lock2 = locker.newLock(lockId)
+    const lock1 = locker.newLock(lockId, abortController.signal)
+    const lock2 = locker.newLock(lockId, abortController.signal)
 
     await lock1.lock(async () => {
       await lock1.unlock()
@@ -32,8 +33,10 @@ describe('MemoryLocker', () => {
     const locker = new MemoryLocker({
       acquireLockTimeout: 500,
     })
+    const abortController = new AbortController()
+
     const lockId = 'upload-id-1'
-    const lock = locker.newLock(lockId)
+    const lock = locker.newLock(lockId, abortController.signal)
 
     const cancel = sinon.spy()
 
@@ -57,8 +60,10 @@ describe('MemoryLocker', () => {
   it('request lock and unlock', async () => {
     const locker = new MemoryLocker()
     const lockId = 'upload-id-1'
-    const lock = locker.newLock(lockId)
-    const lock2 = locker.newLock(lockId)
+    const abortController = new AbortController()
+
+    const lock = locker.newLock(lockId, abortController.signal)
+    const lock2 = locker.newLock(lockId, abortController.signal)
 
     const cancel = sinon.spy()
     await lock.lock(() => {
@@ -79,4 +84,39 @@ describe('MemoryLocker', () => {
       `request released called more times than expected - ${cancel.callCount}`
     )
   })
+
+  it('will stop trying to acquire the lock if the abort signal is aborted', async () => {
+    const locker = new MemoryLocker()
+    const lockId = 'upload-id-1'
+    const abortController = new AbortController()
+
+    const cancel = sinon.spy()
+    const cancel2 = sinon.spy()
+
+    const lock1 = locker.newLock(lockId, abortController.signal)
+    const lock2 = locker.newLock(lockId, abortController.signal)
+
+    await lock1.lock(async () => {
+      // do not unlock when requested
+      cancel()
+    })
+
+    // Abort signal is aborted after lock2 tries to acquire the lock
+    setTimeout(() => {
+      abortController.abort()
+    }, 100)
+
+    try {
+      await lock2.lock(async () => {
+        cancel2()
+      })
+      assert(false, 'lock2 should not have been acquired')
+    } catch (e) {
+      assert(e === ERRORS.ERR_LOCK_TIMEOUT, `error returned is not correct ${e}`)
+    }
+
+
+    assert(cancel.callCount > 1, `calls count dont match ${cancel.callCount} !== 1`)
+    assert(cancel2.callCount === 0, `calls count dont match ${cancel.callCount} !== 1`)
+  })
 })

packages/utils/src/models/Locker.ts

@@ -4,7 +4,7 @@ export type RequestRelease = () => Promise<void> | void
  * The Locker interface creates a Lock instance for a given resource identifier.
  */
 export interface Locker {
-  newLock(id: string): Lock
+  newLock(id: string, signal: AbortSignal): Lock
 }
 
 /**