From 4312af5c5c5dde51873a8f02387bea25e23a36ca Mon Sep 17 00:00:00 2001
From: Ben Lesh <ben@benlesh.com>
Date: Sun, 15 Dec 2019 19:07:33 -0600
Subject: [PATCH] fix(asapScheduler): resolved memory leak (#5183)

Registered handlers would sometimes leak in memory, this resolves that issue and adds a test.

Related #5016
---
 spec/util/Immediate-spec.ts    | 14 ++++++++++++-
 src/internal/util/Immediate.ts | 36 +++++++++++++++++++++++++---------
 2 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/spec/util/Immediate-spec.ts b/spec/util/Immediate-spec.ts
index 932d3226a9..9b1616af96 100644
--- a/spec/util/Immediate-spec.ts
+++ b/spec/util/Immediate-spec.ts
@@ -1,5 +1,5 @@
 import { expect } from 'chai';
-import { Immediate } from 'rxjs/util/Immediate';
+import { Immediate, TestTools } from 'rxjs/internal/util/Immediate';
 
 describe('Immediate', () => {
   it('should schedule on the next microtask', (done) => {
@@ -30,4 +30,16 @@ describe('Immediate', () => {
       done();
     });
   });
+
+  it('should clear the task after execution', (done) => {
+    const results: number[] = [];
+    Immediate.setImmediate(() => results.push(1));
+    Immediate.setImmediate(() => results.push(2));
+
+    setTimeout(() => {
+      const number = TestTools.pending();
+      expect(number).to.equal(0);
+      done();
+    });
+  });
 });
diff --git a/src/internal/util/Immediate.ts b/src/internal/util/Immediate.ts
index 161270706d..ba296613f2 100644
--- a/src/internal/util/Immediate.ts
+++ b/src/internal/util/Immediate.ts
@@ -1,23 +1,41 @@
 let nextHandle = 1;
+const RESOLVED = (() => Promise.resolve())();
+const activeHandles: { [key: number]: any } = {};
 
-const tasksByHandle: { [handle: string]: () => void } = {};
-
-function runIfPresent(handle: number) {
-  const cb = tasksByHandle[handle];
-  if (cb) {
-    cb();
+/**
+ * Finds the handle in the list of active handles, and removes it.
+ * Returns `true` if found, `false` otherwise. Used both to clear
+ * Immediate scheduled tasks, and to identify if a task should be scheduled.
+ */
+function findAndClearHandle(handle: number): boolean {
+  if (handle in activeHandles) {
+    delete activeHandles[handle];
+    return true;
   }
+  return false;
 }
 
+/**
+ * Helper functions to schedule and unschedule microtasks.
+ */
 export const Immediate = {
   setImmediate(cb: () => void): number {
     const handle = nextHandle++;
-    tasksByHandle[handle] = cb;
-    Promise.resolve().then(() => runIfPresent(handle));
+    activeHandles[handle] = true;
+    RESOLVED.then(() => findAndClearHandle(handle) && cb());
     return handle;
   },
 
   clearImmediate(handle: number): void {
-    delete tasksByHandle[handle];
+    findAndClearHandle(handle);
   },
 };
+
+/**
+ * Used for internal testing purposes only. Do not export from library.
+ */
+export const TestTools = {
+  pending() {
+    return Object.keys(activeHandles).length;
+  }
+};