Add ability to asynchronously hold onto a resource to prevent removal while clearing cache

submodules/MediaPlayer/Sources/FFMpegMediaFrameSourceContext.swift

@@ -287,6 +287,7 @@ final class FFMpegMediaFrameSourceContext: NSObject {
     
     fileprivate var requestedDataOffset: Int?
     fileprivate let fetchedDataDisposable = MetaDisposable()
+    fileprivate let keepDataDisposable = MetaDisposable()
     fileprivate let fetchedFullDataDisposable = MetaDisposable()
     fileprivate var requestedCompleteFetch = false
     
@@ -294,6 +295,7 @@ final class FFMpegMediaFrameSourceContext: NSObject {
         didSet {
             self.fetchedDataDisposable.dispose()
             self.fetchedFullDataDisposable.dispose()
+            self.keepDataDisposable.dispose()
         }
     }
     
@@ -316,6 +318,7 @@ final class FFMpegMediaFrameSourceContext: NSObject {
         
         self.fetchedDataDisposable.dispose()
         self.fetchedFullDataDisposable.dispose()
+        self.keepDataDisposable.dispose()
     }
     
     func initializeState(postbox: Postbox, resourceReference: MediaResourceReference, tempFilePath: String?, streamable: Bool, video: Bool, preferSoftwareDecoding: Bool, fetchAutomatically: Bool, maximumFetchSize: Int?) {
@@ -341,6 +344,10 @@ final class FFMpegMediaFrameSourceContext: NSObject {
             }
         }
         
+        if self.tempFilePath == nil {
+            self.keepDataDisposable.set(postbox.mediaBox.keepResource(id: resourceReference.resource.id).start())
+        }
+        
         if streamable {
             if self.tempFilePath == nil {
                 self.fetchedDataDisposable.set(fetchedMediaResource(mediaBox: postbox.mediaBox, reference: resourceReference, range: (0 ..< Int(Int32.max), .elevated), statsCategory: self.statsCategory ?? .generic, preferBackgroundReferenceRevalidation: streamable).start())

submodules/Postbox/Sources/MediaBox.swift

@@ -132,6 +132,14 @@ public enum ResourceDataRequestOption {
     case incremental(waitUntilFetchStatus: Bool)
 }
 
+private final class MediaBoxKeepResourceContext {
+    let subscribers = Bag<Void>()
+    
+    var isEmpty: Bool {
+        return self.subscribers.isEmpty
+    }
+}
+
 public final class MediaBox {
     public let basePath: String
     
@@ -145,6 +153,7 @@ public final class MediaBox {
     private var cachedRepresentationContexts: [CachedMediaResourceRepresentationKey: CachedMediaResourceRepresentationContext] = [:]
     
     private var fileContexts: [WrappedMediaResourceId: MediaBoxFileContext] = [:]
+    private var keepResourceContexts: [WrappedMediaResourceId: MediaBoxKeepResourceContext] = [:]
     
     private var wrappedFetchResource = Promise<(MediaResource, Signal<[(Range<Int>, MediaBoxFetchPriority)], NoError>, MediaResourceFetchParameters?) -> Signal<MediaResourceDataFetchResult, MediaResourceDataFetchError>>()
     public var preFetchedResourcePath: (MediaResource) -> String? = { _ in return nil }
@@ -204,6 +213,10 @@ public final class MediaBox {
         return ResourceStorePaths(partial: "\(self.basePath)/\(fileNameForId(id))_partial", complete: "\(self.basePath)/\(fileNameForId(id))")
     }
     
+    private func fileNamesForId(_ id: MediaResourceId) -> ResourceStorePaths {
+        return ResourceStorePaths(partial: "\(fileNameForId(id))_partial", complete: "\(fileNameForId(id))")
+    }
+    
     private func cachedRepresentationPathsForId(_ id: MediaResourceId, representation: CachedMediaResourceRepresentation) -> ResourceStorePaths {
         let cacheString: String
         switch representation.keepDuration {
@@ -697,6 +710,38 @@ public final class MediaBox {
         }
     }
     
+    public func keepResource(id: MediaResourceId) -> Signal<Never, NoError> {
+        return Signal { subscriber in
+            let disposable = MetaDisposable()
+            
+            let dataQueue = self.dataQueue
+            self.dataQueue.async {
+                let context: MediaBoxKeepResourceContext
+                if let current = self.keepResourceContexts[WrappedMediaResourceId(id)] {
+                    context = current
+                } else {
+                    context = MediaBoxKeepResourceContext()
+                    self.keepResourceContexts[WrappedMediaResourceId(id)] = context
+                }
+                let index = context.subscribers.add(Void())
+                
+                disposable.set(ActionDisposable { [weak self, weak context] in
+                    dataQueue.async {
+                        guard let strongSelf = self, let context = context, let currentContext = strongSelf.keepResourceContexts[WrappedMediaResourceId(id)], currentContext === context else {
+                            return
+                        }
+                        currentContext.subscribers.remove(index)
+                        if currentContext.isEmpty {
+                            strongSelf.keepResourceContexts.removeValue(forKey: WrappedMediaResourceId(id))
+                        }
+                    }
+                })
+            }
+            
+            return disposable
+        }
+    }
+    
     public func cancelInteractiveResourceFetch(_ resource: MediaResource) {
         self.dataQueue.async {
             if let (fileContext, releaseContext) = self.fileContext(for: resource) {
@@ -991,7 +1036,20 @@ public final class MediaBox {
     public func removeOtherCachedResources(paths: [String]) -> Signal<Void, NoError> {
         return Signal { subscriber in
             self.dataQueue.async {
-                for path in paths {
+                var keepPrefixes: [String] = []
+                for id in self.keepResourceContexts.keys {
+                    let resourcePaths = self.fileNamesForId(id.id)
+                    keepPrefixes.append(resourcePaths.partial)
+                    keepPrefixes.append(resourcePaths.complete)
+                }
+                
+                outer: for path in paths {
+                    for prefix in keepPrefixes {
+                        if path.starts(with: prefix) {
+                            continue outer
+                        }
+                    }
+                    
                     unlink(self.basePath + "/" + path)
                 }
                 subscriber.putCompletion()
@@ -1007,6 +1065,9 @@ public final class MediaBox {
                     if self.fileContexts[id] != nil {
                         continue
                     }
+                    if self.keepResourceContexts[id] != nil {
+                        continue
+                    }
                     let paths = self.storePathsForId(id.id)
                     unlink(paths.complete)
                     unlink(paths.partial)
@@ -1044,35 +1105,4 @@ public final class MediaBox {
             return EmptyDisposable
         }
     }
-    
-    public func clearFileContexts() -> Signal<Void, NoError> {
-        return Signal { subscriber in
-            self.dataQueue.async {
-                for (id, _) in self.fileContexts {
-                    let paths = self.storePathsForId(id.id)
-                    unlink(paths.complete)
-                    unlink(paths.partial)
-                    unlink(paths.partial + ".meta")
-                }
-                self.fileContexts.removeAll()
-                subscriber.putCompletion()
-            }
-            return EmptyDisposable
-        }
-    }
-    
-    public func fileConxtets() -> Signal<[(partial: String, complete: String)], NoError> {
-        return Signal { subscriber in
-            self.dataQueue.async {
-                var result: [(partial: String, complete: String)] = []
-                for (id, _) in self.fileContexts {
-                    let paths = self.storePathsForId(id.id)
-                    result.append((partial: paths.partial, complete: paths.complete))
-                }
-                subscriber.putNext(result)
-                subscriber.putCompletion()
-            }
-            return EmptyDisposable
-        }
-    }
 }