[Core ML] Improve asset management (#13560)

## Fix asset management issues in CoreML delegate

- Do not use the temporary directory for asset storage.
- Use the staging directory for storing temporary assets, and ensure it
is cleaned up after use.
- Perform aggressive cleanup of the temporary directory to avoid
leftovers.
- Guarantee proper cleanup of staging directories after operations to
prevent stale files.

---------

Co-authored-by: Scott Roy <[email protected]>

Author: cymbalrush

backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.h

@@ -99,6 +99,17 @@ NS_ASSUME_NONNULL_BEGIN
 - (NSUInteger)compact:(NSUInteger)sizeInBytes error:(NSError* __autoreleasing*)error;
 
 
+/// Executes a block with a unique temporary directory.
+///
+/// A new temporary subdirectory URL is created inside the receiver’s designated
+/// base directory. The directory is passed to the block, which can use it to
+/// perform temporary file operations. After the block finishes executing,
+/// the directory and its contents are removed.
+///
+/// @param block A block to execute. The block receives a unique URL.
+- (void)withTemporaryDirectory:(void (^)(NSURL* directoryURL))block;
+
+
 /// Purges the assets storage. The assets are moved to the trash directory and are asynchronously
 /// deleted.
 ///
@@ -117,6 +128,12 @@ NS_ASSUME_NONNULL_BEGIN
 /// contents are deleted asynchronously.
 @property (copy, readonly, nonatomic) NSURL* trashDirectoryURL;
 
+
+/// The staging directory URL, used to hold assets that are being prepared or processed
+/// before they are moved into their final location. The contents of this directory
+/// are temporary and may be cleared when no longer needed.
+@property (copy, readonly, nonatomic) NSURL* stagingDirectoryURL;
+
 /// The file manager.
 @property (strong, readonly, nonatomic) NSFileManager* fileManager;
 

backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.mm

@@ -254,6 +254,29 @@ BOOL is_asset_alive(NSMapTable<NSString *, ETCoreMLAsset *> *assets_in_use_map,
 
     return assets;
 }
+
+NSURL * _Nullable move_to_directory(NSURL *url,
+                                    NSURL *directoryURL,
+                                    NSFileManager *fileManager,
+                                    NSError * __autoreleasing *error) {
+    if (!url) {
+        ETCoreMLLogErrorAndSetNSError(error, ETCoreMLErrorInternalError, "Move operation failed: source URL is nil.");
+        return nil;
+    }
+
+    if (!directoryURL) {
+        ETCoreMLLogErrorAndSetNSError(error, ETCoreMLErrorInternalError, "Move operation failed: destination URL is nil.");
+        return nil;
+    }
+
+    NSURL *dstURL = [directoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString];
+    if (![fileManager moveItemAtURL:url toURL:dstURL error:error]) {
+        return nil;
+    }
+
+    return dstURL;
+}
+
 } //namespace
 
 @interface ETCoreMLAssetManager () <NSFileManagerDelegate> {
@@ -299,12 +322,17 @@ - (nullable instancetype)initWithDatabase:(const std::shared_ptr<Database>&)data
     if (!managedAssetsDirectoryURL) {
         return nil;
     }
-    
+
     NSURL *managedTrashDirectoryURL = ::create_directory_if_needed(trashDirectoryURL, @"models", fileManager, error);
     if (!managedTrashDirectoryURL) {
         return nil;
     }
-    
+
+    NSURL *managedStagingDirectoryURL = ::create_directory_if_needed(assetsDirectoryURL, @"staging", fileManager, error);
+    if (!managedStagingDirectoryURL) {
+        return nil;
+    }
+
     // If directory is empty then purge the stores
     if (::is_directory_empty(managedAssetsDirectoryURL, fileManager, nil)) {
         assetsMetaStore.impl()->purge(ec);
@@ -315,6 +343,7 @@ - (nullable instancetype)initWithDatabase:(const std::shared_ptr<Database>&)data
         _assetsStore = std::move(assetsStore);
         _assetsMetaStore = std::move(assetsMetaStore);
         _assetsDirectoryURL = managedAssetsDirectoryURL;
+        _stagingDirectoryURL = managedStagingDirectoryURL;
         _trashDirectoryURL = managedTrashDirectoryURL;
         _estimatedSizeInBytes = sizeInBytes.value();
         _maxAssetsSizeInBytes = maxAssetsSizeInBytes;
@@ -346,15 +375,15 @@ - (nullable instancetype)initWithDatabaseURL:(NSURL *)databaseURL
                             error:error];
 }
 
