Draft: [RTDB] Use Firebase Android Executors

firebase-database/firebase-database.gradle

@@ -82,6 +82,7 @@ dependencies {
     implementation('com.google.firebase:firebase-auth-interop:19.0.2') {
         exclude group: "com.google.firebase", module: "firebase-common"
     }
+    androidTestImplementation project(path: ':firebase-database')
 
     javadocClasspath 'com.google.code.findbugs:jsr305:3.0.2'
     javadocClasspath 'com.google.auto.value:auto-value-annotations:1.6.6'

firebase-database/src/androidTest/java/com/google/firebase/database/FirebaseDatabaseTest.java

@@ -256,7 +256,10 @@ public void getReferenceThrowsWithBadUrl() {
 
   @Test
   public void getReferenceWithCustomDatabaseUrl() {
+    FirebaseDatabase db2 = FirebaseDatabase.getInstance();
     FirebaseDatabase db = FirebaseDatabase.getInstance(IntegrationTestValues.getAltNamespace());
+    db.getReference();
+    db2.getReference();
     db.getReferenceFromUrl(IntegrationTestValues.getAltNamespace());
   }
 

firebase-database/src/androidTest/java/com/google/firebase/database/IntegrationTestHelpers.java

@@ -32,6 +32,7 @@
 import com.google.firebase.database.core.DatabaseConfig;
 import com.google.firebase.database.core.EventTarget;
 import com.google.firebase.database.core.Path;
+import com.google.firebase.database.core.ThreadInitializer;
 import com.google.firebase.database.core.persistence.PersistenceManager;
 import com.google.firebase.database.core.utilities.DefaultRunLoop;
 import com.google.firebase.database.core.view.QuerySpec;
@@ -52,9 +53,13 @@
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.ThreadPoolExecutor;
@@ -186,9 +191,56 @@ public void shutdown() {}
     @Override
     public void restart() {}
   }
-
-  private static class TestRunLoop extends DefaultRunLoop {
+  // TODO(mtewani): Refactor this into a common package.
+  public static class TestRunLoop extends DefaultRunLoop {
     AtomicReference<Throwable> caughtException = new AtomicReference<Throwable>(null);
+    private class FirebaseThreadFactory implements ThreadFactory {
+
+      @Override
+      public Thread newThread(Runnable r) {
+        Thread thread = getThreadFactory().newThread(r);
+        ThreadInitializer initializer = getThreadInitializer();
+        initializer.setName(thread, "FirebaseDatabaseWorker");
+        initializer.setDaemon(thread, true);
+        initializer.setUncaughtExceptionHandler(
+                thread,
+                new Thread.UncaughtExceptionHandler() {
+                  @Override
+                  public void uncaughtException(Thread t, Throwable e) {
+                    handleException(e);
+                  }
+                });
+        return thread;
+      }
+    }
+
+    public TestRunLoop() {
+      super();
+      this.setExecutor(new ScheduledThreadPoolExecutor(1, new FirebaseThreadFactory()) {
+        @Override
+        protected void afterExecute(Runnable r, Throwable t) {
+          super.afterExecute(r, t);
+          if (t == null && r instanceof Future<?>) {
+            Future<?> future = (Future<?>) r;
+            try {
+              // Not all Futures will be done, e.g. when used with scheduledAtFixedRate
+              if (future.isDone()) {
+                future.get();
+              }
+            } catch (CancellationException ce) {
+              // Cancellation exceptions are okay, we expect them to happen sometimes
+            } catch (ExecutionException ee) {
+              t = ee.getCause();
+            } catch (InterruptedException e) {
+              Thread.currentThread().interrupt();
+            }
+          }
+          if (t != null) {
+            handleException(t);
+          }
+        }
+      });
+    }
 
     @Override
     public void handleException(Throwable e) {
@@ -262,7 +314,9 @@ public static DatabaseConfig newTestConfig() {
     DatabaseConfig config = new DatabaseConfig();
     config.setLogLevel(Logger.Level.DEBUG);
     config.setEventTarget(new TestEventTarget());
+    // TODO(mtewani): Figure out a way to better instantiate this.
     config.setRunLoop(runLoop);
+    config.setExecutorService(runLoop.getExecutorService());
     config.setFirebaseApp(FirebaseApp.getInstance());
     config.setAuthTokenProvider(
         new AndroidAuthTokenProvider(

firebase-database/src/main/java/com/google/firebase/database/DatabaseReference.java

@@ -283,6 +283,7 @@ private Task<Void> setValueInternal(Object value, Node priority, CompletionListe
     Validation.validateWritableObject(bouncedValue);
     final Node node = NodeUtilities.NodeFromJSON(bouncedValue, priority);
     final Pair<Task<Void>, CompletionListener> wrapped = Utilities.wrapOnComplete(optListener);
+    // Network and Disk I/O
     repo.scheduleNow(
         new Runnable() {
           @Override

firebase-database/src/main/java/com/google/firebase/database/DatabaseRegistrar.java

@@ -17,35 +17,45 @@
 import androidx.annotation.Keep;
 import androidx.annotation.RestrictTo;
 import com.google.firebase.FirebaseApp;
+import com.google.firebase.annotations.concurrent.Background;
+import com.google.firebase.annotations.concurrent.Blocking;
 import com.google.firebase.appcheck.interop.InternalAppCheckTokenProvider;
 import com.google.firebase.auth.internal.InternalAuthProvider;
 import com.google.firebase.components.Component;
 import com.google.firebase.components.ComponentRegistrar;
 import com.google.firebase.components.Dependency;
+import com.google.firebase.components.Qualified;
+import com.google.firebase.concurrent.FirebaseExecutors;
 import com.google.firebase.platforminfo.LibraryVersionComponent;
 import java.util.Arrays;
 import java.util.List;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ScheduledExecutorService;
 
 /** @hide */
 @Keep
 @RestrictTo(RestrictTo.Scope.LIBRARY_GROUP)
 public class DatabaseRegistrar implements ComponentRegistrar {
   private static final String LIBRARY_NAME = "fire-rtdb";
+  Qualified<Executor> blockingExecutor = Qualified.qualified(Blocking.class, Executor.class);
 
   @Override
   public List<Component<?>> getComponents() {
+
+    Qualified<ScheduledExecutorService> qualified = Qualified.qualified(Background.class, ScheduledExecutorService.class);
     return Arrays.asList(
         Component.builder(FirebaseDatabaseComponent.class)
             .name(LIBRARY_NAME)
             .add(Dependency.required(FirebaseApp.class))
             .add(Dependency.deferred(InternalAuthProvider.class))
             .add(Dependency.deferred(InternalAppCheckTokenProvider.class))
+            .add(Dependency.required(blockingExecutor))
             .factory(
                 c ->
                     new FirebaseDatabaseComponent(
                         c.get(FirebaseApp.class),
                         c.getDeferred(InternalAuthProvider.class),
-                        c.getDeferred(InternalAppCheckTokenProvider.class)))
+                        c.getDeferred(InternalAppCheckTokenProvider.class), c.get(qualified)))
             .build(),
         LibraryVersionComponent.create(LIBRARY_NAME, BuildConfig.VERSION_NAME));
   }

firebase-database/src/main/java/com/google/firebase/database/FirebaseDatabase.java

@@ -31,6 +31,8 @@
 import com.google.firebase.database.core.utilities.Validation;
 import com.google.firebase.emulators.EmulatedServiceSettings;
 
+import java.util.concurrent.Executor;
+
 /**
  * The entry point for accessing a Firebase Database. You can get an instance by calling {@link
  * FirebaseDatabase#getInstance()}. To access a location in the database and read or write data, use

firebase-database/src/main/java/com/google/firebase/database/FirebaseDatabaseComponent.java

@@ -26,6 +26,7 @@
 import com.google.firebase.inject.Deferred;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.ScheduledExecutorService;
 
 class FirebaseDatabaseComponent {
   /**
@@ -40,12 +41,15 @@ class FirebaseDatabaseComponent {
   private final FirebaseApp app;
   private final TokenProvider authProvider;
   private final TokenProvider appCheckProvider;
+  private final ScheduledExecutorService scheduledExecutorService;
 
   FirebaseDatabaseComponent(
-      @NonNull FirebaseApp app,
-      Deferred<InternalAuthProvider> authProvider,
-      Deferred<InternalAppCheckTokenProvider> appCheckProvider) {
+          @NonNull FirebaseApp app,
+          Deferred<InternalAuthProvider> authProvider,
+          Deferred<InternalAppCheckTokenProvider> appCheckProvider,
+          ScheduledExecutorService scheduledExecutorService) {
     this.app = app;
+    this.scheduledExecutorService = scheduledExecutorService;
     this.authProvider = new AndroidAuthTokenProvider(authProvider);
     this.appCheckProvider = new AndroidAppCheckTokenProvider(appCheckProvider);
   }
@@ -65,6 +69,7 @@ synchronized FirebaseDatabase get(RepoInfo repo) {
       config.setFirebaseApp(app);
       config.setAuthTokenProvider(authProvider);
       config.setAppCheckTokenProvider(appCheckProvider);
+      config.setExecutorService(scheduledExecutorService);
 
       database = new FirebaseDatabase(app, repo, config);
       instances.put(repo, database);

firebase-database/src/main/java/com/google/firebase/database/android/AndroidPlatform.java

@@ -42,6 +42,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.ScheduledExecutorService;
 
 public class AndroidPlatform implements Platform {
 
@@ -79,7 +80,7 @@ public EventTarget newEventTarget(com.google.firebase.database.core.Context cont
   @Override
   public RunLoop newRunLoop(com.google.firebase.database.core.Context ctx) {
     final LogWrapper logger = ctx.getLogger("RunLoop");
-    return new DefaultRunLoop() {
+    return new DefaultRunLoop(ctx.getExecutorService()) {
       @Override
       public void handleException(final Throwable e) {
         final String message = DefaultRunLoop.messageForException(e);

firebase-database/src/main/java/com/google/firebase/database/connection/WebsocketConnection.java

@@ -190,8 +190,7 @@ private WSClient createConnection(
     extraHeaders.put("X-Firebase-GMPID", connectionContext.getApplicationId());
     extraHeaders.put("X-Firebase-AppCheck", appCheckToken);
     WebSocket ws = new WebSocket(connectionContext, uri, /*protocol=*/ null, extraHeaders);
-    WSClientTubesock client = new WSClientTubesock(ws);
-    return client;
+    return new WSClientTubesock(ws);
   }
 
   public void open() {

firebase-database/src/main/java/com/google/firebase/database/core/Context.java

@@ -51,6 +51,7 @@ public class Context {
   private PersistenceManager forcedPersistenceManager;
   private boolean frozen = false;
   private boolean stopped = false;
+  private ScheduledExecutorService scheduledExecutorService;
 
   private Platform platform;
 
@@ -216,14 +217,16 @@ public PersistentConnection newPersistentConnection(
     return getPlatform().newPersistentConnection(this, this.getConnectionContext(), info, delegate);
   }
 
-  private ScheduledExecutorService getExecutorService() {
-    RunLoop loop = this.getRunLoop();
-    if (!(loop instanceof DefaultRunLoop)) {
-      // TODO: We really need to remove this option from the public DatabaseConfig
-      // object
-      throw new RuntimeException("Custom run loops are not supported!");
-    }
-    return ((DefaultRunLoop) loop).getExecutorService();
+  /**
+   * Sets the current executorService
+   * @param executorService
+   */
+  public void setExecutorService(ScheduledExecutorService executorService) {
+    this.scheduledExecutorService = executorService;
+  }
+
+  public ScheduledExecutorService getExecutorService() {
+    return this.scheduledExecutorService;
   }
 
   private void ensureLogger() {

firebase-database/src/main/java/com/google/firebase/database/core/DatabaseConfig.java

@@ -18,6 +18,7 @@
 import com.google.firebase.database.DatabaseException;
 import com.google.firebase.database.Logger;
 import java.util.List;
+import java.util.concurrent.ScheduledExecutorService;
 
 /**
  * TODO: Merge this class with Context and clean this up. Some methods may need to be re-added to

firebase-database/src/main/java/com/google/firebase/database/core/Repo.java

@@ -16,6 +16,8 @@
 
 import static com.google.firebase.database.core.utilities.Utilities.hardAssert;
 
+import android.util.Log;
+
 import androidx.annotation.NonNull;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
@@ -474,6 +476,7 @@ public void setValue(
             path, newValueUnresolved, newValue, writeId, /*visible=*/ true, /*persist=*/ true);
     this.postEvents(events);
 
+    Log.d(this.getClass().toString(), "before");
     connection.put(
         path.asList(),
         newValueUnresolved.getValue(true),
@@ -484,8 +487,10 @@ public void onRequestResult(String optErrorCode, String optErrorMessage) {
             warnIfWriteFailed("setValue", path, error);
             ackWriteAndRerunTransactions(writeId, path, error);
             callOnComplete(onComplete, error, path);
+            Log.d(this.getClass().toString(), "callback");
           }
         });
+    Log.d(this.getClass().toString(), "after");
 
     Path affectedPath = abortTransactions(path, DatabaseError.OVERRIDDEN_BY_SET);
     this.rerunTransactions(affectedPath);
@@ -616,6 +621,7 @@ public void updateChildren(
 
     // TODO: DatabaseReference.CompleteionListener isn't really appropriate (the DatabaseReference
     // param is meaningless).
+    Log.d(this.getClass().toString(), "started merge");
     connection.merge(
         path.asList(),
         unParsedUpdates,
@@ -626,8 +632,10 @@ public void onRequestResult(String optErrorCode, String optErrorMessage) {
             warnIfWriteFailed("updateChildren", path, error);
             ackWriteAndRerunTransactions(writeId, path, error);
             callOnComplete(onComplete, error, path);
+            Log.d(this.getClass().toString(), "done with callback");
           }
         });
+    Log.d(this.getClass().toString(), "ended merge");
 
     for (Entry<Path, Node> update : updates) {
       Path pathFromRoot = path.child(update.getKey());