From 4810d13094f388ea2240555a8774ee6093e407a7 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Wed, 22 Nov 2017 11:41:55 -0800
Subject: [PATCH] Fix memory leak in .maintainLanguageMode

---
 spec/grammar-registry-spec.js | 67 ++++++++++++++++++++++++++---------
 src/grammar-registry.js       | 16 +++++++--
 2 files changed, 64 insertions(+), 19 deletions(-)

diff --git a/spec/grammar-registry-spec.js b/spec/grammar-registry-spec.js
index d64c0343f..215417d10 100644
--- a/spec/grammar-registry-spec.js
+++ b/spec/grammar-registry-spec.js
@@ -63,7 +63,7 @@ describe('GrammarRegistry', () => {
     })
   })
 
-  describe('.maintainLanguageMode', () => {
+  describe('.maintainLanguageMode(buffer)', () => {
     it('assigns a grammar to the buffer based on its path', async () => {
       const buffer = new TextBuffer()
 
@@ -128,24 +128,55 @@ describe('GrammarRegistry', () => {
       expect(retainedBufferCount(grammarRegistry)).toBe(0)
     })
 
-    describe('when called twice with a given buffer', () => {
-      it('does nothing the second time', async () => {
-        const buffer = new TextBuffer()
-        grammarRegistry.loadGrammarSync(require.resolve('language-javascript/grammars/javascript.cson'))
-        const disposable1 = grammarRegistry.maintainLanguageMode(buffer)
-        const disposable2 = grammarRegistry.maintainLanguageMode(buffer)
+    it('doesn\'t do anything when called a second time with the same buffer', async () => {
+      const buffer = new TextBuffer()
+      grammarRegistry.loadGrammarSync(require.resolve('language-javascript/grammars/javascript.cson'))
+      const disposable1 = grammarRegistry.maintainLanguageMode(buffer)
+      const disposable2 = grammarRegistry.maintainLanguageMode(buffer)
 
-        buffer.setPath('test.js')
-        expect(buffer.getLanguageMode().getLanguageName()).toBe('JavaScript')
+      buffer.setPath('test.js')
+      expect(buffer.getLanguageMode().getLanguageName()).toBe('JavaScript')
 
-        disposable2.dispose()
-        buffer.setPath('test.txt')
-        expect(buffer.getLanguageMode().getLanguageName()).toBe('Null Grammar')
+      disposable2.dispose()
+      buffer.setPath('test.txt')
+      expect(buffer.getLanguageMode().getLanguageName()).toBe('Null Grammar')
 
-        disposable1.dispose()
-        buffer.setPath('test.js')
-        expect(buffer.getLanguageMode().getLanguageName()).toBe('Null Grammar')
-      })
+      disposable1.dispose()
+      buffer.setPath('test.js')
+      expect(buffer.getLanguageMode().getLanguageName()).toBe('Null Grammar')
+    })
+
+    it('does not retain the buffer after the buffer is destroyed', () => {
+      const buffer = new TextBuffer()
+      grammarRegistry.loadGrammarSync(require.resolve('language-javascript/grammars/javascript.cson'))
+
+      const disposable = grammarRegistry.maintainLanguageMode(buffer)
+      expect(retainedBufferCount(grammarRegistry)).toBe(1)
+      expect(subscriptionCount(grammarRegistry)).toBe(2)
+
+      buffer.destroy()
+      expect(retainedBufferCount(grammarRegistry)).toBe(0)
+      expect(subscriptionCount(grammarRegistry)).toBe(0)
+      expect(buffer.emitter.getTotalListenerCount()).toBe(0)
+
+      disposable.dispose()
+      expect(retainedBufferCount(grammarRegistry)).toBe(0)
+      expect(subscriptionCount(grammarRegistry)).toBe(0)
+    })
+
+    it('does not retain the buffer when the grammar registry is destroyed', () => {
+      const buffer = new TextBuffer()
+      grammarRegistry.loadGrammarSync(require.resolve('language-javascript/grammars/javascript.cson'))
+
+      const disposable = grammarRegistry.maintainLanguageMode(buffer)
+      expect(retainedBufferCount(grammarRegistry)).toBe(1)
+      expect(subscriptionCount(grammarRegistry)).toBe(2)
+
+      grammarRegistry.clear()
+
+      expect(retainedBufferCount(grammarRegistry)).toBe(0)
+      expect(subscriptionCount(grammarRegistry)).toBe(0)
+      expect(buffer.emitter.getTotalListenerCount()).toBe(0)
     })
   })
 
@@ -357,3 +388,7 @@ describe('GrammarRegistry', () => {
 function retainedBufferCount (grammarRegistry) {
   return grammarRegistry.grammarScoresByBuffer.size
 }
+
+function subscriptionCount (grammarRegistry) {
+  return grammarRegistry.subscriptions.disposables.size
+}
diff --git a/src/grammar-registry.js b/src/grammar-registry.js
index 361ee0011..dd467cfb9 100644
--- a/src/grammar-registry.js
+++ b/src/grammar-registry.js
@@ -71,12 +71,22 @@ class GrammarRegistry extends FirstMate.GrammarRegistry {
       }
     })
 
-    this.subscriptions.add(pathChangeSubscription)
+    const destroySubscription = buffer.onDidDestroy(() => {
+      this.grammarScoresByBuffer.delete(buffer)
+      this.languageNameOverridesByBufferId.delete(buffer.id)
+      this.subscriptions.remove(destroySubscription)
+      this.subscriptions.remove(pathChangeSubscription)
+    })
+
+    this.subscriptions.add(pathChangeSubscription, destroySubscription)
 
     return new Disposable(() => {
-      this.subscriptions.remove(pathChangeSubscription)
-      this.grammarScoresByBuffer.delete(buffer)
+      destroySubscription.dispose()
       pathChangeSubscription.dispose()
+      this.subscriptions.remove(pathChangeSubscription)
+      this.subscriptions.remove(destroySubscription)
+      this.grammarScoresByBuffer.delete(buffer)
+      this.languageNameOverridesByBufferId.delete(buffer.id)
     })
   }