path: root/java/src/com/android/textclassifier/downloader/ModelDownloadWorker.java

/*
 * Copyright (C) 2018 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.android.textclassifier.downloader;

import static java.lang.Math.min;

import android.content.Context;
import android.os.LocaleList;
import android.util.ArrayMap;
import android.util.Pair;
import androidx.work.ListenableWorker;
import androidx.work.WorkerParameters;
import com.android.textclassifier.common.ModelType;
import com.android.textclassifier.common.ModelType.ModelTypeDef;
import com.android.textclassifier.common.TextClassifierServiceExecutors;
import com.android.textclassifier.common.TextClassifierSettings;
import com.android.textclassifier.common.base.TcLog;
import com.android.textclassifier.downloader.DownloadedModelDatabase.Manifest;
import com.android.textclassifier.downloader.DownloadedModelDatabase.ManifestEnrollment;
import com.android.textclassifier.downloader.DownloadedModelDatabase.Model;
import com.google.auto.value.AutoValue;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.util.concurrent.FluentFuture;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.errorprone.annotations.concurrent.GuardedBy;
import java.time.Clock;
import java.util.ArrayList;
import java.util.Locale;

/** The WorkManager worker to download models for TextClassifierService. */
public final class ModelDownloadWorker extends ListenableWorker {
  private static final String TAG = "ModelDownloadWorker";

  public static final String INPUT_DATA_KEY_WORK_ID = "ModelDownloadWorker_workId";
  public static final String INPUT_DATA_KEY_SCHEDULED_TIMESTAMP =
      "ModelDownloadWorker_scheduledTimestamp";

  private final ListeningExecutorService executorService;
  private final ModelDownloader downloader;
  private final DownloadedModelManager downloadedModelManager;
  private final TextClassifierSettings settings;

  private final long workId;

  private final Clock clock;
  private final long workScheduledTimeMillis;

  private final Object lock = new Object();

  private long workStartedTimeMillis = 0;

  @GuardedBy("lock")
  private final ArrayMap<String, ListenableFuture<Void>> pendingDownloads;

  private ImmutableMap<String, ManifestsToDownloadByType> manifestsToDownload;

  public ModelDownloadWorker(Context context, WorkerParameters workerParams) {
    super(context, workerParams);
    this.executorService = TextClassifierServiceExecutors.getDownloaderExecutor();
    this.downloader = new ModelDownloaderImpl(context, executorService);
    this.downloadedModelManager = DownloadedModelManagerImpl.getInstance(context);
    this.settings = new TextClassifierSettings(context);
    this.pendingDownloads = new ArrayMap<>();
    this.manifestsToDownload = null;

    this.workId = workerParams.getInputData().getLong(INPUT_DATA_KEY_WORK_ID, 0);
    this.workScheduledTimeMillis =
        workerParams.getInputData().getLong(INPUT_DATA_KEY_SCHEDULED_TIMESTAMP, 0);
    this.clock = Clock.systemUTC();
  }

  @VisibleForTesting
  ModelDownloadWorker(
      Context context,
      WorkerParameters workerParams,
      ListeningExecutorService executorService,
      ModelDownloader modelDownloader,
      DownloadedModelManager downloadedModelManager,
      TextClassifierSettings settings,
      long workId,
      Clock clock,
      long workScheduledTimeMillis) {
    super(context, workerParams);
    this.executorService = executorService;
    this.downloader = modelDownloader;
    this.downloadedModelManager = downloadedModelManager;
    this.settings = settings;
    this.pendingDownloads = new ArrayMap<>();
    this.manifestsToDownload = null;
    this.workId = workId;
    this.clock = clock;
    this.workScheduledTimeMillis = workScheduledTimeMillis;
  }

