firebase-ios-sdk/FirebaseMLModelDownloader/Sources/ModelDownloadTask.swift

// Copyright 2021 Google LLC
//
// 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.

import Foundation
import FirebaseCore

/// Possible states of model downloading.
enum ModelDownloadStatus {
  case notStarted
  case inProgress
  case successful
  case failed
}

/// Manager to handle model downloading device and storing downloaded model info to persistent storage.
class ModelDownloadTask: NSObject {
  typealias ProgressHandler = (Float) -> Void
  typealias Completion = (Result<CustomModel, DownloadError>) -> Void

  /// Name of the app associated with this instance of ModelDownloadTask.
  private let appName: String
  /// Model info downloaded from server.
  private(set) var remoteModelInfo: RemoteModelInfo
  /// User defaults to which local model info should ultimately be written.
  private let defaults: UserDefaults
  /// Keeps track of download associated with this model download task.
  private(set) var downloadStatus: ModelDownloadStatus
  /// Downloader instance.
  private let downloader: FileDownloader
  /// Telemetry logger.
  private let telemetryLogger: TelemetryLogger?

  init(remoteModelInfo: RemoteModelInfo,
       appName: String,
       defaults: UserDefaults,
       downloader: FileDownloader,
       telemetryLogger: TelemetryLogger? = nil) {
    self.remoteModelInfo = remoteModelInfo
    self.appName = appName
    self.downloader = downloader
    self.telemetryLogger = telemetryLogger
    self.defaults = defaults
    downloadStatus = .notStarted
  }
}

extension ModelDownloadTask {
  /// Name for model file stored on device.
  var downloadedModelFileName: String {
    return "fbml_model__\(appName)__\(remoteModelInfo.name)"
  }

  func download(progressHandler: ProgressHandler?, completion: @escaping Completion) {
    /// Prevent multiple concurrent downloads.
    guard downloadStatus == .notStarted else {
      completion(.failure(.internalError(description: ModelDownloadTask.ErrorDescription
          .anotherDownloadInProgress)))
      DeviceLogger.logEvent(level: .debug,
                            message: ModelDownloadTask.ErrorDescription.anotherDownloadInProgress,
                            messageCode: .anotherDownloadInProgressError)
      return
    }
    downloader.downloadFile(with: remoteModelInfo.downloadURL,
                            progressHandler: { downloadedBytes, totalBytes in
                              /// Fraction of model file downloaded.
                              let calculatedProgress = Float(downloadedBytes) / Float(totalBytes)
                              progressHandler?(calculatedProgress)
                            }) { result in
      switch result {
      case let .success(response):
        DeviceLogger.logEvent(level: .debug,
                              message: ModelDownloadTask.DebugDescription
                                .receivedServerResponse,
                              messageCode: .validModelDownloadResponse)
        self.handleResponse(
          response: response.urlResponse,
          tempURL: response.fileURL,
          completion: completion
        )
      case let .failure(error):
        var downloadError: DownloadError
        switch error {
        case let FileDownloaderError.networkError(error):
          let description = ModelDownloadTask.ErrorDescription
            .invalidHostName(error.localizedDescription)
          downloadError = .internalError(description: description)
          DeviceLogger.logEvent(level: .debug,
                                message: description,
                                messageCode: .hostnameError)
        // TODO: Handle this case better.
        case FileDownloaderError.sessionInvalidated:
          downloadError = .failedPrecondition
          DeviceLogger.logEvent(level: .debug,
                                message: ModelDownloadTask.ErrorDescription.sessionInvalidated,
                                messageCode: .invalidDownloadSessionError)
        case FileDownloaderError.unexpectedResponseType:
          let description = ModelDownloadTask.ErrorDescription.invalidServerResponse
          downloadError = .internalError(description: description)
          DeviceLogger.logEvent(level: .debug,
                                message: description,
                                messageCode: .invalidResponseError)

        default:
          let description = ModelDownloadTask.ErrorDescription.unknownDownloadError
          downloadError = .internalError(description: description)
          DeviceLogger.logEvent(level: .debug,
                                message: description,
                                messageCode: .modelDownloadError)
        }
        completion(.failure(downloadError))
      }
    }
  }

