Allow headless JS tasks to retry (#23231)

Summary:
`setTimeout` inside a headless JS task does not always works; the function does not get invoked until the user starts an `Activity`.

This was attempted to be used in the context of widgets. When the widget update or user interaction causes the process and React context to be created, the headless JS task may run before other app-specific JS initialisation logic has completed. If it's not possible to change the behaviour of the pre-requisites to be synchronous, then the headless JS task blocks such asynchronous JS work that it may depend on. A primitive solution is the use of `setTimeout` in order to wait for the pre-conditions to be met before continuing with the rest of the headless JS task. But as the function passed to `setTimeout` is not always called, the task will not run to completion.

This PR solves this scenario by allowing the task to be retried again with a delay. If the task returns a promise that resolves to a `{'timeout': number}` object, `AppRegistry.js` will not notify that the task has finished as per master, instead it will tell `HeadlessJsContext` to `startTask` again (cleaning up any posted `Runnable`s beforehand) via a `Handler` within the `HeadlessJsContext`.

Documentation also updated here: https://github.com/facebook/react-native-website/pull/771

### AppRegistry.js
If the task provider does not return any data, or if the data it returns does not contain `timeout` as a number, then it behaves as `master`; notifies that the task has finished. If the response does contain `{timeout: number}`, then it will attempt to queue a retry. If that fails, then it will behaves as if the task provider returned no response i.e. behaves as `master` again. If the retry was successfully queued, then there is nothing to do as we do not want the `Service` to stop itself.

### HeadlessJsTaskSupportModule.java
Similar to notify start/finished, we simply check if the context is running, and if so, pass the request onto `HeadlessJsTaskContext`. The only difference here is that we return a `Promise`, so that `AppRegistry`, as above, knows whether the enqueuing failed and thus needs to perform the usual task clean-up.

### HeadlessJsTaskContext.java
Before retrying, we need to clean-up any timeout `Runnable`'s posted for the first attempt. Then we need to copy the task config so that if this retry (second attempt) also fails, then on the third attempt (second retry) we do not run into a consumed exception. This is also why in `startTask` we copy the config before putting it in the `Map`, so that the initial attempt does leave the config's in the map as consumed. Then we post a `Runnable` to call `startTask` on the main thread's `Handler`. We use the same `taskId` because the `Service` is keeping track of active task IDs in order to calculate whether it needs to `stopSelf`. This negates the need to inform the `Service` of a new task id and us having to remove the old one.

## Changelog
[Android][added] - Allow headless JS tasks to return a promise that will cause the task to be retried again with the specified delay
Pull Request resolved: https://github.com/facebook/react-native/pull/23231

Differential Revision: D15646870

fbshipit-source-id: 4440f4b4392f1fa5c69aab7908b51b7007ba2c40

ReactAndroid/src/main/java/com/facebook/react/bridge/JavaOnlyMap.java

@@ -206,6 +206,13 @@ public class JavaOnlyMap implements ReadableMap, WritableMap {
     mBackingMap.putAll(((JavaOnlyMap) source).mBackingMap);
   }
 
+  @Override
+  public WritableMap copy() {
+    final JavaOnlyMap target = new JavaOnlyMap();
+    target.merge(this);
+    return target;
+  }
+
   @Override
   public void putArray(@Nonnull String key, @Nullable WritableArray value) {
     mBackingMap.put(key, value);

ReactAndroid/src/main/java/com/facebook/react/bridge/UiThreadUtil.java

@@ -44,11 +44,18 @@ public class UiThreadUtil {
    * Runs the given {@code Runnable} on the UI thread.
    */
   public static void runOnUiThread(Runnable runnable) {
+    runOnUiThread(runnable, 0);
+  }
+
+  /**
+   * Runs the given {@code Runnable} on the UI thread with the specified delay.
+   */
+  public static void runOnUiThread(Runnable runnable, long delayInMs) {
     synchronized (UiThreadUtil.class) {
       if (sMainHandler == null) {
         sMainHandler = new Handler(Looper.getMainLooper());
       }
     }
-    sMainHandler.post(runnable);
+    sMainHandler.postDelayed(runnable, delayInMs);
   }
 }

ReactAndroid/src/main/java/com/facebook/react/bridge/WritableMap.java

@@ -24,4 +24,5 @@ public interface WritableMap extends ReadableMap {
   void putMap(@Nonnull String key, @Nullable WritableMap value);
 
   void merge(@Nonnull ReadableMap source);
+  WritableMap copy();
 }

ReactAndroid/src/main/java/com/facebook/react/bridge/WritableNativeMap.java

@@ -61,6 +61,13 @@ public class WritableNativeMap extends ReadableNativeMap implements WritableMap
     mergeNativeMap((ReadableNativeMap) source);
   }
 
+  @Override
+  public WritableMap copy() {
+    final WritableNativeMap target = new WritableNativeMap();
+    target.merge(this);
+    return target;
+  }
+
   public WritableNativeMap() {
     super(initHybrid());
   }

ReactAndroid/src/main/java/com/facebook/react/jstasks/HeadlessJsTaskConfig.java

@@ -15,6 +15,7 @@ public class HeadlessJsTaskConfig {
   private final WritableMap mData;
   private final long mTimeout;
   private final boolean mAllowedInForeground;
+  private final HeadlessJsTaskRetryPolicy mRetryPolicy;
 
   /**
    * Create a HeadlessJsTaskConfig. Equivalent to calling
@@ -55,10 +56,51 @@ public class HeadlessJsTaskConfig {
     WritableMap data,
     long timeout,
     boolean allowedInForeground) {
+    this(taskKey, data, timeout, allowedInForeground, NoRetryPolicy.INSTANCE);
+  }
+
+  /**
+   * Create a HeadlessJsTaskConfig.
+   *
+   * @param taskKey the key for the JS task to execute. This is the same key that you call {@code
+   * AppRegistry.registerTask} with in JS.
+   * @param data a map of parameters passed to the JS task executor.
+   * @param timeout the amount of time (in ms) after which the React instance should be terminated
+   * regardless of whether the task has completed or not. This is meant as a safeguard against
+   * accidentally keeping the device awake for long periods of time because JS crashed or some
+   * request timed out. A value of 0 means no timeout (should only be used for long-running tasks
+   * such as music playback).
+   * @param allowedInForeground whether to allow this task to run while the app is in the foreground
+   * (i.e. there is a host in resumed mode for the current ReactContext). Only set this to true if
+   * you really need it. Note that tasks run in the same JS thread as UI code, so doing expensive
+   * operations would degrade user experience.
+   * @param retryPolicy the number of times & delays the task should be retried on error.
+   */
+  public HeadlessJsTaskConfig(
+          String taskKey,
+          WritableMap data,
+          long timeout,
+          boolean allowedInForeground,
+          HeadlessJsTaskRetryPolicy retryPolicy) {
     mTaskKey = taskKey;
     mData = data;
     mTimeout = timeout;
     mAllowedInForeground = allowedInForeground;
+    mRetryPolicy = retryPolicy;
+  }
+
+  public HeadlessJsTaskConfig(HeadlessJsTaskConfig source) {
+    mTaskKey = source.mTaskKey;
+    mData = source.mData.copy();
+    mTimeout = source.mTimeout;
+    mAllowedInForeground = source.mAllowedInForeground;
+
+    final HeadlessJsTaskRetryPolicy retryPolicy = source.mRetryPolicy;
+    if (retryPolicy != null) {
+      mRetryPolicy = retryPolicy.copy();
+    } else {
+      mRetryPolicy = null;
+    }
   }
 
   /* package */ String getTaskKey() {
@@ -76,4 +118,8 @@ public class HeadlessJsTaskConfig {
   /* package */ boolean isAllowedInForeground() {
     return mAllowedInForeground;
   }
+
+  /* package */ HeadlessJsTaskRetryPolicy getRetryPolicy() {
+    return mRetryPolicy;
+  }
 }

ReactAndroid/src/main/java/com/facebook/react/jstasks/HeadlessJsTaskContext.java

@@ -5,12 +5,6 @@
 
 package com.facebook.react.jstasks;
 
-import java.lang.ref.WeakReference;
-import java.util.Set;
-import java.util.WeakHashMap;
-import java.util.concurrent.CopyOnWriteArraySet;
-import java.util.concurrent.atomic.AtomicInteger;
-
 import android.os.Handler;
 import android.util.SparseArray;
 
@@ -20,6 +14,14 @@ import com.facebook.react.bridge.UiThreadUtil;
 import com.facebook.react.common.LifecycleState;
 import com.facebook.react.modules.appregistry.AppRegistry;
 
+import java.lang.ref.WeakReference;
+import java.util.Map;
+import java.util.Set;
+import java.util.WeakHashMap;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CopyOnWriteArraySet;
+import java.util.concurrent.atomic.AtomicInteger;
+
 /**
  * Helper class for dealing with JS tasks. Handles per-ReactContext active task tracking, starting /
  * stopping tasks and notifying listeners.
@@ -51,6 +53,7 @@ public class HeadlessJsTaskContext {
   private final AtomicInteger mLastTaskId = new AtomicInteger(0);
   private final Handler mHandler = new Handler();
   private final Set<Integer> mActiveTasks = new CopyOnWriteArraySet<>();
+  private final Map<Integer, HeadlessJsTaskConfig> mActiveTaskConfigs = new ConcurrentHashMap<>();
   private final SparseArray<Runnable> mTaskTimeouts = new SparseArray<>();
 
   private HeadlessJsTaskContext(ReactContext reactContext) {
@@ -85,6 +88,16 @@ public class HeadlessJsTaskContext {
    * @return a unique id representing this task instance.
    */
   public synchronized int startTask(final HeadlessJsTaskConfig taskConfig) {
+    final int taskId = mLastTaskId.incrementAndGet();
+    startTask(taskConfig, taskId);
+    return taskId;
+  }
+
+  /**
+   * Start a JS task the provided task id. Handles invoking {@link AppRegistry#startHeadlessTask}
+   * and notifying listeners.
+   */
+  private synchronized void startTask(final HeadlessJsTaskConfig taskConfig, int taskId) {
     UiThreadUtil.assertOnUiThread();
     ReactContext reactContext = Assertions.assertNotNull(
       mReactContext.get(),
@@ -95,8 +108,8 @@ public class HeadlessJsTaskContext {
         "Tried to start task " + taskConfig.getTaskKey() +
           " while in foreground, but this is not allowed.");
     }
-    final int taskId = mLastTaskId.incrementAndGet();
     mActiveTasks.add(taskId);
+    mActiveTaskConfigs.put(taskId, new HeadlessJsTaskConfig(taskConfig));
     reactContext.getJSModule(AppRegistry.class)
       .startHeadlessTask(taskId, taskConfig.getTaskKey(), taskConfig.getData());
     if (taskConfig.getTimeout() > 0) {
@@ -105,7 +118,44 @@ public class HeadlessJsTaskContext {
     for (HeadlessJsTaskEventListener listener : mHeadlessJsTaskEventListeners) {
       listener.onHeadlessJsTaskStart(taskId);
     }
-    return taskId;
+  }
+
+  /**
+   * Retry a running JS task with a delay. Invokes
+   * {@link HeadlessJsTaskContext#startTask(HeadlessJsTaskConfig, int)} as long as the process does
+   * not get killed.
+   *
+   * @return true if a retry attempt has been posted.
+   */
+  public synchronized boolean retryTask(final int taskId) {
+    final HeadlessJsTaskConfig sourceTaskConfig = mActiveTaskConfigs.get(taskId);
+    Assertions.assertCondition(
+      sourceTaskConfig != null,
+      "Tried to retrieve non-existent task config with id " + taskId + ".");
+
+    final HeadlessJsTaskRetryPolicy retryPolicy = sourceTaskConfig.getRetryPolicy();
+    if (!retryPolicy.canRetry()) {
+      return false;
+    }
+
+    removeTimeout(taskId);
+    final HeadlessJsTaskConfig taskConfig = new HeadlessJsTaskConfig(
+            sourceTaskConfig.getTaskKey(),
+            sourceTaskConfig.getData(),
+            sourceTaskConfig.getTimeout(),
+            sourceTaskConfig.isAllowedInForeground(),
+            retryPolicy.update()
+    );
+
+    final Runnable retryAttempt = new Runnable() {
+      @Override
+      public void run() {
+        startTask(taskConfig, taskId);
+      }
+    };
+
+    UiThreadUtil.runOnUiThread(retryAttempt, retryPolicy.getDelay());
+    return true;
   }
 
   /**
@@ -118,11 +168,10 @@ public class HeadlessJsTaskContext {
     Assertions.assertCondition(
       mActiveTasks.remove(taskId),
       "Tried to finish non-existent task with id " + taskId + ".");
-    Runnable timeout = mTaskTimeouts.get(taskId);
-    if (timeout != null) {
-      mHandler.removeCallbacks(timeout);
-      mTaskTimeouts.remove(taskId);
-    }
+    Assertions.assertCondition(
+     mActiveTaskConfigs.remove(taskId) != null,
+     "Tried to remove non-existent task config with id " + taskId + ".");
+    removeTimeout(taskId);
     UiThreadUtil.runOnUiThread(new Runnable() {
       @Override
       public void run() {
@@ -133,6 +182,14 @@ public class HeadlessJsTaskContext {
     });
   }
 
+  private void removeTimeout(int taskId) {
+    Runnable timeout = mTaskTimeouts.get(taskId);
+    if (timeout != null) {
+      mHandler.removeCallbacks(timeout);
+      mTaskTimeouts.remove(taskId);
+    }
+  }
+
   /**
    * Check if a given task is currently running. A task is stopped if either {@link #finishTask} is
    * called or it times out.

ReactAndroid/src/main/java/com/facebook/react/jstasks/HeadlessJsTaskRetryPolicy.java

@@ -0,0 +1,16 @@
+package com.facebook.react.jstasks;
+
+import javax.annotation.CheckReturnValue;
+
+public interface HeadlessJsTaskRetryPolicy {
+
+    boolean canRetry();
+
+    int getDelay();
+
+    @CheckReturnValue
+    HeadlessJsTaskRetryPolicy update();
+
+    HeadlessJsTaskRetryPolicy copy();
+
+}

ReactAndroid/src/main/java/com/facebook/react/jstasks/LinearCountingRetryPolicy.java

@@ -0,0 +1,38 @@
+package com.facebook.react.jstasks;
+
+public class LinearCountingRetryPolicy implements HeadlessJsTaskRetryPolicy {
+
+    private final int mRetryAttempts;
+    private final int mDelayBetweenAttemptsInMs;
+
+    public LinearCountingRetryPolicy(int retryAttempts, int delayBetweenAttemptsInMs) {
+        mRetryAttempts = retryAttempts;
+        mDelayBetweenAttemptsInMs = delayBetweenAttemptsInMs;
+    }
+
+    @Override
+    public boolean canRetry() {
+        return mRetryAttempts > 0;
+    }
+
+    @Override
+    public int getDelay() {
+        return mDelayBetweenAttemptsInMs;
+    }
+
+    @Override
+    public HeadlessJsTaskRetryPolicy update() {
+        final int remainingRetryAttempts = mRetryAttempts - 1;
+
+        if (remainingRetryAttempts > 0) {
+            return new LinearCountingRetryPolicy(remainingRetryAttempts, mDelayBetweenAttemptsInMs);
+        } else {
+            return NoRetryPolicy.INSTANCE;
+        }
+    }
+
+    @Override
+    public HeadlessJsTaskRetryPolicy copy() {
+        return new LinearCountingRetryPolicy(mRetryAttempts, mDelayBetweenAttemptsInMs);
+    }
+}

ReactAndroid/src/main/java/com/facebook/react/jstasks/NoRetryPolicy.java

@@ -0,0 +1,30 @@
+package com.facebook.react.jstasks;
+
+public class NoRetryPolicy implements HeadlessJsTaskRetryPolicy {
+
+    public static final NoRetryPolicy INSTANCE = new NoRetryPolicy();
+
+    private NoRetryPolicy() {
+    }
+
+    @Override
+    public boolean canRetry() {
+        return false;
+    }
+
+    @Override
+    public int getDelay() {
+        throw new IllegalStateException("Should not retrieve delay as canRetry is: " + canRetry());
+    }
+
+    @Override
+    public HeadlessJsTaskRetryPolicy update() {
+        throw new IllegalStateException("Should not update as canRetry is: " + canRetry());
+    }
+
+    @Override
+    public HeadlessJsTaskRetryPolicy copy() {
+        // Class is immutable so no need to copy
+        return this;
+    }
+}

ReactAndroid/src/main/java/com/facebook/react/modules/core/HeadlessJsTaskSupportModule.java

@@ -8,6 +8,7 @@
 package com.facebook.react.modules.core;
 
 import com.facebook.common.logging.FLog;
+import com.facebook.react.bridge.Promise;
 import com.facebook.react.bridge.ReactApplicationContext;
 import com.facebook.react.bridge.ReactContextBaseJavaModule;
 import com.facebook.react.bridge.ReactMethod;
@@ -32,6 +33,22 @@ public class HeadlessJsTaskSupportModule extends ReactContextBaseJavaModule {
     return NAME;
   }
 
+  @ReactMethod
+  public void notifyTaskRetry(int taskId, Promise promise) {
+    HeadlessJsTaskContext headlessJsTaskContext =
+      HeadlessJsTaskContext.getInstance(getReactApplicationContext());
+    if (headlessJsTaskContext.isTaskRunning(taskId)) {
+      final boolean retryPosted = headlessJsTaskContext.retryTask(taskId);
+      promise.resolve(retryPosted);
+    } else {
+      FLog.w(
+        HeadlessJsTaskSupportModule.class,
+        "Tried to retry non-active task with id %d. Did it time out?",
+        taskId);
+      promise.resolve(false);
+    }
+  }
+
   @ReactMethod
   public void notifyTaskFinished(int taskId) {
     HeadlessJsTaskContext headlessJsTaskContext =