Code dump with latest updates. am: f8fef1787d am: 9058c22d3f

Original change: https://googleplex-android-review.googlesource.com/c/platform/external/mobile-data-download/+/22462863

Change-Id: I9daa3d37544ad17f0ce401f7d181c5606b666ca0
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>

5 files changed, 297 insertions, 172 deletions

diff --git a/java/com/google/android/libraries/mobiledatadownload/lite/BUILD b/java/com/google/android/libraries/mobiledatadownload/lite/BUILD
index c1eb8fb..4f8c1f5 100644
--- a/java/com/google/android/libraries/mobiledatadownload/lite/BUILD
+++ b/java/com/google/android/libraries/mobiledatadownload/lite/BUILD
@@ -16,6 +16,7 @@ load("@build_bazel_rules_android//android:rules.bzl", "android_library")
 # MDD Lite visibility is restricted to the following set of packages. Any
 # new clients must be added to this list in order to grant build visibility.
 package(
+    default_applicable_licenses = ["//:license"],
     default_visibility = [
         "//visibility:public",
     ],
@@ -37,11 +38,15 @@ android_library(
         ":DownloadProgressMonitor",
         "//java/com/google/android/libraries/mobiledatadownload:DownloadException",
         "//java/com/google/android/libraries/mobiledatadownload/downloader:FileDownloader",
+        "//java/com/google/android/libraries/mobiledatadownload/foreground:ForegroundDownloadKey",
         "//java/com/google/android/libraries/mobiledatadownload/foreground:NotificationUtil",
         "//java/com/google/android/libraries/mobiledatadownload/internal/logging:LogUtil",
+        "//java/com/google/android/libraries/mobiledatadownload/internal/util:DownloadFutureMap",
+        "//java/com/google/android/libraries/mobiledatadownload/tracing",
+        "//java/com/google/android/libraries/mobiledatadownload/tracing:concurrent",
         "@androidx_core_core",
         "@com_google_auto_value",
-        "@com_google_dagger",
+        "@com_google_errorprone_error_prone_annotations",
         "@com_google_guava_guava",
         "@org_checkerframework_qual",
     ],
@@ -66,9 +71,11 @@ android_library(
         ":DownloadListener",
         "//java/com/google/android/libraries/mobiledatadownload:TimeSource",
         "//java/com/google/android/libraries/mobiledatadownload/file/spi",
+        "//java/com/google/android/libraries/mobiledatadownload/internal:AndroidTimeSource",
         "//java/com/google/android/libraries/mobiledatadownload/internal/logging:LogUtil",
         "@androidx_annotation_annotation",
         "@com_google_code_findbugs_jsr305",
+        "@com_google_errorprone_error_prone_annotations",
         "@com_google_guava_guava",
     ],
 )
diff --git a/java/com/google/android/libraries/mobiledatadownload/lite/DownloadProgressMonitor.java b/java/com/google/android/libraries/mobiledatadownload/lite/DownloadProgressMonitor.java
index d0fd6fa..4c5fbd6 100644
--- a/java/com/google/android/libraries/mobiledatadownload/lite/DownloadProgressMonitor.java
+++ b/java/com/google/android/libraries/mobiledatadownload/lite/DownloadProgressMonitor.java
@@ -21,11 +21,11 @@ import com.google.android.libraries.mobiledatadownload.TimeSource;
 import com.google.android.libraries.mobiledatadownload.file.spi.Monitor;
 import com.google.android.libraries.mobiledatadownload.internal.logging.LogUtil;
 import com.google.common.util.concurrent.MoreExecutors;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
 import java.util.HashMap;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicLong;
 import javax.annotation.Nullable;
-import javax.annotation.concurrent.GuardedBy;
 import javax.annotation.concurrent.ThreadSafe;
 
 /** A Download Progress Monitor to support {@link DownloadListener}. */
@@ -37,7 +37,6 @@ public class DownloadProgressMonitor implements Monitor, SingleFileDownloadProgr
   private final TimeSource timeSource;
   private final Executor sequentialControlExecutor;
 
