fix: dont mark recycled workers as available (#61)

* test: add failing test case for recycled workers * fix: dont mark recycled workers as available
tinylibs · Jun 28, 2023 · f3b4afa · f3b4afa
1 parent 28ba15f
commit f3b4afa
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 13 deletions.
diff --git a/src/index.ts b/src/index.ts
@@ -705,7 +705,10 @@ class ThreadPool {
       const taskInfo = workerInfo.taskInfos.get(taskId)
       workerInfo.taskInfos.delete(taskId)
 
-      if (!this.options.isolateWorkers) pool.workers.maybeAvailable(workerInfo)
+      // Mark worker as available if it's not about to be removed
+      if (!this.shouldRecycleWorker(taskInfo)) {
+        pool.workers.maybeAvailable(workerInfo)
+      }
 
       /* istanbul ignore if */
       if (taskInfo === undefined) {
@@ -894,18 +897,7 @@ class ThreadPool {
           reject(err)
         }
 
-        // When `isolateWorkers` is enabled, remove the worker after task is finished
-        const shouldIsolateWorker =
-          this.options.isolateWorkers && taskInfo.workerInfo
-
-        // When `maxMemoryLimitBeforeRecycle` is enabled, remove workers that have exceeded the memory limit
-        const shouldRecycleWorker =
-          !this.options.isolateWorkers &&
-          this.options.maxMemoryLimitBeforeRecycle !== undefined &&
-          (taskInfo.workerInfo?.usedMemory || 0) >
-            this.options.maxMemoryLimitBeforeRecycle
-
-        if (shouldIsolateWorker || shouldRecycleWorker) {
+        if (this.shouldRecycleWorker(taskInfo)) {
           this._removeWorker(taskInfo.workerInfo!)
             .then(() => this._ensureMinimumWorkers())
             .then(() => this._ensureEnoughWorkersForTaskQueue())
@@ -1002,6 +994,20 @@ class ThreadPool {
     return ret
   }
 
+  shouldRecycleWorker(taskInfo?: TaskInfo): boolean {
+    // When `isolateWorkers` is enabled, remove the worker after task is finished
+    const isWorkerIsolated = this.options.isolateWorkers && taskInfo?.workerInfo
+
+    // When `maxMemoryLimitBeforeRecycle` is enabled, remove workers that have exceeded the memory limit
+    const isWorkersMemoryLimitReached =
+      !this.options.isolateWorkers &&
+      this.options.maxMemoryLimitBeforeRecycle !== undefined &&
+      (taskInfo?.workerInfo?.usedMemory || 0) >
+        this.options.maxMemoryLimitBeforeRecycle
+
+    return Boolean(isWorkerIsolated || isWorkersMemoryLimitReached)
+  }
+
   pendingCapacity(): number {
     return (
       this.workers.pendingItems.size * this.options.concurrentTasksPerWorker

diff --git a/test/resource-limits.test.ts b/test/resource-limits.test.ts
@@ -74,3 +74,20 @@ test('worker is recycled after reaching maxMemoryLimitBeforeRecycle', async () =
   // Thread should have been recycled
   expect(finalThreadId).not.toBe(originalWorkerId)
 })
+
+test('recycled workers should not crash pool (regression)', async () => {
+  const pool = new Tinypool({
+    filename: resolve(__dirname, 'fixtures/leak-memory.js'),
+    maxMemoryLimitBeforeRecycle: 10,
+    isolateWorkers: false,
+    minThreads: 2,
+    maxThreads: 2,
+  })
+
+  // This should not crash the pool
+  await Promise.all(
+    Array(10)
+      .fill(0)
+      .map(() => pool.run(10_000))
+  )
+})