  /// Handle model download response.
  func handleResponse(response: HTTPURLResponse, tempURL: URL, completion: @escaping Completion) {
    guard (200 ..< 299).contains(response.statusCode) else {
      /// Possible failure due to download URL expiry.
      if response.statusCode == 400 {
        completion(.failure(.expiredDownloadURL))
        return
      }
      let description = ModelDownloadTask.ErrorDescription.modelDownloadFailed(response.statusCode)
      DeviceLogger.logEvent(level: .debug,
                            message: description,
                            messageCode: .modelDownloadError)
      completion(.failure(.internalError(description: description)))
      return
    }

    let modelFileURL = ModelFileManager.getDownloadedModelFilePath(
      appName: appName,
      modelName: remoteModelInfo.name
    )

    do {
      try ModelFileManager.moveFile(at: tempURL, to: modelFileURL)
      DeviceLogger.logEvent(level: .debug,
                            message: ModelDownloadTask.DebugDescription.savedModelFile,
                            messageCode: .downloadedModelFileSaved)
      /// Generate local model info.
      let localModelInfo = LocalModelInfo(from: remoteModelInfo, path: modelFileURL.absoluteString)
      /// Write model to user defaults.
      localModelInfo.writeToDefaults(defaults, appName: appName)
      DeviceLogger.logEvent(level: .debug,
                            message: ModelDownloadTask.DebugDescription.savedLocalModelInfo,
                            messageCode: .downloadedModelInfoSaved)
      /// Build model from model info.
      let model = CustomModel(localModelInfo: localModelInfo)
      downloadStatus = .successful
      telemetryLogger?.logModelDownloadEvent(
        eventName: .modelDownload,
        status: downloadStatus,
        model: model
      )
      completion(.success(model))
    } catch let error as DownloadError {
      downloadStatus = .failed
      telemetryLogger?.logModelDownloadEvent(eventName: .modelDownload, status: downloadStatus)
      DeviceLogger.logEvent(level: .debug,
                            message: ModelDownloadTask.ErrorDescription.saveModel,
                            messageCode: .downloadedModelSaveError)
      completion(.failure(error))
      return
    } catch {
      downloadStatus = .failed
      telemetryLogger?.logModelDownloadEvent(eventName: .modelDownload, status: downloadStatus)
      DeviceLogger.logEvent(level: .debug,
                            message: ModelDownloadTask.ErrorDescription.saveModel,
                            messageCode: .downloadedModelSaveError)
      completion(.failure(.internalError(description: error.localizedDescription)))
      return
    }
  }
}

/// Possible error messages for model downloading.
extension ModelDownloadTask {
  /// Debug descriptions.
  private enum DebugDescription {
    static let savedModelFile = "Model file saved successfully to device."
    static let savedLocalModelInfo = "Downloaded model info saved successfully to user defaults."
    static let receivedServerResponse = "Received a valid response from server."
  }

  /// Error descriptions.
  private enum ErrorDescription {
    static let invalidHostName = { (error: String) in
      "Unable to resolve hostname or connect to host: \(error)"
    }

    static let modelDownloadFailed = { (code: Int) in
      "Model download failed with HTTP error code: \(code)"
    }

    static let sessionInvalidated = "Session invalidated due to failed pre-conditions."
    static let invalidServerResponse =
      "Could not get valid server response for model downloading."
    static let unknownDownloadError = "Unable to download model due to unknown error."
    static let saveModel = "Unable to save downloaded remote model file."
    static let expiredModelInfo = "Unable to update expired model info."
    static let anotherDownloadInProgress = "Download already in progress."
  }
}