-  // NOTE: GuardRails prohibits multiple public constructors
   private DownloadProgressMonitor(TimeSource timeSource, Executor controlExecutor) {
     this.timeSource = timeSource;
 
diff --git a/java/com/google/android/libraries/mobiledatadownload/lite/Downloader.java b/java/com/google/android/libraries/mobiledatadownload/lite/Downloader.java
index 208132c..ea4f450 100644
--- a/java/com/google/android/libraries/mobiledatadownload/lite/Downloader.java
+++ b/java/com/google/android/libraries/mobiledatadownload/lite/Downloader.java
@@ -22,6 +22,7 @@ import com.google.common.base.Preconditions;
 import com.google.common.base.Supplier;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.MoreExecutors;
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.CheckReturnValue;
 import java.util.concurrent.Executor;
 
@@ -73,8 +74,19 @@ public interface Downloader {
   @CheckReturnValue
   ListenableFuture<Void> downloadWithForegroundService(DownloadRequest downloadRequest);
 
-  /** Cancel an on-going foreground download. */
-  void cancelForegroundDownload(String destinationFileUri);
+  /**
+   * Cancel an on-going foreground download.
+   *
+   * <p>Use {@link ForegroundDownloadKey} to construct the unique key.
+   *
+   * <p><b>NOTE:</b> In most cases, clients will not need to call this -- it is meant to allow the
+   * ForegroundDownloadService to cancel a download via the Cancel action registered to a
+   * notification.
+   *
+   * <p>Clients should prefer to cancel the future returned to them from {@link
+   * #downloadWithForegroundService} instead.
+   */
+  void cancelForegroundDownload(String downloadKey);
 
   static Downloader.Builder newBuilder() {
     return new Downloader.Builder();
@@ -91,12 +103,14 @@ public interface Downloader {
     private Optional<SingleFileDownloadProgressMonitor> downloadMonitorOptional = Optional.absent();
     private Optional<Class<?>> foregroundDownloadServiceClassOptional = Optional.absent();
 
+    @CanIgnoreReturnValue
     public Builder setContext(Context context) {
       this.context = context.getApplicationContext();
       return this;
     }
 
     /** Set the Control Executor which will run MDDLite control flow. */
+    @CanIgnoreReturnValue
     public Builder setControlExecutor(Executor controlExecutor) {
       Preconditions.checkNotNull(controlExecutor);
       // Executor that will execute tasks sequentially.
@@ -115,6 +129,7 @@ public interface Downloader {
      * DownloadListener} to {@link Downloader#download}. The DownloadListener's {@code onFailure}
      * and {@code onComplete} will be invoked regardless of whether this is set.
      */
+    @CanIgnoreReturnValue
     public Builder setDownloadMonitor(SingleFileDownloadProgressMonitor downloadMonitor) {
       this.downloadMonitorOptional = Optional.of(downloadMonitor);
       return this;
@@ -127,6 +142,7 @@ public interface Downloader {
      * <p>This is required to use {@link Downloader#downloadWithForegroundService}. Not providing
      * this will result in a failed future when calling downloadWithForegroundService.
      */
+    @CanIgnoreReturnValue
     public Builder setForegroundDownloadService(Class<?> foregroundDownloadServiceClass) {
       this.foregroundDownloadServiceClassOptional = Optional.of(foregroundDownloadServiceClass);
       return this;
@@ -136,6 +152,7 @@ public interface Downloader {
      * Set the FileDownloader Supplier. MDDLite takes in a Supplier of FileDownload to support lazy
      * instantiation of the FileDownloader
      */
+    @CanIgnoreReturnValue
     public Builder setFileDownloaderSupplier(Supplier<FileDownloader> fileDownloaderSupplier) {
       this.fileDownloaderSupplier = fileDownloaderSupplier;
       return this;
diff --git a/java/com/google/android/libraries/mobiledatadownload/lite/DownloaderImpl.java b/java/com/google/android/libraries/mobiledatadownload/lite/DownloaderImpl.java
index 1c0cb49..8472667 100644
--- a/java/com/google/android/libraries/mobiledatadownload/lite/DownloaderImpl.java
+++ b/java/com/google/android/libraries/mobiledatadownload/lite/DownloaderImpl.java
@@ -15,6 +15,9 @@
  */
 package com.google.android.libraries.mobiledatadownload.lite;
 
+import static com.google.common.util.concurrent.Futures.immediateFailedFuture;
+import static com.google.common.util.concurrent.Futures.immediateVoidFuture;
+
 import android.content.Context;
 import androidx.annotation.VisibleForTesting;
 import androidx.core.app.NotificationCompat;
@@ -22,17 +25,18 @@ import androidx.core.app.NotificationManagerCompat;
 import com.google.android.libraries.mobiledatadownload.DownloadException;
 import com.google.android.libraries.mobiledatadownload.DownloadException.DownloadResultCode;
 import com.google.android.libraries.mobiledatadownload.downloader.FileDownloader;
+import com.google.android.libraries.mobiledatadownload.foreground.ForegroundDownloadKey;
 import com.google.android.libraries.mobiledatadownload.foreground.NotificationUtil;
 import com.google.android.libraries.mobiledatadownload.internal.logging.LogUtil;
+import com.google.android.libraries.mobiledatadownload.internal.util.DownloadFutureMap;
+import com.google.android.libraries.mobiledatadownload.tracing.PropagatedFluentFuture;
+import com.google.android.libraries.mobiledatadownload.tracing.PropagatedFutures;
 import com.google.common.base.Optional;
-import com.google.common.base.Preconditions;
 import com.google.common.base.Supplier;
 import com.google.common.util.concurrent.FutureCallback;
-import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.ListenableFutureTask;
 import com.google.common.util.concurrent.MoreExecutors;
-import java.util.HashMap;
-import java.util.Map;
 import java.util.concurrent.Executor;
 import org.checkerframework.checker.nullness.compatqual.NullableDecl;
 
@@ -46,9 +50,8 @@ final class DownloaderImpl implements Downloader {
   private final Optional<SingleFileDownloadProgressMonitor> downloadMonitorOptional;
   private final Supplier<FileDownloader> fileDownloaderSupplier;
 
-  // Synchronization will be done through sequentialControlExecutor
-  @VisibleForTesting
-  final Map<String, ListenableFuture<Void>> keyToListenableFuture = new HashMap<>();
+  @VisibleForTesting final DownloadFutureMap<Void> downloadFutureMap;
+  @VisibleForTesting final DownloadFutureMap<Void> foregroundDownloadFutureMap;
 
   DownloaderImpl(
       Context context,
@@ -61,19 +64,25 @@ final class DownloaderImpl implements Downloader {
     this.foregroundDownloadServiceClassOptional = foregroundDownloadServiceClassOptional;
     this.downloadMonitorOptional = downloadMonitorOptional;
     this.fileDownloaderSupplier = fileDownloaderSupplier;
+    this.downloadFutureMap = DownloadFutureMap.create(sequentialControlExecutor);
+    this.foregroundDownloadFutureMap =
+        DownloadFutureMap.create(
+            sequentialControlExecutor,
+            createCallbacksForForegroundService(context, foregroundDownloadServiceClassOptional));
   }
 
   @Override
   public ListenableFuture<Void> download(DownloadRequest downloadRequest) {
     LogUtil.d("%s: download for Uri = %s", TAG, downloadRequest.destinationFileUri().toString());
-    return Futures.submitAsync(
-        () -> {
+    ForegroundDownloadKey foregroundDownloadKey =
+        ForegroundDownloadKey.ofSingleFile(downloadRequest.destinationFileUri());
+
+    return PropagatedFutures.transformAsync(
+        getInProgressDownloadFuture(foregroundDownloadKey.toString()),
+        (Optional<ListenableFuture<Void>> existingDownloadFuture) -> {
           // if there is the same on-going request, return that one.
-          if (keyToListenableFuture.containsKey(downloadRequest.destinationFileUri().toString())) {
-            // uriToListenableFuture.get must return Non-null since we check the containsKey above.
-            // checkNotNull is to suppress false alarm about @Nullable result.
-            return Preconditions.checkNotNull(
-                keyToListenableFuture.get(downloadRequest.destinationFileUri().toString()));
+          if (existingDownloadFuture.isPresent()) {
+            return existingDownloadFuture.get();
           }
 
           // Register listener with monitor if present
@@ -87,13 +96,19 @@ final class DownloaderImpl implements Downloader {
             } else {
               LogUtil.w(
                   "%s: download request included DownloadListener, but DownloadMonitor is not"
-                      + " present! DownloadListener will only be invoked for complete/failure.");
+                      + " present! DownloadListener will only be invoked for complete/failure.",
+                  TAG);
             }
           }
 
-          ListenableFuture<Void> downloadFuture = startDownload(downloadRequest);
+          // Create a ListenableFutureTask to delay starting the downloadFuture until we can add the
+          // future to our map.
+          ListenableFutureTask<Void> startTask = ListenableFutureTask.create(() -> null);
+          ListenableFuture<Void> downloadFuture =
+              PropagatedFutures.transformAsync(
+                  startTask, unused -> startDownload(downloadRequest), sequentialControlExecutor);
 
-          Futures.addCallback(
+          PropagatedFutures.addCallback(
               downloadFuture,
               new FutureCallback<Void>() {
                 @Override
@@ -104,36 +119,36 @@ final class DownloaderImpl implements Downloader {
                   // Remove download listener and remove download future from map after listener
                   // completes
                   if (downloadRequest.listenerOptional().isPresent()) {
-                    Futures.addCallback(
+                    PropagatedFutures.addCallback(
                         downloadRequest.listenerOptional().get().onComplete(),
                         new FutureCallback<Void>() {
                           @Override
                           public void onSuccess(@NullableDecl Void result) {
-                            keyToListenableFuture.remove(
-                                downloadRequest.destinationFileUri().toString());
                             if (downloadMonitorOptional.isPresent()) {
                               downloadMonitorOptional
                                   .get()
                                   .removeDownloadListener(downloadRequest.destinationFileUri());
                             }
+                            ListenableFuture<Void> unused =
+                                downloadFutureMap.remove(foregroundDownloadKey.toString());
                           }
 
                           @Override
                           public void onFailure(Throwable t) {
                             LogUtil.e(t, "%s: Failed to run client onComplete", TAG);
-                            keyToListenableFuture.remove(
-                                downloadRequest.destinationFileUri().toString());
                             if (downloadMonitorOptional.isPresent()) {
                               downloadMonitorOptional
                                   .get()
                                   .removeDownloadListener(downloadRequest.destinationFileUri());
                             }
+                            ListenableFuture<Void> unused =
+                                downloadFutureMap.remove(foregroundDownloadKey.toString());
                           }
                         },
                         sequentialControlExecutor);
                   } else {
-                    // remove from future map immediately
-                    keyToListenableFuture.remove(downloadRequest.destinationFileUri().toString());
+                    ListenableFuture<Void> unused =
+                        downloadFutureMap.remove(foregroundDownloadKey.toString());
                   }
                 }
 
@@ -151,14 +166,20 @@ final class DownloaderImpl implements Downloader {
                           .removeDownloadListener(downloadRequest.destinationFileUri());
                     }
                   }
-                  keyToListenableFuture.remove(downloadRequest.destinationFileUri().toString());
+                  ListenableFuture<Void> unused =
+                      downloadFutureMap.remove(foregroundDownloadKey.toString());
                 }
               },
               MoreExecutors.directExecutor());
 
-          keyToListenableFuture.put(
-              downloadRequest.destinationFileUri().toString(), downloadFuture);
-          return downloadFuture;
+          return PropagatedFutures.transformAsync(
+              downloadFutureMap.add(foregroundDownloadKey.toString(), downloadFuture),
+              unused -> {
+                // Now that the download future is added, start the task and return the future
+                startTask.run();
+                return downloadFuture;
+              },
+              sequentialControlExecutor);
         },
         sequentialControlExecutor);
   }
@@ -178,7 +199,7 @@ final class DownloaderImpl implements Downloader {
       return fileDownloaderSupplier.get().startDownloading(fileDownloaderRequest);
     } catch (RuntimeException e) {
       // Catch any unchecked exceptions that prevented the download from starting.
-      return Futures.immediateFailedFuture(
+      return immediateFailedFuture(
           DownloadException.builder()
               .setDownloadResultCode(DownloadResultCode.UNKNOWN_ERROR)
               .setCause(e)
@@ -192,23 +213,25 @@ final class DownloaderImpl implements Downloader {
         "%s: downloadWithForegroundService for Uri = %s",
         TAG, downloadRequest.destinationFileUri().toString());
     if (!downloadMonitorOptional.isPresent()) {
-      return Futures.immediateFailedFuture(
+      return immediateFailedFuture(
           new IllegalStateException(
               "downloadWithForegroundService: DownloadMonitor is not provided!"));
     }
     if (!foregroundDownloadServiceClassOptional.isPresent()) {
-      return Futures.immediateFailedFuture(
+      return immediateFailedFuture(
           new IllegalStateException(
               "downloadWithForegroundService: ForegroundDownloadService is not provided!"));
     }
-    return Futures.submitAsync(
-        () -> {
+
+    ForegroundDownloadKey foregroundDownloadKey =
+        ForegroundDownloadKey.ofSingleFile(downloadRequest.destinationFileUri());
+
+    return PropagatedFutures.transformAsync(
+        getInProgressDownloadFuture(foregroundDownloadKey.toString()),
+        (Optional<ListenableFuture<Void>> existingDownloadFuture) -> {
           // if there is the same on-going request, return that one.
-          if (keyToListenableFuture.containsKey(downloadRequest.destinationFileUri().toString())) {
-            // uriToListenableFuture.get must return Non-null since we check the containsKey above.
-            // checkNotNull is to suppress false alarm about @Nullable result.
-            return Preconditions.checkNotNull(
-                keyToListenableFuture.get(downloadRequest.destinationFileUri().toString()));
+          if (existingDownloadFuture.isPresent()) {
+            return existingDownloadFuture.get();
           }
 
           // It's OK to recreate the NotificationChannel since it can also be used to restore a
@@ -216,14 +239,6 @@ final class DownloaderImpl implements Downloader {
           // importance.
           NotificationUtil.createNotificationChannel(context);
 
-          // Only start the foreground download service when there is the first download request.
-          if (keyToListenableFuture.isEmpty()) {
-            NotificationUtil.startForegroundDownloadService(
-                context,
-                foregroundDownloadServiceClassOptional.get(),
-                downloadRequest.destinationFileUri().toString());
-          }
-
           DownloadListener downloadListenerWithNotification =
               createDownloadListenerWithNotification(downloadRequest);
 
@@ -233,9 +248,14 @@ final class DownloaderImpl implements Downloader {
               .addDownloadListener(
                   downloadRequest.destinationFileUri(), downloadListenerWithNotification);
 
-          ListenableFuture<Void> downloadFuture = startDownload(downloadRequest);
+          // Create a ListenableFutureTask to delay starting the downloadFuture until we can add the
+          // future to our map.
+          ListenableFutureTask<Void> startTask = ListenableFutureTask.create(() -> null);
+          ListenableFuture<Void> downloadFuture =
+              PropagatedFutures.transformAsync(
+                  startTask, unused -> startDownload(downloadRequest), sequentialControlExecutor);
 
-          Futures.addCallback(
+          PropagatedFutures.addCallback(
               downloadFuture,
               new FutureCallback<Void>() {
                 @Override
@@ -243,7 +263,7 @@ final class DownloaderImpl implements Downloader {
                   // Currently the MobStore monitor does not support onSuccess so we have to add
                   // callback to the download future here.
 
-                  Futures.addCallback(
+                  PropagatedFutures.addCallback(
                       downloadListenerWithNotification.onComplete(),
                       new FutureCallback<Void>() {
                         @Override
@@ -267,15 +287,25 @@ final class DownloaderImpl implements Downloader {
               },
               MoreExecutors.directExecutor());
 
-          keyToListenableFuture.put(
-              downloadRequest.destinationFileUri().toString(), downloadFuture);
-          return downloadFuture;
+          return PropagatedFutures.transformAsync(
+              foregroundDownloadFutureMap.add(foregroundDownloadKey.toString(), downloadFuture),
+              unused -> {
+                // Now that the download future is added, start the task and return the future
+                startTask.run();
+                return downloadFuture;
+              },
+              sequentialControlExecutor);
         },
         sequentialControlExecutor);
   }
 
   // Assertion: foregroundDownloadService and downloadMonitor are present
   private DownloadListener createDownloadListenerWithNotification(DownloadRequest downloadRequest) {
+    String networkPausedMessage =
+        downloadRequest.downloadConstraints().requireUnmeteredNetwork()
+            ? NotificationUtil.getDownloadPausedWifiMessage(context)
+            : NotificationUtil.getDownloadPausedMessage(context);
+
     NotificationManagerCompat notificationManager = NotificationManagerCompat.from(context);
     NotificationCompat.Builder notification =
         NotificationUtil.createNotificationBuilder(
@@ -284,14 +314,16 @@ final class DownloaderImpl implements Downloader {
             downloadRequest.notificationContentTitle(),
             downloadRequest.notificationContentTextOptional().or(downloadRequest.urlToDownload()));
 
-    int notificationKey =
-        NotificationUtil.notificationKeyForKey(downloadRequest.destinationFileUri().toString());
+    ForegroundDownloadKey foregroundDownloadKey =
+        ForegroundDownloadKey.ofSingleFile(downloadRequest.destinationFileUri());
+
+    int notificationKey = NotificationUtil.notificationKeyForKey(foregroundDownloadKey.toString());
 
     // Attach the Cancel action to the notification.
     NotificationUtil.createCancelAction(
         context,
         foregroundDownloadServiceClassOptional.get(),
-        downloadRequest.destinationFileUri().toString(),
+        foregroundDownloadKey.toString(),
         notification,
         notificationKey);
     notificationManager.notify(notificationKey, notification.build());
@@ -299,49 +331,56 @@ final class DownloaderImpl implements Downloader {
     return new DownloadListener() {
       @Override
       public void onProgress(long currentSize) {
-        sequentialControlExecutor.execute(
-            () -> {
-              // There can be a race condition, where onPausedForConnectivity can be called
-              // after onComplete or onFailure which removes the future and the notification.
-              if (keyToListenableFuture.containsKey(
-                  downloadRequest.destinationFileUri().toString())) {
-                notification
-                    .setCategory(NotificationCompat.CATEGORY_PROGRESS)
-                    .setSmallIcon(android.R.drawable.stat_sys_download)
-                    .setProgress(
-                        downloadRequest.fileSizeBytes(),
-                        (int) currentSize,
-                        /* indeterminate = */ downloadRequest.fileSizeBytes() <= 0);
-                notificationManager.notify(notificationKey, notification.build());
-              }
-              if (downloadRequest.listenerOptional().isPresent()) {
-                downloadRequest.listenerOptional().get().onProgress(currentSize);
-              }
-            });
+        // TODO(b/229123693): return this future once DownloadListener has an async api.
+        ListenableFuture<?> unused =
+            PropagatedFutures.transformAsync(
+                foregroundDownloadFutureMap.containsKey(foregroundDownloadKey.toString()),
+                futureInProgress -> {
+                  if (futureInProgress) {
+                    notification
+                        .setCategory(NotificationCompat.CATEGORY_PROGRESS)
+                        .setContentText(
+                            downloadRequest
+                                .notificationContentTextOptional()
+                                .or(downloadRequest.urlToDownload()))
+                        .setSmallIcon(android.R.drawable.stat_sys_download)
+                        .setProgress(
+                            downloadRequest.fileSizeBytes(),
+                            (int) currentSize,
+                            /* indeterminate= */ downloadRequest.fileSizeBytes() <= 0);
+                    notificationManager.notify(notificationKey, notification.build());
+                  }
+                  if (downloadRequest.listenerOptional().isPresent()) {
+                    downloadRequest.listenerOptional().get().onProgress(currentSize);
+                  }
+                  return immediateVoidFuture();
+                },
+                sequentialControlExecutor);
       }
 
       @Override
       public void onPausedForConnectivity() {
-        sequentialControlExecutor.execute(
-            () -> {
-              // There can be a race condition, where onPausedForConnectivity can be called
-              // after onComplete or onFailure which removes the future and the notification.
-              if (keyToListenableFuture.containsKey(
-                  downloadRequest.destinationFileUri().toString())) {
-                notification
-                    .setCategory(NotificationCompat.CATEGORY_STATUS)
-                    .setContentText(NotificationUtil.getDownloadPausedMessage(context))
-                    .setSmallIcon(android.R.drawable.stat_sys_download)
-                    .setOngoing(true)
-                    // hide progress bar.
-                    .setProgress(0, 0, false);
-                notificationManager.notify(notificationKey, notification.build());
-              }
-
-              if (downloadRequest.listenerOptional().isPresent()) {
-                downloadRequest.listenerOptional().get().onPausedForConnectivity();
-              }
-            });
+        // TODO(b/229123693): return this future once DownloadListener has an async api.
+        ListenableFuture<?> unused =
+            PropagatedFutures.transformAsync(
+                foregroundDownloadFutureMap.containsKey(foregroundDownloadKey.toString()),
+                futureInProgress -> {
+                  if (futureInProgress) {
+                    notification
+                        .setCategory(NotificationCompat.CATEGORY_STATUS)
+                        .setContentText(networkPausedMessage)
+                        .setSmallIcon(android.R.drawable.stat_sys_download)
+                        .setOngoing(true)
+                        // hide progress bar.
+                        .setProgress(0, 0, false);
+                    notificationManager.notify(notificationKey, notification.build());
+                  }
+                  if (downloadRequest.listenerOptional().isPresent()) {
+                    downloadRequest.listenerOptional().get().onPausedForConnectivity();
+                  }
+                  return immediateVoidFuture();
+                },
+                sequentialControlExecutor);
       }
 
       @Override
@@ -350,92 +389,154 @@ final class DownloaderImpl implements Downloader {
         ListenableFuture<Void> clientOnCompleteFuture =
             downloadRequest.listenerOptional().isPresent()
                 ? downloadRequest.listenerOptional().get().onComplete()
-                : Futures.immediateVoidFuture();
+                : immediateVoidFuture();
 
         // Logic to shutdown Foreground Download Service after the client's provided onComplete
         // finished
-        clientOnCompleteFuture.addListener(
-            () -> {
-              // Clear the notification action.
-              notification.mActions.clear();
-
-              if (downloadRequest.showDownloadedNotification()) {
-                notification
-                    .setCategory(NotificationCompat.CATEGORY_STATUS)
-                    .setContentText(NotificationUtil.getDownloadSuccessMessage(context))
-                    .setOngoing(false)
-                    .setSmallIcon(android.R.drawable.stat_sys_download_done)
-                    // hide progress bar.
-                    .setProgress(0, 0, false);
-
-                notificationManager.notify(notificationKey, notification.build());
-              } else {
-                NotificationUtil.cancelNotificationForKey(
-                    context, downloadRequest.destinationFileUri().toString());
-              }
-
-              keyToListenableFuture.remove(downloadRequest.destinationFileUri().toString());
-              // If there is no other on-going foreground download, shutdown the
-              // ForegroundDownloadService
-              if (keyToListenableFuture.isEmpty()) {
-                NotificationUtil.stopForegroundDownloadService(
-                    context, foregroundDownloadServiceClassOptional.get());
-              }
+        return PropagatedFluentFuture.from(clientOnCompleteFuture)
+            .transformAsync(
+                unused -> {
+                  // onComplete succeeded, show a success message
+                  notification.mActions.clear();
+
+                  if (downloadRequest.showDownloadedNotification()) {
+                    notification
+                        .setCategory(NotificationCompat.CATEGORY_STATUS)
+                        .setContentText(NotificationUtil.getDownloadSuccessMessage(context))
+                        .setOngoing(false)
+                        .setSmallIcon(android.R.drawable.stat_sys_download_done)
+                        // hide progress bar.
+                        .setProgress(0, 0, false);
+
+                    notificationManager.notify(notificationKey, notification.build());
+                  } else {
+                    NotificationUtil.cancelNotificationForKey(
+                        context, foregroundDownloadKey.toString());
+                  }
+                  return immediateVoidFuture();
+                },
+                sequentialControlExecutor)
+            .catchingAsync(
+                Exception.class,
+                e -> {
+                  LogUtil.w(
+                      e,
+                      "%s: Delegate onComplete failed for uri: %s, showing failure notification.",
+                      TAG,
+                      downloadRequest.destinationFileUri());
+                  notification.mActions.clear();
+
+                  if (downloadRequest.showDownloadedNotification()) {
+                    notification
+                        .setCategory(NotificationCompat.CATEGORY_STATUS)
+                        .setContentText(NotificationUtil.getDownloadFailedMessage(context))
+                        .setOngoing(false)
+                        .setSmallIcon(android.R.drawable.stat_sys_warning)
+                        // hide progress bar.
+                        .setProgress(0, 0, false);
+
+                    notificationManager.notify(notificationKey, notification.build());
+                  } else {
+                    NotificationUtil.cancelNotificationForKey(
+                        context, downloadRequest.destinationFileUri().toString());
+                  }
 
-              downloadMonitorOptional
-                  .get()
-                  .removeDownloadListener(downloadRequest.destinationFileUri());
-            },
-            sequentialControlExecutor);
-        return clientOnCompleteFuture;
+                  return immediateVoidFuture();
+                },
+                sequentialControlExecutor)
+            .transformAsync(
+                unused -> {
+                  // After success or failure notification is shown, clean up
+                  downloadMonitorOptional
+                      .get()
+                      .removeDownloadListener(downloadRequest.destinationFileUri());
+
+                  return foregroundDownloadFutureMap.remove(foregroundDownloadKey.toString());
+                },
+                sequentialControlExecutor);
       }
 
       @Override
       public void onFailure(Throwable t) {
-        sequentialControlExecutor.execute(
-            () -> {
-              // Clear the notification action.
-              notification.mActions.clear();
-
-              // Show download failed in notification.
-              notification
-                  .setCategory(NotificationCompat.CATEGORY_STATUS)
-                  .setContentText(NotificationUtil.getDownloadFailedMessage(context))
-                  .setOngoing(false)
-                  .setSmallIcon(android.R.drawable.stat_sys_warning)
-                  // hide progress bar.
-                  .setProgress(0, 0, false);
-
-              notificationManager.notify(notificationKey, notification.build());
-
-              keyToListenableFuture.remove(downloadRequest.destinationFileUri().toString());
-
-              // If there is no other on-going foreground download, shutdown the
-              // ForegroundDownloadService
-              if (keyToListenableFuture.isEmpty()) {
-                NotificationUtil.stopForegroundDownloadService(
-                    context, foregroundDownloadServiceClassOptional.get());
-              }
+        // TODO(b/229123693): return this future once DownloadListener has an async api.
+        ListenableFuture<?> unused =
+            PropagatedFutures.submitAsync(
+                () -> {
+                  // Clear the notification action.
+                  notification.mActions.clear();
+
+                  // Show download failed in notification.
+                  notification
+                      .setCategory(NotificationCompat.CATEGORY_STATUS)
+                      .setContentText(NotificationUtil.getDownloadFailedMessage(context))
+                      .setOngoing(false)
+                      .setSmallIcon(android.R.drawable.stat_sys_warning)
+                      // hide progress bar.
+                      .setProgress(0, 0, false);
+
+                  notificationManager.notify(notificationKey, notification.build());
 
-              if (downloadRequest.listenerOptional().isPresent()) {
-                downloadRequest.listenerOptional().get().onFailure(t);
-              }
-              downloadMonitorOptional
-                  .get()
-                  .removeDownloadListener(downloadRequest.destinationFileUri());
-            });
+                  if (downloadRequest.listenerOptional().isPresent()) {
+                    downloadRequest.listenerOptional().get().onFailure(t);
+                  }
+                  downloadMonitorOptional
+                      .get()
+                      .removeDownloadListener(downloadRequest.destinationFileUri());
+
+                  return foregroundDownloadFutureMap.remove(foregroundDownloadKey.toString());
+                },
+                sequentialControlExecutor);
       }
     };
   }
 
   @Override
-  public void cancelForegroundDownload(String destinationFileUri) {
-    LogUtil.d("%s: CancelForegroundDownload for Uri = %s", TAG, destinationFileUri);
-    sequentialControlExecutor.execute(
-        () -> {
-          if (keyToListenableFuture.containsKey(destinationFileUri)) {
-            keyToListenableFuture.get(destinationFileUri).cancel(true);
-          }
-        });
+  public void cancelForegroundDownload(String downloadKey) {
+    LogUtil.d("%s: CancelForegroundDownload for Uri = %s", TAG, downloadKey);
+    ListenableFuture<?> unused =
+        PropagatedFutures.transformAsync(
+            getInProgressDownloadFuture(downloadKey),
+            downloadFuture -> {
+              if (downloadFuture.isPresent()) {
+                LogUtil.v(
+                    "%s: CancelForegroundDownload future found for key = %s, cancelling...",
+                    TAG, downloadKey);
+                downloadFuture.get().cancel(false);
+              }
+              return immediateVoidFuture();
+            },
+            sequentialControlExecutor);
+  }
+
+  private ListenableFuture<Optional<ListenableFuture<Void>>> getInProgressDownloadFuture(
+      String key) {
+    return PropagatedFutures.transformAsync(
+        foregroundDownloadFutureMap.containsKey(key),
+        isInForeground ->
+            isInForeground ? foregroundDownloadFutureMap.get(key) : downloadFutureMap.get(key),
+        sequentialControlExecutor);
+  }
+
+  private static DownloadFutureMap.StateChangeCallbacks createCallbacksForForegroundService(
+      Context context, Optional<Class<?>> foregroundDownloadServiceClassOptional) {
+    return new DownloadFutureMap.StateChangeCallbacks() {
+      @Override
+      public void onAdd(String key, int newSize) {
+        // Only start foreground service if this is the first future we are adding.
+        if (newSize == 1 && foregroundDownloadServiceClassOptional.isPresent()) {
+          NotificationUtil.startForegroundDownloadService(
+              context, foregroundDownloadServiceClassOptional.get(), key);
+        }
+      }
+
+      @Override
+      public void onRemove(String key, int newSize) {
+        // Only stop foreground service if there are no more futures remaining.
+        if (newSize == 0 && foregroundDownloadServiceClassOptional.isPresent()) {
+          NotificationUtil.stopForegroundDownloadService(
+              context, foregroundDownloadServiceClassOptional.get(), key);
+        }
+      }
+    };
   }
 }
diff --git a/java/com/google/android/libraries/mobiledatadownload/lite/annotations/BUILD b/java/com/google/android/libraries/mobiledatadownload/lite/annotations/BUILD
index fd00b3b..17ac54b 100644
--- a/java/com/google/android/libraries/mobiledatadownload/lite/annotations/BUILD
+++ b/java/com/google/android/libraries/mobiledatadownload/lite/annotations/BUILD
@@ -14,6 +14,7 @@
 load("@build_bazel_rules_android//android:rules.bzl", "android_library")
 
 package(
+    default_applicable_licenses = ["//:license"],
     default_visibility = [
         "//visibility:public",
     ],