Some improvements to the task caches

firebase-appdistribution/src/main/java/com/google/firebase/appdistribution/impl/TaskCompletionSourceCache.java

@@ -59,7 +59,8 @@ Task<T> getOrCreateTaskFromCompletionSource(TaskCompletionSourceProducer<T> prod
     TaskCompletionSource<T> taskCompletionSource = new TaskCompletionSource<>();
     sequentialExecutor.execute(
         () -> {
-          if (!isOngoing(cachedTaskCompletionSource)) {
+          if (cachedTaskCompletionSource == null
+              || cachedTaskCompletionSource.getTask().isComplete()) {
             cachedTaskCompletionSource = producer.produce();
           }
           TaskUtils.shadowTask(taskCompletionSource, cachedTaskCompletionSource.getTask());
@@ -68,36 +69,62 @@ Task<T> getOrCreateTaskFromCompletionSource(TaskCompletionSourceProducer<T> prod
   }
 
   /**
-   * Sets the result on the cached {@link TaskCompletionSource}, if there is one and it is not
-   * completed.
+   * Sets the result on the cached {@link TaskCompletionSource}.
+   *
+   * <p>The task must be created using {@link #getOrCreateTaskFromCompletionSource} before calling
+   * this method.
+   *
+   * @return A task that resolves when the change is applied, or fails with {@link
+   *     IllegalStateException} if the task was not created yet or was already completed.
    */
   Task<Void> setResult(T result) {
-    TaskCompletionSource<Void> taskCompletionSource = new TaskCompletionSource<>();
-    sequentialExecutor.execute(
-        () -> {
-          if (isOngoing(cachedTaskCompletionSource)) {
-            cachedTaskCompletionSource.setResult(result);
-          }
-        });
-    return taskCompletionSource.getTask();
+    return executeIfOngoing(() -> cachedTaskCompletionSource.setResult(result));
   }
 
   /**
    * Sets the exception on the cached {@link TaskCompletionSource}, if there is one and it is not
    * completed.
+   *
+   * <p>The task must be created using {@link #getOrCreateTaskFromCompletionSource} before calling
+   * this method.
+   *
+   * @return A task that resolves when the change is applied, or fails with {@link
+   *     IllegalStateException} if the task was not created yet or was already completed.
    */
   Task<Void> setException(Exception e) {
+    return executeIfOngoing(() -> cachedTaskCompletionSource.setException(e));
+  }
+
+  /**
+   * Returns a task indicating whether the cached {@link TaskCompletionSource} has been initialized
+   * and is not yet completed.
+   */
+  Task<Boolean> isOngoing() {
+    TaskCompletionSource<Boolean> taskCompletionSource = new TaskCompletionSource<>();
+    sequentialExecutor.execute(
+        () ->
+            taskCompletionSource.setResult(
+                cachedTaskCompletionSource != null
+                    && !cachedTaskCompletionSource.getTask().isComplete()));
+    return taskCompletionSource.getTask();
+  }
+
+  private Task<Void> executeIfOngoing(Runnable r) {
     TaskCompletionSource<Void> taskCompletionSource = new TaskCompletionSource<>();
     sequentialExecutor.execute(
         () -> {
-          if (isOngoing(cachedTaskCompletionSource)) {
-            cachedTaskCompletionSource.setException(e);
+          if (cachedTaskCompletionSource == null) {
+            taskCompletionSource.setException(
+                new IllegalStateException(
+                    "Tried to set exception before calling getOrCreateTaskFromCompletionSource()"));
+          } else if (cachedTaskCompletionSource.getTask().isComplete()) {
+            taskCompletionSource.setException(
+                new IllegalStateException("Tried to set exception on a completed task"));
+          } else {
+            r.run();
+            taskCompletionSource.setResult(null);
           }
         });
     return taskCompletionSource.getTask();
   }
-
-  private static <T> boolean isOngoing(TaskCompletionSource<T> taskCompletionSource) {
-    return taskCompletionSource != null && !taskCompletionSource.getTask().isComplete();
-  }
 }

firebase-appdistribution/src/main/java/com/google/firebase/appdistribution/impl/UpdateTaskCache.java

@@ -63,7 +63,7 @@ UpdateTask getOrCreateUpdateTask(UpdateTaskProducer producer) {
     UpdateTaskImpl impl = new UpdateTaskImpl();
     sequentialExecutor.execute(
         () -> {
-          if (!isOngoing(cachedUpdateTaskImpl)) {
+          if (cachedUpdateTaskImpl == null || cachedUpdateTaskImpl.isComplete()) {
             cachedUpdateTaskImpl = producer.produce();
           }
           impl.shadow(cachedUpdateTaskImpl);
@@ -72,50 +72,59 @@ UpdateTask getOrCreateUpdateTask(UpdateTaskProducer producer) {
   }
 
   /**
-   * Sets the exception on the cached {@link UpdateTaskImpl}, if there is one and it is not
+   * Sets the progress on the cached {@link UpdateTaskImpl}, if there is one and it is not
    * completed.
+   *
+   * <p>The task must be created using {@link #getOrCreateUpdateTask} before calling this method.
+   *
+   * @return A task that resolves when the change is applied, or fails with {@link
+   *     IllegalStateException} if the task was not created yet or was already completed.
    */
   Task<Void> setProgress(UpdateProgress progress) {
-    TaskCompletionSource<Void> taskCompletionSource = new TaskCompletionSource<>();
-    sequentialExecutor.execute(
-        () -> {
-          if (isOngoing(cachedUpdateTaskImpl)) {
-            cachedUpdateTaskImpl.updateProgress(progress);
-          }
-        });
-    return taskCompletionSource.getTask();
+    return executeIfOngoing(() -> cachedUpdateTaskImpl.updateProgress(progress));
   }
 
   /**
    * Sets the result on the cached {@link UpdateTaskImpl}, if there is one and it is not completed.
+   *
+   * <p>The task must be created using {@link #getOrCreateUpdateTask} before calling this method.
+   *
+   * @return A task that resolves when the change is applied, or fails with {@link
+   *     IllegalStateException} if the task was not created yet or was already completed.
    */
   Task<Void> setResult() {
-    TaskCompletionSource<Void> taskCompletionSource = new TaskCompletionSource<>();
-    sequentialExecutor.execute(
-        () -> {
-          if (isOngoing(cachedUpdateTaskImpl)) {
-            cachedUpdateTaskImpl.setResult();
-          }
-        });
-    return taskCompletionSource.getTask();
+    return executeIfOngoing(() -> cachedUpdateTaskImpl.setResult());
   }
 
   /**
    * Sets the exception on the cached {@link UpdateTaskImpl}, if there is one and it is not
    * completed.
+   *
+   * <p>The task must be created using {@link #getOrCreateUpdateTask} before calling this method.
+   *
+   * @return A task that resolves when the change is applied, or fails with {@link
+   *     IllegalStateException} if the task was not created yet or was already completed.
    */
   Task<Void> setException(FirebaseAppDistributionException e) {
+    return executeIfOngoing(() -> cachedUpdateTaskImpl.setException(e));
+  }
+
+  private Task<Void> executeIfOngoing(Runnable r) {
     TaskCompletionSource<Void> taskCompletionSource = new TaskCompletionSource<>();
     sequentialExecutor.execute(
         () -> {
-          if (isOngoing(cachedUpdateTaskImpl)) {
-            cachedUpdateTaskImpl.setException(e);
+          if (cachedUpdateTaskImpl == null) {
+            taskCompletionSource.setException(
+                new IllegalStateException(
+                    "Tried to set exception before calling getOrCreateUpdateTask()"));
+          } else if (cachedUpdateTaskImpl.isComplete()) {
+            taskCompletionSource.setException(
+                new IllegalStateException("Tried to set exception on a completed task"));
+          } else {
+            r.run();
+            taskCompletionSource.setResult(null);
           }
         });
     return taskCompletionSource.getTask();
   }
-
-  private static boolean isOngoing(UpdateTask updateTask) {
-    return updateTask != null && !updateTask.isComplete();
-  }
 }