1 files changed, 13 insertions, 9 deletions
diff --git a/packages/0x.js/src/order_watcher/expiration_watcher.ts b/packages/0x.js/src/order_watcher/expiration_watcher.ts
index 71199e75f..7d6f8df65 100644
--- a/packages/0x.js/src/order_watcher/expiration_watcher.ts
+++ b/packages/0x.js/src/order_watcher/expiration_watcher.ts
@@ -1,9 +1,9 @@
 import * as _ from 'lodash';
 import {BigNumber} from 'bignumber.js';
+import {RBTree} from 'bintrees';
 import {utils} from '../utils/utils';
 import {intervalUtils} from '../utils/interval_utils';
 import {SignedOrder, ZeroExError} from '../types';
-import {Heap} from '../utils/heap';
 import {ZeroEx} from '../0x';
 
 const DEFAULT_EXPIRATION_MARGIN_MS = 0;
@@ -14,7 +14,7 @@ const DEFAULT_ORDER_EXPIRATION_CHECKING_INTERVAL_MS = 50;
  * It stores them in a min heap by expiration time and checks for expired ones every `orderExpirationCheckingIntervalMs`
  */
 export class ExpirationWatcher {
-    private orderHashHeapByExpiration: Heap<string>;
+    private orderHashRBTreeByExpiration: RBTree<string>;
     private expiration: {[orderHash: string]: BigNumber} = {};
     private callbackIfExists?: (orderHash: string) => void;
     private orderExpirationCheckingIntervalMs: number;
@@ -27,7 +27,8 @@ export class ExpirationWatcher {
         this.orderExpirationCheckingIntervalMs = expirationMarginIfExistsMs ||
                                                  DEFAULT_ORDER_EXPIRATION_CHECKING_INTERVAL_MS;
         const scoreFunction = (orderHash: string) => this.expiration[orderHash].toNumber();
-        this.orderHashHeapByExpiration = new Heap(scoreFunction);
+        const comparator = (lhs: string, rhs: string) => scoreFunction(lhs) - scoreFunction(rhs);
+        this.orderHashRBTreeByExpiration = new RBTree(comparator);
     }
     public subscribe(callback: (orderHash: string) => void): void {
         if (!_.isUndefined(this.callbackIfExists)) {
@@ -48,20 +49,23 @@ export class ExpirationWatcher {
     }
     public addOrder(orderHash: string, expirationUnixTimestampMs: BigNumber): void {
         this.expiration[orderHash] = expirationUnixTimestampMs;
-        // We don't remove hashes from the heap on order remove because it's slow (linear).
-        // We just skip them later if the order was already removed from the order watcher.
-        this.orderHashHeapByExpiration.push(orderHash);
+        this.orderHashRBTreeByExpiration.insert(orderHash);
+    }
+    public removeOrder(orderHash: string): void {
+        this.orderHashRBTreeByExpiration.remove(orderHash);
+        delete this.expiration[orderHash];
     }
     private pruneExpiredOrders(): void {
         const currentUnixTimestampMs = utils.getCurrentUnixTimestampMs();
         while (
-            this.orderHashHeapByExpiration.size() !== 0 &&
-            this.expiration[this.orderHashHeapByExpiration.head()].lessThan(
+            this.orderHashRBTreeByExpiration.size !== 0 &&
+            this.expiration[this.orderHashRBTreeByExpiration.min()].lessThan(
                 currentUnixTimestampMs.plus(this.expirationMarginMs),
             ) &&
             !_.isUndefined(this.callbackIfExists)
         ) {
-            const orderHash = this.orderHashHeapByExpiration.pop();
+            const orderHash = this.orderHashRBTreeByExpiration.min();
+            this.orderHashRBTreeByExpiration.remove(orderHash);
             delete this.expiration[orderHash];
             this.callbackIfExists(orderHash);
         }