From 3af104fbd3ba61160ddb1868283ff4e3debe533e Mon Sep 17 00:00:00 2001
From: "Andrew Y. Chen" <ayc@fb.com>
Date: Tue, 25 Oct 2016 07:12:34 -0700
Subject: [PATCH] Fix memory leak in HeadlessJsTaskContext

Reviewed By: foghina, AaaChiuuu

Differential Revision: D4068078

fbshipit-source-id: a45ad83e9ecd8455558968089d80f94ec104c2ef
---
 .../react/jstasks/HeadlessJsTaskContext.java         | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/ReactAndroid/src/main/java/com/facebook/react/jstasks/HeadlessJsTaskContext.java b/ReactAndroid/src/main/java/com/facebook/react/jstasks/HeadlessJsTaskContext.java
index 5dfb919aa..1304ef07e 100644
--- a/ReactAndroid/src/main/java/com/facebook/react/jstasks/HeadlessJsTaskContext.java
+++ b/ReactAndroid/src/main/java/com/facebook/react/jstasks/HeadlessJsTaskContext.java
@@ -2,6 +2,7 @@
 
 package com.facebook.react.jstasks;
 
+import java.lang.ref.WeakReference;
 import java.util.Set;
 import java.util.WeakHashMap;
 import java.util.concurrent.CopyOnWriteArraySet;
@@ -41,7 +42,7 @@ public class HeadlessJsTaskContext {
     return helper;
   }
 
-  private final ReactContext mReactContext;
+  private final WeakReference<ReactContext> mReactContext;
   private final Set<HeadlessJsTaskEventListener> mHeadlessJsTaskEventListeners =
     new CopyOnWriteArraySet<>();
   private final AtomicInteger mLastTaskId = new AtomicInteger(0);
@@ -50,7 +51,7 @@ public class HeadlessJsTaskContext {
   private final SparseArray<Runnable> mTaskTimeouts = new SparseArray<>();
 
   private HeadlessJsTaskContext(ReactContext reactContext) {
-    mReactContext = reactContext;
+    mReactContext = new WeakReference<ReactContext>(reactContext);
   }
 
   /**
@@ -82,14 +83,17 @@ public class HeadlessJsTaskContext {
    */
   public synchronized int startTask(final HeadlessJsTaskConfig taskConfig) {
     UiThreadUtil.assertOnUiThread();
-    if (mReactContext.getLifecycleState() == LifecycleState.RESUMED &&
+    ReactContext reactContext = Assertions.assertNotNull(
+      mReactContext.get(),
+      "Tried to start a task on a react context that has already been destroyed");
+    if (reactContext.getLifecycleState() == LifecycleState.RESUMED &&
       !taskConfig.isAllowedInForeground()) {
       throw new IllegalStateException(
         "Tried to start task " + taskConfig.getTaskKey() +
           " while in foreground, but this is not allowed.");
     }
     final int taskId = mLastTaskId.incrementAndGet();
-    mReactContext.getJSModule(AppRegistry.class)
+    reactContext.getJSModule(AppRegistry.class)
       .startHeadlessTask(taskId, taskConfig.getTaskKey(), taskConfig.getData());
     if (taskConfig.getTimeout() > 0) {
       scheduleTaskTimeout(taskId, taskConfig.getTimeout());