  @Override
  public final ListenableFuture<ListenableWorker.Result> startWork() {
    TcLog.d(TAG, "Start download work...");
    workStartedTimeMillis = getCurrentTimeMillis();
    // Notice: startWork() is invoked on the main thread
    if (!settings.isModelDownloadManagerEnabled()) {
      TcLog.e(TAG, "Model Downloader is disabled. Abort the work.");
      logDownloadWorkCompleted(
          TextClassifierDownloadLogger.WORK_RESULT_FAILURE_MODEL_DOWNLOADER_DISABLED);
      return Futures.immediateFuture(ListenableWorker.Result.failure());
    }
    if (getRunAttemptCount() >= settings.getModelDownloadWorkerMaxAttempts()) {
      TcLog.d(TAG, "Max attempt reached. Abort download work.");
      logDownloadWorkCompleted(
          TextClassifierDownloadLogger.WORK_RESULT_FAILURE_MAX_RUN_ATTEMPT_REACHED);
      return Futures.immediateFuture(ListenableWorker.Result.failure());
    }

    return FluentFuture.from(Futures.submitAsync(this::checkAndDownloadModels, executorService))
        .transform(
            downloadResult -> {
              Preconditions.checkNotNull(manifestsToDownload);
              downloadedModelManager.onDownloadCompleted(manifestsToDownload);
              TcLog.d(TAG, "Download work completed: " + downloadResult);
              if (downloadResult.failureCount() == 0) {
                logDownloadWorkCompleted(
                    downloadResult.successCount() > 0
                        ? TextClassifierDownloadLogger.WORK_RESULT_SUCCESS_MODEL_DOWNLOADED
                        : TextClassifierDownloadLogger.WORK_RESULT_SUCCESS_NO_UPDATE_AVAILABLE);
                return ListenableWorker.Result.success();
              } else {
                logDownloadWorkCompleted(
                    TextClassifierDownloadLogger.WORK_RESULT_RETRY_MODEL_DOWNLOAD_FAILED);
                return ListenableWorker.Result.retry();
              }
            },
            executorService)
        .catching(
            Throwable.class,
            t -> {
              TcLog.e(TAG, "Unexpected Exception during downloading: ", t);
              logDownloadWorkCompleted(
                  TextClassifierDownloadLogger.WORK_RESULT_RETRY_RUNTIME_EXCEPTION);
              return ListenableWorker.Result.retry();
            },
            executorService);
  }

  /**
   * Checks device settings and returns the list of locales to download according to multi language
   * support settings. Guarantees that the primary locale goes first.
   */
  private ImmutableList<Locale> getLocalesToDownload() {
    LocaleList localeList = LocaleList.getAdjustedDefault();
    Locale primaryLocale = localeList.get(0);
    if (!settings.isMultiLanguageSupportEnabled()) {
      return ImmutableList.of(primaryLocale);
    }
    ImmutableList.Builder<Locale> localesToDownloadBuilder = ImmutableList.builder();
    int size = min(settings.getMultiLanguageModelsLimit(), localeList.size());
    for (int i = 0; i < size; i++) {
      localesToDownloadBuilder.add(localeList.get(i));
    }
    return localesToDownloadBuilder.build();
  }

  /**
   * Returns list of locales to download from {@code localeList} for the given {@code modelType}.
   */
  private ImmutableList<Locale> getLocalesToDownloadByType(
      ImmutableList<Locale> localeList, @ModelTypeDef String modelType) {
    if (!settings.getEnabledModelTypesForMultiLanguageSupport().contains(modelType)) {
      return ImmutableList.of(Locale.getDefault());
    }
    return localeList;
  }

