Signal-iOS/SignalServiceKit/Jobs/TSAttachmentMultisendJob.swift

324 lines
13 KiB
Swift

//
// Copyright 2019 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only
//
public struct TSAttachmentMultisendResult {
/// Resolved when the attachments are uploaded and sending is enqueued.
public let enqueuedPromise: Promise<Void>
/// Resolved when the message is sent.
public let sentPromise: Promise<Void>
}
public class TSAttachmentMultisendJobQueue {
private let jobQueueRunner: JobQueueRunner<
JobRecordFinderImpl<TSAttachmentMultisendJobRecord>,
TSAttachmentMultisendJobRunnerFactory
>
private let jobRunnerFactory: TSAttachmentMultisendJobRunnerFactory
public init(db: DB, reachabilityManager: SSKReachabilityManager) {
self.jobRunnerFactory = TSAttachmentMultisendJobRunnerFactory()
self.jobQueueRunner = JobQueueRunner(
canExecuteJobsConcurrently: false,
db: db,
jobFinder: JobRecordFinderImpl(db: db),
jobRunnerFactory: jobRunnerFactory
)
self.jobQueueRunner.listenForReachabilityChanges(reachabilityManager: reachabilityManager)
}
public func start(appContext: AppContext) {
if appContext.isNSE { return }
jobQueueRunner.start(shouldRestartExistingJobs: appContext.isMainApp)
}
@discardableResult
public func add(
attachmentIdMap: [String: [String]],
storyMessagesToSend: [OutgoingStoryMessage],
transaction: SDSAnyWriteTransaction
) -> TSAttachmentMultisendResult {
let jobRecord = TSAttachmentMultisendJobRecord(
attachmentIdMap: attachmentIdMap,
storyMessagesToSend: storyMessagesToSend
)
jobRecord.anyInsert(transaction: transaction)
let enqueuePromise = Promise<Void>.pending()
let sendPromise = Promise<Void>.pending()
let jobFutures = TSAttachmentMultisendFutures(
enqueuedFuture: enqueuePromise.1,
sentFuture: sendPromise.1
)
transaction.addSyncCompletion {
let runner = self.jobRunnerFactory.buildRunner(jobFutures)
self.jobQueueRunner.addPersistedJob(jobRecord, runner: runner)
}
return .init(
enqueuedPromise: enqueuePromise.0,
sentPromise: sendPromise.0
)
}
}
private class TSAttachmentMultisendFutures {
/// Resolved when the attachments are uploaded and sending is enqueued.
public let enqueuedFuture: Future<Void>
/// Resolved when the message is sent.
public let sentFuture: Future<Void>
init(enqueuedFuture: Future<Void>, sentFuture: Future<Void>) {
self.enqueuedFuture = enqueuedFuture
self.sentFuture = sentFuture
}
}
private class TSAttachmentMultisendJobRunnerFactory: JobRunnerFactory {
func buildRunner() -> TSAttachmentMultisendJobRunner {
TSAttachmentMultisendJobRunner(jobFutures: nil)
}
func buildRunner(_ jobFutures: TSAttachmentMultisendFutures?) -> TSAttachmentMultisendJobRunner {
TSAttachmentMultisendJobRunner(jobFutures: jobFutures)
}
}
private class TSAttachmentMultisendJobRunner: JobRunner, Dependencies {
private let jobFutures: TSAttachmentMultisendFutures?
init(jobFutures: TSAttachmentMultisendFutures?) {
self.jobFutures = jobFutures
}
private enum Constants {
static let maxRetries: UInt = 4
}
func runJobAttempt(_ jobRecord: TSAttachmentMultisendJobRecord) async -> JobAttemptResult {
return await JobAttemptResult.executeBlockWithDefaultErrorHandler(
jobRecord: jobRecord,
retryLimit: Constants.maxRetries,
db: DependenciesBridge.shared.db,
block: { try await _runJobAttempt(jobRecord) }
)
}
func didFinishJob(_ jobRecordId: JobRecord.RowId, result: JobResult) async {
switch result.ranSuccessfullyOrError {
case .success:
// When this job finishes, the send is enqueued.
// Send future resolution is handled within the job.
jobFutures?.enqueuedFuture.resolve(())
case .failure(let error):
jobFutures?.enqueuedFuture.reject(error)
}
}
private func _runJobAttempt(_ jobRecord: TSAttachmentMultisendJobRecord) async throws {
try await TSAttachmentMultisendUploader.uploadAttachments(
attachmentIdMap: jobRecord.attachmentIdMap,
sendMessages: { uploadedMessages, tx in
let preparedStoryMessages: [PreparedOutgoingMessage] = jobRecord.storyMessagesToSend?.map {
return PreparedOutgoingMessage.preprepared(
outgoingStoryMessage: $0
)
} ?? []
for preparedMessage in uploadedMessages + preparedStoryMessages {
let sendPromise = SSKEnvironment.shared.messageSenderJobQueueRef.add(
.promise,
message: preparedMessage,
transaction: tx
)
self.jobFutures?.sentFuture.resolve(on: SyncScheduler(), with: sendPromise)
}
jobRecord.anyRemove(transaction: tx)
}
)
}
}
// MARK: -
public enum TSAttachmentMultisendUploader: Dependencies {
public static func uploadAttachments<T>(
attachmentIdMap: [String: [String]],
sendMessages: @escaping (_ messages: [PreparedOutgoingMessage], _ tx: SDSAnyWriteTransaction) -> T
) async throws -> T {
let observer = NotificationCenter.default.addObserver(
forName: Upload.Constants.resourceUploadProgressNotification,
object: nil,
queue: nil
) { notification in
guard let notificationResourceId = notification.userInfo?[Upload.Constants.uploadResourceIDKey] as? TSResourceId else {
owsFailDebug("Missing notificationAttachmentId.")
return
}
guard case let .legacy(notificationAttachmentId) = notificationResourceId else {
// Ignore v2 attachments in a multisend context.
return
}
guard let progress = notification.userInfo?[Upload.Constants.uploadProgressKey] as? NSNumber else {
owsFailDebug("Missing progress.")
return
}
guard let correspondingAttachments = attachmentIdMap[notificationAttachmentId] else {
return
}
// Forward upload progress notifications to the corresponding attachments.
for correspondingId in correspondingAttachments {
guard correspondingId != notificationAttachmentId else {
owsFailDebug("Unexpected attachment id.")
continue
}
NotificationCenter.default.post(
name: Upload.Constants.resourceUploadProgressNotification,
object: nil,
userInfo: [
Upload.Constants.uploadResourceIDKey: TSResourceId.legacy(uniqueId: correspondingId),
Upload.Constants.uploadProgressKey: progress
]
)
}
}
defer { NotificationCenter.default.removeObserver(observer) }
let uploadOperations: [AsyncBlockOperation] = databaseStorage.read { tx in
return attachmentIdMap.map { (attachmentId, correspondingAttachmentIds) in
let messageIds: [String] = correspondingAttachmentIds.compactMap { correspondingId in
let attachment = TSAttachmentStream.anyFetchAttachmentStream(uniqueId: correspondingId, transaction: tx)
guard let attachment else {
Logger.warn("correspondingAttachment is missing. User has since deleted?")
return nil
}
guard let messageId = attachment.albumMessageId else {
return nil
}
return messageId
}
return AsyncBlockOperation {
try await DependenciesBridge.shared.tsResourceUploadManager.uploadAttachment(
attachmentId: .legacy(uniqueId: attachmentId),
legacyMessageOwnerIds: messageIds
)
}
}
}
Logger.info("Starting \(uploadOperations.count) uploads")
try? await withCheckedThrowingContinuation { continuation in
let waitingOperation = AwaitableAsyncBlockOperation(completionContinuation: continuation, asyncBlock: {})
uploadOperations.forEach { waitingOperation.addDependency($0) }
Upload.uploadQueue.addOperations(uploadOperations, waitUntilFinished: false)
Upload.uploadQueue.addOperation(waitingOperation)
}
if let error = (uploadOperations.compactMap { $0.failingError }).first { throw error }
return await databaseStorage.awaitableWrite { transaction in
var messageIdsToSend: Set<String> = Set()
// The attachments we've uploaded don't appear in any thread. Once they're
// uploaded, update the potentially many corresponding attachments (one per
// thread that the attachment was uploaded to) with the details of that
// upload.
let uploadedAttachments = attachmentIdMap.keys.compactMap { attachmentId in
// Probably should foreach to check for missing and warn.
TSAttachmentStream.anyFetchAttachmentStream(
uniqueId: attachmentId,
transaction: transaction
)
}
for uploadedAttachment in uploadedAttachments {
guard let correspondingAttachments = attachmentIdMap[uploadedAttachment.uniqueId] else {
owsFailDebug("correspondingAttachments was unexpectedly nil")
continue
}
let serverId = uploadedAttachment.serverId
let cdnKey = uploadedAttachment.cdnKey
let cdnNumber = uploadedAttachment.cdnNumber
let uploadTimestamp = uploadedAttachment.uploadTimestamp
guard
let encryptionKey = uploadedAttachment.encryptionKey,
let digest = uploadedAttachment.digest,
cdnKey.isEmpty.negated
else {
owsFailDebug("uploaded attachment was incomplete")
continue
}
if uploadTimestamp < 1 {
owsFailDebug("Missing uploadTimestamp.")
}
for correspondingId in correspondingAttachments {
guard let correspondingAttachment = TSAttachmentStream.anyFetchAttachmentStream(
uniqueId: correspondingId,
transaction: transaction
) else {
Logger.warn("correspondingAttachment is missing. User has since deleted?")
continue
}
correspondingAttachment.updateAsUploaded(
withEncryptionKey: encryptionKey,
digest: digest,
serverId: serverId,
cdnKey: cdnKey,
cdnNumber: cdnNumber,
uploadTimestamp: uploadTimestamp,
transaction: transaction
)
uploadedAttachment.blurHash.map { blurHash in
correspondingAttachment.update(withBlurHash: blurHash, transaction: transaction)
}
guard let albumMessageId = correspondingAttachment.albumMessageId else {
continue
}
messageIdsToSend.insert(albumMessageId)
}
}
let messagesToSend: [PreparedOutgoingMessage] = messageIdsToSend.compactMap { messageId in
guard let message = TSOutgoingMessage.anyFetchOutgoingMessage(
uniqueId: messageId,
transaction: transaction
) else {
owsFailDebug("outgoingMessage was unexpectedly nil")
return nil
}
return PreparedOutgoingMessage.preprepared(
forMultisendOf: message,
messageRowId: message.sqliteRowId!
)
}
// The attachment we uploaded should not be associated with any actual
// messages/threads, and is effectively orphaned.
owsAssertDebug(uploadedAttachments.allSatisfy { $0.albumMessageId == nil })
#if DEBUG
for uploadedAttachment in uploadedAttachments {
guard let uploadedAttachmentInDb = TSAttachmentStream.anyFetchAttachmentStream(
uniqueId: uploadedAttachment.uniqueId,
transaction: transaction
) else {
owsFailDebug("Unexpectedly missing uploaded attachment from DB")
continue
}
owsAssertDebug(uploadedAttachmentInDb.albumMessageId == nil)
}
#endif
// TODO: should we delete the orphaned attachments from the DB and disk here, now that we're done with them?
return sendMessages(messagesToSend, transaction)
}
}
}