xref: /aosp_15_r20/external/executorch/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.mm (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1//
2// ETCoreMLAssetManager.mm
3//
4// Copyright © 2024 Apple Inc. All rights reserved.
5//
6// Please refer to the license found in the LICENSE file in the root directory of the source tree.
7
8#import "ETCoreMLAssetManager.h"
9#import <ETCoreMLAsset.h>
10#import <ETCoreMLLogging.h>
11#import <database.hpp>
12#import <iostream>
13#import <json_key_value_store.hpp>
14#import <serde_json.h>
15#import <sstream>
16
17namespace  {
18
19using namespace executorchcoreml;
20using namespace executorchcoreml::sqlite;
21
22constexpr size_t kBusyTimeIntervalInMS = 100;
23
24constexpr std::string_view kModelAssetsStoreName = "MODEL_ASSETS_STORE";
25constexpr std::string_view kModelAssetsMetaStoreName = "MODEL_ASSETS_STORE_META";
26
27class ModelAssetsStore {
28public:
29    using StoreType = JSONKeyValueStore<std::string, Asset>;
30
31    ModelAssetsStore(std::unique_ptr<StoreType> impl) noexcept
32    :impl_(std::move(impl))
33    {}
34
35    ModelAssetsStore() noexcept
36    :impl_(nullptr)
37    {}
38
39    ModelAssetsStore(const ModelAssetsStore &) = delete;
40    ModelAssetsStore &operator=(const ModelAssetsStore &) = delete;
41
42    ModelAssetsStore& operator=(ModelAssetsStore&& rhs) noexcept {
43        rhs.impl_.swap(impl_);
44        return *this;
45    }
46
47    ModelAssetsStore(ModelAssetsStore&& rhs) noexcept
48    :impl_(std::move(rhs.impl_))
49    {}
50
51    inline StoreType *impl() {
52        return impl_.get();
53    }
54
55private:
56    std::unique_ptr<StoreType> impl_;
57};
58
59class ModelAssetsMetaStore {
60public:
61    using StoreType = KeyValueStore<std::string, size_t>;
62
63    ModelAssetsMetaStore() noexcept
64    :impl_(nullptr)
65    {}
66
67    ModelAssetsMetaStore(std::unique_ptr<StoreType> impl) noexcept
68    :impl_(std::move(impl))
69    {}
70
71    ModelAssetsMetaStore(const ModelAssetsMetaStore &) = delete;
72    ModelAssetsMetaStore &operator=(const ModelAssetsMetaStore &) = delete;
73
74    ModelAssetsMetaStore& operator=(ModelAssetsMetaStore&& rhs) noexcept {
75        rhs.impl_.swap(impl_);
76        return *this;
77    }
78
79    ModelAssetsMetaStore(ModelAssetsMetaStore&& rhs) noexcept
80    :impl_(std::move(rhs.impl_))
81    {}
82
83    inline StoreType *impl() {
84        return impl_.get();
85    }
86
87private:
88    std::unique_ptr<StoreType> impl_;
89};
90
91void set_error_from_error_code(const std::error_code& cppError, NSError * __autoreleasing *error) {
92    if (!error || !cppError) {
93        return;
94    }
95
96    NSString *message = @(cppError.message().c_str());
97    NSString *domain =  @(cppError.category().name());
98    NSInteger code = cppError.value();
99    NSError *localError = [NSError errorWithDomain:domain code:code userInfo:@{NSLocalizedDescriptionKey : message}];
100    *error = localError;
101}
102
103std::shared_ptr<Database> make_database(NSURL *database_url,
104                                        NSTimeInterval busy_time_interval,
105                                        NSError * __autoreleasing *error) {
106    Database::OpenOptions options;
107    options.set_read_write_option(true);
108    options.set_create_option(true);
109
110    std::error_code ec;
111    auto database = Database::make(database_url.path.UTF8String,
112                                   options,
113                                   Database::SynchronousMode::Normal,
114                                   busy_time_interval,
115                                   ec);
116    if (!database) {
117        ::set_error_from_error_code(ec, error);
118        return nullptr;
119    }
120
121    return database;
122}
123
124ModelAssetsStore make_assets_store(const std::shared_ptr<Database>& database,
125                                   NSError * __autoreleasing *error) {
126    std::error_code ec;
127    auto store = ModelAssetsStore::StoreType::make(std::move(database), std::string(kModelAssetsStoreName), ec);
128    if (!store) {
129        ::set_error_from_error_code(ec, error);
130        return ModelAssetsStore(nullptr);
131    }
132
133    return ModelAssetsStore(std::move(store));
134}
135
136ModelAssetsMetaStore make_assets_meta_store(const std::shared_ptr<Database>& database,
137                                            NSError * __autoreleasing *error) {
138    std::error_code ec;
139    auto store = ModelAssetsMetaStore::StoreType::make(std::move(database), std::string(kModelAssetsMetaStoreName), ec);
140    if (!store) {
141        ::set_error_from_error_code(ec, error);
142        return ModelAssetsMetaStore(nullptr);
143    }
144
145    return ModelAssetsMetaStore(std::move(store));
146}
147
148std::optional<size_t> get_total_assets_size(ModelAssetsMetaStore& store,
149                                            std::error_code& ec) {
150    std::string name = std::string(kModelAssetsStoreName);
151    auto total_size = store.impl()->get(name, ec);
152
153    if (!total_size && !store.impl()->put(name, size_t(0), ec)) {
154        return std::nullopt;
155    }
156
157    size_t result = total_size.has_value() ? total_size.value() : size_t(0);
158    return result;
159}
160
161bool set_total_assets_size(size_t total_size,
162                           ModelAssetsMetaStore& store,
163                           std::error_code& ec) {
164    if (!store.impl()->put(std::string(kModelAssetsStoreName), total_size, ec)) {
165        return false;
166    }
167
168    return true;
169}
170
171bool exclude_item_from_backup(NSURL *url, NSError * __autoreleasing *error) {
172    return [url setResourceValue:@(YES) forKey:NSURLIsExcludedFromBackupKey error:error];
173}
174
175NSURL * _Nullable create_directory_if_needed(NSURL *url,
176                                             NSString *name,
177                                             NSFileManager *fm,
178                                             NSError * __autoreleasing *error) {
179    NSURL *directory_url = [url URLByAppendingPathComponent:name];
180    if (![fm fileExistsAtPath:directory_url.path] &&
181        ![fm createDirectoryAtURL:directory_url withIntermediateDirectories:NO attributes:@{} error:error]) {
182        return nil;
183    }
184
185    ::exclude_item_from_backup(directory_url, nil);
186
187    return directory_url;
188}
189
190bool is_directory_empty(NSURL *url, NSFileManager *fm, NSError * __autoreleasing *error) {
191    BOOL is_directory = NO;
192    if (![fm fileExistsAtPath:url.path isDirectory:&is_directory] && !is_directory) {
193        return true;
194    }
195
196    __block NSError *local_error = nil;
197    BOOL (^errorHandler)(NSURL *url, NSError *error) = ^BOOL(NSURL *url, NSError *enumeration_error) {
198        local_error = enumeration_error;
199        return NO;
200    };
201
202    NSDirectoryEnumerator *enumerator = [fm enumeratorAtURL:url
203                                 includingPropertiesForKeys:@[]
204                                                    options:NSDirectoryEnumerationProducesRelativePathURLs
205                                               errorHandler:errorHandler];
206    if (local_error && error) {
207        *error = local_error;
208    }
209
210    return [enumerator nextObject] == nil;
211}
212
213NSURL * _Nullable get_asset_url(const Asset& asset) {
214    return [NSURL fileURLWithPath:@(asset.path.c_str())];
215}
216
217BOOL is_asset_alive(NSMapTable<NSString *, ETCoreMLAsset *> *assets_in_use_map, NSString *identifier) {
218    ETCoreMLAsset *asset = [assets_in_use_map objectForKey:identifier];
219    return asset && asset.isAlive;
220}
221
222std::vector<executorchcoreml::Asset>
223get_assets_to_remove(ModelAssetsStore& store,
224                     ssize_t bytes_to_remove,
225                     NSMapTable<NSString *, ETCoreMLAsset *> *assets_in_use_map,
226                     std::error_code &error) {
227    std::vector<Asset> assets;
228    store.impl()->get_keys_sorted_by_access_count([store = store.impl(),
229                                                   &bytes_to_remove,
230                                                   &assets,
231                                                   assets_in_use_map,
232                                                   &error](const std::string& key) {
233        if (bytes_to_remove <= 0) {
234            return false;
235        }
236
237        NSString *identifier = @(key.c_str());
238        // Asset is in use, we can't remove it
239        if (::is_asset_alive(assets_in_use_map, identifier)) {
240            return true;
241        }
242
243        auto asset = store->get(key, error);
244        if (asset) {
245            auto& asset_value = asset.value();
246            bytes_to_remove -= static_cast<ssize_t>(asset_value.total_size_in_bytes());
247            assets.emplace_back(std::move(asset_value));
248        }
249
250        return true;
251    }, SortOrder::Ascending, error);
252
253    return assets;
254}
255} //namespace
256
257@interface ETCoreMLAssetManager () <NSFileManagerDelegate> {
258    ModelAssetsStore _assetsStore;
259    ModelAssetsMetaStore _assetsMetaStore;
260}
261
262@property (assign, readwrite, atomic) NSInteger estimatedSizeInBytes;
263@property (copy, readonly, nonatomic) NSURL *assetsDirectoryURL;
264@property (strong, readonly, nonatomic) dispatch_queue_t syncQueue;
265@property (strong, readonly, nonatomic) dispatch_queue_t trashQueue;
266@property (strong, readonly, nonatomic) NSMapTable<NSString *, ETCoreMLAsset *> *assetsInUseMap;
267
268@end
269
270@implementation ETCoreMLAssetManager
271
272- (nullable instancetype)initWithDatabase:(const std::shared_ptr<Database>&)database
273                       assetsDirectoryURL:(NSURL *)assetsDirectoryURL
274                        trashDirectoryURL:(NSURL *)trashDirectoryURL
275                     maxAssetsSizeInBytes:(NSInteger)maxAssetsSizeInBytes
276                                    error:(NSError * __autoreleasing *)error {
277
278    auto assetsStore = ::make_assets_store(database, error);
279    if (assetsStore.impl() == nullptr) {
280        return nil;
281    }
282
283    auto assetsMetaStore = ::make_assets_meta_store(database, error);
284    if (assetsMetaStore.impl() == nullptr) {
285        return nil;
286    }
287
288    std::error_code ec;
289    auto sizeInBytes = ::get_total_assets_size(assetsMetaStore, ec);
290    if (!sizeInBytes) {
291        ::set_error_from_error_code(ec, error);
292        return nil;
293    }
294
295    NSFileManager *fileManager = [[NSFileManager alloc] init];
296    NSURL *managedAssetsDirectoryURL = ::create_directory_if_needed(assetsDirectoryURL, @"models", fileManager, error);
297    if (!managedAssetsDirectoryURL) {
298        return nil;
299    }
300
301    NSURL *managedTrashDirectoryURL = ::create_directory_if_needed(trashDirectoryURL, @"models", fileManager, error);
302    if (!managedTrashDirectoryURL) {
303        return nil;
304    }
305
306    // If directory is empty then purge the stores
307    if (::is_directory_empty(managedAssetsDirectoryURL, fileManager, nil)) {
308        assetsMetaStore.impl()->purge(ec);
309        assetsStore.impl()->purge(ec);
310    }
311
312    if (self = [super init]) {
313        _assetsStore = std::move(assetsStore);
314        _assetsMetaStore = std::move(assetsMetaStore);
315        _assetsDirectoryURL = managedAssetsDirectoryURL;
316        _trashDirectoryURL = managedTrashDirectoryURL;
317        _estimatedSizeInBytes = sizeInBytes.value();
318        _maxAssetsSizeInBytes = maxAssetsSizeInBytes;
319
320        _fileManager = fileManager;
321        _trashQueue = dispatch_queue_create("com.executorchcoreml.assetmanager.trash", DISPATCH_QUEUE_SERIAL_WITH_AUTORELEASE_POOL);
322        _syncQueue = dispatch_queue_create("com.executorchcoreml.assetmanager.sync", DISPATCH_QUEUE_SERIAL_WITH_AUTORELEASE_POOL);
323        _assetsInUseMap = [NSMapTable strongToWeakObjectsMapTable];
324    }
325
326    [self triggerCompaction];
327    return self;
328}
329
330- (nullable instancetype)initWithDatabaseURL:(NSURL *)databaseURL
331                          assetsDirectoryURL:(NSURL *)assetsDirectoryURL
332                           trashDirectoryURL:(NSURL *)trashDirectoryURL
333                        maxAssetsSizeInBytes:(NSInteger)maxAssetsSizeInBytes
334                                       error:(NSError * __autoreleasing *)error {
335    auto database = make_database(databaseURL, kBusyTimeIntervalInMS, error);
336    if (!database) {
337        return nil;
338    }
339
340    return [self initWithDatabase:database
341               assetsDirectoryURL:assetsDirectoryURL
342                trashDirectoryURL:trashDirectoryURL
343             maxAssetsSizeInBytes:maxAssetsSizeInBytes
344                            error:error];
345}
346
347- (nullable NSURL *)moveURL:(NSURL *)url
348     toUniqueURLInDirectory:(NSURL *)directoryURL
349                      error:(NSError * __autoreleasing *)error {
350    NSURL *dstURL = [directoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString];
351    if (![self.fileManager moveItemAtURL:url toURL:dstURL error:error]) {
352        return nil;
353    }
354
355    return dstURL;
356}
357
358- (void)cleanupAssetIfNeeded:(ETCoreMLAsset *)asset {
359    if (!asset || asset.isValid) {
360        return;
361    }
362
363    NSString *identifier = asset.identifier;
364    dispatch_async(self.syncQueue, ^{
365        NSError *cleanupError = nil;
366        if (![self _removeAssetWithIdentifier:asset.identifier error:&cleanupError]) {
367            ETCoreMLLogError(cleanupError,
368                             "%@: Failed to remove asset with identifier = %@",
369                             NSStringFromClass(ETCoreMLAssetManager.class),
370                             identifier);
371        }
372    });
373}
374
375- (nullable ETCoreMLAsset *)_storeAssetAtURL:(NSURL *)srcURL
376                              withIdentifier:(NSString *)identifier
377                                       error:(NSError * __autoreleasing *)error {
378    dispatch_assert_queue(self.syncQueue);
379    NSString *extension = srcURL.lastPathComponent.pathExtension;
380    NSURL *dstURL = [self.assetsDirectoryURL URLByAppendingPathComponent:[NSString stringWithFormat:@"%@.%@", identifier, extension]];
381    auto asset = Asset::make(srcURL, identifier, self.fileManager, error);
382    if (!asset) {
383        return nil;
384    }
385
386    auto& assetValue = asset.value();
387    size_t assetSizeInBytes = assetValue.total_size_in_bytes();
388    std::error_code ec;
389    // Update the stores inside a transaction, if anything fails it will automatically rollback to the previous state.
390    bool status = _assetsStore.impl()->transaction([self, &assetValue, assetSizeInBytes, srcURL, dstURL, &ec, error]() {
391        const std::string& assetIdentifier = assetValue.identifier;
392        // If an asset exists with the same identifier then remove it.
393        if (![self _removeAssetWithIdentifier:@(assetIdentifier.c_str()) error:error]) {
394            return false;
395        }
396
397        // Update asset path.
398        assetValue.path = dstURL.path.UTF8String;
399        // Store the asset.
400        if (!_assetsStore.impl()->put(assetIdentifier, assetValue, ec)) {
401            return false;
402        }
403
404        // Update the size of the store.
405        if (!::set_total_assets_size(_estimatedSizeInBytes + assetSizeInBytes, _assetsMetaStore, ec)) {
406            return false;
407        }
408
409        // If an asset exists move it
410        [self moveURL:dstURL toUniqueURLInDirectory:self.trashDirectoryURL error:nil];
411
412        // Move the asset to assets directory.
413        if (![self.fileManager moveItemAtURL:srcURL toURL:dstURL error:error]) {
414            return false;
415        }
416
417        return true;
418    }, Database::TransactionBehavior::Immediate, ec);
419
420    // Update the estimated size if the transaction succeeded.
421    _estimatedSizeInBytes += status ? assetSizeInBytes : 0;
422    ::set_error_from_error_code(ec, error);
423
424    ETCoreMLAsset *result = status ? [[ETCoreMLAsset alloc] initWithBackingAsset:assetValue] : nil;
425    if ([result keepAliveAndReturnError:error]) {
426        [self.assetsInUseMap setObject:result forKey:identifier];
427    } else {
428        [self cleanupAssetIfNeeded:result];
429    }
430
431    return result;
432}
433
434- (void)triggerCompaction {
435    if (self.estimatedSizeInBytes < self.maxAssetsSizeInBytes) {
436        return;
437    }
438
439    __weak __typeof(self) weakSelf = self;
440    dispatch_async(self.syncQueue, ^{
441        NSError *localError = nil;
442        if (![weakSelf _compact:self.maxAssetsSizeInBytes error:&localError]) {
443            ETCoreMLLogError(localError,
444                             "%@: Failed to compact asset store.",
445                             NSStringFromClass(ETCoreMLAssetManager.class));
446        }
447    });
448}
449
450- (nullable ETCoreMLAsset *)storeAssetAtURL:(NSURL *)url
451                             withIdentifier:(NSString *)identifier
452                                      error:(NSError * __autoreleasing *)error {
453    __block ETCoreMLAsset *result = nil;
454    dispatch_sync(self.syncQueue, ^{
455        result = [self _storeAssetAtURL:url withIdentifier:identifier error:error];
456    });
457
458    [self triggerCompaction];
459    return result;
460}
461
462- (nullable ETCoreMLAsset *)_assetWithIdentifier:(NSString *)identifier
463                                           error:(NSError * __autoreleasing *)error {
464    dispatch_assert_queue(self.syncQueue);
465    std::string assetIdentifier(identifier.UTF8String);
466    std::error_code ec;
467    auto asset = _assetsStore.impl()->get(assetIdentifier, ec);
468    if (!asset) {
469        ::set_error_from_error_code(ec, error);
470        return nil;
471    }
472
473    const auto& assetValue = asset.value();
474    ETCoreMLAsset *modelAsset = [[ETCoreMLAsset alloc] initWithBackingAsset:assetValue];
475    [self.assetsInUseMap setObject:modelAsset forKey:identifier];
476
477    return modelAsset;
478}
479
480- (nullable ETCoreMLAsset *)assetWithIdentifier:(NSString *)identifier
481                                          error:(NSError * __autoreleasing *)error {
482    __block ETCoreMLAsset *result = nil;
483    dispatch_sync(self.syncQueue, ^{
484        result = [self _assetWithIdentifier:identifier error:error];
485    });
486
487    if ([result keepAliveAndReturnError:error]) {
488        [self.assetsInUseMap setObject:result forKey:identifier];
489    } else {
490        [self cleanupAssetIfNeeded:result];
491    }
492
493    return result;
494}
495
496- (BOOL)_containsAssetWithIdentifier:(NSString *)identifier
497                               error:(NSError * __autoreleasing *)error {
498    dispatch_assert_queue(self.syncQueue);
499    std::error_code ec;
500    BOOL result = static_cast<BOOL>(_assetsStore.impl()->exists(std::string(identifier.UTF8String), ec));
501    ::set_error_from_error_code(ec, error);
502
503    return result;
504}
505
506- (BOOL)hasAssetWithIdentifier:(NSString *)identifier
507                         error:(NSError * __autoreleasing *)error {
508    __block BOOL result = NO;
509    dispatch_sync(self.syncQueue, ^{
510        result = [self _containsAssetWithIdentifier:identifier error:error];
511    });
512
513    return result;
514}
515
516- (BOOL)_removeAssetWithIdentifier:(NSString *)identifier
517                             error:(NSError * __autoreleasing *)error {
518    dispatch_assert_queue(self.syncQueue);
519    // Asset is alive we can't delete it.
520    if (is_asset_alive(self.assetsInUseMap, identifier)) {
521        return NO;
522    }
523
524    std::error_code ec;
525    std::string assetIdentifier(identifier.UTF8String);
526    auto asset = _assetsStore.impl()->get(assetIdentifier, ec);
527    // If it's an error then we can't proceed.
528    if (ec) {
529        ::set_error_from_error_code(ec, error);
530        return NO;
531    }
532
533    // Asset doesn't exists, we are good.
534    if (!asset) {
535        return YES;
536    }
537
538    const auto& assetValue = asset.value();
539    size_t assetSizeInBytes = std::min(_estimatedSizeInBytes, static_cast<NSInteger>(assetValue.total_size_in_bytes()));
540    // Update the stores inside a transaction, if anything fails it will automatically rollback to the previous state.
541    bool status = _assetsStore.impl()->transaction([self, &assetValue, assetSizeInBytes, &ec, error]() {
542        if (!self->_assetsStore.impl()->remove(assetValue.identifier, ec)) {
543            return false;
544        }
545
546        if (!::set_total_assets_size(_estimatedSizeInBytes - assetSizeInBytes, _assetsMetaStore, ec)) {
547            return false;
548        }
549
550        NSURL *assetURL = ::get_asset_url(assetValue);
551        if ([self.fileManager fileExistsAtPath:assetURL.path] &&
552            ![self moveURL:assetURL toUniqueURLInDirectory:self.trashDirectoryURL error:error]) {
553            return false;
554        }
555
556        return true;
557    }, Database::TransactionBehavior::Immediate, ec);
558
559    // Update the estimated size if the transaction succeeded.
560    _estimatedSizeInBytes -= status ? assetSizeInBytes : 0;
561    ::set_error_from_error_code(ec, error);
562    return static_cast<BOOL>(status);
563}
564
565- (BOOL)removeAssetWithIdentifier:(NSString *)identifier
566                            error:(NSError * __autoreleasing *)error {
567    __block BOOL result = NO;
568    dispatch_sync(self.syncQueue, ^{
569        result = [self _removeAssetWithIdentifier:identifier error:error];
570    });
571
572    return result;
573}
574
575- (nullable NSArray<ETCoreMLAsset *> *)_recentlyUsedAssetsWithMaxCount:(NSUInteger)maxCount
576                                                                 error:(NSError * __autoreleasing *)error {
577    dispatch_assert_queue(self.syncQueue);
578
579    NSMutableArray<ETCoreMLAsset *> *assets = [NSMutableArray arrayWithCapacity:maxCount];
580    std::error_code ec;
581    bool status = _assetsStore.impl()->get_keys_sorted_by_access_time([self, maxCount, assets](const std::string& key) {
582        NSError *localError = nil;
583        NSString *identifier = @(key.c_str());
584        ETCoreMLAsset *asset = [self _assetWithIdentifier:identifier error:&localError];
585
586        if (asset) {
587            [assets addObject:asset];
588        } else if (localError) {
589            ETCoreMLLogError(localError,
590                             "%@: Failed to retrieve asset with identifier = %@",
591                             NSStringFromClass(ETCoreMLAssetManager.class),
592                             identifier);
593        }
594
595        return assets.count < maxCount;
596    }, SortOrder::Descending, ec);
597
598    ::set_error_from_error_code(ec, error);
599    return status ? assets : nil;
600}
601
602- (nullable NSArray<ETCoreMLAsset *> *)mostRecentlyUsedAssetsWithMaxCount:(NSUInteger)maxCount
603                                                                    error:(NSError * __autoreleasing *)error {
604    __block NSArray<ETCoreMLAsset *> *result = nil;
605    dispatch_sync(self.syncQueue, ^{
606        result = [self _recentlyUsedAssetsWithMaxCount:maxCount error:error];
607    });
608
609    return result;
610}
611
612- (BOOL)_canPurgeStore {
613    dispatch_assert_queue(self.syncQueue);
614
615    NSEnumerator *keyEnumerator = self.assetsInUseMap.keyEnumerator;
616    for (NSString *identifier in keyEnumerator) {
617        if (is_asset_alive(self.assetsInUseMap, identifier)) {
618            return NO;
619        }
620    }
621
622    return YES;
623}
624
625- (NSUInteger)_compact:(NSUInteger)sizeInBytes error:(NSError * __autoreleasing *)error {
626    dispatch_assert_queue(self.syncQueue);
627
628    if (sizeInBytes == 0 && [self _canPurgeStore]) {
629        return [self _purge:error] ? 0 : _estimatedSizeInBytes;
630    }
631
632    if (_estimatedSizeInBytes <= sizeInBytes) {
633        return _estimatedSizeInBytes;
634    }
635
636    std::error_code ec;
637    ssize_t bytesToRemove = _estimatedSizeInBytes - sizeInBytes;
638    const auto& assets = ::get_assets_to_remove(_assetsStore, bytesToRemove, self.assetsInUseMap, ec);
639
640    if (ec) {
641        ::set_error_from_error_code(ec, error);
642        return _estimatedSizeInBytes;
643    }
644
645    for (const auto& asset : assets) {
646        NSError *cleanupError = nil;
647        NSString *identifier = @(asset.identifier.c_str());
648        if (![self _removeAssetWithIdentifier:identifier error:&cleanupError] && cleanupError) {
649            ETCoreMLLogError(cleanupError,
650                             "%@: Failed to remove asset with identifier = %@",
651                             NSStringFromClass(ETCoreMLAssetManager.class),
652                             identifier);
653        }
654    }
655
656    // Trigger cleanup.
657    __weak __typeof(self) weakSelf = self;
658    dispatch_async(self.trashQueue, ^{
659        [weakSelf removeFilesInTrashDirectory];
660    });
661
662    return _estimatedSizeInBytes;
663}
664
665- (NSUInteger)compact:(NSUInteger)sizeInBytes error:(NSError * __autoreleasing *)error {
666    __block NSUInteger result = 0;
667    dispatch_sync(self.syncQueue, ^{
668        result = [self _compact:sizeInBytes error:error];
669    });
670
671    return result;
672}
673
674- (void)removeFilesInTrashDirectory {
675    dispatch_assert_queue(self.trashQueue);
676
677    NSFileManager *fileManager = [[NSFileManager alloc] init];
678    fileManager.delegate = self;
679    __block NSError *localError = nil;
680    BOOL (^errorHandler)(NSURL *url, NSError *error) = ^BOOL(NSURL *url, NSError *enumerationError) {
681        localError = enumerationError;
682        return YES;
683    };
684
685    NSDirectoryEnumerator *enumerator = [fileManager enumeratorAtURL:self.trashDirectoryURL
686                                          includingPropertiesForKeys:@[]
687                                                             options:NSDirectoryEnumerationSkipsSubdirectoryDescendants
688                                                        errorHandler:errorHandler];
689    for (NSURL *itemURL in enumerator) {
690        if (![fileManager removeItemAtURL:itemURL error:&localError]) {
691            ETCoreMLLogError(localError,
692                             "%@: Failed to remove item in trash directory with name = %@",
693                             NSStringFromClass(ETCoreMLAssetManager.class),
694                             itemURL.lastPathComponent);
695        }
696    }
697}
698
699- (BOOL)_purge:(NSError * __autoreleasing *)error {
700    dispatch_assert_queue(self.syncQueue);
701
702    std::error_code ec;
703    bool status = _assetsStore.impl()->transaction([self, &ec, error]() {
704        // Purge the assets store.
705        if (!self->_assetsStore.impl()->purge(ec)) {
706            return false;
707        }
708
709        // Purge the assets size store.
710        if (!self->_assetsMetaStore.impl()->purge(ec)) {
711            return false;
712        }
713
714        // Move the the whole assets directory to the temp directory.
715        if (![self moveURL:self.assetsDirectoryURL toUniqueURLInDirectory:self.trashDirectoryURL error:error]) {
716            return false;
717        }
718
719        self->_estimatedSizeInBytes = 0;
720        NSError *localError = nil;
721        // Create the assets directory, if we fail here it's okay.
722        if (![self.fileManager createDirectoryAtURL:self.assetsDirectoryURL withIntermediateDirectories:NO attributes:@{} error:&localError]) {
723            ETCoreMLLogError(localError,
724                             "%@: Failed to create assets directory",
725                             NSStringFromClass(ETCoreMLAssetManager.class));
726        }
727
728        return true;
729    }, Database::TransactionBehavior::Immediate, ec);
730
731    ::set_error_from_error_code(ec, error);
732    // Trigger cleanup
733    if (status) {
734        __weak __typeof(self) weakSelf = self;
735        dispatch_async(self.trashQueue, ^{
736            [weakSelf removeFilesInTrashDirectory];
737        });
738    }
739
740    return static_cast<BOOL>(status);
741}
742
743- (BOOL)purgeAndReturnError:(NSError * __autoreleasing *)error {
744    __block BOOL result = 0;
745    dispatch_sync(self.syncQueue, ^{
746        result = [self _purge:error];
747    });
748
749    return result;
750}
751
752- (BOOL)fileManager:(NSFileManager *)fileManager shouldProceedAfterError:(NSError *)error removingItemAtURL:(NSURL *)URL {
753    return YES;
754}
755
756@end
757