  /**
   * Check device config and dispatch download tasks for all modelTypes.
   *
   * <p>Download tasks will be combined and logged after completion. Return true if all tasks
   * succeeded
   */
  private ListenableFuture<DownloadResult> checkAndDownloadModels() {
    ImmutableList<Locale> localesToDownload = getLocalesToDownload();
    ArrayList<ListenableFuture<Boolean>> downloadResultFutures = new ArrayList<>();
    ImmutableMap.Builder<String, ManifestsToDownloadByType> manifestsToDownloadBuilder =
        ImmutableMap.builder();
    for (String modelType : ModelType.values()) {
      ImmutableList<Locale> localesToDownloadByType =
          getLocalesToDownloadByType(localesToDownload, modelType);
      ImmutableMap.Builder<String, String> localeTagToManifestUrlBuilder = ImmutableMap.builder();
      for (Locale locale : localesToDownloadByType) {
        Pair<String, String> bestLocaleTagAndManifestUrl =
            LocaleUtils.lookupBestLocaleTagAndManifestUrl(modelType, locale, settings);
        if (bestLocaleTagAndManifestUrl == null) {
          TcLog.w(
              TAG,
              String.format(
                  Locale.US, "No suitable manifest for %s, %s", modelType, locale.toLanguageTag()));
          continue;
        }
        String bestLocaleTag = bestLocaleTagAndManifestUrl.first;
        String manifestUrl = bestLocaleTagAndManifestUrl.second;
        localeTagToManifestUrlBuilder.put(bestLocaleTag, manifestUrl);
        TcLog.d(
            TAG,
            String.format(
                Locale.US,
                "model type: %s, current locale tag: %s, best locale tag: %s, manifest url: %s",
                modelType,
                locale.toLanguageTag(),
                bestLocaleTag,
                manifestUrl));
        if (!shouldDownloadManifest(modelType, bestLocaleTag, manifestUrl)) {
          continue;
        }
        downloadResultFutures.add(
            downloadManifestAndRegister(modelType, bestLocaleTag, manifestUrl));
      }
      manifestsToDownloadBuilder.put(
          modelType,
          ManifestsToDownloadByType.create(localeTagToManifestUrlBuilder.buildOrThrow()));
    }
    manifestsToDownload = manifestsToDownloadBuilder.buildOrThrow();

    return Futures.whenAllComplete(downloadResultFutures)
        .call(
            () -> {
              TcLog.d(TAG, "All Download Tasks Completed");
              int successCount = 0;
              int failureCount = 0;
              for (ListenableFuture<Boolean> downloadResultFuture : downloadResultFutures) {
                if (Futures.getDone(downloadResultFuture)) {
                  successCount += 1;
                } else {
                  failureCount += 1;
                }
              }
              return DownloadResult.create(successCount, failureCount);
            },
            executorService);
  }

  private boolean shouldDownloadManifest(
      @ModelTypeDef String modelType, String localeTag, String manifestUrl) {
    Manifest downloadedManifest = downloadedModelManager.getManifest(manifestUrl);
    if (downloadedManifest == null) {
      return true;
    }
    if (downloadedManifest.getStatus() == Manifest.STATUS_FAILED) {
      if (downloadedManifest.getFailureCounts() >= settings.getManifestDownloadMaxAttempts()) {
        TcLog.w(
            TAG,
            String.format(
                Locale.US,
                "Manifest failed too many times, stop retrying: %s %d",
                manifestUrl,
                downloadedManifest.getFailureCounts()));
        return false;
      } else {
        return true;
      }
    }
    ManifestEnrollment manifestEnrollment =
        downloadedModelManager.getManifestEnrollment(modelType, localeTag);
    return manifestEnrollment == null || !manifestUrl.equals(manifestEnrollment.getManifestUrl());
  }

  /**
   * Downloads a single manifest and models configured inside it.
   *
   * <p>The returned future should always resolve to a ManifestDownloadResult as we catch all
   * exceptions.
   */
  private ListenableFuture<Boolean> downloadManifestAndRegister(
      @ModelTypeDef String modelType, String localeTag, String manifestUrl) {
    long downloadStartTimestamp = getCurrentTimeMillis();
    return FluentFuture.from(downloadManifest(manifestUrl))
        .transform(
            unused -> {
              downloadedModelManager.registerManifestEnrollment(modelType, localeTag, manifestUrl);
              TextClassifierDownloadLogger.downloadSucceeded(
                  workId,
                  modelType,
                  manifestUrl,
                  getRunAttemptCount(),
                  getCurrentTimeMillis() - downloadStartTimestamp);
              TcLog.d(TAG, "Manifest downloaded and registered: " + manifestUrl);
              return true;
            },
            executorService)
        .catching(
            Throwable.class,
            t -> {
              downloadedModelManager.registerManifestDownloadFailure(manifestUrl);
              int errorCode = ModelDownloadException.UNKNOWN_FAILURE_REASON;
              int downloaderLibErrorCode = 0;
              if (t instanceof ModelDownloadException) {
                ModelDownloadException mde = (ModelDownloadException) t;
                errorCode = mde.getErrorCode();
                downloaderLibErrorCode = mde.getDownloaderLibErrorCode();
              }
              TcLog.e(TAG, "Failed to download manfiest: " + manifestUrl, t);
              TextClassifierDownloadLogger.downloadFailed(
                  workId,
                  modelType,
                  manifestUrl,
                  errorCode,
                  getRunAttemptCount(),
                  downloaderLibErrorCode,
                  getCurrentTimeMillis() - downloadStartTimestamp);
              return false;
            },
            executorService);
  }

