Path: blob/main/SignalServiceKit/Backups/Attachments/BackupAttachmentUploadScheduler.swift
1 views
//
// Copyright 2024 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only
//
import Foundation
public struct BackupAttachmentUploadEnqueueMode: OptionSet {
public let rawValue: Int
public init(rawValue: Int) {
self.rawValue = rawValue
}
public static let fullsize = Self(rawValue: 1 << 0)
public static let thumbnail = Self(rawValue: 1 << 1)
static let all: BackupAttachmentUploadEnqueueMode = [.fullsize, .thumbnail]
}
public class BackupAttachmentUploadScheduler {
private let attachmentStore: AttachmentStore
private let backupAttachmentUploadStore: BackupAttachmentUploadStore
private let backupAttachmentUploadEraStore: BackupAttachmentUploadEraStore
private let dateProvider: DateProvider
private let interactionStore: InteractionStore
private let remoteConfigProvider: any RemoteConfigProvider
private let tsAccountManager: TSAccountManager
public init(
attachmentStore: AttachmentStore,
backupAttachmentUploadStore: BackupAttachmentUploadStore,
backupAttachmentUploadEraStore: BackupAttachmentUploadEraStore,
dateProvider: @escaping DateProvider,
interactionStore: InteractionStore,
remoteConfigProvider: any RemoteConfigProvider,
tsAccountManager: TSAccountManager,
) {
self.attachmentStore = attachmentStore
self.backupAttachmentUploadStore = backupAttachmentUploadStore
self.backupAttachmentUploadEraStore = backupAttachmentUploadEraStore
self.dateProvider = dateProvider
self.interactionStore = interactionStore
self.remoteConfigProvider = remoteConfigProvider
self.tsAccountManager = tsAccountManager
}
/// Returns true if the provided attachment is eligible to be uploaded
/// to backup/media tier, independently of the current backupPlan state.
///
/// - parameter mode: OptionSet instructing what eligibility type to check
public func isEligibleToUpload(
_ attachment: Attachment,
mode: BackupAttachmentUploadEnqueueMode,
currentUploadEra: String,
tx: DBReadTransaction,
) -> Bool {
guard
tsAccountManager.registrationState(tx: tx).isRegisteredPrimaryDevice,
let stream = attachment.asStream()
else {
return false
}
let eligibility = Eligibility.forAttachment(
stream,
currentUploadEra: currentUploadEra,
remoteConfig: remoteConfigProvider.currentConfig(),
)
if mode.contains(.fullsize), !eligibility.needsUploadFullsize {
return false
}
if mode.contains(.thumbnail), !eligibility.needsUploadThumbnail {
return false
}
let highestPriorityEligibleOwner = self.highestPriorityEligibleOwner(
attachment,
tx: tx,
)
return highestPriorityEligibleOwner != nil
}
/// "Enqueue" an attachment from a backup for upload, if needed and eligible, otherwise do nothing.
///
/// Fetches all attachment owners and uses the highest priority one available.
///
/// Doesn't actually trigger an upload; callers must later call
/// ``BackupAttachmentUploadQueueRunner.backUpAllAttachments()`` to upload.
public func enqueueUsingHighestPriorityOwnerIfNeeded(
_ attachment: Attachment,
mode: BackupAttachmentUploadEnqueueMode,
tx: DBWriteTransaction,
file: StaticString? = #file,
function: StaticString? = #function,
line: UInt? = #line,
) {
guard tsAccountManager.registrationState(tx: tx).isRegisteredPrimaryDevice else {
return
}
// Before we fetch references, check if the attachment is
// eligible to begin with.
guard let stream = attachment.asStream() else {
if let file, let function, let line {
Logger.info("Skipping enqueue of non-stream \(attachment.id) from \(file) \(line): \(function)")
}
return
}
let currentUploadEra = backupAttachmentUploadEraStore.currentUploadEra(tx: tx)
let eligibility = Eligibility.forAttachment(
stream,
currentUploadEra: currentUploadEra,
remoteConfig: remoteConfigProvider.currentConfig(),
)
guard eligibility.needsUploadFullsize || eligibility.needsUploadThumbnail else {
if let file, let function, let line {
Logger.info("Skipping enqueue of fullsize+thumbnail \(attachment.id) from \(file) \(line): \(function)")
}
return
}
guard
let uploadOwnerType = highestPriorityEligibleOwner(attachment, tx: tx)
else {
if let file, let function, let line {
Logger.info("No eligible owners; skipping enqueue of \(attachment.id) from \(file) \(line): \(function)")
}
return
}
if mode.contains(.fullsize) {
if eligibility.needsUploadFullsize {
backupAttachmentUploadStore.enqueue(
stream,
owner: uploadOwnerType,
fullsize: true,
tx: tx,
file: file,
function: function,
line: line,
)
} else if let file, let function, let line {
Logger.info("Skipping enqueue of fullsize \(attachment.id) from \(file) \(line): \(function)")
}
}
if mode.contains(.thumbnail) {
if eligibility.needsUploadThumbnail {
backupAttachmentUploadStore.enqueue(
stream,
owner: uploadOwnerType,
fullsize: false,
tx: tx,
file: file,
function: function,
line: line,
)
} else if let file, let function, let line {
Logger.info("Skipping enqueue of thumbnail \(attachment.id) from \(file) \(line): \(function)")
}
}
}
/// "Enqueue" an attachment from a backup for upload, if needed and eligible via the provided
/// owner, otherwise do nothing.
///
/// The attachment may or may not already be enqueued for upload using via other owners;
/// if so the provided owner (if eligible) may increase the priority.
///
/// Doesn't actually trigger an upload; callers must later call
/// ``BackupAttachmentUploadQueueRunner.backUpAllAttachments()`` to upload.
public func enqueueIfNeededWithOwner(
_ attachment: Attachment,
owner: AttachmentReference.Owner,
tx: DBWriteTransaction,
file: StaticString? = #file,
function: StaticString? = #function,
line: UInt? = #line,
) {
guard tsAccountManager.registrationState(tx: tx).isRegisteredPrimaryDevice else {
return
}
guard let stream = attachment.asStream() else {
if let file, let function, let line {
Logger.info("Skipping enqueue of non-stream \(attachment.id) from \(file) \(line): \(function)")
}
return
}
let currentUploadEra = backupAttachmentUploadEraStore.currentUploadEra(tx: tx)
let eligibility = Eligibility.forAttachment(
stream,
currentUploadEra: currentUploadEra,
remoteConfig: remoteConfigProvider.currentConfig(),
)
guard eligibility.needsUploadFullsize || eligibility.needsUploadThumbnail else {
if let file, let function, let line {
Logger.info("Skipping enqueue of fullsize+thumbnail \(attachment.id) from \(file) \(line): \(function)")
}
return
}
// We only include the provided owner because this is an incremental check;
// if some other owner made the attachment eligible for upload, it'd already
// be enqueued. We only care if this particular owner makes it newly eligible
// (or it was eligible both before and now, but the enqueuing it idempotent).
guard let uploadOwnerType = self.asEligibleUploadOwnerType(owner, tx: tx) else {
if let file, let function, let line {
Logger.info(
"Passed in owner not eligible (may be eligible with other owners);"
+ " skipping enqueue of \(attachment.id) from \(file) \(line): \(function)",
)
}
return
}
if eligibility.needsUploadFullsize {
backupAttachmentUploadStore.enqueue(
stream,
owner: uploadOwnerType,
fullsize: true,
tx: tx,
)
} else if let file, let function, let line {
Logger.info("Skipping enqueue of fullsize \(attachment.id) from \(file) \(line): \(function)")
}
if eligibility.needsUploadThumbnail {
backupAttachmentUploadStore.enqueue(
stream,
owner: uploadOwnerType,
fullsize: false,
tx: tx,
)
} else if let file, let function, let line {
Logger.info("Skipping enqueue of thumbnail \(attachment.id) from \(file) \(line): \(function)")
}
}
// MARK: - Eligibility
private struct Eligibility {
let needsUploadFullsize: Bool
let needsUploadThumbnail: Bool
static func forAttachment(
_ attachment: AttachmentStream,
currentUploadEra: String,
remoteConfig: RemoteConfig,
) -> Self {
let needsUploadFullsize = { () -> Bool in
if attachment.encryptedByteCount > remoteConfig.backupAttachmentMaxEncryptedBytes {
Logger.info("Skipping upload of too-large attachment \(attachment.id), \(attachment.encryptedByteCount) bytes")
return false
}
if let mediaTierInfo = attachment.attachment.mediaTierInfo {
return !mediaTierInfo.isUploaded(currentUploadEra: currentUploadEra)
} else {
return true
}
}()
let needsUploadThumbnail = { () -> Bool in
if let thumbnailMediaTierInfo = attachment.attachment.thumbnailMediaTierInfo {
return !thumbnailMediaTierInfo.isUploaded(currentUploadEra: currentUploadEra)
} else {
return AttachmentBackupThumbnail.canBeThumbnailed(attachment.attachment)
}
}()
return Self(
needsUploadFullsize: needsUploadFullsize,
needsUploadThumbnail: needsUploadThumbnail,
)
}
}
private func asEligibleUploadOwnerType(
_ owner: AttachmentReference.Owner,
tx: DBReadTransaction,
) -> QueuedBackupAttachmentUpload.OwnerType? {
switch owner {
case .message(let messageSource):
switch messageSource {
case .oversizeText:
// We inline oversize text in the backup, and don't back
// up the corresponding attachment.
Logger.info("Skip: oversized text")
return nil
case
.bodyAttachment,
.contactAvatar,
.linkPreview,
.quotedReply,
.sticker:
break
}
guard
let message = interactionStore.fetchInteraction(
rowId: messageSource.messageRowId,
tx: tx,
) as? TSMessage
else {
owsFailDebug("Missing message!")
return nil
}
// For every owning reference, check that we _would_ include the
// attachment for it in a remote backup. If we wouldn't, that
// reference shouldn't be used as the anchor for upload (and if
// it is the only reference, we shouldn't upload at all!)
let includedContentFilter = BackupArchive.IncludedContentFilter(
backupPurpose: .remoteBackup,
)
if
includedContentFilter.shouldSkipAttachment(
owningMessage: message,
currentTimestamp: dateProvider().ows_millisecondsSince1970,
)
{
return nil
}
return .message(timestamp: messageSource.receivedAtTimestamp)
case .thread(let threadSource):
switch threadSource {
case .threadWallpaperImage, .globalThreadWallpaperImage:
return .threadWallpaper
}
case .storyMessage:
Logger.info("Skip: story message")
return nil
}
}
private func highestPriorityEligibleOwner(
_ attachment: Attachment,
tx: DBReadTransaction,
) -> QueuedBackupAttachmentUpload.OwnerType? {
var ineligibleOwners = [AttachmentReference.Owner]()
var eligibleOwner: AttachmentReference.Owner?
// Backup uploads are prioritized by attachment owner. Find the highest
// priority owner to use.
var uploadOwnerType: QueuedBackupAttachmentUpload.OwnerType?
attachmentStore.enumerateAllReferences(
toAttachmentId: attachment.id,
tx: tx,
) { reference, _ in
guard
let ownerType = self.asEligibleUploadOwnerType(
reference.owner,
tx: tx,
)
else {
ineligibleOwners.append(reference.owner)
return
}
if uploadOwnerType?.isHigherPriority(than: ownerType) != true {
eligibleOwner = reference.owner
uploadOwnerType = ownerType
}
}
// If an eligible owner was found, and an ineligible owner also exists,
// log this info for debugging
if ineligibleOwners.isEmpty == false, let eligibleOwner {
func debugString(_ type: AttachmentReference.Owner) -> String {
return switch type {
case .message: "message"
case .thread: "thread"
case .storyMessage: "story"
}
}
Logger.info("Attachment \(attachment.id): eligible owner: \(debugString(eligibleOwner)) [\(eligibleOwner.id)]")
ineligibleOwners.forEach { Logger.info("Attachment \(attachment.id): inegible owner: \(debugString($0)) [\($0.id)]") }
}
return uploadOwnerType
}
}
// MARK: -
extension Attachment.MediaTierInfo {
public func isUploaded(currentUploadEra: String) -> Bool {
return self.uploadEra == currentUploadEra && cdnNumber != nil
}
}
extension Attachment.ThumbnailMediaTierInfo {
public func isUploaded(currentUploadEra: String) -> Bool {
return self.uploadEra == currentUploadEra
&& cdnNumber != nil
}
}