Fix notification service media download

Telegram/NotificationService/Sources/NotificationService.swift

@@ -460,6 +460,7 @@ private struct NotificationContent: CustomStringConvertible {
         string += " userInfo: \(String(describing: self.userInfo)),\n"
         string += " senderImage: \(self.senderImage != nil ? "non-empty" : "empty"),\n"
         string += " isLockedMessage: \(String(describing: self.isLockedMessage)),\n"
+        string += " attachments: \(self.attachments),\n"
         string += "}"
         return string
     }
@@ -1164,7 +1165,13 @@ private final class NotificationServiceHandler {
                                                 } else {
                                                     let intervals: Signal<[(Range<Int64>, MediaBoxFetchPriority)], NoError> = .single([(0 ..< Int64.max, MediaBoxFetchPriority.maximum)])
                                                     fetchMediaSignal = Signal { subscriber in
-                                                        let collectedData = Atomic<Data>(value: Data())
+                                                        final class DataValue {
+                                                            var data = Data()
+                                                            var totalSize: Int64?
+                                                        }
+                                                        
+                                                        let collectedData = Atomic<DataValue>(value: DataValue())
+                                                        
                                                         return standaloneMultipartFetch(
                                                             postbox: stateManager.postbox,
                                                             network: stateManager.network,
@@ -1191,11 +1198,33 @@ private final class NotificationServiceHandler {
                                                         ).start(next: { result in
                                                             switch result {
                                                             case let .dataPart(_, data, _, _):
+                                                                var isCompleted = false
                                                                 let _ = collectedData.modify { current in
-                                                                    var current = current
+                                                                    let current = current
-                                                                    current.append(data)
+                                                                    current.data.append(data)
+                                                                    if let totalSize = current.totalSize, Int64(current.data.count) >= totalSize {
+                                                                        isCompleted = true
+                                                                    }
                                                                     return current
                                                                 }
+                                                                if isCompleted {
+                                                                    subscriber.putNext(collectedData.with({ $0.data }))
+                                                                    subscriber.putCompletion()
+                                                                }
+                                                            case let .resourceSizeUpdated(size):
+                                                                var isCompleted = false
+                                                                let _ = collectedData.modify { current in
+                                                                    let current = current
+                                                                    current.totalSize = size
+                                                                    if Int64(current.data.count) >= size {
+                                                                        isCompleted = true
+                                                                    }
+                                                                    return current
+                                                                }
+                                                                if isCompleted {
+                                                                    subscriber.putNext(collectedData.with({ $0.data }))
+                                                                    subscriber.putCompletion()
+                                                                }
                                                             default:
                                                                 break
                                                             }
@@ -1203,7 +1232,7 @@ private final class NotificationServiceHandler {
                                                             subscriber.putNext(nil)
                                                             subscriber.putCompletion()
                                                         }, completed: {
-                                                            subscriber.putNext(collectedData.with({ $0 }))
+                                                            subscriber.putNext(collectedData.with({ $0.data }))
                                                             subscriber.putCompletion()
                                                         })
                                                     }
@@ -1304,9 +1333,14 @@ private final class NotificationServiceHandler {
 
                                                 Logger.shared.log("NotificationService \(episode)", "Unread count: \(value.0), isCurrentAccount: \(isCurrentAccount)")
                                                 
+                                                Logger.shared.log("NotificationService \(episode)", "mediaAttachment: \(String(describing: mediaAttachment)), mediaData: \(String(describing: mediaData?.count))")
+
                                                 if let image = mediaAttachment as? TelegramMediaImage, let resource = largestImageRepresentation(image.representations)?.resource {
                                                     if let mediaData = mediaData {
                                                         stateManager.postbox.mediaBox.storeResourceData(resource.id, data: mediaData, synchronous: true)
+                                                        if let messageId {
+                                                            let _ = addSynchronizeAutosaveItemOperation(postbox: stateManager.postbox, messageId: messageId, mediaId: image.imageId).start()
+                                                        }
                                                     }
                                                     if let storedPath = stateManager.postbox.mediaBox.completedResourcePath(resource, pathExtension: "jpg") {
                                                         if let attachment = try? UNNotificationAttachment(identifier: "image", url: URL(fileURLWithPath: storedPath), options: nil) {

submodules/TelegramCore/Sources/Account/AccountManager.swift

@@ -193,6 +193,7 @@ private var declaredEncodables: Void = {
     declareEncodable(TelegramPeerUsername.self, f: { TelegramPeerUsername(decoder: $0) })
     declareEncodable(MediaSpoilerMessageAttribute.self, f: { MediaSpoilerMessageAttribute(decoder: $0) })
     declareEncodable(TranslationMessageAttribute.self, f: { TranslationMessageAttribute(decoder: $0) })
+    declareEncodable(SynchronizeAutosaveItemOperation.self, f: { SynchronizeAutosaveItemOperation(decoder: $0) })
     return
 }()
 

submodules/TelegramCore/Sources/State/SynchronizeAutosaveItemOperations.swift

@@ -0,0 +1,78 @@
+import Foundation
+import Postbox
+import SwiftSignalKit
+
+public final class SynchronizeAutosaveItemOperation: PostboxCoding {
+    public struct Content: Codable {
+        public var messageId: MessageId
+        public var mediaId: MediaId
+        
+        public init(messageId: MessageId, mediaId: MediaId) {
+            self.messageId = messageId
+            self.mediaId = mediaId
+        }
+    }
+    
+    public let messageId: MessageId
+    public let mediaId: MediaId
+    
+    public init(messageId: MessageId, mediaId: MediaId) {
+        self.messageId = messageId
+        self.mediaId = mediaId
+    }
+    
+    public init(decoder: PostboxDecoder) {
+        if let content = decoder.decode(Content.self, forKey: "c") {
+            self.messageId = content.messageId
+            self.mediaId = content.mediaId
+        } else {
+            self.messageId = MessageId(peerId: PeerId(0), namespace: 0, id: 0)
+            self.mediaId = MediaId(namespace: 0, id: 0)
+        }
+    }
+    
+    public func encode(_ encoder: PostboxEncoder) {
+        encoder.encode(Content(messageId: self.messageId, mediaId: self.mediaId), forKey: "c")
+    }
+}
+
+public func addSynchronizeAutosaveItemOperation(transaction: Transaction, messageId: MessageId, mediaId: MediaId) {
+    let tag: PeerOperationLogTag = OperationLogTags.SynchronizeAutosaveItems
+    let peerId = PeerId(0)
+    
+    transaction.operationLogAddEntry(peerId: peerId, tag: tag, tagLocalIndex: .automatic, tagMergedIndex: .automatic, contents: SynchronizeAutosaveItemOperation(messageId: messageId, mediaId: mediaId))
+}
+
+public func addSynchronizeAutosaveItemOperation(postbox: Postbox, messageId: MessageId, mediaId: MediaId) -> Signal<Never, NoError> {
+    return postbox.transaction { transaction -> Void in
+        addSynchronizeAutosaveItemOperation(transaction: transaction, messageId: messageId, mediaId: mediaId)
+    }
+    |> ignoreValues
+}
+
+public func _internal_getSynchronizeAutosaveItemOperations(transaction: Transaction) -> [(index: Int32, message: Message, mediaId: MediaId)] {
+    let peerId = PeerId(0)
+    var result: [(index: Int32, message: Message, mediaId: MediaId)] = []
+    var removeIndices: [Int32] = []
+    transaction.operationLogEnumerateEntries(peerId: peerId, tag: OperationLogTags.SynchronizeAutosaveItems, { entry in
+        if let operation = entry.contents as? SynchronizeAutosaveItemOperation {
+            if let message = transaction.getMessage(operation.messageId) {
+                result.append((index: entry.tagLocalIndex, message: message, mediaId: operation.mediaId))
+            } else {
+                removeIndices.append(entry.tagLocalIndex)
+            }
+        }
+        return true
+    })
+    for index in removeIndices {
+        let _ = transaction.operationLogRemoveEntry(peerId: PeerId(0), tag: OperationLogTags.SynchronizeAutosaveItems, tagLocalIndex: index)
+    }
+    
+    return result
+}
+
+public func _internal_removeSyncrhonizeAutosaveItemOperations(transaction: Transaction, indices: [Int32]) {
+    for index in indices {
+        let _ = transaction.operationLogRemoveEntry(peerId: PeerId(0), tag: OperationLogTags.SynchronizeAutosaveItems, tagLocalIndex: index)
+    }
+}

submodules/TelegramCore/Sources/SyncCore/SyncCore_Namespaces.swift

@@ -180,6 +180,7 @@ public struct OperationLogTags {
     public static let SynchronizeChatListFilters = PeerOperationLogTag(value: 20)
     public static let SynchronizeMarkAllUnseenReactions = PeerOperationLogTag(value: 21)
     public static let SynchronizeInstalledEmoji = PeerOperationLogTag(value: 22)
+    public static let SynchronizeAutosaveItems = PeerOperationLogTag(value: 23)
 }
 
 public struct LegacyPeerSummaryCounterTags: OptionSet, Sequence, Hashable {

submodules/TelegramCore/Sources/TelegramEngine/Messages/TelegramEngineMessages.swift

@@ -507,5 +507,17 @@ public extension TelegramEngine {
             return managedSynchronizeMessageHistoryTagSummaries(postbox: self.account.postbox, network: self.account.network, stateManager: self.account.stateManager, peerId: peerId, threadId: threadId)
             |> ignoreValues
         }
+        
+        public func getSynchronizeAutosaveItemOperations() -> Signal<[(index: Int32, message: Message, mediaId: MediaId)], NoError> {
+            return self.account.postbox.transaction { transaction -> [(index: Int32, message: Message, mediaId: MediaId)] in
+                return _internal_getSynchronizeAutosaveItemOperations(transaction: transaction)
+            }
+        }
+
+        func removeSyncrhonizeAutosaveItemOperations(indices: [Int32]) {
+            let _ = (self.account.postbox.transaction { transaction -> Void in
+                _internal_removeSyncrhonizeAutosaveItemOperations(transaction: transaction, indices: indices)
+            }).start()
+        }
     }
 }

submodules/TelegramUI/Sources/AppDelegate.swift

@@ -1310,6 +1310,8 @@ private func extractAccountManagerState(records: AccountRecordsView<TelegramAcco
             self.isInForegroundPromise.set(true)
             self.isActiveValue = true
             self.isActivePromise.set(true)
+            
+            self.runForegroundTasks()
         }
         
         if UIApplication.shared.isStatusBarHidden {
@@ -1519,6 +1521,22 @@ private func extractAccountManagerState(records: AccountRecordsView<TelegramAcco
                 }
             }
         }
+        
+        self.runForegroundTasks()
+    }
+    
+    func runForegroundTasks() {
+        let _ = (self.sharedContextPromise.get()
+        |> take(1)
+        |> deliverOnMainQueue).start(next: { sharedApplicationContext in
+            let _ = (sharedApplicationContext.sharedContext.activeAccountContexts
+             |> take(1)
+             |> deliverOnMainQueue).start(next: { activeAccounts in
+                for (_, context, _) in activeAccounts.accounts {
+                    (context.downloadedMediaStoreManager as? DownloadedMediaStoreManagerImpl)?.runTasks()
+                }
+            })
+        })
     }
 
     func applicationDidBecomeActive(_ application: UIApplication) {

submodules/TelegramUI/Sources/StoreDownloadedMedia.swift

@@ -266,10 +266,12 @@ private final class DownloadedMediaStoreManagerPrivateImpl {
 
 final class DownloadedMediaStoreManagerImpl: DownloadedMediaStoreManager {
     private let queue = Queue()
+    private let postbox: Postbox
     private let impl: QueueLocalObject<DownloadedMediaStoreManagerPrivateImpl>
     
     init(postbox: Postbox, accountManager: AccountManager<TelegramAccountManagerTypes>) {
         let queue = self.queue
+        self.postbox = postbox
         self.impl = QueueLocalObject(queue: queue, generate: {
             return DownloadedMediaStoreManagerPrivateImpl(queue: queue, postbox: postbox, accountManager: accountManager)
         })
@@ -280,4 +282,27 @@ final class DownloadedMediaStoreManagerImpl: DownloadedMediaStoreManager {
             impl.store(media, timestamp: timestamp, peerId: peerId)
         }
     }
+    
+    func runTasks() {
+        let _ = (self.postbox.transaction({ transaction -> [(index: Int32, message: Message, mediaId: MediaId)] in
+            return _internal_getSynchronizeAutosaveItemOperations(transaction: transaction)
+        })
+        |> deliverOnMainQueue).start(next: { [weak self] items in
+            guard let self else {
+                return
+            }
+            for item in items {
+                for media in item.message.media {
+                    if let id = media.id, id == item.mediaId {
+                        self.store(.standalone(media: media), timestamp: item.message.timestamp, peerId: item.message.id.peerId)
+                        break
+                    }
+                }
+            }
+            
+            let _ = self.postbox.transaction({ transaction -> Void in
+                return _internal_removeSyncrhonizeAutosaveItemOperations(transaction: transaction, indices: items.map(\.index))
+            }).start()
+        })
+    }
 }