  // Download a manifest and its models, and register it to Manifest table.
  private ListenableFuture<Void> downloadManifest(String manifestUrl) {
    synchronized (lock) {
      Manifest downloadedManifest = downloadedModelManager.getManifest(manifestUrl);
      if (downloadedManifest != null
          && downloadedManifest.getStatus() == Manifest.STATUS_SUCCEEDED) {
        TcLog.d(TAG, "Manifest already downloaded: " + manifestUrl);
        return Futures.immediateVoidFuture();
      }
      if (pendingDownloads.containsKey(manifestUrl)) {
        return pendingDownloads.get(manifestUrl);
      }
      ListenableFuture<Void> manfiestDownloadFuture =
          FluentFuture.from(downloader.downloadManifest(manifestUrl))
              .transformAsync(
                  manifest -> {
                    ModelManifest.Model modelInfo = manifest.getModels(0);
                    return Futures.transform(
                        downloadModel(modelInfo), unused -> modelInfo, executorService);
                  },
                  executorService)
              .transform(
                  modelInfo -> {
                    downloadedModelManager.registerManifest(manifestUrl, modelInfo.getUrl());
                    return null;
                  },
                  executorService);
      pendingDownloads.put(manifestUrl, manfiestDownloadFuture);
      return manfiestDownloadFuture;
    }
  }

  // Download a model and register it into Model table.
  private ListenableFuture<Void> downloadModel(ModelManifest.Model modelInfo) {
    String modelUrl = modelInfo.getUrl();
    synchronized (lock) {
      Model downloadedModel = downloadedModelManager.getModel(modelUrl);
      if (downloadedModel != null) {
        TcLog.d(TAG, "Model file already exists: " + downloadedModel.getModelPath());
        return Futures.immediateVoidFuture();
      }
      if (pendingDownloads.containsKey(modelUrl)) {
        return pendingDownloads.get(modelUrl);
      }
      ListenableFuture<Void> modelDownloadFuture =
          FluentFuture.from(
                  downloader.downloadModel(
                      downloadedModelManager.getModelDownloaderDir(), modelInfo))
              .transform(
                  modelFile -> {
                    downloadedModelManager.registerModel(modelUrl, modelFile.getAbsolutePath());
                    TcLog.d(TAG, "Model File downloaded: " + modelUrl);
                    return null;
                  },
                  executorService);
      pendingDownloads.put(modelUrl, modelDownloadFuture);
      return modelDownloadFuture;
    }
  }

  /**
   * This method will be called when we our work gets interrupted by the system. Result future
   * should have already been cancelled in that case. Unless it's because the REPLACE policy of
   * WorkManager unique queue, the interrupted work will be rescheduled later.
   */
  @Override
  public final void onStopped() {
    TcLog.d(TAG, String.format(Locale.US, "Stop download. Attempt:%d", getRunAttemptCount()));
    logDownloadWorkCompleted(TextClassifierDownloadLogger.WORK_RESULT_RETRY_STOPPED_BY_OS);
  }

  private long getCurrentTimeMillis() {
    return clock.instant().toEpochMilli();
  }

  private void logDownloadWorkCompleted(int workResult) {
    if (workStartedTimeMillis < workScheduledTimeMillis) {
      TcLog.w(
          TAG,
          String.format(
              Locale.US,
              "Bad workStartedTimeMillis: %d, workScheduledTimeMillis: %d",
              workStartedTimeMillis,
              workScheduledTimeMillis));
      workStartedTimeMillis = workScheduledTimeMillis;
    }
    TextClassifierDownloadLogger.downloadWorkCompleted(
        workId,
        workResult,
        getRunAttemptCount(),
        workStartedTimeMillis - workScheduledTimeMillis,
        getCurrentTimeMillis() - workStartedTimeMillis);
  }

  @AutoValue
  abstract static class DownloadResult {
    public abstract int successCount();

    public abstract int failureCount();

    public static DownloadResult create(int successCount, int failureCount) {
      return new AutoValue_ModelDownloadWorker_DownloadResult(successCount, failureCount);
    }
  }
}