-- (nullable NSURL *)moveURL:(NSURL *)url
-     toUniqueURLInDirectory:(NSURL *)directoryURL
-                      error:(NSError * __autoreleasing *)error {
-    NSURL *dstURL = [directoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString];
-    if (![self.fileManager moveItemAtURL:url toURL:dstURL error:error]) {
-        return nil;
+- (void)withTemporaryDirectory:(void (^)(NSURL *directoryURL))block {
+    NSURL *dstURL = [self.stagingDirectoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString];
+    block(dstURL);
+    if (![self.fileManager fileExistsAtPath:dstURL.path]) {
+        return;
     }
-    
-    return dstURL;
+
+    move_to_directory(dstURL, self.trashDirectoryURL, self.fileManager, nil);
+    [self cleanupTrashDirectory];
 }
 
 - (void)cleanupAssetIfNeeded:(ETCoreMLAsset *)asset {
@@ -407,9 +436,8 @@ - (nullable ETCoreMLAsset *)_storeAssetAtURL:(NSURL *)srcURL
             return false;
         }
 
-        // If an asset exists move it
-        [self moveURL:dstURL toUniqueURLInDirectory:self.trashDirectoryURL error:nil];
-        
+        // If a file already exists at `dstURL`, move it to the trash for removal.
+        move_to_directory(dstURL, self.trashDirectoryURL, self.fileManager, nil);
         // Move the asset to assets directory.
         if (![self.fileManager moveItemAtURL:srcURL toURL:dstURL error:error]) {
             return false;
@@ -433,16 +461,25 @@ - (nullable ETCoreMLAsset *)_storeAssetAtURL:(NSURL *)srcURL
 }
 
 - (void)triggerCompaction {
-    if (self.estimatedSizeInBytes < self.maxAssetsSizeInBytes) {
-        return;
+    if (self.estimatedSizeInBytes >= self.maxAssetsSizeInBytes) {
+        __weak __typeof(self) weakSelf = self;
+        dispatch_async(self.syncQueue, ^{
+            NSError *localError = nil;
+            if (![weakSelf _compact:self.maxAssetsSizeInBytes error:&localError]) {
+                ETCoreMLLogError(localError, "Failed to compact asset store.");
+            }
+        });
     }
-    
+
+    // Always clean the trash directory to ensure a minimal footprint.
+    // The `trashQueue` is serialized, so only one cleanup will run at a time.
+    [self cleanupTrashDirectory];
+}
+
+- (void)cleanupTrashDirectory {
     __weak __typeof(self) weakSelf = self;
-    dispatch_async(self.syncQueue, ^{
-        NSError *localError = nil;
-        if (![weakSelf _compact:self.maxAssetsSizeInBytes error:&localError]) {
-            ETCoreMLLogError(localError, "Failed to compact asset store.");
-        }
+    dispatch_async(self.trashQueue, ^{
+        [weakSelf removeFilesInTrashDirectory];
     });
 }
 
@@ -548,7 +585,7 @@ - (BOOL)_removeAssetWithIdentifier:(NSString *)identifier
 
         NSURL *assetURL = ::get_asset_url(assetValue);
         if ([self.fileManager fileExistsAtPath:assetURL.path] &&
-            ![self moveURL:assetURL toUniqueURLInDirectory:self.trashDirectoryURL error:error]) {
+            !move_to_directory(assetURL, self.trashDirectoryURL, self.fileManager, error)) {
             return false;
         }
 
@@ -649,13 +686,7 @@ - (NSUInteger)_compact:(NSUInteger)sizeInBytes error:(NSError * __autoreleasing
                              identifier);
         }
     }
-    
-    // Trigger cleanup.
-    __weak __typeof(self) weakSelf = self;
-    dispatch_async(self.trashQueue, ^{
-        [weakSelf removeFilesInTrashDirectory];
-    });
-    
+
     return _estimatedSizeInBytes;
 }
 
@@ -664,7 +695,10 @@ - (NSUInteger)compact:(NSUInteger)sizeInBytes error:(NSError * __autoreleasing *
     dispatch_sync(self.syncQueue, ^{
         result = [self _compact:sizeInBytes error:error];
     });
-    
+
+    // Always clean the trash directory to ensure a minimal footprint.
+    // The `trashQueue` is serialized, so only one cleanup will run at a time.
+    [self cleanupTrashDirectory];
     return result;
 }
 
@@ -708,7 +742,7 @@ - (BOOL)_purge:(NSError * __autoreleasing *)error {
         }
 
         // Move the the whole assets directory to the temp directory.
-        if (![self moveURL:self.assetsDirectoryURL toUniqueURLInDirectory:self.trashDirectoryURL error:error]) {
+        if (!move_to_directory(self.assetsDirectoryURL, self.trashDirectoryURL, self.fileManager, error)) {
             return false;
         }
 
@@ -724,13 +758,7 @@ - (BOOL)_purge:(NSError * __autoreleasing *)error {
 
     ::set_error_from_error_code(ec, error);
     // Trigger cleanup
-    if (status) {
-        __weak __typeof(self) weakSelf = self;
-        dispatch_async(self.trashQueue, ^{
-            [weakSelf removeFilesInTrashDirectory];
-        });
-    }
-    
+    [self cleanupTrashDirectory];
     return static_cast<BOOL>(status);
 }
 

backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm

@@ -62,21 +62,12 @@ + (nullable ETCoreMLModel *)loadModelWithContentsOfURL:(NSURL *)compiledModelURL
     if (model) {
         return model;
     }
-    
-    if (localError) {
-        ETCoreMLLogError(localError,
-                         "Failed to load model from compiled asset with identifier = %@",
-                         identifier);
-    }
-    
-    // If store failed then we will load the model from compiledURL.
-    auto backingAsset = Asset::make(compiledModelURL, identifier, assetManager.fileManager, error);
-    if (!backingAsset) {
-        return nil;
+
+    if (error) {
+        *error = localError;
     }
-    
-    asset = [[ETCoreMLAsset alloc] initWithBackingAsset:backingAsset.value()];
-    return ::get_model_from_asset(asset, configuration, metadata, error);
+
+    return nil;
 }
 
 @end