From 291dc54d1368962162c20be15693064bb9a783fb Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Sun, 4 Oct 2026 19:51:11 -0700 Subject: [PATCH 1/5] [CoreAI] Add root and PTE cache clearing Summary: Add public APIs to clear Core AI cache state, either for a whole assets root or for the models referenced by one PTE. - `runtime/coreai_cache.{h,mm}`: - `clear_cache(coreai_assets_dir)` clears every keyed entry in an explicit assets root. - `clear_cache_for_pte(loader, dir)` and `clear_cache_for_pte(pte_path, dir)` clear only the entries for the current platform and architecture that the PTE's Core AI delegates reference. They find those delegates with `inspect_coreai_pte`, so no Program/Method is constructed and no Core AI assets are materialized. Every selected manifest is validated before any key is cleared. A null directory selects the backend's default root. - Per key, clearing preflights the staged source tree, evicts the SDK entry, removes staging, then removes the bookmark. The bookmark is kept until staging removal succeeds so a failed clear can be retried. SDK errors other than confirmed absence preserve the key's files. Keys are processed in sorted order and the first error is returned. Each removal syncs the directory it changed. Missing roots and entries are no-ops. - Everything under a keyed `staging/` directory belongs to the backend and is removed with that key. Outside keyed entries, PTEs, other files, lock files and unkeyed temporaries are never removed. - Supporting changes: descriptor-relative keyed staging removal in `coreai_storage`, bookmark inventory and removal in `coreai_bookmarks`, and SDK eviction plus per-key clearing in `coreai_load_coordinator`. - Tests: cache-clearing and eviction suites in `coreai_host_test`. `coreai_runtime_smoke` also links every public cache-clearing overload. - README: "Clearing cached assets". Callers must unload affected models and stop concurrent loads, inference and maintenance until these calls return. Per-key locks are not a whole-root barrier. No changes outside `backends/apple/coreai`. Test Plan: ``` cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release \ -DCMAKE_OSX_DEPLOYMENT_TARGET=27.0 -DEXECUTORCH_BUILD_COREAI=ON \ -DEXECUTORCH_BUILD_TESTS=ON -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON cmake --build build --target backends/apple/coreai/all ctest --test-dir build/backends/apple/coreai --output-on-failure --no-tests=error ``` CTest runs `coreai_host_test` (adding the cache-clearing and eviction suites), `coreai_swift_bridge_test` and `coreai_runtime_smoke`. --- backends/apple/coreai/CMakeLists.txt | 25 +- backends/apple/coreai/README.md | 56 +- .../apple/coreai/runtime/coreai_bookmarks.h | 16 +- .../apple/coreai/runtime/coreai_bookmarks.mm | 52 +- backends/apple/coreai/runtime/coreai_cache.h | 57 ++ backends/apple/coreai/runtime/coreai_cache.mm | 120 +++++ .../coreai/runtime/coreai_load_coordinator.h | 12 + .../coreai/runtime/coreai_load_coordinator.mm | 63 ++- .../apple/coreai/runtime/coreai_storage.h | 6 +- .../apple/coreai/runtime/coreai_storage.mm | 7 +- .../coreai/runtime/test/coreai_cache_test.mm | 500 ++++++++++++++++++ .../runtime/test/coreai_eviction_test.mm | 402 ++++++++++++++ .../coreai/runtime/test/coreai_fake_loader.h | 2 + .../coreai/runtime/test/coreai_fake_loader.mm | 11 +- .../coreai/runtime/test/runtime_smoke.cpp | 21 + 15 files changed, 1322 insertions(+), 28 deletions(-) create mode 100644 backends/apple/coreai/runtime/coreai_cache.h create mode 100644 backends/apple/coreai/runtime/coreai_cache.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_cache_test.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_eviction_test.mm diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index a08ee029635..7c65baa2d0e 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/backends/apple/coreai/CMakeLists.txt @@ -28,9 +28,13 @@ endforeach() enable_language(OBJCXX) set(_coreai_runtime_sources - runtime/coreai_backend.mm runtime/coreai_assets.mm - runtime/coreai_storage.mm runtime/coreai_bookmarks.mm - runtime/coreai_load_coordinator.mm runtime/coreai_pte.cpp + runtime/coreai_backend.mm + runtime/coreai_assets.mm + runtime/coreai_storage.mm + runtime/coreai_bookmarks.mm + runtime/coreai_load_coordinator.mm + runtime/coreai_pte.cpp + runtime/coreai_cache.mm ) # Keep ARC and C++ exception settings away from the private Swift module. @@ -64,6 +68,8 @@ if(EXECUTORCH_BUILD_TESTS AND CMAKE_SYSTEM_NAME STREQUAL "Darwin") runtime/test/coreai_delegate_test.mm runtime/test/coreai_pte_test.mm runtime/test/coreai_pte_fixture.mm + runtime/test/coreai_cache_test.mm + runtime/test/coreai_eviction_test.mm runtime/ETCoreAITensor.mm ${_coreai_runtime_sources} ) @@ -73,12 +79,15 @@ if(EXECUTORCH_BUILD_TESTS AND CMAKE_SYSTEM_NAME STREQUAL "Darwin") # The consolidated runtime target is created after this directory. if(EXECUTORCH_BUILD_SHARED) set(_coreai_host_runtime executorch_shared) + set(_coreai_host_data_loader "") else() set(_coreai_host_runtime executorch) + set(_coreai_host_data_loader extension_data_loader) endif() target_link_libraries( - coreai_host_test PRIVATE ${_coreai_host_runtime} program_schema - GTest::gtest ${COREAI_FOUNDATION_FRAMEWORK} + coreai_host_test + PRIVATE ${_coreai_host_runtime} ${_coreai_host_data_loader} program_schema + GTest::gtest ${COREAI_FOUNDATION_FRAMEWORK} ) add_test(NAME coreai_host_test COMMAND coreai_host_test) set_tests_properties(coreai_host_test PROPERTIES TIMEOUT 120) @@ -434,6 +443,12 @@ target_link_libraries( coreai_runtime_smoke PRIVATE coreaidelegate ${_coreai_runtime} ) +install( + FILES runtime/coreai_cache.h + DESTINATION + ${CMAKE_INSTALL_INCLUDEDIR}/executorch/backends/apple/coreai/runtime +) + install( TARGETS coreaidelegate coreai_swift EXPORT ExecuTorchTargets diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index 6ffa262a339..b2f0e677846 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -225,6 +225,55 @@ left untouched; they can be removed separately during a fully quiescent offline cleanup. The 256-file bound applies to the new stripe pool, not pre-existing legacy files. Bookmarks and staged bundles need no migration. +### Clearing cached assets + +The public C++ maintenance APIs are in `executorch::backends::coreai`: + +```cpp +runtime::Error clear_cache(const char* coreai_assets_dir); +runtime::Error clear_cache_for_pte( + runtime::DataLoader& pte, const char* coreai_assets_dir = nullptr); +runtime::Error clear_cache_for_pte( + const char* pte_path, const char* coreai_assets_dir = nullptr); +``` + +`clear_cache` requires an explicit absolute assets root. The PTE overloads use +user-domain `NSCachesDirectory` plus `executorch_coreai` when the directory is +omitted or null, just like backend initialization. An invalid supplied path is +an error. If loading used a custom root, pass that same root when clearing; +the PTE does not store runtime directory overrides. + +Root-wide clearing visits keyed bookmarks and staged directories. PTE clearing +selects and deduplicates keys for all Core AI delegates in every method, using +the current platform, SDK architecture and default specialization configuration. +Both public `clear_cache_for_pte` overloads use the private `inspect_coreai_pte` +reader and retain its `CoreAIPteData`, including the verified program image and +selected processed buffers, through `clear_keys`. Discovery does not load a +Program or initialize a Method, load SDK models, bind functions, extract Core AI +assets or specialize. All selected manifests are validated before any key is +cleared; root-wide clearing remains +available for existing keyed entries. Backend-local structural and semantic +verification is mandatory even with `ET_ENABLE_PROGRAM_VERIFICATION=0`. + +These synchronous calls require callers to unload affected models and prevent +concurrent loads, inference and maintenance across processes. Per-key locks are +not a whole-root barrier. Clearing first deletes the recorded SDK entry, then +staging, then its bookmark. Busy or unknown SDK errors preserve the files; +SDK-confirmed absence permits cleanup. Filesystem errors may leave partial +cleanup, but the bookmark is retained until staging removal succeeds so cleanup +can be retried. Independent keys are still attempted in sorted order, and the +first key error is returned. Each removal syncs the directory it changed. A +missing root or absent entry is a noncreating no-op. + +Everything under a keyed `staging/` directory belongs to the backend and is +removed with that key. Outside keyed entries, clearing does not delete PTEs, +other files or the stable lock files. Unkeyed interrupted-staging directories +also remain untouched. Bookmarks are SDK eviction handles, not weights: lost +bookmarks can leave SDK entries that these APIs cannot target. There is no +history or ownership registry. Copied or reused exported artifacts can share +keys; clearing them also removes the other copy's warm-cache benefit. Subsequent +loads can specialize again. + ### Runtime options Pass options through the public `Module` API before loading. In an error-returning @@ -309,8 +358,8 @@ artifact. The durable-save message confirms local bookmark publication, including its final directory flush, not independent verification of SDK-private artifact storage. Failed operations do not emit their success messages. -The delegate does not require `EXECUTORCH_BUILD_EXTENSION_DATA_LOADER`. -Enable it for consumers such as `Module` that load program data. +Enable `EXECUTORCH_BUILD_EXTENSION_DATA_LOADER` for the public PTE-file +maintenance API and consumers such as `Module`. Link the CMake target `coreaidelegate`, not just its archive filename. Its transitive dependencies supply Swift and Foundation/CoreAI linkage, and its link interface retains static backend registration. With @@ -323,7 +372,8 @@ Select the Apple SDK/toolchain when configuring the consumer project. SwiftPM and XCFramework distribution are not integrated yet. The `coreai_runtime_smoke` target checks final linkage of a C++ consumer, -including Swift dependencies and backend registration retention. With +including Swift dependencies, backend registration retention, and the public +cache-clearing overloads. With `EXECUTORCH_BUILD_TESTS=ON` it is built; otherwise it is excluded from the default build. SDK-linked tests are registered with CTest only for native macOS builds when the host meets the target's deployment minimum (at least 27.0). diff --git a/backends/apple/coreai/runtime/coreai_bookmarks.h b/backends/apple/coreai/runtime/coreai_bookmarks.h index 26b841a4f8e..d9d1a067fa5 100644 --- a/backends/apple/coreai/runtime/coreai_bookmarks.h +++ b/backends/apple/coreai/runtime/coreai_bookmarks.h @@ -26,6 +26,8 @@ runtime::Result bookmark_lock_name(NSString* key); runtime::Result resolve_bookmark_root( NSString* path, bool create = true); +// Missing roots produce an empty inventory. +runtime::Result*> inventory_bookmark_keys(NSString* root); class BookmarkLock final { public: @@ -37,7 +39,10 @@ class BookmarkLock final { lock_bookmark(NSString* root, NSString* key, bool create); friend runtime::Result read_bookmark(const BookmarkLock& lock); friend runtime::Error write_bookmark(const BookmarkLock& lock, NSData* data); - friend runtime::Error remove_bookmark_staging(const BookmarkLock& lock); + friend runtime::Error remove_bookmark(const BookmarkLock& lock); + friend runtime::Error remove_bookmark_staging( + const BookmarkLock& lock, + bool remove); friend runtime::Result prepare_bookmark_staging( const BookmarkLock& lock); BookmarkLock( @@ -66,9 +71,12 @@ lock_bookmark(NSString* root, NSString* key, bool create = true); runtime::Result read_bookmark(const BookmarkLock& lock); // A staging-only lock has no bookmark directory; writes return InvalidArgument. runtime::Error write_bookmark(const BookmarkLock& lock, NSData* data); -// Requires successful persistent specialization and durable bookmark -// publication. -runtime::Error remove_bookmark_staging(const BookmarkLock& lock); +runtime::Error remove_bookmark(const BookmarkLock& lock); +// The caller must ensure the SDK no longer needs staging. With remove=false, +// only preflight the keyed tree. +runtime::Error remove_bookmark_staging( + const BookmarkLock& lock, + bool remove = true); // Prepared only for cold loads; source preparation appends key/bundle. runtime::Result prepare_bookmark_staging(const BookmarkLock& lock); diff --git a/backends/apple/coreai/runtime/coreai_bookmarks.mm b/backends/apple/coreai/runtime/coreai_bookmarks.mm index 03521ad9e91..f0c1e163f72 100644 --- a/backends/apple/coreai/runtime/coreai_bookmarks.mm +++ b/backends/apple/coreai/runtime/coreai_bookmarks.mm @@ -158,12 +158,35 @@ bool digest(id value) { return path; } -BookmarkLock::BookmarkLock( - int root, - int bookmarks, - int file, - NSString* root_path, - NSString* key) +Result*> inventory_bookmark_keys(NSString* root) { + auto resolved = resolve_bookmark_root(root, false); + if (!resolved.ok()) return resolved.error(); + auto opened = open_root(resolved.get()); + if (!opened.ok()) return opened.error(); + FileDescriptor directory(opened.get()); + if (directory.get() < 0) return @[]; + NSMutableSet* keys = [NSMutableSet set]; + for (NSString* name in @[ @"bookmarks", @"staging" ]) { + auto child_result = open_child(directory.get(), name, false); + if (!child_result.ok()) return child_result.error(); + FileDescriptor child(child_result.get()); + if (child.get() < 0) continue; + auto children = storage_children(child.get(), true); + if (!children.ok()) return children.error(); + for (NSString* entry in children.get()) { + NSString* key = entry; + if ([name isEqual:@"bookmarks"]) { + if (![entry hasSuffix:@".bookmark"]) continue; + key = [entry substringToIndex:entry.length - @".bookmark".length]; + } + if (digest(key)) [keys addObject:key]; + } + } + return [keys.allObjects sortedArrayUsingSelector:@selector(compare:)]; +} + +BookmarkLock::BookmarkLock(int root, int bookmarks, int file, + NSString* root_path, NSString* key) : root_(root), bookmarks_(bookmarks), file_(file), @@ -283,8 +306,21 @@ Error write_bookmark(const BookmarkLock& lock, NSData* data) { lock.bookmarks_.get(), bookmark_name(lock.key_), data); } -Error remove_bookmark_staging(const BookmarkLock& lock) { - return remove_storage_staging(lock.root_.get(), lock.key_); +Error remove_bookmark(const BookmarkLock& lock) { + if (lock.bookmarks_.get() < 0) return Error::Ok; + ET_CHECK_OR_RETURN_ERROR(storage_fault(StorageOperation::Remove) == 0, + AccessFailed, + "Core AI bookmark removal interrupted"); + ET_CHECK_OR_RETURN_ERROR( + unlinkat(lock.bookmarks_.get(), + bookmark_name(lock.key_).fileSystemRepresentation, 0) == 0 || + errno == ENOENT, + AccessFailed, "Cannot remove Core AI bookmark"); + return sync_storage_directory(lock.bookmarks_.get()); +} + +Error remove_bookmark_staging(const BookmarkLock& lock, bool remove) { + return remove_storage_staging(lock.root_.get(), lock.key_, remove); } Result prepare_bookmark_staging(const BookmarkLock& lock) { diff --git a/backends/apple/coreai/runtime/coreai_cache.h b/backends/apple/coreai/runtime/coreai_cache.h new file mode 100644 index 00000000000..c87dd6f9daa --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache.h @@ -0,0 +1,57 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +#pragma once + +#include + +namespace executorch::runtime { +class DataLoader; +} + +namespace executorch::backends::coreai { + +/** + * Clears tracked SDK entries, keyed staging, and bookmarks in an explicit, + * dedicated Core AI assets directory. Missing roots are successful no-ops. + * PTEs, stable lock files, and unkeyed temporaries are not removed. + * SDK entries without saved bookmarks cannot be targeted. + * + * Callers must unload affected models and prevent concurrent loading, + * inference, and maintenance across processes until this call returns. + * Independent keys are attempted in sorted order, returning the first error; + * clearing is not a transaction. SDK failures preserve an entry's files unless + * the SDK confirms that entry is absent. + */ +[[nodiscard]] runtime::Error clear_cache(const char* coreai_assets_dir); + +/** + * Clears current-platform/SDK-architecture cache entries referenced by every + * CoreAI delegate in the PTE, after validating all selected manifests. Does not + * initialize delegates or materialize Core AI assets. Backend-local structural + * and semantic verification is mandatory regardless of runtime build settings. + * Reads the program region, including inline constants and unrelated inline + * blobs, and only selected Core AI external processed-data segments. + * + * A null directory selects the same Caches default as backend initialization. + * Supplied directories must be valid explicit roots; they never fall back to + * the default. The directory is never derived from the PTE's location. + * Copied PTEs can share cache keys. The quiescence and failure contract of + * clear_cache also applies. The loader and its data must remain stable for + * this call; the retained program and selected buffers survive through + * eviction. + */ +[[nodiscard]] runtime::Error clear_cache_for_pte( + runtime::DataLoader& loader, + const char* coreai_assets_dir = nullptr); + +/** FileDataLoader convenience overload with the same contract. */ +[[nodiscard]] runtime::Error clear_cache_for_pte( + const char* pte_path, + const char* coreai_assets_dir = nullptr); + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_cache.mm b/backends/apple/coreai/runtime/coreai_cache.mm new file mode 100644 index 00000000000..20e07827f05 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache.mm @@ -0,0 +1,120 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +#include "coreai_cache.h" + +#import "coreai_assets.h" +#import "coreai_load_coordinator.h" +#include "coreai_pte.h" +#import "coreai_storage.h" + +#include +#include + +namespace executorch::backends::coreai { +namespace { +using runtime::Error; +using runtime::Result; + +Result explicit_root(const char* path) { + NSString* root = path == nullptr ? nil : [NSString stringWithUTF8String:path]; + return resolve_bookmark_root(root, false); +} + +Error clear_keys(NSString* root, NSArray* keys) { + NSArray* ordered = [[[NSSet setWithArray:keys] allObjects] + sortedArrayUsingSelector:@selector(compare:)]; + Error first_error = Error::Ok; + for (NSString* key in ordered) { + Error error = clear_bookmark_assets(root, key, nil); + if (error != Error::Ok) { + ET_LOG(Error, "Core AI cache clear failed for key %s: 0x%x", + key.UTF8String, static_cast(error)); + if (first_error == Error::Ok) first_error = error; + } + } + return first_error; +} + +Error clear_pte(runtime::DataLoader& loader, NSString* root) { + // Keep the verified image and original segment owners alive through eviction. + auto pte = inspect_coreai_pte(loader); + if (!pte.ok()) return pte.error(); + NSMutableArray* keys = [NSMutableArray array]; + NSString* architecture = nil; +#if TARGET_OS_OSX + NSString* platform = @"macOS"; +#elif TARGET_OS_IOS + NSString* platform = @"iOS"; +#else + NSString* platform = @"unsupported"; +#endif + for (const auto& buffer : pte->processed) { + auto data = buffer.data_safe(); + if (!data.ok()) return data.error(); + ET_CHECK_OR_RETURN_ERROR(data.get() != nullptr && buffer.size() > 0, + InvalidProgram, "Missing Core AI manifest"); + NSData* bytes = [NSData dataWithBytesNoCopy:const_cast(data.get()) + length:buffer.size() + freeWhenDone:NO]; + auto manifest = parse_manifest(bytes); + if (!manifest.ok()) return manifest.error(); + ET_CHECK_OR_RETURN_ERROR( + [NSProcessInfo.processInfo + isOperatingSystemAtLeastVersion:manifest->minimum_version], + DelegateInvalidCompatibility, + "Core AI model requires a newer operating system"); + if (architecture == nil) { + ET_CHECK_OR_RETURN_ERROR(ETCoreAIIsAvailable(), NotSupported, + "Core AI requires macOS 27 or iOS 27"); + architecture = ETCoreAIDeviceArchitectureName(); + } + auto selected = select_assets(manifest.get(), architecture, platform); + if (!selected.ok()) return selected.error(); + auto key = bookmark_key(selected.get(), platform, architecture); + if (!key.ok()) return key.error(); + [keys addObject:key.get()]; + } + return clear_keys(root, keys); +} +} // namespace + +Error clear_cache(const char* coreai_assets_dir) { + @autoreleasepool { + auto root = explicit_root(coreai_assets_dir); + if (!root.ok()) return root.error(); + auto keys = inventory_bookmark_keys(root.get()); + if (!keys.ok()) return keys.error(); + return clear_keys(root.get(), keys.get()); + } +} + +Error clear_cache_for_pte(runtime::DataLoader& loader, + const char* coreai_assets_dir) { + @autoreleasepool { + if (coreai_assets_dir != nullptr) { + auto root = explicit_root(coreai_assets_dir); + if (!root.ok()) return root.error(); + return clear_pte(loader, root.get()); + } + auto default_root = default_coreai_assets_root(); + if (!default_root.ok()) return default_root.error(); + auto root = resolve_bookmark_root(default_root.get(), false); + if (!root.ok()) return root.error(); + return clear_pte(loader, root.get()); + } +} + +Error clear_cache_for_pte(const char* pte_path, const char* coreai_assets_dir) { + ET_CHECK_OR_RETURN_ERROR(pte_path != nullptr && pte_path[0] != '\0', + InvalidArgument, "Missing PTE path"); + auto loader = extension::FileDataLoader::from(pte_path); + if (!loader.ok()) return loader.error(); + return clear_cache_for_pte(loader.get(), coreai_assets_dir); +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.h b/backends/apple/coreai/runtime/coreai_load_coordinator.h index 018117a18f9..eb41450c06b 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.h +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.h @@ -22,4 +22,16 @@ runtime::Result> acquire_bookmark_model( NSString* architecture, id loader); +// Deletes only the saved current SDK entry, not historical entries or sources. +runtime::Error evict_bookmark_model( + NSString* root, + NSString* key, + id loader); + +// Callers must keep affected loads and sessions quiescent across processes. +runtime::Error clear_bookmark_assets( + NSString* root, + NSString* key, + id loader); + } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.mm b/backends/apple/coreai/runtime/coreai_load_coordinator.mm index d239ccb0f64..7f77776972a 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.mm +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.mm @@ -78,7 +78,31 @@ void cleanup_source(const BookmarkLock& lock) { static_cast(error)); } } -} // namespace + +Error evict_one(id loader, NSData* bookmark) { + ET_CHECK_OR_RETURN_ERROR(storage_fault(StorageOperation::BeforeEvict) == 0, + AccessFailed, "Core AI eviction interrupted"); + dispatch_semaphore_t ready = dispatch_semaphore_create(0); + __block NSError* failure = nil; + [loader evictModelWithBookmark:bookmark + completion:^(NSError* error) { + failure = error; + dispatch_semaphore_signal(ready); + }]; + dispatch_semaphore_wait(ready, DISPATCH_TIME_FOREVER); + if (failure != nil) { + // Probe after failed deletion so our own probe cannot make the entry busy. + @autoreleasepool { + auto probe = restore(loader, bookmark); + if (probe.ok() && probe.get() == nil) return Error::Ok; + } + return sdk_error(failure); + } + ET_CHECK_OR_RETURN_ERROR(storage_fault(StorageOperation::AfterEvict) == 0, + AccessFailed, "Core AI eviction interrupted"); + return Error::Ok; +} +} // namespace Result> acquire_bookmark_model( const Manifest& selected, @@ -156,4 +180,39 @@ void cleanup_source(const BookmarkLock& lock) { return model; } -} // namespace executorch::backends::coreai +Error evict_bookmark_model(NSString* root, NSString* key, + id loader) { + ET_CHECK_OR_RETURN_ERROR(loader != nil, InvalidArgument, + "Missing Core AI model loader"); + auto locked = lock_bookmark(root, key); + if (!locked.ok()) return locked.error(); + auto bookmark = read_bookmark(*locked.get()); + if (!bookmark.ok()) return bookmark.error(); + if (bookmark.get() == nil) return Error::Ok; + ET_CHECK_OK_OR_RETURN_ERROR(evict_one(loader, bookmark.get())); + return remove_bookmark(*locked.get()); +} + +Error clear_bookmark_assets(NSString* root, NSString* key, + id loader) { + auto locked = lock_bookmark(root, key, false); + if (!locked.ok()) return locked.error(); + if (locked.get() == nullptr) return Error::Ok; + auto bookmark = read_bookmark(*locked.get()); + if (!bookmark.ok()) return bookmark.error(); + ET_CHECK_OK_OR_RETURN_ERROR(remove_bookmark_staging(*locked.get(), false)); + if (bookmark.get() != nil) { + if (loader == nil) { + ET_CHECK_OR_RETURN_ERROR(ETCoreAIIsAvailable(), NotSupported, + "Core AI requires macOS 27 or iOS 27"); + loader = ETCoreAICreateModelLoader(); + ET_CHECK_OR_RETURN_ERROR(loader != nil, Internal, + "Cannot create Core AI model loader"); + } + ET_CHECK_OK_OR_RETURN_ERROR(evict_one(loader, bookmark.get())); + } + ET_CHECK_OK_OR_RETURN_ERROR(remove_bookmark_staging(*locked.get())); + return bookmark.get() == nil ? Error::Ok : remove_bookmark(*locked.get()); +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_storage.h b/backends/apple/coreai/runtime/coreai_storage.h index 9de9a3ced28..3d489a5bf7e 100644 --- a/backends/apple/coreai/runtime/coreai_storage.h +++ b/backends/apple/coreai/runtime/coreai_storage.h @@ -48,9 +48,11 @@ runtime::Result*> storage_children( int fd, bool skip_invalid_names = false); // The caller holds the key's lock. Missing staging is already removed. -runtime::Error remove_storage_staging(int root_fd, NSString* key); +// With remove=false, only preflight the keyed tree without modifying it. +runtime::Error remove_storage_staging( + int root_fd, NSString* key, bool remove = true); // Removes a keyed child of source_root_fd under the same lock requirement. -runtime::Error remove_storage_source(int source_root_fd, NSString* key); +runtime::Error remove_storage_source(int source_root_fd, NSString* key, bool remove = true); enum class StorageOperation { Write, diff --git a/backends/apple/coreai/runtime/coreai_storage.mm b/backends/apple/coreai/runtime/coreai_storage.mm index 2712b3a1253..50c192b26f1 100644 --- a/backends/apple/coreai/runtime/coreai_storage.mm +++ b/backends/apple/coreai/runtime/coreai_storage.mm @@ -373,7 +373,7 @@ bool storage_key(NSString* key) { } } // namespace -Error remove_storage_staging(int root_fd, NSString* key) { +Error remove_storage_staging(int root_fd, NSString* key, bool remove) { ET_CHECK_OR_RETURN_ERROR( storage_key(key), InvalidArgument, "Invalid Core AI staging key"); FileDescriptor staging(retry_eintr([&] { @@ -386,10 +386,10 @@ FileDescriptor staging(retry_eintr([&] { staging.get() >= 0, AccessFailed, "Cannot open Core AI staging directory"); - return remove_storage_source(staging.get(), key); + return remove_storage_source(staging.get(), key, remove); } -Error remove_storage_source(int source_root_fd, NSString* key) { +Error remove_storage_source(int source_root_fd, NSString* key, bool remove) { ET_CHECK_OR_RETURN_ERROR( storage_key(key), InvalidArgument, "Invalid Core AI staging key"); struct stat info; @@ -421,6 +421,7 @@ FileDescriptor source(retry_eintr([&] { AccessFailed, "Cannot open Core AI keyed staging directory"); ET_CHECK_OK_OR_RETURN_ERROR(visit_source_tree(source.get(), false)); + if (!remove) return Error::Ok; ET_CHECK_OK_OR_RETURN_ERROR(visit_source_tree(source.get(), true)); ET_CHECK_OR_RETURN_ERROR( storage_call( diff --git a/backends/apple/coreai/runtime/test/coreai_cache_test.mm b/backends/apple/coreai/runtime/test/coreai_cache_test.mm new file mode 100644 index 00000000000..a6597795d14 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_cache_test.mm @@ -0,0 +1,500 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include "coreai_cache.h" +#include +#include +#include +#include +#include +#include "coreai_bookmark_fixture.h" +#include "coreai_fake_loader.h" +#include "coreai_fault_scope.h" +#include "coreai_filesystem_fixture.h" +#include "coreai_pte_fixture.h" +#include "coreai_source_fixture.h" + +namespace executorch::backends::coreai::testing { +namespace { + +using runtime::Error; +using SegmentType = runtime::DataLoader::SegmentInfo::Type; +using ::testing::AssertionFailure; +using ::testing::AssertionResult; +using ::testing::AssertionSuccess; + +class CoreAICacheTest : public ::testing::Test { + protected: + void SetUp() override { runtime::runtime_init(); } + + void drain_callbacks() { + auto drained = bridge.wait_for_callbacks(dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC)); + if (!drained) { + ADD_FAILURE() << drained.message(); + // Timed-out callbacks may still borrow this test's PTE or observations. + std::_Exit(EXIT_FAILURE); + } + } + + ScopedFakeBridgeState bridge; +}; + +struct CacheEntry { + NSString* key = nil; + NSData* bookmark = nil; + NSURL* stage = nil; +}; + +AssertionResult seed_cache(NSURL* root, NSDictionary* dict, BookmarkFakeLoader* loader, + CacheEntry& output, bool bookmark = true, bool stage = true, + NSString* architecture = @"arch_b") { + if (root == nil || loader == nil) { + return AssertionFailure() << "Cache seed requires a root and fake loader"; + } + auto parsed = parse_manifest(encode(dict)); + if (!parsed.ok()) { + return AssertionFailure() << "Cannot parse cache seed manifest"; + } + auto selected = select_assets(parsed.get(), architecture, @"macOS"); + if (!selected.ok()) { + return AssertionFailure() << "Cannot select cache seed assets"; + } + auto key = bookmark_key(selected.get(), @"macOS", architecture); + if (!key.ok()) { + return AssertionFailure() << "Cannot compute cache seed key"; + } + CacheEntry entry{ + key.get(), nil, + [[root URLByAppendingPathComponent:@"staging"] URLByAppendingPathComponent:key.get()]}; + auto lock = lock_bookmark(root.path, entry.key); + if (!lock.ok()) { + return AssertionFailure() << "Cannot lock cache seed bookmark"; + } + if (stage) { + auto staging = prepare_bookmark_staging(*lock.get()); + if (!staging.ok()) { + return AssertionFailure() << "Cannot prepare cache seed staging"; + } + TestData data; + if (selected->aot_compiled) { + aot_data(data, architecture); + } + auto staged = executorch::backends::coreai::prepare_source_bundle(selected.get(), &data, + staging.get(), entry.key); + if (data.requests != data.releases) { + return AssertionFailure() << "Cache seed source buffers were not released"; + } + if (!staged.ok()) { + return AssertionFailure() << "Cannot stage cache seed source bundle"; + } + } + if (bookmark) { + entry.bookmark = [NSUUID.UUID.UUIDString dataUsingEncoding:NSUTF8StringEncoding]; + BookmarkFakeModel* model = [[BookmarkFakeModel alloc] init]; + if (model == nil || entry.bookmark == nil) { + return AssertionFailure() << "Cannot create cache seed model/bookmark"; + } + model.owner = loader; + model.bookmark = entry.bookmark; + loader.models[entry.bookmark] = model; + if (write_bookmark(*lock.get(), entry.bookmark) != Error::Ok) { + return AssertionFailure() << "Cannot write cache seed bookmark"; + } + } + output = entry; + return AssertionSuccess(); +} + +bool cache_exists(NSURL* url) { return [NSFileManager.defaultManager fileExistsAtPath:url.path]; } + +void cache_entry_removed(NSURL* root, const CacheEntry& entry) { + EXPECT_FALSE(cache_exists(test_bookmark_url(root, entry.key))); + EXPECT_FALSE(cache_exists(entry.stage)); + EXPECT_TRUE(cache_exists([[root URLByAppendingPathComponent:@"locks"] + URLByAppendingPathComponent:[entry.key stringByAppendingString:@".lock"]])); +} + +NSMutableDictionary* cache_manifest(char identity) { + auto dict = manifest_dict(); + dict[@"bundle_digests"] = @{@"model.aimodel" : fixture_identity(identity)}; + return dict; +} + +TEST_F(CoreAICacheTest, EmptyAndMissingRootsAreNoncreatingNoops) { + @autoreleasepool { + BookmarkDirectory root; + ASSERT_NE(root.url, nil); + BookmarkLoaderScope fake; + ASSERT_NE(fake.loader, nil); + FBSyntheticPTE empty({{"forward", {}}}); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(root.url, snapshot)); + const int queries = bridge.state().architecture_queries; + bridge.state().available = false; + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + EXPECT_EQ(clear_cache_for_pte(empty, root.url.path.UTF8String), Error::Ok); + NSURL* missing = [root.url URLByAppendingPathComponent:@"absent/nested/models"]; + EXPECT_EQ(clear_cache(missing.path.UTF8String), Error::Ok); + EXPECT_EQ(clear_cache_for_pte(empty, missing.path.UTF8String), Error::Ok); + EXPECT_FALSE(cache_exists(missing)); + EXPECT_TRUE(snapshot_matches(root.url, snapshot)); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(bridge.state().architecture_queries.load(), queries); + EXPECT_TRUE(empty.check_released()); + } +} + +TEST_F(CoreAICacheTest, ValidatesExplicitRootsAndPtePaths) { + @autoreleasepool { + BookmarkDirectory root; + ASSERT_NE(root.url, nil); + BookmarkLoaderScope fake; + ASSERT_NE(fake.loader, nil); + FBSyntheticPTE empty({{"forward", {}}}); + EXPECT_EQ(clear_cache("relative"), Error::InvalidArgument); + EXPECT_EQ(clear_cache_for_pte(empty, "relative"), Error::InvalidArgument); + EXPECT_EQ(clear_cache(nullptr), Error::InvalidArgument); + EXPECT_EQ(clear_cache_for_pte(static_cast(nullptr), root.url.path.UTF8String), + Error::InvalidArgument); + EXPECT_TRUE(empty.check_released()); + } +} + +TEST_F(CoreAICacheTest, SelectsAllMethodsDeduplicatesAndEvictsInKeyOrder) { + @autoreleasepool { + BookmarkDirectory root; + BookmarkDirectory isolated; + BookmarkLoaderScope fake; + auto source_dict = manifest_dict(); + auto second_dict = cache_manifest('2'); + auto aot_dict = aot_manifest_dict(); + CacheEntry source, second, aot, unselected, unrelated, other_root; + ASSERT_TRUE(seed_cache(root.url, source_dict, fake.loader, source)); + ASSERT_TRUE(seed_cache(root.url, second_dict, fake.loader, second)); + ASSERT_TRUE(seed_cache(root.url, aot_dict, fake.loader, aot)); + ASSERT_TRUE(seed_cache(root.url, aot_dict, fake.loader, unselected, true, true, @"arch_a")); + ASSERT_TRUE(seed_cache(root.url, cache_manifest('3'), fake.loader, unrelated)); + ASSERT_TRUE(seed_cache(isolated.url, source_dict, fake.loader, other_root)); + NSDictionary* isolation = nil; + NSDictionary* unselected_stage = nil; + NSDictionary* unrelated_stage = nil; + ASSERT_TRUE(asset_tree_snapshot(isolated.url, isolation)); + ASSERT_TRUE(asset_tree_snapshot(unselected.stage, unselected_stage)); + ASSERT_TRUE(asset_tree_snapshot(unrelated.stage, unrelated_stage)); + FBSyntheticPTE pte( + {{"empty", {}}, + {"first", + {{"UnregisteredBackend", [@"not JSON" dataUsingEncoding:NSUTF8StringEncoding], true}, + {"CoreAIBackend", encode(source_dict)}, + {"CoreAIBackend", encode(second_dict), true}}}, + {"last", + {{"CoreAIBackend", encode(source_dict), true}, + {"CoreAIBackend", encode(aot_manifest_dict()), true}}}}); + NSMutableArray* evictions = [NSMutableArray array]; + fake.loader.onEvict = ^(NSData* bookmark) { + [evictions addObject:bookmark]; + }; + const int queries = bridge.state().architecture_queries; + EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); + drain_callbacks(); + EXPECT_EQ(bridge.state().architecture_queries.load(), queries + 1); + EXPECT_EQ(fake.loader->evictions.load(), 3); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(fake.loader->bindings.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + NSArray* ordered = + [@[ source.key, second.key, aot.key ] sortedArrayUsingSelector:@selector(compare:)]; + NSDictionary* bookmarks = + @{source.key : source.bookmark, second.key : second.bookmark, aot.key : aot.bookmark}; + ASSERT_EQ(evictions.count, ordered.count); + for (NSUInteger i = 0; i < ordered.count; ++i) { + SCOPED_TRACE(::testing::Message() << "sorted eviction index " << i); + EXPECT_TRUE([evictions[i] isEqual:bookmarks[ordered[i]]]); + } + cache_entry_removed(root.url, source); + cache_entry_removed(root.url, second); + cache_entry_removed(root.url, aot); + EXPECT_NE(fake.loader.models[other_root.bookmark], nil); + EXPECT_NE(fake.loader.models[unselected.bookmark], nil); + EXPECT_TRUE(snapshot_matches(isolated.url, isolation)); + EXPECT_TRUE(snapshot_matches(unselected.stage, unselected_stage)); + EXPECT_TRUE(snapshot_matches(unrelated.stage, unrelated_stage)); + EXPECT_EQ(pte.backend_indices, (std::vector{1, 2, 3})); + EXPECT_EQ(pte.backend_descriptors, + (std::vector{"CoreAIBackend", "CoreAIBackend", "CoreAIBackend"})); + EXPECT_TRUE(pte.check_released()); + } +} + +TEST_F(CoreAICacheTest, FilePathOverloadClearsReferencedEntry) { + @autoreleasepool { + BookmarkDirectory root; + BookmarkLoaderScope fake; + TestDirectory distribution; + ASSERT_NE(distribution.url, nil); + CacheEntry entry; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, entry)); + FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", encode(manifest_dict())}}}}); + NSURL* path = [distribution.url URLByAppendingPathComponent:@"model.pte"]; + ASSERT_TRUE([[NSData dataWithBytes:pte.bytes.data() length:pte.bytes.size()] writeToURL:path + options:0 + error:nil]); + EXPECT_EQ(clear_cache_for_pte(path.path.UTF8String, root.url.path.UTF8String), Error::Ok); + drain_callbacks(); + EXPECT_EQ(fake.loader->evictions.load(), 1); + cache_entry_removed(root.url, entry); + } +} + +TEST_F(CoreAICacheTest, RejectsInvalidLaterDelegatesBeforeMutation) { + @autoreleasepool { + const char* rows[] = {"malformed manifest", "wrong AOT platform", "backend load failure"}; + for (int row = 0; row < 3; ++row) { + SCOPED_TRACE(rows[row]); + BookmarkDirectory root; + BookmarkLoaderScope fake; + CacheEntry source; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, source)); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(root.url, snapshot)); + PteDelegateSpec later{"CoreAIBackend", encode(aot_manifest_dict()), true}; + if (row == 0) later.data = [@"malformed" dataUsingEncoding:NSUTF8StringEncoding]; + if (row == 1) { + auto dict = aot_manifest_dict(); + dict[@"platform"] = @"iOS"; + later.data = encode(dict); + } + FBSyntheticPTE pte( + {{"first", {{"CoreAIBackend", encode(manifest_dict())}}}, {"later", {later}}}); + if (row == 2) { + pte.fault_call = 2; + pte.load_fault = FBSyntheticPTE::LoadFault::Error; + pte.load_error = Error::AccessFailed; + } + NSDictionary* models = [fake.loader.models copy]; + const auto error = clear_cache_for_pte(pte, root.url.path.UTF8String); + drain_callbacks(); + EXPECT_NE(error, Error::Ok); + if (row == 2) EXPECT_EQ(error, Error::AccessFailed); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(fake.loader->bindings.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_TRUE(snapshot_matches(root.url, snapshot)); + EXPECT_TRUE([fake.loader.models isEqual:models]); + EXPECT_TRUE(pte.check_released()); + } + } +} + +TEST_F(CoreAICacheTest, ProcessedBuffersRemainAliveThroughEviction) { + @autoreleasepool { + BookmarkDirectory root; + CacheEntry source; + NSData* manifest = encode(manifest_dict()); + FBSyntheticPTE valid({{"forward", + {{"CoreAIBackend", manifest, true}, + {"CoreAIBackend", manifest, true}, + {"CoreAIBackend", manifest}}}}); + // Destroy the loader scope before the PTE captured by its callback. + BookmarkLoaderScope fake; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, source)); + __block int observations = 0; + __block bool program_alive = false; + __block int backend_releases = -1; + FBSyntheticPTE* observed = &valid; + fake.loader.onEvict = ^(NSData*) { + ++observations; + program_alive = observed->program_loads.load() == observed->program_releases.load() + 1; + backend_releases = observed->backend_releases.load(); + }; + EXPECT_EQ(clear_cache_for_pte(valid, root.url.path.UTF8String), Error::Ok); + drain_callbacks(); + EXPECT_EQ(observations, 1); + EXPECT_TRUE(program_alive); + EXPECT_EQ(backend_releases, 0); + EXPECT_EQ(valid.backend_loads.load(), 2); + EXPECT_EQ(valid.backend_indices, (std::vector{0, 1})); + EXPECT_TRUE(valid.program_alive_at_backend_release); + EXPECT_TRUE(valid.check_released()); + ASSERT_FALSE(valid.release_order.empty()); + EXPECT_EQ(valid.release_order.back(), SegmentType::Program); + } +} + +TEST_F(CoreAICacheTest, InventoriesBookmarkAndStagingUnionPreservingUnownedFiles) { + @autoreleasepool { + BookmarkDirectory root; + BookmarkLoaderScope fake; + CacheEntry both, bookmark_only, stage_only; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, both)); + ASSERT_TRUE(seed_cache(root.url, cache_manifest('2'), fake.loader, bookmark_only, true, false)); + ASSERT_TRUE(seed_cache(root.url, cache_manifest('3'), fake.loader, stage_only, false, true)); + NSURL* temporary = [root.url URLByAppendingPathComponent:@"staging/.staging-unkeyed"]; + ASSERT_TRUE([NSFileManager.defaultManager createDirectoryAtURL:temporary + withIntermediateDirectories:NO + attributes:nil + error:nil]); + NSURL* untouched = [temporary URLByAppendingPathComponent:@"partial"]; + ASSERT_TRUE([[@"do not infer ownership" dataUsingEncoding:NSUTF8StringEncoding] + writeToURL:untouched + options:0 + error:nil]); + NSURL* external_sentinel = [root.url URLByAppendingPathComponent:@"caller.pte"]; + ASSERT_TRUE([[@"PTE" dataUsingEncoding:NSUTF8StringEncoding] writeToURL:external_sentinel + options:0 + error:nil]); + NSDictionary* temp_snapshot = nil; + NSDictionary* locks = nil; + ASSERT_TRUE(asset_tree_snapshot(temporary, temp_snapshot)); + ASSERT_TRUE(asset_tree_snapshot([root.url URLByAppendingPathComponent:@"locks"], locks)); + const int queries = bridge.state().architecture_queries; + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + EXPECT_EQ(fake.loader->evictions.load(), 2); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(bridge.state().architecture_queries.load(), queries); + cache_entry_removed(root.url, both); + cache_entry_removed(root.url, bookmark_only); + cache_entry_removed(root.url, stage_only); + EXPECT_TRUE(snapshot_matches([root.url URLByAppendingPathComponent:@"locks"], locks)); + EXPECT_TRUE(snapshot_matches(temporary, temp_snapshot)); + EXPECT_TRUE(cache_exists(external_sentinel)); + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + EXPECT_EQ(fake.loader->evictions.load(), 2); + } +} + +TEST_F(CoreAICacheTest, RemovesStagingOnlyEntriesWithoutSdkAvailability) { + @autoreleasepool { + BookmarkDirectory root; + BookmarkLoaderScope fake; + CacheEntry stage_only; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, stage_only, false, true)); + bridge.state().available = false; + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + cache_entry_removed(root.url, stage_only); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + } +} + +TEST_F(CoreAICacheTest, NilLoaderFactoryPreservesEntriesAndAllowsRetry) { + @autoreleasepool { + ASSERT_TRUE(bridge.state().available); + ASSERT_FALSE(bridge.state().fail_loader_factory); + BookmarkDirectory root; + BookmarkLoaderScope fake; + CacheEntry entry; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, entry)); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(root.url, snapshot)); + bridge.state().fail_loader_factory = true; + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Internal); + bridge.state().fail_loader_factory = false; + EXPECT_TRUE(snapshot_matches(root.url, snapshot)); + EXPECT_NE(fake.loader.models[entry.bookmark], nil); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(fake.loader->bindings.load(), 0); + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + cache_entry_removed(root.url, entry); + } +} + +TEST_F(CoreAICacheTest, UnavailableSdkPreservesTrackedEntriesAndAllowsRetry) { + @autoreleasepool { + BookmarkDirectory root; + BookmarkLoaderScope fake; + CacheEntry entry; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, entry)); + NSDictionary* snapshot = nil; + ASSERT_TRUE(asset_tree_snapshot(root.url, snapshot)); + bridge.state().available = false; + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::NotSupported); + bridge.state().available = true; + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_TRUE(snapshot_matches(root.url, snapshot)); + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + cache_entry_removed(root.url, entry); + } +} + +TEST_F(CoreAICacheTest, SourceRemovalFailuresKeepBookmarkForRetry) { + @autoreleasepool { + const struct { + int match_index; + bool bookmark_removal; + const char* name; + } rows[] = {{2, false, "staged-tree removal"}, {6, true, "bookmark removal"}}; + for (const auto& row : rows) { + SCOPED_TRACE(row.name); + BookmarkDirectory root; + BookmarkLoaderScope fake; + CacheEntry entry; + ASSERT_TRUE(seed_cache(root.url, aot_manifest_dict(), fake.loader, entry)); + { + StorageFaultScope fault(StorageOperation::Remove); + fault.match_index = row.match_index; + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + } + drain_callbacks(); + auto saved = saved_bookmark(root.url, entry.key); + ASSERT_TRUE(saved.ok()); + EXPECT_TRUE([saved.get() isEqual:entry.bookmark]); + EXPECT_EQ(cache_exists(entry.stage), !row.bookmark_removal); + if (fake.loader.models[entry.bookmark] == nil) { + fake.loader.evictError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:ENOENT + userInfo:nil]; + } + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + cache_entry_removed(root.url, entry); + } + } +} + +TEST_F(CoreAICacheTest, IndependentKeysContinueInSortedOrderAfterFailures) { + @autoreleasepool { + BookmarkDirectory root; + BookmarkLoaderScope fake; + CacheEntry first, second, third; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, first)); + ASSERT_TRUE(seed_cache(root.url, cache_manifest('2'), fake.loader, second)); + ASSERT_TRUE(seed_cache(root.url, cache_manifest('3'), fake.loader, third)); + NSArray* ordered = + [@[ first.key, second.key, third.key ] sortedArrayUsingSelector:@selector(compare:)]; + NSDictionary* bookmarks = + @{first.key : first.bookmark, second.key : second.bookmark, third.key : third.bookmark}; + NSMutableArray* evicted = [NSMutableArray array]; + BookmarkFakeLoader* loader = fake.loader; + fake.loader.onEvict = ^(NSData* bookmark) { + [evicted addObject:bookmark]; + loader.evictError = [bookmark isEqual:bookmarks[ordered[1]]] + ? [NSError errorWithDomain:NSPOSIXErrorDomain code:EBUSY userInfo:nil] + : nil; + }; + StorageFaultScope fault(StorageOperation::BeforeEvict); + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::AccessFailed); + drain_callbacks(); + EXPECT_EQ(fault.hits, 1); + ASSERT_EQ(evicted.count, 2u); + EXPECT_TRUE([evicted[0] isEqual:bookmarks[ordered[1]]]); + EXPECT_TRUE([evicted[1] isEqual:bookmarks[ordered[2]]]); + EXPECT_TRUE(cache_exists(test_bookmark_url(root.url, ordered[0]))); + EXPECT_TRUE(cache_exists(test_bookmark_url(root.url, ordered[1]))); + EXPECT_FALSE(cache_exists(test_bookmark_url(root.url, ordered[2]))); + } +} + +} // namespace +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_eviction_test.mm b/backends/apple/coreai/runtime/test/coreai_eviction_test.mm new file mode 100644 index 00000000000..619da5e1754 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_eviction_test.mm @@ -0,0 +1,402 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include "coreai_bookmark_fixture.h" +#include "coreai_fake_loader.h" +#include "coreai_fault_scope.h" +#include "coreai_load_coordinator.h" +#include "coreai_source_fixture.h" + +using namespace executorch::runtime; +using namespace executorch::backends::coreai; +using namespace executorch::backends::coreai::testing; + +namespace { +NSError* bookmark_test_error(NSInteger code = EIO) { + return [NSError errorWithDomain:NSPOSIXErrorDomain code:code userInfo:nil]; +} + +// These cases never retain a delegate handle beyond initialization. +Error initialize_and_destroy(NSDictionary* dict, const NamedDataMap* data, NSString* root) { + if (root == nil) return Error::InvalidArgument; + BackendOptions<1> options; + auto error = options.set_option("coreai_assets_dir", root.UTF8String); + if (error != Error::Ok) return error; + auto* backend = get_backend_class("CoreAIBackend"); + if (backend == nullptr) return Error::Internal; + BackendInitContext context(nullptr, nullptr, "forward", data, + {options.view().data(), options.view().size()}); + NSData* json = encode(dict); + if (json == nil) return Error::InvalidArgument; + FreeableBuffer processed(json.bytes, json.length, nullptr); + auto handle = backend->init(context, &processed, {}); + if (!handle.ok()) return handle.error(); + backend->destroy(handle.get()); + return Error::Ok; +} + +dispatch_time_t deadline() { return dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC); } + +class CoreAIEvictionTest : public ::testing::Test { + protected: + ScopedFakeBridgeState bridge; + void SetUp() override { runtime_init(); } +}; + +// Release the callback before waiting; never unwind storage borrowed by a hang. +struct ReleaseAndJoin final { + dispatch_semaphore_t gate; + dispatch_semaphore_t eviction_done; + dispatch_semaphore_t loading_done; + ScopedFakeBridgeState& bridge; + pthread_t eviction{}; + pthread_t loading{}; + bool eviction_started = false; + bool loading_started = false; + bool released = false; + bool joined = false; + long eviction_wait = -1; + long loading_wait = -1; + + void release() { + if (!released) { + dispatch_semaphore_signal(gate); + released = true; + } + } + void join() { + release(); + if (joined) return; + const auto completion_deadline = deadline(); + if (eviction_started) + eviction_wait = dispatch_semaphore_wait(eviction_done, completion_deadline); + if (loading_started) loading_wait = dispatch_semaphore_wait(loading_done, completion_deadline); + if ((eviction_started && eviction_wait != 0) || (loading_started && loading_wait != 0)) { + ADD_FAILURE() << "Eviction/load workers did not return after releasing the gate"; + std::_Exit(EXIT_FAILURE); + } + const int eviction_join = eviction_started ? pthread_join(eviction, nullptr) : 0; + const int loading_join = loading_started ? pthread_join(loading, nullptr) : 0; + if (eviction_join != 0 || loading_join != 0) { + ADD_FAILURE() << "Cannot join eviction/load workers: " << eviction_join << ", " + << loading_join; + std::_Exit(EXIT_FAILURE); + } + auto drained = bridge.wait_for_callbacks(deadline()); + if (!drained) { + ADD_FAILURE() << drained.message(); + std::_Exit(EXIT_FAILURE); + } + joined = true; + } + ~ReleaseAndJoin() { join(); } +}; +} // namespace + +@interface PinBookmarkLoader : BookmarkFakeLoader +@property(nonatomic, weak) BookmarkFakeModel* temporary; +@property(nonatomic, copy) ETCoreAIAcquisitionCompletion retainedCompletion; +@end + +@implementation PinBookmarkLoader +- (void)specializeModelAtURL:(NSURL*)url completion:(ETCoreAIAcquisitionCompletion)completion { + ++specializations; + self.retainedCompletion = completion; + @autoreleasepool { + BookmarkFakeModel* model = [[BookmarkFakeModel alloc] init]; + model.owner = self; + model.bookmark = [@"pinned-bookmark" dataUsingEncoding:NSUTF8StringEncoding]; + self.temporary = model; + completion(ETCoreAIAcquisitionStatusHit, model, nil); + model = nil; + } +} +- (void)restoreModelFromBookmark:(NSData*)bookmark + completion:(ETCoreAIAcquisitionCompletion)completion { + ++restores; + self.retainedCompletion = completion; + @autoreleasepool { + BookmarkFakeModel* model = self.temporary; + completion(model == nil ? ETCoreAIAcquisitionStatusMiss : ETCoreAIAcquisitionStatusHit, model, + nil); + model = nil; + } +} +- (void)evictModelWithBookmark:(NSData*)bookmark completion:(void (^)(NSError*))completion { + ++evictions; + bool busy; + @autoreleasepool { + busy = self.temporary != nil; + } + completion(busy ? bookmark_test_error(EBUSY) : nil); +} +- (void)clearRetainedState { + self.retainedCompletion = nil; + [super clearRetainedState]; +} +@end + +namespace { +TEST_F(CoreAIEvictionTest, ConfirmedMissRecoveryEvictsOnlyCurrentBookmark) { + @autoreleasepool { + BookmarkDirectory root; + ASSERT_NE(root.url, nil); + BookmarkLoaderScope fake; + ASSERT_NE(fake.loader, nil); + TestData data; + auto selected = bookmark_manifest(); + ASSERT_TRUE(selected.ok()); + auto key_result = test_bookmark_key(selected.get()); + ASSERT_TRUE(key_result.ok()); + NSString* key = key_result.get(); + NSMutableDictionary* dict = manifest_dict(); + ASSERT_EQ(initialize_and_destroy(dict, &data, root.url.path), Error::Ok); + auto original_result = saved_bookmark(root.url, key); + ASSERT_TRUE(original_result.ok()); + NSData* original = original_result.get(); + ASSERT_NE(original, nil); + NSString* source_path = fake.loader.lastSpecializeURL; + ASSERT_NE(source_path, nil); + NSURL* source_url = [NSURL fileURLWithPath:source_path]; + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + + fake.loader.restoreMiss = YES; + fake.loader.specializeError = bookmark_test_error(); + EXPECT_NE(initialize_and_destroy(dict, &data, root.url.path), Error::Ok); + auto after_failure = saved_bookmark(root.url, key); + ASSERT_TRUE(after_failure.ok()); + EXPECT_TRUE([after_failure.get() isEqual:original]); + EXPECT_TRUE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + fake.loader.specializeError = nil; + ASSERT_EQ(initialize_and_destroy(dict, &data, root.url.path), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 3); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_TRUE([fake.loader.lastSpecializeURL isEqual:source_path]); + auto saved_result = saved_bookmark(root.url, key); + ASSERT_TRUE(saved_result.ok()); + NSData* saved = saved_result.get(); + ASSERT_NE(saved, nil); + EXPECT_FALSE([saved isEqual:original]); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + ASSERT_EQ(evict_bookmark_model(root.url.path, key, fake.loader), Error::Ok); + EXPECT_NE(fake.loader.models[original], nil); + EXPECT_EQ(fake.loader.models[saved], nil); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + fake.loader.restoreMiss = NO; + fake.loader.emptyBookmark = YES; + EXPECT_NE(initialize_and_destroy(dict, &data, root.url.path), Error::Ok); + auto empty = saved_bookmark(root.url, key); + ASSERT_TRUE(empty.ok()); + EXPECT_EQ(empty.get(), nil); + fake.loader.emptyBookmark = NO; + ASSERT_EQ(initialize_and_destroy(dict, &data, root.url.path), Error::Ok); + EXPECT_EQ(fake.loader->specializations.load(), 5); + EXPECT_EQ(data.requests.load(), data.releases.load()); + } +} + +TEST_F(CoreAIEvictionTest, BusyAndUnknownEvictionRetainBookmarkUntilConfirmedAbsent) { + @autoreleasepool { + BookmarkDirectory root; + ASSERT_NE(root.url, nil); + BookmarkLoaderScope fake; + ASSERT_NE(fake.loader, nil); + TestData data; + auto selected = bookmark_manifest(); + ASSERT_TRUE(selected.ok()); + auto key_result = test_bookmark_key(selected.get()); + ASSERT_TRUE(key_result.ok()); + NSString* key = key_result.get(); + EXPECT_EQ(evict_bookmark_model(root.url.path, key, fake.loader), Error::Ok); + EXPECT_EQ(fake.loader->evictions.load(), 0); + ASSERT_EQ(initialize_and_destroy(manifest_dict(), &data, root.url.path), Error::Ok); + auto saved_result = saved_bookmark(root.url, key); + ASSERT_TRUE(saved_result.ok()); + NSData* saved = saved_result.get(); + ASSERT_NE(saved, nil); + NSURL* source_url = + [[[root.url URLByAppendingPathComponent:@"staging"] URLByAppendingPathComponent:key] + URLByAppendingPathComponent:selected->path.lastPathComponent]; + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + for (bool probe_error : {false, true}) { + SCOPED_TRACE(probe_error ? "unknown restore failure" : "restore hit/busy"); + fake.loader.evictError = bookmark_test_error(EBUSY); + fake.loader.restoreError = probe_error ? bookmark_test_error() : nil; + EXPECT_NE(evict_bookmark_model(root.url.path, key, fake.loader), Error::Ok); + auto retained = saved_bookmark(root.url, key); + ASSERT_TRUE(retained.ok()); + EXPECT_TRUE([retained.get() isEqual:saved]); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + } + fake.loader.restoreError = nil; + fake.loader.evictError = nil; + { + StorageFaultScope fault(StorageOperation::AfterEvict); + EXPECT_NE(evict_bookmark_model(root.url.path, key, fake.loader), Error::Ok); + EXPECT_EQ(fault.hits, 1); + } + auto retained = saved_bookmark(root.url, key); + ASSERT_TRUE(retained.ok()); + EXPECT_TRUE([retained.get() isEqual:saved]); + fake.loader.evictError = bookmark_test_error(ENOENT); + ASSERT_EQ(evict_bookmark_model(root.url.path, key, fake.loader), Error::Ok); + auto absent = saved_bookmark(root.url, key); + ASSERT_TRUE(absent.ok()); + EXPECT_EQ(absent.get(), nil); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + fake.loader.evictError = nil; + ASSERT_EQ(initialize_and_destroy(manifest_dict(), &data, root.url.path), Error::Ok); + { + StorageFaultScope fault(StorageOperation::Remove); + EXPECT_NE(evict_bookmark_model(root.url.path, key, fake.loader), Error::Ok); + EXPECT_EQ(fault.hits, 1); + auto before_remove = saved_bookmark(root.url, key); + ASSERT_TRUE(before_remove.ok()); + EXPECT_NE(before_remove.get(), nil); + } + ASSERT_EQ(evict_bookmark_model(root.url.path, key, fake.loader), Error::Ok); + auto removed = saved_bookmark(root.url, key); + ASSERT_TRUE(removed.ok()); + EXPECT_EQ(removed.get(), nil); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + EXPECT_EQ(data.requests.load(), data.releases.load()); + } +} + +TEST_F(CoreAIEvictionTest, PreparedPinsBlockEvictionWithoutRetainedCompletionLeak) { + @autoreleasepool { + BookmarkDirectory root; + ASSERT_NE(root.url, nil); + TestData data; + auto selected = bookmark_manifest(); + ASSERT_TRUE(selected.ok()); + auto key_result = test_bookmark_key(selected.get()); + ASSERT_TRUE(key_result.ok()); + NSString* key = key_result.get(); + PinBookmarkLoader* loader = [[PinBookmarkLoader alloc] init]; + ASSERT_NE(loader, nil); + // Scope teardown clears retained completions even after fatal test + // failures. + bridge.state().bookmark_loader = loader; + @autoreleasepool { + auto pinned = + acquire_bookmark_model(selected.get(), &data, root.url.path, @"macOS", @"arch_b", loader); + ASSERT_TRUE(pinned.ok()); + EXPECT_NE(loader.temporary, nil); + EXPECT_NE(evict_bookmark_model(root.url.path, key, loader), Error::Ok); + auto saved = saved_bookmark(root.url, key); + ASSERT_TRUE(saved.ok()); + EXPECT_NE(saved.get(), nil); + } + EXPECT_EQ(loader.temporary, nil); + EXPECT_NE(loader.retainedCompletion, nil); + ASSERT_EQ(evict_bookmark_model(root.url.path, key, loader), Error::Ok); + auto removed = saved_bookmark(root.url, key); + ASSERT_TRUE(removed.ok()); + EXPECT_EQ(removed.get(), nil); + EXPECT_EQ(loader->evictions.load(), 2); + NSURL* source_url = + [[[root.url URLByAppendingPathComponent:@"staging"] URLByAppendingPathComponent:key] + URLByAppendingPathComponent:selected->path.lastPathComponent]; + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + [loader clearRetainedState]; + EXPECT_EQ(data.requests.load(), data.releases.load()); + } +} + +TEST_F(CoreAIEvictionTest, LoadContendsOnActualKeyLockUntilEvictionFinishes) { + @autoreleasepool { + BookmarkDirectory root; + ASSERT_NE(root.url, nil); + BookmarkLoaderScope fake; + ASSERT_NE(fake.loader, nil); + TestData data; + auto selected = bookmark_manifest(); + ASSERT_TRUE(selected.ok()); + auto key_result = test_bookmark_key(selected.get()); + ASSERT_TRUE(key_result.ok()); + NSString* key = key_result.get(); + ASSERT_EQ(initialize_and_destroy(manifest_dict(), &data, root.url.path), Error::Ok); + dispatch_semaphore_t deleting = dispatch_semaphore_create(0); + dispatch_semaphore_t proceed = dispatch_semaphore_create(0); + dispatch_semaphore_t loading_observed = dispatch_semaphore_create(0); + dispatch_semaphore_t eviction_done = dispatch_semaphore_create(0); + dispatch_semaphore_t loading_done = dispatch_semaphore_create(0); + __block long gate_wait = -1; + fake.loader.onEvict = ^(NSData*) { + dispatch_semaphore_signal(deleting); + gate_wait = dispatch_semaphore_wait(proceed, deadline()); + }; + Error eviction_error = Error::Internal; + bool load_ok = false; + auto evict = [&] { + @autoreleasepool { + eviction_error = evict_bookmark_model(root.url.path, key, fake.loader); + } + dispatch_semaphore_signal(eviction_done); + }; + auto load = [&] { + @autoreleasepool { + auto acquired = acquire_bookmark_model(selected.get(), &data, root.url.path, @"macOS", + @"arch_b", fake.loader); + load_ok = acquired.ok() && acquired.get() != nil; + } + dispatch_semaphore_signal(loading_observed); + dispatch_semaphore_signal(loading_done); + }; + ReleaseAndJoin cleanup{proceed, eviction_done, loading_done, bridge}; + const int eviction_start = pthread_create( + &cleanup.eviction, nullptr, + [](void* context) -> void* { + (*static_cast(context))(); + return nullptr; + }, + &evict); + cleanup.eviction_started = eviction_start == 0; + const long entered = + cleanup.eviction_started ? dispatch_semaphore_wait(deleting, deadline()) : -1; + int loading_start = -1; + long premature_completion = -1; + if (entered == 0) { + loading_start = pthread_create( + &cleanup.loading, nullptr, + [](void* context) -> void* { + (*static_cast(context))(); + return nullptr; + }, + &load); + cleanup.loading_started = loading_start == 0; + if (cleanup.loading_started) { + premature_completion = dispatch_semaphore_wait( + loading_observed, dispatch_time(DISPATCH_TIME_NOW, 50 * NSEC_PER_MSEC)); + } + } + cleanup.join(); + fake.loader.onEvict = nil; + + EXPECT_EQ(eviction_start, 0); + EXPECT_EQ(loading_start, 0); + EXPECT_EQ(entered, 0); + EXPECT_NE(premature_completion, 0); + EXPECT_EQ(gate_wait, 0); + EXPECT_EQ(cleanup.eviction_wait, 0); + EXPECT_EQ(cleanup.loading_wait, 0); + EXPECT_EQ(eviction_error, Error::Ok); + EXPECT_TRUE(load_ok); + EXPECT_EQ(fake.loader->specializations.load(), 2); + auto saved = saved_bookmark(root.url, key); + ASSERT_TRUE(saved.ok()); + EXPECT_NE(saved.get(), nil); + EXPECT_EQ(data.requests.load(), data.releases.load()); + } +} +} // namespace diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.h b/backends/apple/coreai/runtime/test/coreai_fake_loader.h index 935fa036179..faefbf15f94 100644 --- a/backends/apple/coreai/runtime/test/coreai_fake_loader.h +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.h @@ -94,9 +94,11 @@ class ScopedFakeBridgeState final { @property(nonatomic, strong) NSError* restoreError; @property(nonatomic, strong) NSError* specializeError; @property(nonatomic, strong) NSError* bindError; +@property(nonatomic, strong) NSError* evictError; @property(nonatomic) BOOL restoreMiss; @property(nonatomic) BOOL emptyBookmark; @property(nonatomic, copy) void (^onSpecialize)(NSURL*); +@property(nonatomic, copy) void (^onEvict)(NSData*); @property(nonatomic, strong) BookmarkFakeModel* lastAcquired; @property(nonatomic, strong) BookmarkFakeModel* lastBound; @property(nonatomic, copy) NSString* lastSpecializeURL; diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm index 50878e915b9..d53b6542748 100644 --- a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm @@ -394,11 +394,20 @@ - (void)evictModelWithBookmark:(NSData*)bookmark completion:(void (^)(NSError*))completion { ++evictions; run_async(self.context, ^{ - completion(fake_error(ETCoreAIErrorUnsupported)); + if (self.onEvict != nil) { + self.onEvict(bookmark); + } + @synchronized(self) { + if (self.evictError == nil) { + [self.models removeObjectForKey:bookmark]; + } + completion(self.evictError); + } }); } - (void)clearRetainedState { self.onSpecialize = nil; + self.onEvict = nil; self.lastAcquired = nil; self.lastBound = nil; [self.models removeAllObjects]; diff --git a/backends/apple/coreai/runtime/test/runtime_smoke.cpp b/backends/apple/coreai/runtime/test/runtime_smoke.cpp index 3d91a760b4b..a136dceb1df 100644 --- a/backends/apple/coreai/runtime/test/runtime_smoke.cpp +++ b/backends/apple/coreai/runtime/test/runtime_smoke.cpp @@ -6,13 +6,34 @@ * LICENSE file in the root directory of this source tree. */ +#include #include #include #include +namespace { +using executorch::runtime::DataLoader; +using executorch::runtime::Error; +using RootClear = Error (*)(const char*); +using FileClear = Error (*)(const char*, const char*); +using LoaderClear = Error (*)(DataLoader&, const char*); + +// Keep every overload in the link check without touching cache storage. +RootClear volatile clear_root = &executorch::backends::coreai::clear_cache; +FileClear volatile clear_file = + &executorch::backends::coreai::clear_cache_for_pte; +LoaderClear volatile clear_loader = + &executorch::backends::coreai::clear_cache_for_pte; +} // namespace + int main() { executorch::runtime::runtime_init(); + if (clear_root == nullptr || clear_file == nullptr || + clear_loader == nullptr) { + std::fputs("Core AI cache APIs were not linked\n", stderr); + return 1; + } const auto* backend = executorch::runtime::get_backend_class("CoreAIBackend"); if (backend == nullptr) { std::fputs("CoreAIBackend was not retained by the linker\n", stderr); From b18e831b3f4c5835777fe9d9e8d6e1510637f690 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Mon, 26 Oct 2026 17:54:33 -0700 Subject: [PATCH 2/5] [CoreAI] Align cache clearing with loading and cover production eviction Summary: Prevent loading and maintenance from drifting in how they select model assets, and let applications clear the default cache without reconstructing a private path. Exercise real clearing behavior for eviction failures and loaded models while preserving strict compatibility checks and storage validation. Simplify the public maintenance documentation. Test Plan: Built the CoreAI targets and ran all four CTest suites successfully, including native host tests, Swift bridge tests, SDK source-lifetime tests, and runtime smoke coverage. New tests verify portable/AOT load-to-clear identity, isolated default-root clearing, rejection before mutation, and production eviction failure/retry behavior. Ran the complete CoreAI Python suite: 113 passed. Whitespace checks passed. The generic validator could not complete its Buck ownership query; native and Swift validation passed through CMake. # Conflicts: # backends/apple/coreai/runtime/coreai_storage.h # backends/apple/coreai/runtime/coreai_storage.mm # Please enter the commit message for your changes. Lines starting # with '#' will be kept; you may remove them yourself if you want to. # An empty message aborts the commit. # # interactive rebase in progress; onto f13d818a91 # Last commands done (2 commands done): # pick 2fef739086 # [CoreAI] Add root and PTE cache clearing # pick 5b46516cd7 # [CoreAI] Align cache clearing with loading and cover production eviction # No commands remaining. # You are currently rebasing branch 'coreai-v2/07-cache-clear' on 'f13d818a91'. # # Changes to be committed: # modified: backends/apple/coreai/README.md # modified: backends/apple/coreai/runtime/coreai_backend.mm # modified: backends/apple/coreai/runtime/coreai_cache.h # modified: backends/apple/coreai/runtime/coreai_cache.mm # modified: backends/apple/coreai/runtime/coreai_load_coordinator.h # modified: backends/apple/coreai/runtime/coreai_load_coordinator.mm # modified: backends/apple/coreai/runtime/coreai_storage.h # modified: backends/apple/coreai/runtime/coreai_storage.mm # modified: backends/apple/coreai/runtime/test/coreai_cache_test.mm # modified: backends/apple/coreai/runtime/test/coreai_delegate_test.mm # modified: backends/apple/coreai/runtime/test/coreai_eviction_test.mm # # Changes not staged for commit: # modified: backends/mlx/third-party/mlx (modified content) # --- backends/apple/coreai/README.md | 23 ++-- .../apple/coreai/runtime/coreai_backend.mm | 22 +--- backends/apple/coreai/runtime/coreai_cache.h | 21 ++- backends/apple/coreai/runtime/coreai_cache.mm | 47 +++---- .../coreai/runtime/coreai_load_coordinator.h | 17 ++- .../coreai/runtime/coreai_load_coordinator.mm | 44 +++++-- .../apple/coreai/runtime/coreai_storage.h | 2 + .../apple/coreai/runtime/coreai_storage.mm | 18 ++- .../coreai/runtime/test/coreai_cache_test.mm | 21 ++- .../runtime/test/coreai_delegate_test.mm | 122 ++++++++++++++++++ .../runtime/test/coreai_eviction_test.mm | 42 ++++-- 11 files changed, 264 insertions(+), 115 deletions(-) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index b2f0e677846..e73c1bb4bf9 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -230,30 +230,25 @@ legacy files. Bookmarks and staged bundles need no migration. The public C++ maintenance APIs are in `executorch::backends::coreai`: ```cpp -runtime::Error clear_cache(const char* coreai_assets_dir); +runtime::Error clear_cache(const char* coreai_assets_dir = nullptr); runtime::Error clear_cache_for_pte( runtime::DataLoader& pte, const char* coreai_assets_dir = nullptr); runtime::Error clear_cache_for_pte( const char* pte_path, const char* coreai_assets_dir = nullptr); ``` -`clear_cache` requires an explicit absolute assets root. The PTE overloads use -user-domain `NSCachesDirectory` plus `executorch_coreai` when the directory is -omitted or null, just like backend initialization. An invalid supplied path is -an error. If loading used a custom root, pass that same root when clearing; -the PTE does not store runtime directory overrides. +All three APIs use the same default assets directory as backend initialization +when the directory is omitted or null. A supplied path must be absolute; invalid +paths fail without falling back to the default. If loading used a custom root, +pass that same root when clearing; the PTE does not store runtime overrides. Root-wide clearing visits keyed bookmarks and staged directories. PTE clearing selects and deduplicates keys for all Core AI delegates in every method, using the current platform, SDK architecture and default specialization configuration. -Both public `clear_cache_for_pte` overloads use the private `inspect_coreai_pte` -reader and retain its `CoreAIPteData`, including the verified program image and -selected processed buffers, through `clear_keys`. Discovery does not load a -Program or initialize a Method, load SDK models, bind functions, extract Core AI -assets or specialize. All selected manifests are validated before any key is -cleared; root-wide clearing remains -available for existing keyed entries. Backend-local structural and semantic -verification is mandatory even with `ET_ENABLE_PROGRAM_VERIFICATION=0`. +All selected manifests must be valid and compatible before any entry is cleared. +Incompatible delegates are not silently skipped. Root-wide clearing remains +available independently of the PTE. Discovery does not load models or materialize +assets. These synchronous calls require callers to unload affected models and prevent concurrent loads, inference and maintenance across processes. Per-key locks are diff --git a/backends/apple/coreai/runtime/coreai_backend.mm b/backends/apple/coreai/runtime/coreai_backend.mm index 15e651564a1..9295130d041 100644 --- a/backends/apple/coreai/runtime/coreai_backend.mm +++ b/backends/apple/coreai/runtime/coreai_backend.mm @@ -12,7 +12,6 @@ #import "coreai_load_coordinator.h" #import "coreai_storage.h" -#include #include #include #include @@ -356,24 +355,11 @@ bool is_available() const override { if (!manifest.ok()) { return manifest.error(); } - ET_CHECK_OR_RETURN_ERROR( - [NSProcessInfo.processInfo - isOperatingSystemAtLeastVersion:manifest->minimum_version], - DelegateInvalidCompatibility, - "Core AI model requires a newer operating system"); -#if TARGET_OS_OSX - NSString* platform = @"macOS"; -#elif TARGET_OS_IOS - NSString* platform = @"iOS"; -#else - NSString* platform = @"unsupported"; -#endif - NSString* architecture = ETCoreAIDeviceArchitectureName(); - auto selected = select_assets(manifest.get(), architecture, platform); + auto selected = select_runtime_assets(manifest.get()); if (!selected.ok()) { return selected.error(); } - manifest.get() = std::move(selected.get()); + manifest.get() = std::move(selected->manifest); auto assets_option = path_option(context, "coreai_assets_dir"); if (!assets_option.ok()) { return assets_option.error(); @@ -396,8 +382,8 @@ bool is_available() const override { manifest.get(), context.get_named_data_map(), coreai_assets_dir, - platform, - architecture, + selected->platform, + selected->architecture, loader); if (!prepared.ok()) { return prepared.error(); diff --git a/backends/apple/coreai/runtime/coreai_cache.h b/backends/apple/coreai/runtime/coreai_cache.h index c87dd6f9daa..2e1d28bbf28 100644 --- a/backends/apple/coreai/runtime/coreai_cache.h +++ b/backends/apple/coreai/runtime/coreai_cache.h @@ -16,8 +16,9 @@ class DataLoader; namespace executorch::backends::coreai { /** - * Clears tracked SDK entries, keyed staging, and bookmarks in an explicit, - * dedicated Core AI assets directory. Missing roots are successful no-ops. + * Clears tracked SDK entries, keyed staging, and bookmarks in a dedicated + * Core AI assets directory. Null selects the same default as loading; supplied + * paths must be absolute. Missing roots are successful no-ops. * PTEs, stable lock files, and unkeyed temporaries are not removed. * SDK entries without saved bookmarks cannot be targeted. * @@ -27,23 +28,21 @@ namespace executorch::backends::coreai { * clearing is not a transaction. SDK failures preserve an entry's files unless * the SDK confirms that entry is absent. */ -[[nodiscard]] runtime::Error clear_cache(const char* coreai_assets_dir); +[[nodiscard]] runtime::Error clear_cache( + const char* coreai_assets_dir = nullptr); /** * Clears current-platform/SDK-architecture cache entries referenced by every - * CoreAI delegate in the PTE, after validating all selected manifests. Does not - * initialize delegates or materialize Core AI assets. Backend-local structural - * and semantic verification is mandatory regardless of runtime build settings. - * Reads the program region, including inline constants and unrelated inline - * blobs, and only selected Core AI external processed-data segments. + * CoreAI delegate in the PTE, after validating all selected manifests. Invalid + * or incompatible delegates fail before any entries are cleared. Does not + * initialize delegates or materialize Core AI assets. * * A null directory selects the same Caches default as backend initialization. * Supplied directories must be valid explicit roots; they never fall back to * the default. The directory is never derived from the PTE's location. * Copied PTEs can share cache keys. The quiescence and failure contract of - * clear_cache also applies. The loader and its data must remain stable for - * this call; the retained program and selected buffers survive through - * eviction. + * clear_cache also applies. Keep the loader alive and its data unchanged until + * the call returns. */ [[nodiscard]] runtime::Error clear_cache_for_pte( runtime::DataLoader& loader, diff --git a/backends/apple/coreai/runtime/coreai_cache.mm b/backends/apple/coreai/runtime/coreai_cache.mm index 20e07827f05..06c23a118d8 100644 --- a/backends/apple/coreai/runtime/coreai_cache.mm +++ b/backends/apple/coreai/runtime/coreai_cache.mm @@ -12,7 +12,6 @@ #include "coreai_pte.h" #import "coreai_storage.h" -#include #include namespace executorch::backends::coreai { @@ -20,8 +19,16 @@ using runtime::Error; using runtime::Result; -Result explicit_root(const char* path) { - NSString* root = path == nullptr ? nil : [NSString stringWithUTF8String:path]; +Result cache_root(const char* path) { + NSString* root = nil; + if (path == nullptr) { + auto default_root = default_coreai_assets_root(); + if (!default_root.ok()) + return default_root.error(); + root = default_root.get(); + } else { + root = [NSString stringWithUTF8String:path]; + } return resolve_bookmark_root(root, false); } @@ -46,13 +53,6 @@ Error clear_pte(runtime::DataLoader& loader, NSString* root) { if (!pte.ok()) return pte.error(); NSMutableArray* keys = [NSMutableArray array]; NSString* architecture = nil; -#if TARGET_OS_OSX - NSString* platform = @"macOS"; -#elif TARGET_OS_IOS - NSString* platform = @"iOS"; -#else - NSString* platform = @"unsupported"; -#endif for (const auto& buffer : pte->processed) { auto data = buffer.data_safe(); if (!data.ok()) return data.error(); @@ -63,19 +63,11 @@ Error clear_pte(runtime::DataLoader& loader, NSString* root) { freeWhenDone:NO]; auto manifest = parse_manifest(bytes); if (!manifest.ok()) return manifest.error(); - ET_CHECK_OR_RETURN_ERROR( - [NSProcessInfo.processInfo - isOperatingSystemAtLeastVersion:manifest->minimum_version], - DelegateInvalidCompatibility, - "Core AI model requires a newer operating system"); - if (architecture == nil) { - ET_CHECK_OR_RETURN_ERROR(ETCoreAIIsAvailable(), NotSupported, - "Core AI requires macOS 27 or iOS 27"); - architecture = ETCoreAIDeviceArchitectureName(); - } - auto selected = select_assets(manifest.get(), architecture, platform); + auto selected = select_runtime_assets(manifest.get(), architecture); if (!selected.ok()) return selected.error(); - auto key = bookmark_key(selected.get(), platform, architecture); + architecture = selected->architecture; + auto key = + bookmark_key(selected->manifest, selected->platform, architecture); if (!key.ok()) return key.error(); [keys addObject:key.get()]; } @@ -85,7 +77,7 @@ Error clear_pte(runtime::DataLoader& loader, NSString* root) { Error clear_cache(const char* coreai_assets_dir) { @autoreleasepool { - auto root = explicit_root(coreai_assets_dir); + auto root = cache_root(coreai_assets_dir); if (!root.ok()) return root.error(); auto keys = inventory_bookmark_keys(root.get()); if (!keys.ok()) return keys.error(); @@ -96,14 +88,7 @@ Error clear_cache(const char* coreai_assets_dir) { Error clear_cache_for_pte(runtime::DataLoader& loader, const char* coreai_assets_dir) { @autoreleasepool { - if (coreai_assets_dir != nullptr) { - auto root = explicit_root(coreai_assets_dir); - if (!root.ok()) return root.error(); - return clear_pte(loader, root.get()); - } - auto default_root = default_coreai_assets_root(); - if (!default_root.ok()) return default_root.error(); - auto root = resolve_bookmark_root(default_root.get(), false); + auto root = cache_root(coreai_assets_dir); if (!root.ok()) return root.error(); return clear_pte(loader, root.get()); } diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.h b/backends/apple/coreai/runtime/coreai_load_coordinator.h index eb41450c06b..cfd44efdf00 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.h +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.h @@ -12,6 +12,17 @@ namespace executorch::backends::coreai { +struct RuntimeAssets { + Manifest manifest; + NSString* platform; + NSString* architecture; +}; + +// Reuse a resolved architecture when selecting multiple delegates in one PTE. +runtime::Result select_runtime_assets( + const Manifest& manifest, + NSString* architecture = nil); + // The key lock outlives all SDK callbacks and raw bookmark publication. // Cleanup after a restore hit or durable publication is best-effort. runtime::Result> acquire_bookmark_model( @@ -22,12 +33,6 @@ runtime::Result> acquire_bookmark_model( NSString* architecture, id loader); -// Deletes only the saved current SDK entry, not historical entries or sources. -runtime::Error evict_bookmark_model( - NSString* root, - NSString* key, - id loader); - // Callers must keep affected loads and sessions quiescent across processes. runtime::Error clear_bookmark_assets( NSString* root, diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.mm b/backends/apple/coreai/runtime/coreai_load_coordinator.mm index 7f77776972a..35810caffff 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.mm +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.mm @@ -8,6 +8,9 @@ #import "coreai_load_coordinator.h" #import "coreai_storage.h" +#include +#include + namespace executorch::backends::coreai { namespace { using runtime::Error; @@ -104,6 +107,34 @@ Error evict_one(id loader, NSData* bookmark) { } } // namespace +Result select_runtime_assets( + const Manifest& manifest, + NSString* architecture) { + ET_CHECK_OR_RETURN_ERROR( + [NSProcessInfo.processInfo + isOperatingSystemAtLeastVersion:manifest.minimum_version], + DelegateInvalidCompatibility, + "Core AI model requires a newer operating system"); + ET_CHECK_OR_RETURN_ERROR( + ETCoreAIIsAvailable(), + NotSupported, + "Core AI requires macOS 27 or iOS 27"); +#if TARGET_OS_OSX + NSString* platform = @"macOS"; +#elif TARGET_OS_IOS + NSString* platform = @"iOS"; +#else + NSString* platform = @"unsupported"; +#endif + if (architecture == nil) { + architecture = ETCoreAIDeviceArchitectureName(); + } + auto selected = select_assets(manifest, architecture, platform); + if (!selected.ok()) + return selected.error(); + return RuntimeAssets{std::move(selected.get()), platform, architecture}; +} + Result> acquire_bookmark_model( const Manifest& selected, const runtime::NamedDataMap* named_data, @@ -180,19 +211,6 @@ Error evict_one(id loader, NSData* bookmark) { return model; } -Error evict_bookmark_model(NSString* root, NSString* key, - id loader) { - ET_CHECK_OR_RETURN_ERROR(loader != nil, InvalidArgument, - "Missing Core AI model loader"); - auto locked = lock_bookmark(root, key); - if (!locked.ok()) return locked.error(); - auto bookmark = read_bookmark(*locked.get()); - if (!bookmark.ok()) return bookmark.error(); - if (bookmark.get() == nil) return Error::Ok; - ET_CHECK_OK_OR_RETURN_ERROR(evict_one(loader, bookmark.get())); - return remove_bookmark(*locked.get()); -} - Error clear_bookmark_assets(NSString* root, NSString* key, id loader) { auto locked = lock_bookmark(root, key, false); diff --git a/backends/apple/coreai/runtime/coreai_storage.h b/backends/apple/coreai/runtime/coreai_storage.h index 3d489a5bf7e..902b993a0a3 100644 --- a/backends/apple/coreai/runtime/coreai_storage.h +++ b/backends/apple/coreai/runtime/coreai_storage.h @@ -75,6 +75,8 @@ int storage_fault(StorageOperation operation); namespace testing { using StorageFaultCallback = int (*)(StorageOperation); extern thread_local StorageFaultCallback storage_fault_callback; +// Borrowed for the scope of tests that must not access the real default cache. +extern thread_local const char* default_assets_root_override; } // namespace testing #else constexpr int storage_fault(StorageOperation) { diff --git a/backends/apple/coreai/runtime/coreai_storage.mm b/backends/apple/coreai/runtime/coreai_storage.mm index 50c192b26f1..57eb333576b 100644 --- a/backends/apple/coreai/runtime/coreai_storage.mm +++ b/backends/apple/coreai/runtime/coreai_storage.mm @@ -49,12 +49,17 @@ Error ensure_excluded_from_backup(NSURL* directory) { } Result default_coreai_assets_root() { - NSURL* caches = - [NSFileManager.defaultManager URLsForDirectory:NSCachesDirectory - inDomains:NSUserDomainMask] - .firstObject; - ET_CHECK_OR_RETURN_ERROR( - caches.isFileURL, AccessFailed, "Cannot resolve Core AI cache directory"); +#if defined(COREAI_ASSETS_TESTING) && COREAI_ASSETS_TESTING + if (testing::default_assets_root_override != nullptr) { + return + [NSString stringWithUTF8String:testing::default_assets_root_override]; + } +#endif + NSURL* caches = [NSFileManager.defaultManager + URLsForDirectory:NSCachesDirectory + inDomains:NSUserDomainMask].firstObject; + ET_CHECK_OR_RETURN_ERROR(caches.isFileURL, AccessFailed, + "Cannot resolve Core AI cache directory"); return [caches.path stringByAppendingPathComponent:@"executorch_coreai"]; } @@ -104,6 +109,7 @@ Error ensure_excluded_from_backup(NSURL* directory) { #if defined(COREAI_ASSETS_TESTING) && COREAI_ASSETS_TESTING namespace testing { thread_local StorageFaultCallback storage_fault_callback = nullptr; +thread_local const char* default_assets_root_override = nullptr; } int storage_fault(StorageOperation operation) { diff --git a/backends/apple/coreai/runtime/test/coreai_cache_test.mm b/backends/apple/coreai/runtime/test/coreai_cache_test.mm index a6597795d14..5878c574e07 100644 --- a/backends/apple/coreai/runtime/test/coreai_cache_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_cache_test.mm @@ -160,7 +160,8 @@ void cache_entry_removed(NSURL* root, const CacheEntry& entry) { FBSyntheticPTE empty({{"forward", {}}}); EXPECT_EQ(clear_cache("relative"), Error::InvalidArgument); EXPECT_EQ(clear_cache_for_pte(empty, "relative"), Error::InvalidArgument); - EXPECT_EQ(clear_cache(nullptr), Error::InvalidArgument); + EXPECT_EQ(clear_cache(""), Error::InvalidArgument); + EXPECT_EQ(clear_cache("\xff"), Error::InvalidArgument); EXPECT_EQ(clear_cache_for_pte(static_cast(nullptr), root.url.path.UTF8String), Error::InvalidArgument); EXPECT_TRUE(empty.check_released()); @@ -255,8 +256,13 @@ FBSyntheticPTE pte( TEST_F(CoreAICacheTest, RejectsInvalidLaterDelegatesBeforeMutation) { @autoreleasepool { - const char* rows[] = {"malformed manifest", "wrong AOT platform", "backend load failure"}; - for (int row = 0; row < 3; ++row) { + const char* rows[] = { + "malformed manifest", + "wrong AOT platform", + "backend load failure", + "newer OS", + "missing architecture"}; + for (int row = 0; row < 5; ++row) { SCOPED_TRACE(rows[row]); BookmarkDirectory root; BookmarkLoaderScope fake; @@ -271,6 +277,13 @@ FBSyntheticPTE pte( dict[@"platform"] = @"iOS"; later.data = encode(dict); } + if (row == 3) { + auto dict = aot_manifest_dict(); + dict[@"min_deployment_version"] = @"999999.0"; + later.data = encode(dict); + } + if (row == 4) + bridge.state().device_architecture = @"missing_arch"; FBSyntheticPTE pte( {{"first", {{"CoreAIBackend", encode(manifest_dict())}}}, {"later", {later}}}); if (row == 2) { @@ -283,6 +296,8 @@ FBSyntheticPTE pte( drain_callbacks(); EXPECT_NE(error, Error::Ok); if (row == 2) EXPECT_EQ(error, Error::AccessFailed); + if (row == 1 || row >= 3) + EXPECT_EQ(error, Error::DelegateInvalidCompatibility); EXPECT_EQ(fake.loader->evictions.load(), 0); EXPECT_EQ(fake.loader->specializations.load(), 0); EXPECT_EQ(fake.loader->bindings.load(), 0); diff --git a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm index 6dac8179962..5088bef0a6c 100644 --- a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm @@ -16,8 +16,11 @@ #include #include "coreai_bookmark_fixture.h" +#include "coreai_cache.h" #import "coreai_fake_loader.h" +#include "coreai_pte_fixture.h" #include "coreai_source_fixture.h" +#include "coreai_storage.h" namespace executorch::backends::coreai::testing { namespace { @@ -110,6 +113,125 @@ Error execute(Span args) { ScopedFakeBridgeState& bridge_; }; +class DefaultAssetsRootScope final { + public: + explicit DefaultAssetsRootScope(NSString* root) + : root_([root copy]), previous_(default_assets_root_override) { + default_assets_root_override = root_.UTF8String; + } + ~DefaultAssetsRootScope() { + default_assets_root_override = previous_; + } + DefaultAssetsRootScope(const DefaultAssetsRootScope&) = delete; + DefaultAssetsRootScope& operator=(const DefaultAssetsRootScope&) = delete; + + private: + NSString* root_; + const char* previous_; +}; + +TEST_F(CoreAIDelegateTest, PteClearingFindsLoadedSourceAndAotEntries) { + @autoreleasepool { + for (bool aot : {false, true}) { + for (NSString* arch in @[ @"arch_a", @"arch_b" ]) { + SCOPED_TRACE(aot ? "AOT" : "source"); + SCOPED_TRACE(arch.UTF8String); + bridge.state().device_architecture = arch; + BookmarkLoaderScope fake; + DelegateFixture fixture( + backend, bridge, aot ? aot_manifest_dict() : manifest_dict()); + if (aot) + aot_data(fixture.data, arch); + const int queries = bridge.state().architecture_queries.load(); + ASSERT_EQ(fixture.init(), Error::Ok); + fixture.reset(); + auto keys = inventory_bookmark_keys(fixture.storage.url.path); + ASSERT_TRUE(keys.ok()); + ASSERT_EQ(keys.get().count, 1); + auto bookmark = saved_bookmark(fixture.storage.url, keys.get()[0]); + ASSERT_TRUE(bookmark.ok()); + ASSERT_NE(bookmark.get(), nil); + ASSERT_NE(fake.loader.models[bookmark.get()], nil); + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", fixture.json}}}, + {"other", {{"CoreAIBackend", fixture.json, true}}}}); + ASSERT_EQ( + clear_cache_for_pte(pte, fixture.storage.url.path.UTF8String), + Error::Ok); + EXPECT_EQ(bridge.state().architecture_queries.load(), queries + 2); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(fake.loader->evictions.load(), 1); + EXPECT_EQ(fake.loader.models[bookmark.get()], nil); + auto remaining = inventory_bookmark_keys(fixture.storage.url.path); + ASSERT_TRUE(remaining.ok()); + EXPECT_EQ(remaining.get().count, 0); + EXPECT_TRUE(pte.check_released()); + } + } + } +} + +TEST_F(CoreAIDelegateTest, DefaultRootClearingMatchesDefaultInitialization) { + @autoreleasepool { + for (int api = 0; api < 4; ++api) { + SCOPED_TRACE(api); + BookmarkLoaderScope fake; + DelegateFixture fixture(backend, bridge); + ASSERT_NE(fixture.storage.url, nil); + DefaultAssetsRootScope default_root(fixture.storage.url.path); + ASSERT_EQ(fixture.init_with_options({}), Error::Ok); + fixture.reset(); + auto keys = inventory_bookmark_keys(fixture.storage.url.path); + ASSERT_TRUE(keys.ok()); + ASSERT_EQ(keys.get().count, 1); + FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", fixture.json}}}}); + for (const char* invalid : {"", "relative", "\xff"}) { + EXPECT_EQ(clear_cache(invalid), Error::InvalidArgument); + EXPECT_EQ(clear_cache_for_pte(pte, invalid), Error::InvalidArgument); + } + EXPECT_EQ(fake.loader->evictions.load(), 0); + switch (api) { + case 0: + ASSERT_EQ(clear_cache(), Error::Ok); + break; + case 1: + ASSERT_EQ(clear_cache(nullptr), Error::Ok); + break; + case 2: + ASSERT_EQ(clear_cache_for_pte(pte), Error::Ok); + break; + case 3: + ASSERT_EQ(clear_cache_for_pte(pte, nullptr), Error::Ok); + break; + } + EXPECT_EQ(fake.loader->evictions.load(), 1); + EXPECT_EQ(fake.loader.models.count, 0); + auto remaining = inventory_bookmark_keys(fixture.storage.url.path); + ASSERT_TRUE(remaining.ok()); + EXPECT_EQ(remaining.get().count, 0); + EXPECT_TRUE(pte.check_released()); + } + } +} + +TEST_F(CoreAIDelegateTest, MissingDefaultRootClearingDoesNotCreateStorage) { + @autoreleasepool { + BookmarkDirectory root; + ASSERT_NE(root.url, nil); + NSString* missing = + [root.url.path stringByAppendingPathComponent:@"missing"]; + DefaultAssetsRootScope default_root(missing); + bridge.state().available = false; + EXPECT_EQ(clear_cache(), Error::Ok); + FBSyntheticPTE empty({}); + EXPECT_EQ(clear_cache_for_pte(empty), Error::Ok); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:missing]); + EXPECT_EQ(bridge.state().architecture_queries.load(), 0); + EXPECT_EQ(bridge.state().loaders.load(), 0); + EXPECT_TRUE(empty.check_released()); + } +} + struct FloatTensors { float input[2] = {3, 7}; float output[2] = {-1, -1}; diff --git a/backends/apple/coreai/runtime/test/coreai_eviction_test.mm b/backends/apple/coreai/runtime/test/coreai_eviction_test.mm index 619da5e1754..9cf065e9fc8 100644 --- a/backends/apple/coreai/runtime/test/coreai_eviction_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_eviction_test.mm @@ -187,7 +187,8 @@ - (void)clearRetainedState { ASSERT_NE(saved, nil); EXPECT_FALSE([saved isEqual:original]); EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); - ASSERT_EQ(evict_bookmark_model(root.url.path, key, fake.loader), Error::Ok); + ASSERT_EQ( + clear_bookmark_assets(root.url.path, key, fake.loader), Error::Ok); EXPECT_NE(fake.loader.models[original], nil); EXPECT_EQ(fake.loader.models[saved], nil); EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); @@ -216,9 +217,17 @@ - (void)clearRetainedState { auto key_result = test_bookmark_key(selected.get()); ASSERT_TRUE(key_result.ok()); NSString* key = key_result.get(); - EXPECT_EQ(evict_bookmark_model(root.url.path, key, fake.loader), Error::Ok); + EXPECT_EQ( + clear_bookmark_assets(root.url.path, key, fake.loader), Error::Ok); EXPECT_EQ(fake.loader->evictions.load(), 0); - ASSERT_EQ(initialize_and_destroy(manifest_dict(), &data, root.url.path), Error::Ok); + // Retain staging after publication to exercise clearing's failure paths. + { + StorageFaultScope fault(StorageOperation::Remove); + ASSERT_EQ( + initialize_and_destroy(manifest_dict(), &data, root.url.path), + Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + } auto saved_result = saved_bookmark(root.url, key); ASSERT_TRUE(saved_result.ok()); NSData* saved = saved_result.get(); @@ -226,29 +235,34 @@ - (void)clearRetainedState { NSURL* source_url = [[[root.url URLByAppendingPathComponent:@"staging"] URLByAppendingPathComponent:key] URLByAppendingPathComponent:selected->path.lastPathComponent]; - EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + EXPECT_TRUE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); for (bool probe_error : {false, true}) { SCOPED_TRACE(probe_error ? "unknown restore failure" : "restore hit/busy"); fake.loader.evictError = bookmark_test_error(EBUSY); fake.loader.restoreError = probe_error ? bookmark_test_error() : nil; - EXPECT_NE(evict_bookmark_model(root.url.path, key, fake.loader), Error::Ok); + EXPECT_NE( + clear_bookmark_assets(root.url.path, key, fake.loader), Error::Ok); auto retained = saved_bookmark(root.url, key); ASSERT_TRUE(retained.ok()); EXPECT_TRUE([retained.get() isEqual:saved]); - EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + EXPECT_TRUE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); } fake.loader.restoreError = nil; fake.loader.evictError = nil; { StorageFaultScope fault(StorageOperation::AfterEvict); - EXPECT_NE(evict_bookmark_model(root.url.path, key, fake.loader), Error::Ok); + EXPECT_NE( + clear_bookmark_assets(root.url.path, key, fake.loader), Error::Ok); EXPECT_EQ(fault.hits, 1); } auto retained = saved_bookmark(root.url, key); ASSERT_TRUE(retained.ok()); EXPECT_TRUE([retained.get() isEqual:saved]); fake.loader.evictError = bookmark_test_error(ENOENT); - ASSERT_EQ(evict_bookmark_model(root.url.path, key, fake.loader), Error::Ok); + ASSERT_EQ( + clear_bookmark_assets(root.url.path, key, fake.loader), Error::Ok); auto absent = saved_bookmark(root.url, key); ASSERT_TRUE(absent.ok()); EXPECT_EQ(absent.get(), nil); @@ -257,13 +271,15 @@ - (void)clearRetainedState { ASSERT_EQ(initialize_and_destroy(manifest_dict(), &data, root.url.path), Error::Ok); { StorageFaultScope fault(StorageOperation::Remove); - EXPECT_NE(evict_bookmark_model(root.url.path, key, fake.loader), Error::Ok); + EXPECT_NE( + clear_bookmark_assets(root.url.path, key, fake.loader), Error::Ok); EXPECT_EQ(fault.hits, 1); auto before_remove = saved_bookmark(root.url, key); ASSERT_TRUE(before_remove.ok()); EXPECT_NE(before_remove.get(), nil); } - ASSERT_EQ(evict_bookmark_model(root.url.path, key, fake.loader), Error::Ok); + ASSERT_EQ( + clear_bookmark_assets(root.url.path, key, fake.loader), Error::Ok); auto removed = saved_bookmark(root.url, key); ASSERT_TRUE(removed.ok()); EXPECT_EQ(removed.get(), nil); @@ -292,14 +308,14 @@ - (void)clearRetainedState { acquire_bookmark_model(selected.get(), &data, root.url.path, @"macOS", @"arch_b", loader); ASSERT_TRUE(pinned.ok()); EXPECT_NE(loader.temporary, nil); - EXPECT_NE(evict_bookmark_model(root.url.path, key, loader), Error::Ok); + EXPECT_NE(clear_bookmark_assets(root.url.path, key, loader), Error::Ok); auto saved = saved_bookmark(root.url, key); ASSERT_TRUE(saved.ok()); EXPECT_NE(saved.get(), nil); } EXPECT_EQ(loader.temporary, nil); EXPECT_NE(loader.retainedCompletion, nil); - ASSERT_EQ(evict_bookmark_model(root.url.path, key, loader), Error::Ok); + ASSERT_EQ(clear_bookmark_assets(root.url.path, key, loader), Error::Ok); auto removed = saved_bookmark(root.url, key); ASSERT_TRUE(removed.ok()); EXPECT_EQ(removed.get(), nil); @@ -340,7 +356,7 @@ - (void)clearRetainedState { bool load_ok = false; auto evict = [&] { @autoreleasepool { - eviction_error = evict_bookmark_model(root.url.path, key, fake.loader); + eviction_error = clear_bookmark_assets(root.url.path, key, fake.loader); } dispatch_semaphore_signal(eviction_done); }; From 0ba28828464f0e8fce0b3d92c536825f00cc5ea7 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Tue, 27 Oct 2026 15:17:44 -0700 Subject: [PATCH 3/5] [CoreAI] Clear all recorded PTE cache variants safely Clear every recorded variant associated with a PTE regardless of host loading compatibility. Persist durable per-key attribution without changing existing cache keys, and report incomplete cleanup when older entries cannot be safely identified. Recover malformed local entries without following links, while preserving uncertain SDK eviction state for retry. Test Plan: Ran all five CoreAI CTest suites: 236 native tests, 18 Swift tests, PTE export, real SDK source-lifetime checks, and runtime smoke all passed. Coverage includes cross-platform and architecture variants, legacy attribution, malformed entries, publication/removal failures, retry durability, and preservation of outside files. Ran the complete CoreAI Python suite: 124 tests passed. Formatting and whitespace checks passed. The generic Buck Swift owner lookup still fails without diagnostics; CMake build and functional validation passed. --- backends/apple/coreai/CMakeLists.txt | 6 +- backends/apple/coreai/README.md | 53 +- .../apple/coreai/runtime/coreai_bookmarks.h | 35 +- .../apple/coreai/runtime/coreai_bookmarks.mm | 283 +++++++-- backends/apple/coreai/runtime/coreai_cache.h | 11 +- backends/apple/coreai/runtime/coreai_cache.mm | 128 +++- .../coreai/runtime/coreai_load_coordinator.mm | 73 ++- .../apple/coreai/runtime/coreai_storage.h | 24 +- .../apple/coreai/runtime/coreai_storage.mm | 119 +++- .../runtime/test/coreai_acquisition_test.mm | 106 +++- .../runtime/test/coreai_bookmark_test.mm | 447 ++++++++++++-- .../coreai/runtime/test/coreai_cache_test.mm | 583 ++++++++++++++---- .../runtime/test/coreai_delegate_test.mm | 4 +- .../runtime/test/coreai_eviction_test.mm | 174 ++++-- .../runtime/test/coreai_storage_test.mm | 225 +++++++ tools/cmake/preset/default.cmake | 4 + 16 files changed, 1909 insertions(+), 366 deletions(-) diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index 7c65baa2d0e..481790ff8a4 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/backends/apple/coreai/CMakeLists.txt @@ -402,9 +402,11 @@ endif() if(EXECUTORCH_BUILD_SHARED) set(_coreai_library_type SHARED) set(_coreai_runtime executorch_shared) + set(_coreai_data_loader "") else() set(_coreai_library_type STATIC) set(_coreai_runtime executorch) + set(_coreai_data_loader extension_data_loader) endif() add_library( coreaidelegate ${_coreai_library_type} ${_coreai_runtime_sources} @@ -417,8 +419,8 @@ coreai_require_os27(coreaidelegate) target_include_directories(coreaidelegate PUBLIC ${_common_include_directories}) add_dependencies(coreaidelegate coreai_swift) target_link_libraries( - coreaidelegate PRIVATE coreai_swift ${_coreai_runtime} program_schema - executorch::coreai_dependencies + coreaidelegate PRIVATE coreai_swift ${_coreai_runtime} ${_coreai_data_loader} + program_schema executorch::coreai_dependencies ) executorch_target_link_options_shared_lib(coreaidelegate) if(EXECUTORCH_BUILD_SHARED) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index e73c1bb4bf9..a6afb391bb9 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -242,13 +242,14 @@ when the directory is omitted or null. A supplied path must be absolute; invalid paths fail without falling back to the default. If loading used a custom root, pass that same root when clearing; the PTE does not store runtime overrides. -Root-wide clearing visits keyed bookmarks and staged directories. PTE clearing -selects and deduplicates keys for all Core AI delegates in every method, using -the current platform, SDK architecture and default specialization configuration. -All selected manifests must be valid and compatible before any entry is cleared. -Incompatible delegates are not silently skipped. Root-wide clearing remains -available independently of the PTE. Discovery does not load models or materialize -assets. +Root-wide clearing visits keyed bookmarks, identity records, and staged +directories. PTE clearing finds every Core AI delegate in every method and clears +all recorded variants of its bundles, including other platforms and SDK +architectures. Manifests must be structurally valid before any entry is cleared; +their minimum OS and runtime compatibility do not restrict maintenance. +Discovery requires neither SDK availability nor an architecture query, and does +not load models or materialize assets. Actual SDK eviction still requires a +usable bookmark and an available loader. These synchronous calls require callers to unload affected models and prevent concurrent loads, inference and maintenance across processes. Per-key locks are @@ -263,11 +264,39 @@ missing root or absent entry is a noncreating no-op. Everything under a keyed `staging/` directory belongs to the backend and is removed with that key. Outside keyed entries, clearing does not delete PTEs, other files or the stable lock files. Unkeyed interrupted-staging directories -also remain untouched. Bookmarks are SDK eviction handles, not weights: lost -bookmarks can leave SDK entries that these APIs cannot target. There is no -history or ownership registry. Copied or reused exported artifacts can share -keys; clearing them also removes the other copy's warm-cache benefit. Subsequent -loads can specialize again. +also remain untouched. Bookmarks are SDK eviction handles, not weights. For an +existing entry with a missing or structurally invalid bookmark, clearing logs an +Info message with the cache key, attempts local cleanup, and continues to other keys. +The message states that SDK eviction could not be confirmed; bookmark contents +are not logged. Absent entries remain quiet no-ops. I/O failures and SDK eviction +errors still preserve bookmarks for retry. Lost bookmarks can leave SDK entries +that these APIs cannot target. No history of removed entries is retained. +Copied or reused exported artifacts can share keys; clearing them also removes +the other copy's warm-cache benefit. Subsequent loads can specialize again. + +Identity records are small JSON sidecars beside bookmarks. They retain bundle +identity and cache context, and are checked against their hashed key. Loading +publishes the record before staging or SDK acquisition, and clearing removes it +last. Existing keys and opaque bookmark bytes are unchanged. Older AOT entries +can be found directly from the PTE's architecture list. Older source entries +without attribution cannot always be identified: PTE clearing removes known +matches, logs unresolved keys, and reports incomplete cleanup rather than +removing unrelated entries. A one-time root-wide `clear_cache` handles these +older entries. Missing or corrupt attribution has the same conservative behavior. + +### Persistent SDK Rejection + +SDK errors are not classified as transient or permanent. +Busy errors and all unclassified eviction/restore errors retain the local +bookmark for retry; only successful eviction or a confirmed SDK miss permits +normal removal. If retries keep failing, first stop every process using the +assets root. Back up or quarantine only the affected key's bookmark, identity +sidecar, and keyed staging outside that root, leaving the stable lock files in +place. A later load can attempt cold acquisition and create new tracking data. +This manual recovery does not confirm SDK eviction or guarantee disk-space +reclamation: keep the backup because it may contain the only remaining eviction +handle. Do not perform manual recovery while loading, inference, or maintenance +is running. ### Runtime options diff --git a/backends/apple/coreai/runtime/coreai_bookmarks.h b/backends/apple/coreai/runtime/coreai_bookmarks.h index d9d1a067fa5..3fdda5d1467 100644 --- a/backends/apple/coreai/runtime/coreai_bookmarks.h +++ b/backends/apple/coreai/runtime/coreai_bookmarks.h @@ -10,7 +10,7 @@ #include #include #import "coreai_assets.h" -#include "coreai_file.h" +#include "coreai_storage.h" namespace executorch::backends::coreai { @@ -18,6 +18,9 @@ runtime::Result bookmark_key( const Manifest& selected, NSString* platform, NSString* architecture); +runtime::Result bookmark_identity( + const Manifest& manifest, + NSString* bundle); // The first key byte selects a stripe in a fixed, versioned lock namespace. constexpr size_t kBookmarkLockStripeCount = 256; runtime::Result bookmark_lock_name(NSString* key); @@ -37,12 +40,23 @@ class BookmarkLock final { private: friend runtime::Result> lock_bookmark(NSString* root, NSString* key, bool create); - friend runtime::Result read_bookmark(const BookmarkLock& lock); + friend runtime::Result read_bookmark( + const BookmarkLock& lock, + bool maintenance); + friend runtime::Result read_bookmark_identity( + const BookmarkLock& lock); + friend runtime::Error ensure_bookmark_identity( + const BookmarkLock& lock, + const Manifest& selected, + NSString* platform, + NSString* architecture); + friend runtime::Error remove_bookmark_identity(const BookmarkLock& lock); friend runtime::Error write_bookmark(const BookmarkLock& lock, NSData* data); friend runtime::Error remove_bookmark(const BookmarkLock& lock); friend runtime::Error remove_bookmark_staging( const BookmarkLock& lock, - bool remove); + bool remove, + StorageRemovalMode mode); friend runtime::Result prepare_bookmark_staging( const BookmarkLock& lock); BookmarkLock( @@ -68,7 +82,17 @@ runtime::Result> lock_bookmark(NSString* root, NSString* key, bool create = true); // Missing is nil. Empty, oversized, inaccessible or nonregular files fail // closed. -runtime::Result read_bookmark(const BookmarkLock& lock); +runtime::Result read_bookmark( + const BookmarkLock& lock, + bool maintenance = false); +// Missing metadata is nil; malformed or misattributed records fail closed. +runtime::Result read_bookmark_identity(const BookmarkLock& lock); +runtime::Error ensure_bookmark_identity( + const BookmarkLock& lock, + const Manifest& selected, + NSString* platform, + NSString* architecture); +runtime::Error remove_bookmark_identity(const BookmarkLock& lock); // A staging-only lock has no bookmark directory; writes return InvalidArgument. runtime::Error write_bookmark(const BookmarkLock& lock, NSData* data); runtime::Error remove_bookmark(const BookmarkLock& lock); @@ -76,7 +100,8 @@ runtime::Error remove_bookmark(const BookmarkLock& lock); // only preflight the keyed tree. runtime::Error remove_bookmark_staging( const BookmarkLock& lock, - bool remove = true); + bool remove = true, + StorageRemovalMode mode = StorageRemovalMode::Strict); // Prepared only for cold loads; source preparation appends key/bundle. runtime::Result prepare_bookmark_staging(const BookmarkLock& lock); diff --git a/backends/apple/coreai/runtime/coreai_bookmarks.mm b/backends/apple/coreai/runtime/coreai_bookmarks.mm index f0c1e163f72..16d153d6926 100644 --- a/backends/apple/coreai/runtime/coreai_bookmarks.mm +++ b/backends/apple/coreai/runtime/coreai_bookmarks.mm @@ -42,30 +42,53 @@ bool digest(id value) { NSString* bookmark_name(NSString* key) { return [key stringByAppendingString:@".bookmark"]; } -} // namespace -Result bookmark_key( +constexpr size_t kIdentityByteLimit = 128 * 1024; +NSString* identity_name(NSString* key) { + return [key stringByAppendingString:@".metadata"]; +} + +Result*> identity_fields( const Manifest& selected, - NSString* platform, - NSString* architecture) { - NSString* bundle = selected.path.lastPathComponent; + NSString* bundle) { + ET_CHECK_OR_RETURN_ERROR( + component(bundle), + DelegateInvalidCompatibility, + "Invalid Core AI bundle name"); NSString* exported_digest = selected.bundle_digests[bundle]; ET_CHECK_OR_RETURN_ERROR( - digest(exported_digest) && component(bundle) && text(platform) && - text(architecture), + digest(exported_digest), DelegateInvalidCompatibility, - "Core AI bookmark key requires a selected bundle and SDK architecture"); - NSArray* fields = @[ + "Missing Core AI bundle digest"); + return @[ @"1", exported_digest, selected.aot_compiled ? @"aot" : @"source", - bundle, + bundle + ]; +} + +Result*> key_fields( + const Manifest& selected, + NSString* platform, + NSString* architecture) { + auto identity = identity_fields(selected, selected.path.lastPathComponent); + if (!identity.ok()) + return identity.error(); + ET_CHECK_OR_RETURN_ERROR( + text(platform) && text(architecture), + DelegateInvalidCompatibility, + "Missing Core AI cache context"); + return [identity.get() arrayByAddingObjectsFromArray:@[ platform, architecture, @"AIModelCache.default", @"default-options-v1", @"persistent" - ]; + ]]; +} + +Result hash_fields(NSArray* fields) { NSData* encoded = [NSJSONSerialization dataWithJSONObject:fields options:0 error:nil]; @@ -79,6 +102,27 @@ bool digest(id value) { return [key copy]; } +} // namespace + +Result bookmark_identity( + const Manifest& manifest, + NSString* bundle) { + auto fields = identity_fields(manifest, bundle); + if (!fields.ok()) + return fields.error(); + return hash_fields(fields.get()); +} + +Result bookmark_key( + const Manifest& selected, + NSString* platform, + NSString* architecture) { + auto fields = key_fields(selected, platform, architecture); + if (!fields.ok()) + return fields.error(); + return hash_fields(fields.get()); +} + Result bookmark_lock_name(NSString* key) { ET_CHECK_OR_RETURN_ERROR( digest(key), InvalidArgument, "Invalid Core AI bookmark key"); @@ -120,19 +164,23 @@ bool digest(id value) { FileDescriptor child(opened.get()); if (child.get() < 0) continue; - NSString* entry = [name isEqual:@"bookmarks"] ? bookmark_name(key) : key; - struct stat info; - const int result = retry_eintr([&] { - return fstatat( - child.get(), - entry.fileSystemRepresentation, - &info, - AT_SYMLINK_NOFOLLOW); - }); - if (result == 0) - return true; - ET_CHECK_OR_RETURN_ERROR( - errno == ENOENT, AccessFailed, "Cannot inspect Core AI cache entry"); + NSArray* entries = [name isEqual:@"bookmarks"] + ? @[ bookmark_name(key), identity_name(key) ] + : @[ key ]; + for (NSString* entry in entries) { + struct stat info; + const int result = retry_eintr([&] { + return fstatat( + child.get(), + entry.fileSystemRepresentation, + &info, + AT_SYMLINK_NOFOLLOW); + }); + if (result == 0) + return true; + ET_CHECK_OR_RETURN_ERROR( + errno == ENOENT, AccessFailed, "Cannot inspect Core AI cache entry"); + } } return false; } @@ -160,33 +208,47 @@ bool digest(id value) { Result*> inventory_bookmark_keys(NSString* root) { auto resolved = resolve_bookmark_root(root, false); - if (!resolved.ok()) return resolved.error(); + if (!resolved.ok()) + return resolved.error(); auto opened = open_root(resolved.get()); - if (!opened.ok()) return opened.error(); + if (!opened.ok()) + return opened.error(); FileDescriptor directory(opened.get()); - if (directory.get() < 0) return @[]; + if (directory.get() < 0) + return @[]; NSMutableSet* keys = [NSMutableSet set]; for (NSString* name in @[ @"bookmarks", @"staging" ]) { auto child_result = open_child(directory.get(), name, false); - if (!child_result.ok()) return child_result.error(); + if (!child_result.ok()) + return child_result.error(); FileDescriptor child(child_result.get()); - if (child.get() < 0) continue; + if (child.get() < 0) + continue; auto children = storage_children(child.get(), true); - if (!children.ok()) return children.error(); + if (!children.ok()) + return children.error(); for (NSString* entry in children.get()) { NSString* key = entry; if ([name isEqual:@"bookmarks"]) { - if (![entry hasSuffix:@".bookmark"]) continue; - key = [entry substringToIndex:entry.length - @".bookmark".length]; + NSString* suffix = + [entry hasSuffix:@".bookmark"] ? @".bookmark" : @".metadata"; + if (![entry hasSuffix:suffix]) + continue; + key = [entry substringToIndex:entry.length - suffix.length]; } - if (digest(key)) [keys addObject:key]; + if (digest(key)) + [keys addObject:key]; } } return [keys.allObjects sortedArrayUsingSelector:@selector(compare:)]; } -BookmarkLock::BookmarkLock(int root, int bookmarks, int file, - NSString* root_path, NSString* key) +BookmarkLock::BookmarkLock( + int root, + int bookmarks, + int file, + NSString* root_path, + NSString* key) : root_(root), bookmarks_(bookmarks), file_(file), @@ -256,24 +318,47 @@ FileDescriptor file(open_storage_shared_file( [key copy])); } -Result read_bookmark(const BookmarkLock& lock) { - if (lock.bookmarks_.get() < 0) +namespace { +Result +read_record(int directory, NSString* name, size_t limit, bool maintenance) { + if (directory < 0) return static_cast(nil); + if (maintenance) { + struct stat entry; + const int inspected = retry_eintr([&] { + return fstatat( + directory, + name.fileSystemRepresentation, + &entry, + AT_SYMLINK_NOFOLLOW); + }); + if (inspected < 0 && errno == ENOENT) + return static_cast(nil); + ET_CHECK_OR_RETURN_ERROR( + inspected == 0, AccessFailed, "Cannot inspect Core AI cache record"); + if (!S_ISREG(entry.st_mode)) + return Error::InvalidExternalData; + } FileDescriptor file(retry_eintr([&] { return openat( - lock.bookmarks_.get(), - bookmark_name(lock.key_).fileSystemRepresentation, + directory, + name.fileSystemRepresentation, O_RDONLY | O_NOFOLLOW | O_NONBLOCK | O_CLOEXEC); })); if (file.get() < 0 && errno == ENOENT) return static_cast(nil); + if (file.get() < 0 && errno == ELOOP && maintenance) + return Error::InvalidExternalData; ET_CHECK_OR_RETURN_ERROR( file.get() >= 0, AccessFailed, "Cannot open Core AI bookmark"); struct stat info; ET_CHECK_OR_RETURN_ERROR( - fstat(file.get(), &info) == 0 && S_ISREG(info.st_mode) && - info.st_size > 0 && - static_cast(info.st_size) <= kBookmarkByteLimit, + fstat(file.get(), &info) == 0, + AccessFailed, + "Cannot inspect Core AI bookmark"); + ET_CHECK_OR_RETURN_ERROR( + S_ISREG(info.st_mode) && info.st_size > 0 && + static_cast(info.st_size) <= limit, InvalidExternalData, "Invalid Core AI bookmark file"); NSMutableData* bytes = [NSMutableData dataWithLength:info.st_size]; @@ -292,6 +377,100 @@ FileDescriptor file(retry_eintr([&] { return bytes; } +} // namespace + +Result read_bookmark(const BookmarkLock& lock, bool maintenance) { + return read_record( + lock.bookmarks_.get(), + bookmark_name(lock.key_), + kBookmarkByteLimit, + maintenance); +} + +Result read_bookmark_identity(const BookmarkLock& lock) { + auto bytes = read_record( + lock.bookmarks_.get(), + identity_name(lock.key_), + kIdentityByteLimit, + true); + if (!bytes.ok()) + return bytes.error(); + if (bytes.get() == nil) + return static_cast(nil); + id parsed = [NSJSONSerialization JSONObjectWithData:bytes.get() + options:0 + error:nil]; + if (![parsed isKindOfClass:NSArray.class] || [parsed count] != 9) + return Error::InvalidExternalData; + NSArray* fields = parsed; + for (id field in fields) { + if (!text(field)) + return Error::InvalidExternalData; + } + if (![fields[2] isEqual:@"source"] && ![fields[2] isEqual:@"aot"]) + return Error::InvalidExternalData; + Manifest selected; + selected.path = fields[3]; + selected.bundle_digests = @{fields[3] : fields[1]}; + selected.aot_compiled = [fields[2] isEqual:@"aot"]; + auto expected = key_fields(selected, fields[4], fields[5]); + if (!expected.ok() || ![expected.get() isEqual:fields]) + return Error::InvalidExternalData; + auto key = hash_fields(fields); + if (!key.ok()) + return key.error(); + if (![key.get() isEqual:lock.key_]) + return Error::InvalidExternalData; + return bookmark_identity(selected, fields[3]); +} + +Error ensure_bookmark_identity( + const BookmarkLock& lock, + const Manifest& selected, + NSString* platform, + NSString* architecture) { + ET_CHECK_OR_RETURN_ERROR( + lock.bookmarks_.get() >= 0, + InvalidArgument, + "Cannot attribute a staging-only lock"); + auto fields = key_fields(selected, platform, architecture); + if (!fields.ok()) + return fields.error(); + auto key = hash_fields(fields.get()); + if (!key.ok()) + return key.error(); + ET_CHECK_OR_RETURN_ERROR( + [key.get() isEqual:lock.key_], + InvalidArgument, + "Core AI cache identity does not match key"); + NSData* encoded = [NSJSONSerialization dataWithJSONObject:fields.get() + options:0 + error:nil]; + ET_CHECK_OR_RETURN_ERROR( + encoded != nil && encoded.length <= kIdentityByteLimit, + InvalidExternalData, + "Invalid Core AI cache identity"); + auto existing = read_record( + lock.bookmarks_.get(), + identity_name(lock.key_), + kIdentityByteLimit, + true); + if (!existing.ok() && existing.error() != Error::InvalidExternalData) + return existing.error(); + if (existing.ok() && [existing.get() isEqual:encoded]) { + // A prior publication may have failed its post-rename flush. + return sync_storage_directory(lock.bookmarks_.get(), true); + } + return publish_storage_data( + lock.bookmarks_.get(), identity_name(lock.key_), encoded); +} + +Error remove_bookmark_identity(const BookmarkLock& lock) { + if (lock.bookmarks_.get() < 0) + return Error::Ok; + return remove_storage_entry(lock.bookmarks_.get(), identity_name(lock.key_)); +} + Error write_bookmark(const BookmarkLock& lock, NSData* data) { ET_CHECK_OR_RETURN_ERROR( [data isKindOfClass:NSData.class] && data.length > 0 && @@ -307,20 +486,16 @@ Error write_bookmark(const BookmarkLock& lock, NSData* data) { } Error remove_bookmark(const BookmarkLock& lock) { - if (lock.bookmarks_.get() < 0) return Error::Ok; - ET_CHECK_OR_RETURN_ERROR(storage_fault(StorageOperation::Remove) == 0, - AccessFailed, - "Core AI bookmark removal interrupted"); - ET_CHECK_OR_RETURN_ERROR( - unlinkat(lock.bookmarks_.get(), - bookmark_name(lock.key_).fileSystemRepresentation, 0) == 0 || - errno == ENOENT, - AccessFailed, "Cannot remove Core AI bookmark"); - return sync_storage_directory(lock.bookmarks_.get()); + if (lock.bookmarks_.get() < 0) + return Error::Ok; + return remove_storage_entry(lock.bookmarks_.get(), bookmark_name(lock.key_)); } -Error remove_bookmark_staging(const BookmarkLock& lock, bool remove) { - return remove_storage_staging(lock.root_.get(), lock.key_, remove); +Error remove_bookmark_staging( + const BookmarkLock& lock, + bool remove, + StorageRemovalMode mode) { + return remove_storage_staging(lock.root_.get(), lock.key_, remove, mode); } Result prepare_bookmark_staging(const BookmarkLock& lock) { diff --git a/backends/apple/coreai/runtime/coreai_cache.h b/backends/apple/coreai/runtime/coreai_cache.h index 2e1d28bbf28..b6ce1c06258 100644 --- a/backends/apple/coreai/runtime/coreai_cache.h +++ b/backends/apple/coreai/runtime/coreai_cache.h @@ -32,10 +32,13 @@ namespace executorch::backends::coreai { const char* coreai_assets_dir = nullptr); /** - * Clears current-platform/SDK-architecture cache entries referenced by every - * CoreAI delegate in the PTE, after validating all selected manifests. Invalid - * or incompatible delegates fail before any entries are cleared. Does not - * initialize delegates or materialize Core AI assets. + * Clears all recorded cache variants referenced by every CoreAI delegate in + * the PTE after validating its manifests. Discovery does not require host + * compatibility or SDK availability. Actual SDK eviction requires a loader. + * Missing/corrupt attribution can leave legacy source entries unidentified: + * known matches are cleared, then an error reports incomplete cleanup. A + * root-wide clear can recover these entries. Does not initialize delegates or + * materialize Core AI assets. * * A null directory selects the same Caches default as backend initialization. * Supplied directories must be valid explicit roots; they never fall back to diff --git a/backends/apple/coreai/runtime/coreai_cache.mm b/backends/apple/coreai/runtime/coreai_cache.mm index 06c23a118d8..f20155dd843 100644 --- a/backends/apple/coreai/runtime/coreai_cache.mm +++ b/backends/apple/coreai/runtime/coreai_cache.mm @@ -39,9 +39,13 @@ Error clear_keys(NSString* root, NSArray* keys) { for (NSString* key in ordered) { Error error = clear_bookmark_assets(root, key, nil); if (error != Error::Ok) { - ET_LOG(Error, "Core AI cache clear failed for key %s: 0x%x", - key.UTF8String, static_cast(error)); - if (first_error == Error::Ok) first_error = error; + ET_LOG( + Error, + "Core AI cache clear failed for key %s: 0x%x", + key.UTF8String, + static_cast(error)); + if (first_error == Error::Ok) + first_error = error; } } return first_error; @@ -50,56 +54,120 @@ Error clear_keys(NSString* root, NSArray* keys) { Error clear_pte(runtime::DataLoader& loader, NSString* root) { // Keep the verified image and original segment owners alive through eviction. auto pte = inspect_coreai_pte(loader); - if (!pte.ok()) return pte.error(); - NSMutableArray* keys = [NSMutableArray array]; - NSString* architecture = nil; + if (!pte.ok()) + return pte.error(); + NSMutableSet* identities = [NSMutableSet set]; + NSMutableSet* legacy_aot_keys = [NSMutableSet set]; + bool has_source = false; for (const auto& buffer : pte->processed) { auto data = buffer.data_safe(); - if (!data.ok()) return data.error(); - ET_CHECK_OR_RETURN_ERROR(data.get() != nullptr && buffer.size() > 0, - InvalidProgram, "Missing Core AI manifest"); + if (!data.ok()) + return data.error(); + ET_CHECK_OR_RETURN_ERROR( + data.get() != nullptr && buffer.size() > 0, + InvalidProgram, + "Missing Core AI manifest"); NSData* bytes = [NSData dataWithBytesNoCopy:const_cast(data.get()) - length:buffer.size() - freeWhenDone:NO]; + length:buffer.size() + freeWhenDone:NO]; auto manifest = parse_manifest(bytes); - if (!manifest.ok()) return manifest.error(); - auto selected = select_runtime_assets(manifest.get(), architecture); - if (!selected.ok()) return selected.error(); - architecture = selected->architecture; - auto key = - bookmark_key(selected->manifest, selected->platform, architecture); - if (!key.ok()) return key.error(); - [keys addObject:key.get()]; + if (!manifest.ok()) + return manifest.error(); + for (NSString* bundle in manifest->bundle_digests) { + auto identity = bookmark_identity(manifest.get(), bundle); + if (!identity.ok()) + return identity.error(); + [identities addObject:identity.get()]; + } + if (manifest->aot_compiled) { + for (NSString* architecture in manifest->archs) { + auto selected = + select_assets(manifest.get(), architecture, manifest->platform); + if (!selected.ok()) + return selected.error(); + auto key = + bookmark_key(selected.get(), manifest->platform, architecture); + if (!key.ok()) + return key.error(); + [legacy_aot_keys addObject:key.get()]; + } + } else { + has_source = true; + } + } + if (identities.count == 0) + return Error::Ok; + auto inventory = inventory_bookmark_keys(root); + if (!inventory.ok()) + return inventory.error(); + NSMutableArray* keys = [NSMutableArray array]; + Error lookup_error = Error::Ok; + for (NSString* key in inventory.get()) { + if ([legacy_aot_keys containsObject:key]) { + [keys addObject:key]; + continue; + } + auto locked = lock_bookmark(root, key, false); + if (!locked.ok()) { + if (lookup_error == Error::Ok) + lookup_error = locked.error(); + continue; + } + if (locked.get() == nullptr) + continue; + auto identity = read_bookmark_identity(*locked.get()); + if (!identity.ok() || identity.get() == nil) { + if (has_source) { + ET_LOG( + Info, + "Core AI cache key %s has missing or invalid attribution; PTE cleanup is incomplete", + key.UTF8String); + if (lookup_error == Error::Ok) + lookup_error = + identity.ok() ? Error::InvalidExternalData : identity.error(); + } + continue; + } + if ([identities containsObject:identity.get()]) + [keys addObject:key]; } - return clear_keys(root, keys); + const auto cleared = clear_keys(root, keys); + return lookup_error != Error::Ok ? lookup_error : cleared; } -} // namespace +} // namespace Error clear_cache(const char* coreai_assets_dir) { @autoreleasepool { auto root = cache_root(coreai_assets_dir); - if (!root.ok()) return root.error(); + if (!root.ok()) + return root.error(); auto keys = inventory_bookmark_keys(root.get()); - if (!keys.ok()) return keys.error(); + if (!keys.ok()) + return keys.error(); return clear_keys(root.get(), keys.get()); } } -Error clear_cache_for_pte(runtime::DataLoader& loader, - const char* coreai_assets_dir) { +Error clear_cache_for_pte( + runtime::DataLoader& loader, + const char* coreai_assets_dir) { @autoreleasepool { auto root = cache_root(coreai_assets_dir); - if (!root.ok()) return root.error(); + if (!root.ok()) + return root.error(); return clear_pte(loader, root.get()); } } Error clear_cache_for_pte(const char* pte_path, const char* coreai_assets_dir) { - ET_CHECK_OR_RETURN_ERROR(pte_path != nullptr && pte_path[0] != '\0', - InvalidArgument, "Missing PTE path"); + ET_CHECK_OR_RETURN_ERROR( + pte_path != nullptr && pte_path[0] != '\0', + InvalidArgument, + "Missing PTE path"); auto loader = extension::FileDataLoader::from(pte_path); - if (!loader.ok()) return loader.error(); + if (!loader.ok()) + return loader.error(); return clear_cache_for_pte(loader.get(), coreai_assets_dir); } -} // namespace executorch::backends::coreai +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.mm b/backends/apple/coreai/runtime/coreai_load_coordinator.mm index 35810caffff..7dfde306521 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.mm +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.mm @@ -83,8 +83,10 @@ void cleanup_source(const BookmarkLock& lock) { } Error evict_one(id loader, NSData* bookmark) { - ET_CHECK_OR_RETURN_ERROR(storage_fault(StorageOperation::BeforeEvict) == 0, - AccessFailed, "Core AI eviction interrupted"); + ET_CHECK_OR_RETURN_ERROR( + storage_fault(StorageOperation::BeforeEvict) == 0, + AccessFailed, + "Core AI eviction interrupted"); dispatch_semaphore_t ready = dispatch_semaphore_create(0); __block NSError* failure = nil; [loader evictModelWithBookmark:bookmark @@ -97,15 +99,18 @@ Error evict_one(id loader, NSData* bookmark) { // Probe after failed deletion so our own probe cannot make the entry busy. @autoreleasepool { auto probe = restore(loader, bookmark); - if (probe.ok() && probe.get() == nil) return Error::Ok; + if (probe.ok() && probe.get() == nil) + return Error::Ok; } return sdk_error(failure); } - ET_CHECK_OR_RETURN_ERROR(storage_fault(StorageOperation::AfterEvict) == 0, - AccessFailed, "Core AI eviction interrupted"); + ET_CHECK_OR_RETURN_ERROR( + storage_fault(StorageOperation::AfterEvict) == 0, + AccessFailed, + "Core AI eviction interrupted"); return Error::Ok; } -} // namespace +} // namespace Result select_runtime_assets( const Manifest& manifest, @@ -152,6 +157,8 @@ Error evict_one(id loader, NSData* bookmark) { auto locked = lock_bookmark(root, cache_key); if (!locked.ok()) return locked.error(); + ET_CHECK_OK_OR_RETURN_ERROR(ensure_bookmark_identity( + *locked.get(), selected, platform, architecture)); auto bookmark = read_bookmark(*locked.get()); if (!bookmark.ok()) return bookmark.error(); @@ -211,26 +218,48 @@ Error evict_one(id loader, NSData* bookmark) { return model; } -Error clear_bookmark_assets(NSString* root, NSString* key, - id loader) { +Error clear_bookmark_assets( + NSString* root, + NSString* key, + id loader) { auto locked = lock_bookmark(root, key, false); - if (!locked.ok()) return locked.error(); - if (locked.get() == nullptr) return Error::Ok; - auto bookmark = read_bookmark(*locked.get()); - if (!bookmark.ok()) return bookmark.error(); - ET_CHECK_OK_OR_RETURN_ERROR(remove_bookmark_staging(*locked.get(), false)); - if (bookmark.get() != nil) { + if (!locked.ok()) + return locked.error(); + if (locked.get() == nullptr) + return Error::Ok; + auto bookmark = read_bookmark(*locked.get(), /*maintenance=*/true); + if (!bookmark.ok() && bookmark.error() != Error::InvalidExternalData) { + return bookmark.error(); + } + NSData* bytes = bookmark.ok() ? bookmark.get() : nil; + if (bytes == nil) { + ET_LOG( + Info, + "Core AI cache key %s has a %s bookmark; SDK eviction could not be " + "confirmed; attempting local cleanup", + key.UTF8String, + bookmark.ok() ? "missing" : "corrupt"); + } + ET_CHECK_OK_OR_RETURN_ERROR(remove_bookmark_staging( + *locked.get(), false, StorageRemovalMode::Maintenance)); + if (bytes != nil) { if (loader == nil) { - ET_CHECK_OR_RETURN_ERROR(ETCoreAIIsAvailable(), NotSupported, - "Core AI requires macOS 27 or iOS 27"); + ET_CHECK_OR_RETURN_ERROR( + ETCoreAIIsAvailable(), + NotSupported, + "Core AI requires macOS 27 or iOS 27"); loader = ETCoreAICreateModelLoader(); - ET_CHECK_OR_RETURN_ERROR(loader != nil, Internal, - "Cannot create Core AI model loader"); + ET_CHECK_OR_RETURN_ERROR( + loader != nil, Internal, "Cannot create Core AI model loader"); } - ET_CHECK_OK_OR_RETURN_ERROR(evict_one(loader, bookmark.get())); + ET_CHECK_OK_OR_RETURN_ERROR(evict_one(loader, bytes)); + } + ET_CHECK_OK_OR_RETURN_ERROR(remove_bookmark_staging( + *locked.get(), true, StorageRemovalMode::Maintenance)); + if (!bookmark.ok() || bytes != nil) { + ET_CHECK_OK_OR_RETURN_ERROR(remove_bookmark(*locked.get())); } - ET_CHECK_OK_OR_RETURN_ERROR(remove_bookmark_staging(*locked.get())); - return bookmark.get() == nil ? Error::Ok : remove_bookmark(*locked.get()); + return remove_bookmark_identity(*locked.get()); } -} // namespace executorch::backends::coreai +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_storage.h b/backends/apple/coreai/runtime/coreai_storage.h index 902b993a0a3..ea317757994 100644 --- a/backends/apple/coreai/runtime/coreai_storage.h +++ b/backends/apple/coreai/runtime/coreai_storage.h @@ -47,12 +47,32 @@ runtime::Error publish_storage_data(int root_fd, NSString* name, NSData* data); runtime::Result*> storage_children( int fd, bool skip_invalid_names = false); +enum class StorageRemovalMode { + Strict, + Maintenance, +}; + // The caller holds the key's lock. Missing staging is already removed. // With remove=false, only preflight the keyed tree without modifying it. +// Maintenance recovers malformed backend-owned entries without following +// symlinks. Strict keeps ordinary acquisition validation. runtime::Error remove_storage_staging( - int root_fd, NSString* key, bool remove = true); + int root_fd, + NSString* key, + bool remove = true, + StorageRemovalMode mode = StorageRemovalMode::Strict); // Removes a keyed child of source_root_fd under the same lock requirement. -runtime::Error remove_storage_source(int source_root_fd, NSString* key, bool remove = true); +// Maintenance also unlinks a non-directory keyed entry. +runtime::Error remove_storage_source( + int source_root_fd, + NSString* key, + bool remove = true, + StorageRemovalMode mode = StorageRemovalMode::Strict); +// Removes one validated entry without following symlinks, covering directory +// trees as well as bookmark-style files, links, and FIFOs. Missing entries +// are already removed. +runtime::Error +remove_storage_entry(int parent_fd, NSString* name, bool remove = true); enum class StorageOperation { Write, diff --git a/backends/apple/coreai/runtime/coreai_storage.mm b/backends/apple/coreai/runtime/coreai_storage.mm index 57eb333576b..2c246e0623f 100644 --- a/backends/apple/coreai/runtime/coreai_storage.mm +++ b/backends/apple/coreai/runtime/coreai_storage.mm @@ -55,11 +55,12 @@ Error ensure_excluded_from_backup(NSURL* directory) { [NSString stringWithUTF8String:testing::default_assets_root_override]; } #endif - NSURL* caches = [NSFileManager.defaultManager - URLsForDirectory:NSCachesDirectory - inDomains:NSUserDomainMask].firstObject; - ET_CHECK_OR_RETURN_ERROR(caches.isFileURL, AccessFailed, - "Cannot resolve Core AI cache directory"); + NSURL* caches = + [NSFileManager.defaultManager URLsForDirectory:NSCachesDirectory + inDomains:NSUserDomainMask] + .firstObject; + ET_CHECK_OR_RETURN_ERROR( + caches.isFileURL, AccessFailed, "Cannot resolve Core AI cache directory"); return [caches.path stringByAppendingPathComponent:@"executorch_coreai"]; } @@ -110,7 +111,7 @@ Error ensure_excluded_from_backup(NSURL* directory) { namespace testing { thread_local StorageFaultCallback storage_fault_callback = nullptr; thread_local const char* default_assets_root_override = nullptr; -} +} // namespace testing int storage_fault(StorageOperation operation) { return testing::storage_fault_callback == nullptr @@ -205,8 +206,9 @@ FileDescriptor file(retry_eintr([&] { } // Preflight the whole tree before unlinking anything; never follow source -// links. -Error visit_source_tree(int fd, bool remove) { +// links. Maintenance also accepts symlinks, FIFOs, and other non-directory +// entries so malformed backend-owned state can be recovered. +Error visit_source_tree(int fd, bool remove, StorageRemovalMode mode) { auto names = storage_children(fd); if (!names.ok()) return names.error(); @@ -219,11 +221,14 @@ Error visit_source_tree(int fd, bool remove) { }) == 0, AccessFailed, "Cannot inspect Core AI source entry"); - ET_CHECK_OR_RETURN_ERROR( - S_ISDIR(info.st_mode) || S_ISREG(info.st_mode), - InvalidExternalData, - "Unexpected Core AI source entry type"); - if (S_ISDIR(info.st_mode)) { + const bool directory = S_ISDIR(info.st_mode); + if (mode == StorageRemovalMode::Strict) { + ET_CHECK_OR_RETURN_ERROR( + directory || S_ISREG(info.st_mode), + InvalidExternalData, + "Unexpected Core AI source entry type"); + } + if (directory) { FileDescriptor child(retry_eintr([&] { return openat( fd, @@ -234,7 +239,7 @@ FileDescriptor child(retry_eintr([&] { child.get() >= 0, AccessFailed, "Cannot open Core AI source directory"); - ET_CHECK_OK_OR_RETURN_ERROR(visit_source_tree(child.get(), remove)); + ET_CHECK_OK_OR_RETURN_ERROR(visit_source_tree(child.get(), remove, mode)); } if (remove) { ET_CHECK_OR_RETURN_ERROR( @@ -244,7 +249,7 @@ FileDescriptor child(retry_eintr([&] { return unlinkat( fd, name.fileSystemRepresentation, - S_ISDIR(info.st_mode) ? AT_REMOVEDIR : 0); + directory ? AT_REMOVEDIR : 0); }) == 0, AccessFailed, "Cannot remove Core AI source entry"); @@ -379,7 +384,63 @@ bool storage_key(NSString* key) { } } // namespace -Error remove_storage_staging(int root_fd, NSString* key, bool remove) { +Error remove_storage_entry(int parent_fd, NSString* name, bool remove) { + ET_CHECK_OR_RETURN_ERROR( + storage_component(name), InvalidArgument, "Invalid Core AI entry name"); + struct stat info; + const int inspected = retry_eintr([&] { + return fstatat( + parent_fd, name.fileSystemRepresentation, &info, AT_SYMLINK_NOFOLLOW); + }); + if (inspected < 0 && errno == ENOENT) + return Error::Ok; + ET_CHECK_OR_RETURN_ERROR( + inspected == 0, AccessFailed, "Cannot inspect Core AI storage entry"); + if (!S_ISDIR(info.st_mode)) { + if (!remove) + return Error::Ok; + ET_CHECK_OR_RETURN_ERROR( + storage_call( + StorageOperation::Remove, + [&] { + return unlinkat(parent_fd, name.fileSystemRepresentation, 0); + }) == 0 || + errno == ENOENT, + AccessFailed, + "Cannot remove Core AI storage entry"); + return sync_storage_directory(parent_fd); + } + FileDescriptor entry(retry_eintr([&] { + return openat( + parent_fd, + name.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); + })); + ET_CHECK_OR_RETURN_ERROR( + entry.get() >= 0, AccessFailed, "Cannot open Core AI storage directory"); + ET_CHECK_OK_OR_RETURN_ERROR( + visit_source_tree(entry.get(), false, StorageRemovalMode::Maintenance)); + if (!remove) + return Error::Ok; + ET_CHECK_OK_OR_RETURN_ERROR( + visit_source_tree(entry.get(), true, StorageRemovalMode::Maintenance)); + ET_CHECK_OR_RETURN_ERROR( + storage_call( + StorageOperation::Remove, + [&] { + return unlinkat( + parent_fd, name.fileSystemRepresentation, AT_REMOVEDIR); + }) == 0, + AccessFailed, + "Cannot remove Core AI storage directory"); + return sync_storage_directory(parent_fd); +} + +Error remove_storage_staging( + int root_fd, + NSString* key, + bool remove, + StorageRemovalMode mode) { ET_CHECK_OR_RETURN_ERROR( storage_key(key), InvalidArgument, "Invalid Core AI staging key"); FileDescriptor staging(retry_eintr([&] { @@ -392,10 +453,14 @@ FileDescriptor staging(retry_eintr([&] { staging.get() >= 0, AccessFailed, "Cannot open Core AI staging directory"); - return remove_storage_source(staging.get(), key, remove); + return remove_storage_source(staging.get(), key, remove, mode); } -Error remove_storage_source(int source_root_fd, NSString* key, bool remove) { +Error remove_storage_source( + int source_root_fd, + NSString* key, + bool remove, + StorageRemovalMode mode) { ET_CHECK_OR_RETURN_ERROR( storage_key(key), InvalidArgument, "Invalid Core AI staging key"); struct stat info; @@ -412,10 +477,13 @@ Error remove_storage_source(int source_root_fd, NSString* key, bool remove) { inspected == 0, AccessFailed, "Cannot inspect Core AI keyed staging directory"); - ET_CHECK_OR_RETURN_ERROR( - S_ISDIR(info.st_mode), - InvalidExternalData, - "Core AI keyed staging entry is not a directory"); + if (!S_ISDIR(info.st_mode)) { + ET_CHECK_OR_RETURN_ERROR( + mode == StorageRemovalMode::Maintenance, + InvalidExternalData, + "Core AI keyed staging entry is not a directory"); + return remove_storage_entry(source_root_fd, key, remove); + } FileDescriptor source(retry_eintr([&] { return openat( source_root_fd, @@ -426,9 +494,10 @@ FileDescriptor source(retry_eintr([&] { source.get() >= 0, AccessFailed, "Cannot open Core AI keyed staging directory"); - ET_CHECK_OK_OR_RETURN_ERROR(visit_source_tree(source.get(), false)); - if (!remove) return Error::Ok; - ET_CHECK_OK_OR_RETURN_ERROR(visit_source_tree(source.get(), true)); + ET_CHECK_OK_OR_RETURN_ERROR(visit_source_tree(source.get(), false, mode)); + if (!remove) + return Error::Ok; + ET_CHECK_OK_OR_RETURN_ERROR(visit_source_tree(source.get(), true, mode)); ET_CHECK_OR_RETURN_ERROR( storage_call( StorageOperation::Remove, diff --git a/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm index b70ae5380d6..99735b4ce1f 100644 --- a/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm @@ -99,6 +99,89 @@ ino_t lock_inode() { } }; +TEST_F( + CoreAIAcquisitionTest, + AttributionPublicationFailurePrecedesSdkAndSourceAccess) { + BookmarkLoaderScope fake; + { + StorageFaultScope fault(StorageOperation::Rename); + EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + } + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(data.attempts.load(), 0); + EXPECT_EQ(saved(), nil); + EXPECT_FALSE(staging_key_exists()); + EXPECT_EQ(acquire(fake.loader, &data), Error::Ok); + auto locked = lock_bookmark(root.url.path, key, false); + ASSERT_TRUE(locked.ok()); + ASSERT_NE(locked.get(), nullptr); + auto identity = read_bookmark_identity(*locked.get()); + ASSERT_TRUE(identity.ok()); + EXPECT_NE(identity.get(), nil); +} + +TEST_F( + CoreAIAcquisitionTest, + AttributionFlushFailureMustBeRetriedBeforeSdkAccess) { + BookmarkLoaderScope fake; + { + StorageFaultScope fault(StorageOperation::FullSync); + fault.match_index = 2; + EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + } + { + auto locked = lock_bookmark(root.url.path, key, false); + ASSERT_TRUE(locked.ok()); + ASSERT_NE(locked.get(), nullptr); + auto identity = read_bookmark_identity(*locked.get()); + ASSERT_TRUE(identity.ok()); + EXPECT_NE(identity.get(), nil); + } + { + StorageFaultScope fault(StorageOperation::FullSync); + EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + } + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(data.attempts.load(), 0); + EXPECT_EQ(saved(), nil); + EXPECT_FALSE(staging_key_exists()); + EXPECT_EQ(acquire(fake.loader, &data), Error::Ok); +} + +TEST_F(CoreAIAcquisitionTest, WarmLegacyEntryGetsAttributionWithoutRestaging) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + NSData* original = saved(); + { + auto locked = lock_bookmark(root.url.path, key, false); + ASSERT_TRUE(locked.ok()); + ASSERT_NE(locked.get(), nullptr); + ASSERT_EQ(remove_bookmark_identity(*locked.get()), Error::Ok); + } + { + StorageFaultScope fault(StorageOperation::Write); + EXPECT_EQ(acquire(fake.loader, nullptr), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + } + EXPECT_TRUE([saved() isEqual:original]); + EXPECT_EQ(fake.loader->restores.load(), 0); + ASSERT_EQ(acquire(fake.loader, nullptr), Error::Ok); + EXPECT_EQ(fake.loader->restores.load(), 1); + EXPECT_EQ(fake.loader->specializations.load(), 1); + EXPECT_EQ(data.attempts.load(), 1); + auto locked = lock_bookmark(root.url.path, key, false); + ASSERT_TRUE(locked.ok()); + ASSERT_NE(locked.get(), nullptr); + auto identity = read_bookmark_identity(*locked.get()); + ASSERT_TRUE(identity.ok()); + EXPECT_NE(identity.get(), nil); +} + TEST_F(CoreAIAcquisitionTest, ColdSpecializesStagedSourceAndPersistsBookmark) { BookmarkLoaderScope fake; ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); @@ -365,13 +448,19 @@ ino_t lock_inode() { TEST_F(CoreAIAcquisitionTest, BookmarkFinalFlushFailurePerformsNoCleanup) { BookmarkLoaderScope fake; + { + auto locked = lock_bookmark(root.url.path, key); + ASSERT_TRUE(locked.ok()); + ASSERT_EQ( + ensure_bookmark_identity(*locked.get(), manifest, @"macOS", @"arch_b"), + Error::Ok); + } __block std::vector ops; Error error = Error::Internal; { - // Cold source contributes two FullSyncs and the bookmark two more, so - // index four is the post-rename bookmark flush. + // Attribution reuse flushes once, then source and bookmark twice each. StorageFaultScope fault(StorageOperation::FullSync); - fault.match_index = 4; + fault.match_index = 5; fault.observe = ^(StorageOperation operation) { ops.push_back(operation); }; @@ -590,9 +679,14 @@ ino_t lock_inode() { fake.loader.restoreMiss = YES; } { - // Cold and rematerialized sources publish before the bookmark, so the - // bookmark rename is the second rename once the first success has - // cleaned the source. + auto locked = lock_bookmark(root.url.path, key); + ASSERT_TRUE(locked.ok()); + ASSERT_EQ( + ensure_bookmark_identity(*locked.get(), manifest, @"macOS", @"arch_b"), + Error::Ok); + } + { + // With attribution present, the bookmark follows source publication. StorageFaultScope fault(StorageOperation::Rename); fault.match_index = 2; EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); diff --git a/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm b/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm index 7d9fc96feb6..9e2801e0ae6 100644 --- a/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm @@ -144,6 +144,11 @@ FileDescriptor fd(open( return ::testing::AssertionSuccess(); } +NSURL* identity_url(NSURL* root, NSString* key) { + return [[root URLByAppendingPathComponent:@"bookmarks"] + URLByAppendingPathComponent:[key stringByAppendingString:@".metadata"]]; +} + TEST_F(CoreAIBookmarkIdentityTest, IgnoresFunctionBindingsAndDeliveryLocation) { EXPECT_EQ(key.length, 64u); manifest.function = @"another_function"; @@ -575,6 +580,371 @@ FileDescriptor file(open( EXPECT_EQ(before.st_ino, after.st_ino); } +TEST_F(CoreAIBookmarkIdentityTest, IdentityExcludesContextAndFunction) { + NSString* bundle = manifest.path.lastPathComponent; + auto identity = bookmark_identity(manifest, bundle); + ASSERT_TRUE(identity.ok()); + EXPECT_EQ(identity.get().length, 64u); + manifest.function = @"another_function"; + manifest.inputs = @[ @"different_input" ]; + manifest.outputs = @[ @"different_output" ]; + manifest.hash = @"another_delivery_location"; + manifest.minimum_version = (NSOperatingSystemVersion){27, 0, 0}; + auto unchanged = bookmark_identity(manifest, bundle); + ASSERT_TRUE(unchanged.ok()); + EXPECT_TRUE([unchanged.get() isEqual:identity.get()]); + auto context = bookmark_key(manifest, @"iOS", @"arch_a"); + ASSERT_TRUE(context.ok()); + EXPECT_FALSE([context.get() isEqual:key]); +} + +TEST_F(CoreAIBookmarkIdentityTest, IdentitySeparatesDigestBundleAndKind) { + NSString* bundle = manifest.path.lastPathComponent; + auto identity = bookmark_identity(manifest, bundle); + ASSERT_TRUE(identity.ok()); + Manifest reweighted = manifest; + reweighted.bundle_digests = @{bundle : fixture_identity('2')}; + auto weights = bookmark_identity(reweighted, bundle); + ASSERT_TRUE(weights.ok()); + EXPECT_FALSE([weights.get() isEqual:identity.get()]); + Manifest renamed = manifest; + renamed.bundle_digests = + @{@"other.aimodel" : manifest.bundle_digests[bundle]}; + auto bundle_name = bookmark_identity(renamed, @"other.aimodel"); + ASSERT_TRUE(bundle_name.ok()); + EXPECT_FALSE([bundle_name.get() isEqual:identity.get()]); + Manifest compiled = manifest; + compiled.aot_compiled = true; + auto kind = bookmark_identity(compiled, bundle); + ASSERT_TRUE(kind.ok()); + EXPECT_FALSE([kind.get() isEqual:identity.get()]); +} + +TEST_F( + CoreAIBookmarkIdentityTest, + IdentityRejectsUnselectedAndEscapingBundles) { + EXPECT_FALSE(bookmark_identity(manifest, @"other.aimodel").ok()); + Manifest escaped = manifest; + escaped.bundle_digests = @{@"../escape" : fixture_identity('1')}; + EXPECT_FALSE(bookmark_identity(escaped, @"../escape").ok()); +} + +TEST_F(CoreAIBookmarkRawTest, EnsurePublishesFullKeyFieldsReadableAsIdentity) { + ASSERT_EQ( + ensure_bookmark_identity(*locked, manifest, @"macOS", @"arch_b"), + Error::Ok); + auto identity = read_bookmark_identity(*locked); + ASSERT_TRUE(identity.ok()); + auto expected = bookmark_identity(manifest, manifest.path.lastPathComponent); + ASSERT_TRUE(expected.ok()); + EXPECT_TRUE([identity.get() isEqual:expected.get()]); + NSData* bytes = [NSData dataWithContentsOfURL:identity_url(root.url, key)]; + ASSERT_NE(bytes, nil); + id parsed = [NSJSONSerialization JSONObjectWithData:bytes + options:0 + error:nil]; + ASSERT_TRUE([parsed isKindOfClass:NSArray.class]); + NSArray* fields = parsed; + ASSERT_EQ(fields.count, 9u); + EXPECT_TRUE([fields[0] isEqual:@"1"]); + EXPECT_TRUE([fields[1] + isEqual:manifest.bundle_digests[manifest.path.lastPathComponent]]); + EXPECT_TRUE([fields[2] isEqual:@"source"]); + EXPECT_TRUE([fields[3] isEqual:manifest.path.lastPathComponent]); + EXPECT_TRUE([fields[4] isEqual:@"macOS"]); + EXPECT_TRUE([fields[5] isEqual:@"arch_b"]); + for (id field in fields) { + EXPECT_TRUE([field isKindOfClass:NSString.class] && [field length] > 0); + } + EXPECT_TRUE(no_temporary_bookmarks(root.url)); +} + +TEST_F(CoreAIBookmarkRawTest, EnsureReusesMatchingRecordWithoutRepublishing) { + ASSERT_EQ( + ensure_bookmark_identity(*locked, manifest, @"macOS", @"arch_b"), + Error::Ok); + NSData* published = + [NSData dataWithContentsOfURL:identity_url(root.url, key)]; + ASSERT_NE(published, nil); + for (StorageOperation operation : + {StorageOperation::Write, StorageOperation::Rename}) { + SCOPED_TRACE(static_cast(operation)); + StorageFaultScope fault(operation); + EXPECT_EQ( + ensure_bookmark_identity(*locked, manifest, @"macOS", @"arch_b"), + Error::Ok); + EXPECT_EQ(fault.hits, 0); + } + EXPECT_TRUE([[NSData dataWithContentsOfURL:identity_url(root.url, key)] + isEqual:published]); + EXPECT_TRUE(no_temporary_bookmarks(root.url)); +} + +TEST_F(CoreAIBookmarkRawTest, MissingIdentityReadsAsNil) { + auto identity = read_bookmark_identity(*locked); + ASSERT_TRUE(identity.ok()); + EXPECT_EQ(identity.get(), nil); +} + +TEST_F(CoreAIBookmarkRawTest, EnsureRepairsCorruptAndMisattributedRecords) { + NSURL* metadata = identity_url(root.url, key); + NSData* garbage = [@"not metadata" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([garbage writeToURL:metadata atomically:NO]); + ASSERT_EQ( + ensure_bookmark_identity(*locked, manifest, @"macOS", @"arch_b"), + Error::Ok); + auto repaired = read_bookmark_identity(*locked); + ASSERT_TRUE(repaired.ok()); + auto expected = bookmark_identity(manifest, manifest.path.lastPathComponent); + ASSERT_TRUE(expected.ok()); + EXPECT_TRUE([repaired.get() isEqual:expected.get()]); + Manifest reweighted = manifest; + reweighted.bundle_digests = + @{manifest.path.lastPathComponent : fixture_identity('2')}; + auto foreign_key = test_bookmark_key(reweighted); + ASSERT_TRUE(foreign_key.ok()); + auto foreign = lock_bookmark(root.url.path, foreign_key.get()); + ASSERT_TRUE(foreign.ok()); + ASSERT_EQ( + ensure_bookmark_identity(*foreign.get(), reweighted, @"macOS", @"arch_b"), + Error::Ok); + NSData* foreign_bytes = + [NSData dataWithContentsOfURL:identity_url(root.url, foreign_key.get())]; + ASSERT_NE(foreign_bytes, nil); + ASSERT_TRUE([foreign_bytes writeToURL:metadata atomically:NO]); + ASSERT_EQ( + ensure_bookmark_identity(*locked, manifest, @"macOS", @"arch_b"), + Error::Ok); + auto restored = read_bookmark_identity(*locked); + ASSERT_TRUE(restored.ok()); + EXPECT_TRUE([restored.get() isEqual:expected.get()]); +} + +TEST_F(CoreAIBookmarkRawTest, RejectsMalformedIdentityRecords) { + NSString* bundle = manifest.path.lastPathComponent; + NSString* digest = manifest.bundle_digests[bundle]; + NSArray* base = @[ + @"1", + digest, + @"source", + bundle, + @"macOS", + @"arch_b", + @"AIModelCache.default", + @"default-options-v1", + @"persistent" + ]; + NSMutableArray* payloads = [NSMutableArray array]; + [payloads addObject:[@"not json" dataUsingEncoding:NSUTF8StringEncoding]]; + for (id document in @[ + @{@"version" : @"1"}, + @"metadata", + [base subarrayWithRange:NSMakeRange(0, 8)], + [base arrayByAddingObject:@"extra"] + ]) { + [payloads addObject:[NSJSONSerialization + dataWithJSONObject:document + options:NSJSONWritingFragmentsAllowed + error:nil]]; + } + NSMutableArray* numeric = [base mutableCopy]; + numeric[1] = @1; + NSMutableArray* empty = [base mutableCopy]; + empty[4] = @""; + NSMutableArray* kind = [base mutableCopy]; + kind[2] = @"compiled"; + NSMutableArray* undigestable = [base mutableCopy]; + undigestable[1] = [@"" stringByPaddingToLength:64 + withString:@"z" + startingAtIndex:0]; + NSMutableArray* nested = [base mutableCopy]; + nested[3] = @"a/b"; + NSMutableArray* retargeted = [base mutableCopy]; + retargeted[8] = @"other"; + for (NSArray* variant in + @[ numeric, empty, kind, undigestable, nested, retargeted ]) { + [payloads addObject:[NSJSONSerialization dataWithJSONObject:variant + options:0 + error:nil]]; + } + NSURL* metadata = identity_url(root.url, key); + for (NSData* payload in payloads) { + SCOPED_TRACE(payload.length); + ASSERT_NE(payload, nil); + ASSERT_TRUE([payload writeToURL:metadata atomically:NO]); + EXPECT_EQ( + read_bookmark_identity(*locked).error(), Error::InvalidExternalData); + EXPECT_TRUE([[NSData dataWithContentsOfURL:metadata] isEqual:payload]); + } +} + +TEST_F(CoreAIBookmarkRawTest, RejectsOversizedIdentityRecord) { + NSURL* metadata = identity_url(root.url, key); + FileDescriptor file(open( + metadata.fileSystemRepresentation, + O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, + 0600)); + ASSERT_GE(file.get(), 0); + ASSERT_EQ(ftruncate(file.get(), 128 * 1024 + 1), 0); + ASSERT_EQ(close(file.release()), 0); + EXPECT_EQ( + read_bookmark_identity(*locked).error(), Error::InvalidExternalData); + struct stat info {}; + ASSERT_EQ(lstat(metadata.fileSystemRepresentation, &info), 0); + EXPECT_EQ(info.st_size, static_cast(128 * 1024 + 1)); +} + +TEST_F(CoreAIBookmarkRawTest, RejectsLinkedIdentityRecordWithoutReadingTarget) { + NSURL* metadata = identity_url(root.url, key); + NSURL* target = [root.url URLByAppendingPathComponent:@"target"]; + NSData* bytes = [@"target bytes" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([bytes writeToURL:target atomically:NO]); + ASSERT_EQ( + symlink( + target.fileSystemRepresentation, metadata.fileSystemRepresentation), + 0); + EXPECT_EQ( + read_bookmark_identity(*locked).error(), Error::InvalidExternalData); + EXPECT_TRUE(([[NSData dataWithContentsOfURL:target] isEqual:bytes])); + struct stat info {}; + ASSERT_EQ(lstat(metadata.fileSystemRepresentation, &info), 0); + EXPECT_TRUE(S_ISLNK(info.st_mode)); +} + +TEST_F(CoreAIBookmarkRawTest, RejectsDirectoryAndFifoIdentityRecords) { + NSURL* metadata = identity_url(root.url, key); + ASSERT_EQ(mkdir(metadata.fileSystemRepresentation, 0700), 0); + EXPECT_EQ( + read_bookmark_identity(*locked).error(), Error::InvalidExternalData); + ASSERT_EQ(rmdir(metadata.fileSystemRepresentation), 0); + ASSERT_EQ(mkfifo(metadata.fileSystemRepresentation, 0600), 0); + EXPECT_EQ( + read_bookmark_identity(*locked).error(), Error::InvalidExternalData); +} + +TEST_F(CoreAIBookmarkRawTest, RemoveIdentityPreservesBookmarkAndStaging) { + NSData* bytes = [@"bookmark" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_EQ(write_bookmark(*locked, bytes), Error::Ok); + ASSERT_EQ( + ensure_bookmark_identity(*locked, manifest, @"macOS", @"arch_b"), + Error::Ok); + auto staging = prepare_bookmark_staging(*locked); + ASSERT_TRUE(staging.ok()); + ASSERT_EQ(remove_bookmark_identity(*locked), Error::Ok); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:identity_url(root.url, key).path]); + EXPECT_TRUE(bookmark_bytes_equal(*locked, bytes)); + EXPECT_TRUE([NSFileManager.defaultManager fileExistsAtPath:staging.get()]); + auto reread = read_bookmark_identity(*locked); + ASSERT_TRUE(reread.ok()); + EXPECT_EQ(reread.get(), nil); + EXPECT_EQ(remove_bookmark_identity(*locked), Error::Ok); +} + +TEST_F(CoreAIBookmarkPathTest, EnsureRejectsSelectedOrContextOutsideHeldKey) { + auto locked = lock_bookmark(root.url.path, key); + ASSERT_TRUE(locked.ok()); + Manifest reweighted = manifest; + reweighted.bundle_digests = + @{manifest.path.lastPathComponent : fixture_identity('2')}; + EXPECT_EQ( + ensure_bookmark_identity(*locked.get(), reweighted, @"macOS", @"arch_b"), + Error::InvalidArgument); + EXPECT_EQ( + ensure_bookmark_identity(*locked.get(), manifest, @"iOS", @"arch_b"), + Error::InvalidArgument); + EXPECT_EQ( + ensure_bookmark_identity(*locked.get(), manifest, @"macOS", @"arch_a"), + Error::InvalidArgument); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:identity_url(root.url, key).path]); +} + +TEST_F(CoreAIBookmarkPathTest, EnsureRejectsStagingOnlyLock) { + NSURL* staged = [[root.url URLByAppendingPathComponent:@"staging"] + URLByAppendingPathComponent:key]; + ASSERT_TRUE([NSFileManager.defaultManager createDirectoryAtURL:staged + withIntermediateDirectories:YES + attributes:nil + error:nil]); + auto existing = lock_bookmark(root.url.path, key, false); + ASSERT_TRUE(existing.ok()); + ASSERT_NE(existing.get(), nullptr); + EXPECT_EQ( + ensure_bookmark_identity(*existing.get(), manifest, @"macOS", @"arch_b"), + Error::InvalidArgument); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:identity_url(root.url, key).path]); +} + +TEST_F(CoreAIBookmarkPathTest, RemoveIdentityThroughStagingOnlyLockSucceeds) { + NSURL* staged = [[root.url URLByAppendingPathComponent:@"staging"] + URLByAppendingPathComponent:key]; + ASSERT_TRUE([NSFileManager.defaultManager createDirectoryAtURL:staged + withIntermediateDirectories:YES + attributes:nil + error:nil]); + auto existing = lock_bookmark(root.url.path, key, false); + ASSERT_TRUE(existing.ok()); + ASSERT_NE(existing.get(), nullptr); + EXPECT_EQ(remove_bookmark_identity(*existing.get()), Error::Ok); +} + +TEST_F(CoreAIBookmarkPathTest, RejectsMetadataWrittenForAnotherKey) { + Manifest reweighted = manifest; + reweighted.bundle_digests = + @{manifest.path.lastPathComponent : fixture_identity('2')}; + auto foreign_key = test_bookmark_key(reweighted); + ASSERT_TRUE(foreign_key.ok()); + auto foreign = lock_bookmark(root.url.path, foreign_key.get()); + ASSERT_TRUE(foreign.ok()); + ASSERT_EQ( + ensure_bookmark_identity(*foreign.get(), reweighted, @"macOS", @"arch_b"), + Error::Ok); + NSData* foreign_bytes = + [NSData dataWithContentsOfURL:identity_url(root.url, foreign_key.get())]; + ASSERT_NE(foreign_bytes, nil); + auto locked = lock_bookmark(root.url.path, key); + ASSERT_TRUE(locked.ok()); + ASSERT_TRUE([foreign_bytes writeToURL:identity_url(root.url, key) + atomically:NO]); + EXPECT_EQ( + read_bookmark_identity(*locked.get()).error(), + Error::InvalidExternalData); + EXPECT_TRUE([[NSData dataWithContentsOfURL:identity_url(root.url, key)] + isEqual:foreign_bytes]); +} + +TEST_F(CoreAIBookmarkPathTest, MetadataOnlyEntriesListLockAndRemove) { + auto locked = lock_bookmark(root.url.path, key); + ASSERT_TRUE(locked.ok()); + ASSERT_EQ( + ensure_bookmark_identity(*locked.get(), manifest, @"macOS", @"arch_b"), + Error::Ok); + locked.get().reset(); + auto inventory = inventory_bookmark_keys(root.url.path); + ASSERT_TRUE(inventory.ok()); + EXPECT_TRUE([inventory.get() containsObject:key]); + auto existing = lock_bookmark(root.url.path, key, false); + ASSERT_TRUE(existing.ok()); + ASSERT_NE(existing.get(), nullptr); + auto identity = read_bookmark_identity(*existing.get()); + ASSERT_TRUE(identity.ok()); + auto expected = bookmark_identity(manifest, manifest.path.lastPathComponent); + ASSERT_TRUE(expected.ok()); + EXPECT_TRUE([identity.get() isEqual:expected.get()]); + EXPECT_TRUE(bookmark_bytes_equal(*existing.get(), nil)); + ASSERT_EQ(remove_bookmark_identity(*existing.get()), Error::Ok); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:identity_url(root.url, key).path]); + auto relisted = inventory_bookmark_keys(root.url.path); + ASSERT_TRUE(relisted.ok()); + EXPECT_FALSE([relisted.get() containsObject:key]); + auto absent = lock_bookmark(root.url.path, key, false); + ASSERT_TRUE(absent.ok()); + EXPECT_EQ(absent.get(), nullptr); +} + TEST_F(CoreAIBookmarkPathTest, LockNamesUseOnlyValidatedFirstByte) { EXPECT_EQ(bookmark_lock_name(nil).error(), Error::InvalidArgument); for (NSString* invalid in @@ -591,6 +961,44 @@ FileDescriptor file(open( isEqual:test_bookmark_lock_url(root.url, changed_key(key, true))]); } +TEST_F(CoreAIBookmarkPathTest, LockPoolIsBoundedAcrossManyKeys) { + std::array inodes{}; + NSURL* locks = [root.url URLByAppendingPathComponent:@"locks"]; + for (size_t generation = 0; generation < 4; ++generation) { + for (size_t stripe = 0; stripe < kBookmarkLockStripeCount; ++stripe) { + @autoreleasepool { + NSString* entry = + [NSString stringWithFormat:@"%02zx%062zx", stripe, generation]; + auto locked = lock_bookmark(root.url.path, entry); + ASSERT_TRUE(locked.ok()); + struct stat info {}; + NSURL* file = test_bookmark_lock_url(root.url, entry); + ASSERT_EQ(lstat(file.fileSystemRepresentation, &info), 0); + ASSERT_TRUE(S_ISREG(info.st_mode)); + EXPECT_EQ(info.st_size, 0); + if (generation == 0) + inodes[stripe] = info.st_ino; + else + EXPECT_EQ(info.st_ino, inodes[stripe]); + locked.get().reset(); + auto reused = lock_bookmark(root.url.path, entry); + ASSERT_TRUE(reused.ok()); + ASSERT_EQ(lstat(file.fileSystemRepresentation, &info), 0); + EXPECT_EQ(info.st_ino, inodes[stripe]); + } + } + NSArray* children = + [NSFileManager.defaultManager contentsOfDirectoryAtPath:locks.path + error:nil]; + ASSERT_NE(children, nil); + ASSERT_EQ(children.count, kBookmarkLockStripeCount); + for (NSString* name in children) { + EXPECT_TRUE([name hasPrefix:@"stripe-v1-"]); + EXPECT_EQ(name.length, 17u); + } + } +} + TEST_F(CoreAIBookmarkPathTest, CollidingKeysSerializeAcrossThreads) { auto holder = lock_bookmark(root.url.path, key); ASSERT_TRUE(holder.ok()); @@ -630,6 +1038,7 @@ FileDescriptor file(open( EXPECT_EQ(names.count, 1u); } + TEST_F(CoreAIBookmarkPathTest, RejectsSymlinkDirectoryAndFifoStripeFiles) { NSData* bytes = [@"untouched target" dataUsingEncoding:NSUTF8StringEncoding]; NSURL* target = [root.url URLByAppendingPathComponent:@"target"]; @@ -693,42 +1102,4 @@ FileDescriptor file(open( EXPECT_EQ(before.st_ino, after.st_ino); } -TEST_F(CoreAIBookmarkPathTest, LockPoolIsBoundedAcrossManyKeys) { - std::array inodes{}; - NSURL* locks = [root.url URLByAppendingPathComponent:@"locks"]; - for (size_t generation = 0; generation < 4; ++generation) { - for (size_t stripe = 0; stripe < kBookmarkLockStripeCount; ++stripe) { - @autoreleasepool { - NSString* entry = - [NSString stringWithFormat:@"%02zx%062zx", stripe, generation]; - auto locked = lock_bookmark(root.url.path, entry); - ASSERT_TRUE(locked.ok()); - struct stat info {}; - NSURL* file = test_bookmark_lock_url(root.url, entry); - ASSERT_EQ(lstat(file.fileSystemRepresentation, &info), 0); - ASSERT_TRUE(S_ISREG(info.st_mode)); - EXPECT_EQ(info.st_size, 0); - if (generation == 0) - inodes[stripe] = info.st_ino; - else - EXPECT_EQ(info.st_ino, inodes[stripe]); - locked.get().reset(); - auto reused = lock_bookmark(root.url.path, entry); - ASSERT_TRUE(reused.ok()); - ASSERT_EQ(lstat(file.fileSystemRepresentation, &info), 0); - EXPECT_EQ(info.st_ino, inodes[stripe]); - } - } - NSArray* children = - [NSFileManager.defaultManager contentsOfDirectoryAtPath:locks.path - error:nil]; - ASSERT_NE(children, nil); - ASSERT_EQ(children.count, kBookmarkLockStripeCount); - for (NSString* name in children) { - EXPECT_TRUE([name hasPrefix:@"stripe-v1-"]); - EXPECT_EQ(name.length, 17u); - } - } -} - } // namespace diff --git a/backends/apple/coreai/runtime/test/coreai_cache_test.mm b/backends/apple/coreai/runtime/test/coreai_cache_test.mm index 5878c574e07..c9e55500c03 100644 --- a/backends/apple/coreai/runtime/test/coreai_cache_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_cache_test.mm @@ -6,13 +6,19 @@ * LICENSE file in the root directory of this source tree. */ -#include "coreai_cache.h" +#include #include +#include #include +#include +#include #include #include +#include +#include #include #include "coreai_bookmark_fixture.h" +#include "coreai_cache.h" #include "coreai_fake_loader.h" #include "coreai_fault_scope.h" #include "coreai_filesystem_fixture.h" @@ -30,10 +36,13 @@ class CoreAICacheTest : public ::testing::Test { protected: - void SetUp() override { runtime::runtime_init(); } + void SetUp() override { + runtime::runtime_init(); + } void drain_callbacks() { - auto drained = bridge.wait_for_callbacks(dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC)); + auto drained = bridge.wait_for_callbacks( + dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC)); if (!drained) { ADD_FAILURE() << drained.message(); // Timed-out callbacks may still borrow this test's PTE or observations. @@ -50,9 +59,16 @@ void drain_callbacks() { NSURL* stage = nil; }; -AssertionResult seed_cache(NSURL* root, NSDictionary* dict, BookmarkFakeLoader* loader, - CacheEntry& output, bool bookmark = true, bool stage = true, - NSString* architecture = @"arch_b") { +AssertionResult seed_cache( + NSURL* root, + NSDictionary* dict, + BookmarkFakeLoader* loader, + CacheEntry& output, + bool bookmark = true, + bool stage = true, + NSString* architecture = @"arch_b", + NSString* platform = @"macOS", + bool metadata = true) { if (root == nil || loader == nil) { return AssertionFailure() << "Cache seed requires a root and fake loader"; } @@ -60,21 +76,28 @@ AssertionResult seed_cache(NSURL* root, NSDictionary* dict, BookmarkFakeLoader* if (!parsed.ok()) { return AssertionFailure() << "Cannot parse cache seed manifest"; } - auto selected = select_assets(parsed.get(), architecture, @"macOS"); + auto selected = select_assets(parsed.get(), architecture, platform); if (!selected.ok()) { return AssertionFailure() << "Cannot select cache seed assets"; } - auto key = bookmark_key(selected.get(), @"macOS", architecture); + auto key = bookmark_key(selected.get(), platform, architecture); if (!key.ok()) { return AssertionFailure() << "Cannot compute cache seed key"; } CacheEntry entry{ - key.get(), nil, - [[root URLByAppendingPathComponent:@"staging"] URLByAppendingPathComponent:key.get()]}; + key.get(), + nil, + [[root URLByAppendingPathComponent:@"staging"] + URLByAppendingPathComponent:key.get()]}; auto lock = lock_bookmark(root.path, entry.key); if (!lock.ok()) { return AssertionFailure() << "Cannot lock cache seed bookmark"; } + if (metadata && + ensure_bookmark_identity( + *lock.get(), selected.get(), platform, architecture) != Error::Ok) { + return AssertionFailure() << "Cannot attribute cache seed"; + } if (stage) { auto staging = prepare_bookmark_staging(*lock.get()); if (!staging.ok()) { @@ -84,17 +107,19 @@ AssertionResult seed_cache(NSURL* root, NSDictionary* dict, BookmarkFakeLoader* if (selected->aot_compiled) { aot_data(data, architecture); } - auto staged = executorch::backends::coreai::prepare_source_bundle(selected.get(), &data, - staging.get(), entry.key); + auto staged = executorch::backends::coreai::prepare_source_bundle( + selected.get(), &data, staging.get(), entry.key); if (data.requests != data.releases) { - return AssertionFailure() << "Cache seed source buffers were not released"; + return AssertionFailure() + << "Cache seed source buffers were not released"; } if (!staged.ok()) { return AssertionFailure() << "Cannot stage cache seed source bundle"; } } if (bookmark) { - entry.bookmark = [NSUUID.UUID.UUIDString dataUsingEncoding:NSUTF8StringEncoding]; + entry.bookmark = + [NSUUID.UUID.UUIDString dataUsingEncoding:NSUTF8StringEncoding]; BookmarkFakeModel* model = [[BookmarkFakeModel alloc] init]; if (model == nil || entry.bookmark == nil) { return AssertionFailure() << "Cannot create cache seed model/bookmark"; @@ -110,13 +135,22 @@ AssertionResult seed_cache(NSURL* root, NSDictionary* dict, BookmarkFakeLoader* return AssertionSuccess(); } -bool cache_exists(NSURL* url) { return [NSFileManager.defaultManager fileExistsAtPath:url.path]; } +bool cache_exists(NSURL* url) { + return [NSFileManager.defaultManager fileExistsAtPath:url.path]; +} + +NSURL* cache_identity_url(NSURL* root, NSString* key) { + return [[test_bookmark_url(root, key) URLByDeletingPathExtension] + URLByAppendingPathExtension:@"metadata"]; +} void cache_entry_removed(NSURL* root, const CacheEntry& entry) { + EXPECT_FALSE(cache_exists(cache_identity_url(root, entry.key))); EXPECT_FALSE(cache_exists(test_bookmark_url(root, entry.key))); EXPECT_FALSE(cache_exists(entry.stage)); EXPECT_TRUE(cache_exists([[root URLByAppendingPathComponent:@"locks"] - URLByAppendingPathComponent:[entry.key stringByAppendingString:@".lock"]])); + URLByAppendingPathComponent:[entry.key + stringByAppendingString:@".lock"]])); } NSMutableDictionary* cache_manifest(char identity) { @@ -127,7 +161,7 @@ void cache_entry_removed(NSURL* root, const CacheEntry& entry) { TEST_F(CoreAICacheTest, EmptyAndMissingRootsAreNoncreatingNoops) { @autoreleasepool { - BookmarkDirectory root; + TestDirectory root; ASSERT_NE(root.url, nil); BookmarkLoaderScope fake; ASSERT_NE(fake.loader, nil); @@ -138,9 +172,15 @@ void cache_entry_removed(NSURL* root, const CacheEntry& entry) { bridge.state().available = false; EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); EXPECT_EQ(clear_cache_for_pte(empty, root.url.path.UTF8String), Error::Ok); - NSURL* missing = [root.url URLByAppendingPathComponent:@"absent/nested/models"]; + NSURL* missing = + [root.url URLByAppendingPathComponent:@"absent/nested/models"]; EXPECT_EQ(clear_cache(missing.path.UTF8String), Error::Ok); EXPECT_EQ(clear_cache_for_pte(empty, missing.path.UTF8String), Error::Ok); + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", encode(manifest_dict())}}}}); + EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); + EXPECT_EQ(clear_cache_for_pte(pte, missing.path.UTF8String), Error::Ok); + EXPECT_TRUE(pte.check_released()); EXPECT_FALSE(cache_exists(missing)); EXPECT_TRUE(snapshot_matches(root.url, snapshot)); EXPECT_EQ(fake.loader->evictions.load(), 0); @@ -153,7 +193,7 @@ void cache_entry_removed(NSURL* root, const CacheEntry& entry) { TEST_F(CoreAICacheTest, ValidatesExplicitRootsAndPtePaths) { @autoreleasepool { - BookmarkDirectory root; + TestDirectory root; ASSERT_NE(root.url, nil); BookmarkLoaderScope fake; ASSERT_NE(fake.loader, nil); @@ -162,37 +202,41 @@ void cache_entry_removed(NSURL* root, const CacheEntry& entry) { EXPECT_EQ(clear_cache_for_pte(empty, "relative"), Error::InvalidArgument); EXPECT_EQ(clear_cache(""), Error::InvalidArgument); EXPECT_EQ(clear_cache("\xff"), Error::InvalidArgument); - EXPECT_EQ(clear_cache_for_pte(static_cast(nullptr), root.url.path.UTF8String), - Error::InvalidArgument); + EXPECT_EQ( + clear_cache_for_pte( + static_cast(nullptr), root.url.path.UTF8String), + Error::InvalidArgument); EXPECT_TRUE(empty.check_released()); } } TEST_F(CoreAICacheTest, SelectsAllMethodsDeduplicatesAndEvictsInKeyOrder) { @autoreleasepool { - BookmarkDirectory root; - BookmarkDirectory isolated; + TestDirectory root; + TestDirectory isolated; BookmarkLoaderScope fake; auto source_dict = manifest_dict(); auto second_dict = cache_manifest('2'); auto aot_dict = aot_manifest_dict(); - CacheEntry source, second, aot, unselected, unrelated, other_root; + CacheEntry source, second, aot, other_arch, unrelated, other_root; ASSERT_TRUE(seed_cache(root.url, source_dict, fake.loader, source)); ASSERT_TRUE(seed_cache(root.url, second_dict, fake.loader, second)); ASSERT_TRUE(seed_cache(root.url, aot_dict, fake.loader, aot)); - ASSERT_TRUE(seed_cache(root.url, aot_dict, fake.loader, unselected, true, true, @"arch_a")); - ASSERT_TRUE(seed_cache(root.url, cache_manifest('3'), fake.loader, unrelated)); + ASSERT_TRUE(seed_cache( + root.url, aot_dict, fake.loader, other_arch, true, true, @"arch_a")); + ASSERT_TRUE( + seed_cache(root.url, cache_manifest('3'), fake.loader, unrelated)); ASSERT_TRUE(seed_cache(isolated.url, source_dict, fake.loader, other_root)); NSDictionary* isolation = nil; - NSDictionary* unselected_stage = nil; NSDictionary* unrelated_stage = nil; ASSERT_TRUE(asset_tree_snapshot(isolated.url, isolation)); - ASSERT_TRUE(asset_tree_snapshot(unselected.stage, unselected_stage)); ASSERT_TRUE(asset_tree_snapshot(unrelated.stage, unrelated_stage)); FBSyntheticPTE pte( {{"empty", {}}, {"first", - {{"UnregisteredBackend", [@"not JSON" dataUsingEncoding:NSUTF8StringEncoding], true}, + {{"UnregisteredBackend", + [@"not JSON" dataUsingEncoding:NSUTF8StringEncoding], + true}, {"CoreAIBackend", encode(source_dict)}, {"CoreAIBackend", encode(second_dict), true}}}, {"last", @@ -205,15 +249,20 @@ FBSyntheticPTE pte( const int queries = bridge.state().architecture_queries; EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); drain_callbacks(); - EXPECT_EQ(bridge.state().architecture_queries.load(), queries + 1); - EXPECT_EQ(fake.loader->evictions.load(), 3); + EXPECT_EQ(bridge.state().architecture_queries.load(), queries); + EXPECT_EQ(fake.loader->evictions.load(), 4); EXPECT_EQ(fake.loader->specializations.load(), 0); EXPECT_EQ(fake.loader->bindings.load(), 0); EXPECT_EQ(fake.loader->restores.load(), 0); NSArray* ordered = - [@[ source.key, second.key, aot.key ] sortedArrayUsingSelector:@selector(compare:)]; - NSDictionary* bookmarks = - @{source.key : source.bookmark, second.key : second.bookmark, aot.key : aot.bookmark}; + [@[ source.key, second.key, aot.key, other_arch.key ] + sortedArrayUsingSelector:@selector(compare:)]; + NSDictionary* bookmarks = @{ + source.key : source.bookmark, + second.key : second.bookmark, + aot.key : aot.bookmark, + other_arch.key : other_arch.bookmark + }; ASSERT_EQ(evictions.count, ordered.count); for (NSUInteger i = 0; i < ordered.count; ++i) { SCOPED_TRACE(::testing::Message() << "sorted eviction index " << i); @@ -223,31 +272,37 @@ FBSyntheticPTE pte( cache_entry_removed(root.url, second); cache_entry_removed(root.url, aot); EXPECT_NE(fake.loader.models[other_root.bookmark], nil); - EXPECT_NE(fake.loader.models[unselected.bookmark], nil); + EXPECT_EQ(fake.loader.models[other_arch.bookmark], nil); + cache_entry_removed(root.url, other_arch); EXPECT_TRUE(snapshot_matches(isolated.url, isolation)); - EXPECT_TRUE(snapshot_matches(unselected.stage, unselected_stage)); EXPECT_TRUE(snapshot_matches(unrelated.stage, unrelated_stage)); EXPECT_EQ(pte.backend_indices, (std::vector{1, 2, 3})); - EXPECT_EQ(pte.backend_descriptors, - (std::vector{"CoreAIBackend", "CoreAIBackend", "CoreAIBackend"})); + EXPECT_EQ( + pte.backend_descriptors, + (std::vector{ + "CoreAIBackend", "CoreAIBackend", "CoreAIBackend"})); EXPECT_TRUE(pte.check_released()); } } TEST_F(CoreAICacheTest, FilePathOverloadClearsReferencedEntry) { @autoreleasepool { - BookmarkDirectory root; + TestDirectory root; BookmarkLoaderScope fake; TestDirectory distribution; ASSERT_NE(distribution.url, nil); CacheEntry entry; ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, entry)); - FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", encode(manifest_dict())}}}}); + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", encode(manifest_dict())}}}}); NSURL* path = [distribution.url URLByAppendingPathComponent:@"model.pte"]; - ASSERT_TRUE([[NSData dataWithBytes:pte.bytes.data() length:pte.bytes.size()] writeToURL:path - options:0 - error:nil]); - EXPECT_EQ(clear_cache_for_pte(path.path.UTF8String, root.url.path.UTF8String), Error::Ok); + ASSERT_TRUE([[NSData dataWithBytes:pte.bytes.data() + length:pte.bytes.size()] writeToURL:path + options:0 + error:nil]); + EXPECT_EQ( + clear_cache_for_pte(path.path.UTF8String, root.url.path.UTF8String), + Error::Ok); drain_callbacks(); EXPECT_EQ(fake.loader->evictions.load(), 1); cache_entry_removed(root.url, entry); @@ -256,37 +311,22 @@ FBSyntheticPTE pte( TEST_F(CoreAICacheTest, RejectsInvalidLaterDelegatesBeforeMutation) { @autoreleasepool { - const char* rows[] = { - "malformed manifest", - "wrong AOT platform", - "backend load failure", - "newer OS", - "missing architecture"}; - for (int row = 0; row < 5; ++row) { + const char* rows[] = {"malformed manifest", "backend load failure"}; + for (int row = 0; row < 2; ++row) { SCOPED_TRACE(rows[row]); - BookmarkDirectory root; + TestDirectory root; BookmarkLoaderScope fake; CacheEntry source; ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, source)); NSDictionary* snapshot = nil; ASSERT_TRUE(asset_tree_snapshot(root.url, snapshot)); PteDelegateSpec later{"CoreAIBackend", encode(aot_manifest_dict()), true}; - if (row == 0) later.data = [@"malformed" dataUsingEncoding:NSUTF8StringEncoding]; - if (row == 1) { - auto dict = aot_manifest_dict(); - dict[@"platform"] = @"iOS"; - later.data = encode(dict); - } - if (row == 3) { - auto dict = aot_manifest_dict(); - dict[@"min_deployment_version"] = @"999999.0"; - later.data = encode(dict); - } - if (row == 4) - bridge.state().device_architecture = @"missing_arch"; + if (row == 0) + later.data = [@"malformed" dataUsingEncoding:NSUTF8StringEncoding]; FBSyntheticPTE pte( - {{"first", {{"CoreAIBackend", encode(manifest_dict())}}}, {"later", {later}}}); - if (row == 2) { + {{"first", {{"CoreAIBackend", encode(manifest_dict())}}}, + {"later", {later}}}); + if (row == 1) { pte.fault_call = 2; pte.load_fault = FBSyntheticPTE::LoadFault::Error; pte.load_error = Error::AccessFailed; @@ -295,9 +335,8 @@ FBSyntheticPTE pte( const auto error = clear_cache_for_pte(pte, root.url.path.UTF8String); drain_callbacks(); EXPECT_NE(error, Error::Ok); - if (row == 2) EXPECT_EQ(error, Error::AccessFailed); - if (row == 1 || row >= 3) - EXPECT_EQ(error, Error::DelegateInvalidCompatibility); + if (row == 1) + EXPECT_EQ(error, Error::AccessFailed); EXPECT_EQ(fake.loader->evictions.load(), 0); EXPECT_EQ(fake.loader->specializations.load(), 0); EXPECT_EQ(fake.loader->bindings.load(), 0); @@ -309,15 +348,181 @@ FBSyntheticPTE pte( } } +TEST_F(CoreAICacheTest, ClearsWarmVariantsWithoutHostCompatibility) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + auto source = manifest_dict(); + auto aot = aot_manifest_dict(); + aot[@"platform"] = @"iOS"; + CacheEntry mac, ios, compiled; + ASSERT_TRUE(seed_cache( + root.url, source, fake.loader, mac, true, false, @"unknown_mac_arch")); + ASSERT_TRUE(seed_cache( + root.url, + source, + fake.loader, + ios, + true, + false, + @"unknown_ios_arch", + @"iOS")); + ASSERT_TRUE(seed_cache( + root.url, aot, fake.loader, compiled, true, false, @"arch_a", @"iOS")); + source[@"min_deployment_version"] = @"999999.0"; + aot[@"min_deployment_version"] = @"999999.0"; + bridge.state().device_architecture = nil; + const int queries = bridge.state().architecture_queries; + FBSyntheticPTE pte( + {{"forward", + {{"CoreAIBackend", encode(source)}, + {"CoreAIBackend", encode(aot), true}}}}); + EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); + EXPECT_EQ(fake.loader->evictions.load(), 3); + EXPECT_EQ(bridge.state().architecture_queries.load(), queries); + cache_entry_removed(root.url, mac); + cache_entry_removed(root.url, ios); + cache_entry_removed(root.url, compiled); + } +} + +TEST_F(CoreAICacheTest, ClearsLegacyAotWithoutAttribution) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry first, second; + auto dict = aot_manifest_dict(); + ASSERT_TRUE(seed_cache( + root.url, + dict, + fake.loader, + first, + true, + false, + @"arch_a", + @"macOS", + false)); + ASSERT_TRUE(seed_cache( + root.url, + dict, + fake.loader, + second, + true, + false, + @"arch_b", + @"macOS", + false)); + FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", encode(dict)}}}}); + EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); + EXPECT_EQ(fake.loader->evictions.load(), 2); + cache_entry_removed(root.url, first); + cache_entry_removed(root.url, second); + } +} + +TEST_F( + CoreAICacheTest, + UnattributedEntriesReportIncompleteButClearKnownMatches) { + for (bool corrupt : {false, true}) { + SCOPED_TRACE(corrupt); + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry unknown, known; + ASSERT_TRUE(seed_cache( + root.url, + manifest_dict(), + fake.loader, + unknown, + true, + false, + @"old_arch", + @"macOS", + corrupt)); + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, known)); + if (corrupt) { + ASSERT_TRUE([[@"{}" dataUsingEncoding:NSUTF8StringEncoding] + writeToURL:cache_identity_url(root.url, unknown.key) + options:0 + error:nil]); + } + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", encode(manifest_dict())}}}}); + EXPECT_EQ( + clear_cache_for_pte(pte, root.url.path.UTF8String), + Error::InvalidExternalData); + EXPECT_EQ(fake.loader->evictions.load(), 1); + cache_entry_removed(root.url, known); + EXPECT_TRUE(cache_exists(test_bookmark_url(root.url, unknown.key))); + EXPECT_NE(fake.loader.models[unknown.bookmark], nil); + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + cache_entry_removed(root.url, unknown); + } + } +} + +TEST_F(CoreAICacheTest, ClearsAttributionOnlyAndStagingEntriesWithoutSdk) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry metadata_only, staged; + ASSERT_TRUE(seed_cache( + root.url, + manifest_dict(), + fake.loader, + metadata_only, + false, + false, + @"other_arch")); + ASSERT_TRUE(seed_cache( + root.url, manifest_dict(), fake.loader, staged, false, true)); + bridge.state().available = false; + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", encode(manifest_dict())}}}}); + EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); + EXPECT_EQ(fake.loader->evictions.load(), 0); + cache_entry_removed(root.url, metadata_only); + cache_entry_removed(root.url, staged); + } +} + +TEST_F(CoreAICacheTest, MaintenanceUnlinksStagingRootWithoutFollowingIt) { + @autoreleasepool { + TestDirectory root, outside; + BookmarkLoaderScope fake; + CacheEntry entry; + ASSERT_TRUE( + seed_cache(root.url, manifest_dict(), fake.loader, entry, true, false)); + NSURL* sentinel = [outside.url URLByAppendingPathComponent:@"sentinel"]; + NSData* contents = [@"untouched" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([contents writeToURL:sentinel options:0 error:nil]); + ASSERT_TRUE([NSFileManager.defaultManager + createDirectoryAtURL:[entry.stage URLByDeletingLastPathComponent] + withIntermediateDirectories:NO + attributes:nil + error:nil]); + ASSERT_EQ( + symlink( + outside.url.fileSystemRepresentation, + entry.stage.fileSystemRepresentation), + 0); + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + EXPECT_TRUE([[NSData dataWithContentsOfURL:sentinel] isEqual:contents]); + EXPECT_EQ(fake.loader->evictions.load(), 1); + cache_entry_removed(root.url, entry); + } +} + TEST_F(CoreAICacheTest, ProcessedBuffersRemainAliveThroughEviction) { @autoreleasepool { - BookmarkDirectory root; + TestDirectory root; CacheEntry source; NSData* manifest = encode(manifest_dict()); - FBSyntheticPTE valid({{"forward", - {{"CoreAIBackend", manifest, true}, - {"CoreAIBackend", manifest, true}, - {"CoreAIBackend", manifest}}}}); + FBSyntheticPTE valid( + {{"forward", + {{"CoreAIBackend", manifest, true}, + {"CoreAIBackend", manifest, true}, + {"CoreAIBackend", manifest}}}}); // Destroy the loader scope before the PTE captured by its callback. BookmarkLoaderScope fake; ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, source)); @@ -327,7 +532,8 @@ FBSyntheticPTE valid({{"forward", FBSyntheticPTE* observed = &valid; fake.loader.onEvict = ^(NSData*) { ++observations; - program_alive = observed->program_loads.load() == observed->program_releases.load() + 1; + program_alive = observed->program_loads.load() == + observed->program_releases.load() + 1; backend_releases = observed->backend_releases.load(); }; EXPECT_EQ(clear_cache_for_pte(valid, root.url.path.UTF8String), Error::Ok); @@ -344,32 +550,45 @@ FBSyntheticPTE valid({{"forward", } } -TEST_F(CoreAICacheTest, InventoriesBookmarkAndStagingUnionPreservingUnownedFiles) { +TEST_F( + CoreAICacheTest, + InventoriesBookmarkAndStagingUnionPreservingUnownedFiles) { @autoreleasepool { - BookmarkDirectory root; + TestDirectory root; BookmarkLoaderScope fake; CacheEntry both, bookmark_only, stage_only; ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, both)); - ASSERT_TRUE(seed_cache(root.url, cache_manifest('2'), fake.loader, bookmark_only, true, false)); - ASSERT_TRUE(seed_cache(root.url, cache_manifest('3'), fake.loader, stage_only, false, true)); - NSURL* temporary = [root.url URLByAppendingPathComponent:@"staging/.staging-unkeyed"]; + ASSERT_TRUE(seed_cache( + root.url, + cache_manifest('2'), + fake.loader, + bookmark_only, + true, + false)); + ASSERT_TRUE(seed_cache( + root.url, cache_manifest('3'), fake.loader, stage_only, false, true)); + NSURL* temporary = + [root.url URLByAppendingPathComponent:@"staging/.staging-unkeyed"]; ASSERT_TRUE([NSFileManager.defaultManager createDirectoryAtURL:temporary withIntermediateDirectories:NO attributes:nil error:nil]); NSURL* untouched = [temporary URLByAppendingPathComponent:@"partial"]; - ASSERT_TRUE([[@"do not infer ownership" dataUsingEncoding:NSUTF8StringEncoding] - writeToURL:untouched + ASSERT_TRUE([[@"do not infer ownership" + dataUsingEncoding:NSUTF8StringEncoding] writeToURL:untouched + options:0 + error:nil]); + NSURL* external_sentinel = + [root.url URLByAppendingPathComponent:@"caller.pte"]; + ASSERT_TRUE([[@"PTE" dataUsingEncoding:NSUTF8StringEncoding] + writeToURL:external_sentinel options:0 error:nil]); - NSURL* external_sentinel = [root.url URLByAppendingPathComponent:@"caller.pte"]; - ASSERT_TRUE([[@"PTE" dataUsingEncoding:NSUTF8StringEncoding] writeToURL:external_sentinel - options:0 - error:nil]); NSDictionary* temp_snapshot = nil; NSDictionary* locks = nil; ASSERT_TRUE(asset_tree_snapshot(temporary, temp_snapshot)); - ASSERT_TRUE(asset_tree_snapshot([root.url URLByAppendingPathComponent:@"locks"], locks)); + ASSERT_TRUE(asset_tree_snapshot( + [root.url URLByAppendingPathComponent:@"locks"], locks)); const int queries = bridge.state().architecture_queries; EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); EXPECT_EQ(fake.loader->evictions.load(), 2); @@ -379,7 +598,8 @@ FBSyntheticPTE valid({{"forward", cache_entry_removed(root.url, both); cache_entry_removed(root.url, bookmark_only); cache_entry_removed(root.url, stage_only); - EXPECT_TRUE(snapshot_matches([root.url URLByAppendingPathComponent:@"locks"], locks)); + EXPECT_TRUE(snapshot_matches( + [root.url URLByAppendingPathComponent:@"locks"], locks)); EXPECT_TRUE(snapshot_matches(temporary, temp_snapshot)); EXPECT_TRUE(cache_exists(external_sentinel)); EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); @@ -389,10 +609,11 @@ FBSyntheticPTE valid({{"forward", TEST_F(CoreAICacheTest, RemovesStagingOnlyEntriesWithoutSdkAvailability) { @autoreleasepool { - BookmarkDirectory root; + TestDirectory root; BookmarkLoaderScope fake; CacheEntry stage_only; - ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, stage_only, false, true)); + ASSERT_TRUE(seed_cache( + root.url, manifest_dict(), fake.loader, stage_only, false, true)); bridge.state().available = false; EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); cache_entry_removed(root.url, stage_only); @@ -401,11 +622,131 @@ FBSyntheticPTE valid({{"forward", } } +TEST_F(CoreAICacheTest, WarnsAndContinuesPastMissingOrCorruptBookmarks) { + for (bool pte_clear : {false, true}) { + for (int corruption = 0; corruption < 6; ++corruption) { + SCOPED_TRACE(pte_clear ? "PTE" : "root"); + SCOPED_TRACE(corruption); + @autoreleasepool { + TestDirectory root, outside; + NSURL* sentinel = [outside.url URLByAppendingPathComponent:@"sentinel"]; + NSData* contents = + [@"untouched" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([contents writeToURL:sentinel options:0 error:nil]); + BookmarkLoaderScope fake; + CacheEntry damaged, healthy; + ASSERT_TRUE( + seed_cache(root.url, manifest_dict(), fake.loader, damaged)); + ASSERT_TRUE( + seed_cache(root.url, cache_manifest('2'), fake.loader, healthy)); + if ([damaged.key compare:healthy.key] == NSOrderedDescending) { + std::swap(damaged, healthy); + } + NSURL* path = test_bookmark_url(root.url, damaged.key); + if (corruption == 0 || corruption >= 3) { + ASSERT_EQ(unlink(path.fileSystemRepresentation), 0); + if (corruption == 3) { + ASSERT_EQ(mkfifo(path.fileSystemRepresentation, 0600), 0); + } else if (corruption == 4) { + ASSERT_EQ( + symlink( + sentinel.fileSystemRepresentation, + path.fileSystemRepresentation), + 0); + } else if (corruption == 5) { + ASSERT_EQ(mkdir(path.fileSystemRepresentation, 0700), 0); + ASSERT_TRUE([contents + writeToURL:[path URLByAppendingPathComponent:@"invalid"] + options:0 + error:nil]); + } + } else { + FileDescriptor file( + open(path.fileSystemRepresentation, O_WRONLY | O_NOFOLLOW)); + ASSERT_GE(file.get(), 0); + ASSERT_EQ( + ftruncate(file.get(), corruption == 1 ? 0 : 8 * 1024 * 1024 + 1), + 0); + } + FBSyntheticPTE pte( + {{"forward", + {{"CoreAIBackend", encode(manifest_dict())}, + {"CoreAIBackend", encode(cache_manifest('2')), true}}}}); + auto clear = [&] { + return pte_clear ? clear_cache_for_pte(pte, root.url.path.UTF8String) + : clear_cache(root.url.path.UTF8String); + }; + ::testing::internal::CaptureStderr(); + const auto error = clear(); + const std::string logs = ::testing::internal::GetCapturedStderr(); + EXPECT_EQ(error, Error::Ok); + if constexpr ( + ET_LOG_ENABLED && + runtime::LogLevel::Info >= runtime::LogLevel::ET_MIN_LOG_LEVEL) { + EXPECT_NE(logs.find(damaged.key.UTF8String), std::string::npos); + EXPECT_NE( + logs.find( + corruption == 0 ? "has a missing bookmark" + : "has a corrupt bookmark"), + std::string::npos); + EXPECT_NE( + logs.find("SDK eviction could not be confirmed"), + std::string::npos); + EXPECT_EQ( + logs.find(std::string( + static_cast(damaged.bookmark.bytes), + damaged.bookmark.length)), + std::string::npos); + } + cache_entry_removed(root.url, damaged); + cache_entry_removed(root.url, healthy); + EXPECT_TRUE([[NSData dataWithContentsOfURL:sentinel] isEqual:contents]); + EXPECT_EQ(fake.loader->evictions.load(), 1); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_NE(fake.loader.models[damaged.bookmark], nil); + EXPECT_EQ(fake.loader.models[healthy.bookmark], nil); + ::testing::internal::CaptureStderr(); + const auto repeated = clear(); + const std::string repeated_logs = + ::testing::internal::GetCapturedStderr(); + EXPECT_EQ(repeated, Error::Ok); + EXPECT_TRUE(repeated_logs.empty()); + EXPECT_EQ(fake.loader->evictions.load(), 1); + EXPECT_TRUE(pte.check_released()); + } + } + } +} + +TEST_F(CoreAICacheTest, CorruptBookmarkCleanupDoesNotRequireSdkAndCanRetry) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry damaged; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, damaged)); + NSURL* path = test_bookmark_url(root.url, damaged.key); + ASSERT_TRUE([NSData.data writeToURL:path options:0 error:nil]); + bridge.state().available = false; + { + StorageFaultScope fault(StorageOperation::Remove); + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + EXPECT_TRUE(cache_exists(path)); + EXPECT_TRUE(cache_exists(damaged.stage)); + } + EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); + cache_entry_removed(root.url, damaged); + EXPECT_EQ(fake.loader->evictions.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_NE(fake.loader.models[damaged.bookmark], nil); + } +} + TEST_F(CoreAICacheTest, NilLoaderFactoryPreservesEntriesAndAllowsRetry) { @autoreleasepool { ASSERT_TRUE(bridge.state().available); ASSERT_FALSE(bridge.state().fail_loader_factory); - BookmarkDirectory root; + TestDirectory root; BookmarkLoaderScope fake; CacheEntry entry; ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, entry)); @@ -427,7 +768,7 @@ FBSyntheticPTE valid({{"forward", TEST_F(CoreAICacheTest, UnavailableSdkPreservesTrackedEntriesAndAllowsRetry) { @autoreleasepool { - BookmarkDirectory root; + TestDirectory root; BookmarkLoaderScope fake; CacheEntry entry; ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, entry)); @@ -449,13 +790,15 @@ FBSyntheticPTE valid({{"forward", int match_index; bool bookmark_removal; const char* name; - } rows[] = {{2, false, "staged-tree removal"}, {6, true, "bookmark removal"}}; + } rows[] = { + {2, false, "staged-tree removal"}, {6, true, "bookmark removal"}}; for (const auto& row : rows) { SCOPED_TRACE(row.name); - BookmarkDirectory root; + TestDirectory root; BookmarkLoaderScope fake; CacheEntry entry; - ASSERT_TRUE(seed_cache(root.url, aot_manifest_dict(), fake.loader, entry)); + ASSERT_TRUE( + seed_cache(root.url, aot_manifest_dict(), fake.loader, entry)); { StorageFaultScope fault(StorageOperation::Remove); fault.match_index = row.match_index; @@ -478,25 +821,55 @@ FBSyntheticPTE valid({{"forward", } } +TEST_F(CoreAICacheTest, AttributionRemovalFailureLeavesDiscoverableRetry) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry entry; + ASSERT_TRUE( + seed_cache(root.url, manifest_dict(), fake.loader, entry, true, false)); + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", encode(manifest_dict())}}}}); + { + StorageFaultScope fault(StorageOperation::Remove); + fault.match_index = 2; + EXPECT_EQ( + clear_cache_for_pte(pte, root.url.path.UTF8String), + Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + } + EXPECT_FALSE(cache_exists(test_bookmark_url(root.url, entry.key))); + EXPECT_TRUE(cache_exists(cache_identity_url(root.url, entry.key))); + EXPECT_EQ(fake.loader->evictions.load(), 1); + bridge.state().available = false; + EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); + EXPECT_EQ(fake.loader->evictions.load(), 1); + cache_entry_removed(root.url, entry); + } +} + TEST_F(CoreAICacheTest, IndependentKeysContinueInSortedOrderAfterFailures) { @autoreleasepool { - BookmarkDirectory root; + TestDirectory root; BookmarkLoaderScope fake; CacheEntry first, second, third; ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, first)); ASSERT_TRUE(seed_cache(root.url, cache_manifest('2'), fake.loader, second)); ASSERT_TRUE(seed_cache(root.url, cache_manifest('3'), fake.loader, third)); - NSArray* ordered = - [@[ first.key, second.key, third.key ] sortedArrayUsingSelector:@selector(compare:)]; - NSDictionary* bookmarks = - @{first.key : first.bookmark, second.key : second.bookmark, third.key : third.bookmark}; + NSArray* ordered = [@[ first.key, second.key, third.key ] + sortedArrayUsingSelector:@selector(compare:)]; + NSDictionary* bookmarks = @{ + first.key : first.bookmark, + second.key : second.bookmark, + third.key : third.bookmark + }; NSMutableArray* evicted = [NSMutableArray array]; BookmarkFakeLoader* loader = fake.loader; fake.loader.onEvict = ^(NSData* bookmark) { [evicted addObject:bookmark]; loader.evictError = [bookmark isEqual:bookmarks[ordered[1]]] - ? [NSError errorWithDomain:NSPOSIXErrorDomain code:EBUSY userInfo:nil] - : nil; + ? [NSError errorWithDomain:NSPOSIXErrorDomain code:EBUSY userInfo:nil] + : nil; }; StorageFaultScope fault(StorageOperation::BeforeEvict); EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::AccessFailed); @@ -511,5 +884,5 @@ FBSyntheticPTE valid({{"forward", } } -} // namespace -} // namespace executorch::backends::coreai::testing +} // namespace +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm index 5088bef0a6c..76c8e7b047a 100644 --- a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm @@ -158,7 +158,7 @@ FBSyntheticPTE pte( ASSERT_EQ( clear_cache_for_pte(pte, fixture.storage.url.path.UTF8String), Error::Ok); - EXPECT_EQ(bridge.state().architecture_queries.load(), queries + 2); + EXPECT_EQ(bridge.state().architecture_queries.load(), queries + 1); EXPECT_EQ(fake.loader->specializations.load(), 1); EXPECT_EQ(fake.loader->evictions.load(), 1); EXPECT_EQ(fake.loader.models[bookmark.get()], nil); @@ -216,7 +216,7 @@ FBSyntheticPTE pte( TEST_F(CoreAIDelegateTest, MissingDefaultRootClearingDoesNotCreateStorage) { @autoreleasepool { - BookmarkDirectory root; + TestDirectory root; ASSERT_NE(root.url, nil); NSString* missing = [root.url.path stringByAppendingPathComponent:@"missing"]; diff --git a/backends/apple/coreai/runtime/test/coreai_eviction_test.mm b/backends/apple/coreai/runtime/test/coreai_eviction_test.mm index 9cf065e9fc8..c14447b31c2 100644 --- a/backends/apple/coreai/runtime/test/coreai_eviction_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_eviction_test.mm @@ -26,30 +26,46 @@ } // These cases never retain a delegate handle beyond initialization. -Error initialize_and_destroy(NSDictionary* dict, const NamedDataMap* data, NSString* root) { - if (root == nil) return Error::InvalidArgument; +Error initialize_and_destroy( + NSDictionary* dict, + const NamedDataMap* data, + NSString* root) { + if (root == nil) + return Error::InvalidArgument; BackendOptions<1> options; auto error = options.set_option("coreai_assets_dir", root.UTF8String); - if (error != Error::Ok) return error; + if (error != Error::Ok) + return error; auto* backend = get_backend_class("CoreAIBackend"); - if (backend == nullptr) return Error::Internal; - BackendInitContext context(nullptr, nullptr, "forward", data, - {options.view().data(), options.view().size()}); + if (backend == nullptr) + return Error::Internal; + BackendInitContext context( + nullptr, + nullptr, + "forward", + data, + {options.view().data(), options.view().size()}); NSData* json = encode(dict); - if (json == nil) return Error::InvalidArgument; + if (json == nil) + return Error::InvalidArgument; FreeableBuffer processed(json.bytes, json.length, nullptr); auto handle = backend->init(context, &processed, {}); - if (!handle.ok()) return handle.error(); + if (!handle.ok()) + return handle.error(); backend->destroy(handle.get()); return Error::Ok; } -dispatch_time_t deadline() { return dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC); } +dispatch_time_t deadline() { + return dispatch_time(DISPATCH_TIME_NOW, 10 * NSEC_PER_SEC); +} class CoreAIEvictionTest : public ::testing::Test { protected: ScopedFakeBridgeState bridge; - void SetUp() override { runtime_init(); } + void SetUp() override { + runtime_init(); + } }; // Release the callback before waiting; never unwind storage borrowed by a hang. @@ -75,20 +91,27 @@ void release() { } void join() { release(); - if (joined) return; + if (joined) + return; const auto completion_deadline = deadline(); if (eviction_started) - eviction_wait = dispatch_semaphore_wait(eviction_done, completion_deadline); - if (loading_started) loading_wait = dispatch_semaphore_wait(loading_done, completion_deadline); - if ((eviction_started && eviction_wait != 0) || (loading_started && loading_wait != 0)) { - ADD_FAILURE() << "Eviction/load workers did not return after releasing the gate"; + eviction_wait = + dispatch_semaphore_wait(eviction_done, completion_deadline); + if (loading_started) + loading_wait = dispatch_semaphore_wait(loading_done, completion_deadline); + if ((eviction_started && eviction_wait != 0) || + (loading_started && loading_wait != 0)) { + ADD_FAILURE() + << "Eviction/load workers did not return after releasing the gate"; std::_Exit(EXIT_FAILURE); } - const int eviction_join = eviction_started ? pthread_join(eviction, nullptr) : 0; - const int loading_join = loading_started ? pthread_join(loading, nullptr) : 0; + const int eviction_join = + eviction_started ? pthread_join(eviction, nullptr) : 0; + const int loading_join = + loading_started ? pthread_join(loading, nullptr) : 0; if (eviction_join != 0 || loading_join != 0) { - ADD_FAILURE() << "Cannot join eviction/load workers: " << eviction_join << ", " - << loading_join; + ADD_FAILURE() << "Cannot join eviction/load workers: " << eviction_join + << ", " << loading_join; std::_Exit(EXIT_FAILURE); } auto drained = bridge.wait_for_callbacks(deadline()); @@ -98,9 +121,11 @@ void join() { } joined = true; } - ~ReleaseAndJoin() { join(); } + ~ReleaseAndJoin() { + join(); + } }; -} // namespace +} // namespace @interface PinBookmarkLoader : BookmarkFakeLoader @property(nonatomic, weak) BookmarkFakeModel* temporary; @@ -108,13 +133,15 @@ @interface PinBookmarkLoader : BookmarkFakeLoader @end @implementation PinBookmarkLoader -- (void)specializeModelAtURL:(NSURL*)url completion:(ETCoreAIAcquisitionCompletion)completion { +- (void)specializeModelAtURL:(NSURL*)url + completion:(ETCoreAIAcquisitionCompletion)completion { ++specializations; self.retainedCompletion = completion; @autoreleasepool { BookmarkFakeModel* model = [[BookmarkFakeModel alloc] init]; model.owner = self; - model.bookmark = [@"pinned-bookmark" dataUsingEncoding:NSUTF8StringEncoding]; + model.bookmark = + [@"pinned-bookmark" dataUsingEncoding:NSUTF8StringEncoding]; self.temporary = model; completion(ETCoreAIAcquisitionStatusHit, model, nil); model = nil; @@ -126,12 +153,16 @@ - (void)restoreModelFromBookmark:(NSData*)bookmark self.retainedCompletion = completion; @autoreleasepool { BookmarkFakeModel* model = self.temporary; - completion(model == nil ? ETCoreAIAcquisitionStatusMiss : ETCoreAIAcquisitionStatusHit, model, - nil); + completion( + model == nil ? ETCoreAIAcquisitionStatusMiss + : ETCoreAIAcquisitionStatusHit, + model, + nil); model = nil; } } -- (void)evictModelWithBookmark:(NSData*)bookmark completion:(void (^)(NSError*))completion { +- (void)evictModelWithBookmark:(NSData*)bookmark + completion:(void (^)(NSError*))completion { ++evictions; bool busy; @autoreleasepool { @@ -148,7 +179,7 @@ - (void)clearRetainedState { namespace { TEST_F(CoreAIEvictionTest, ConfirmedMissRecoveryEvictsOnlyCurrentBookmark) { @autoreleasepool { - BookmarkDirectory root; + TestDirectory root; ASSERT_NE(root.url, nil); BookmarkLoaderScope fake; ASSERT_NE(fake.loader, nil); @@ -167,7 +198,8 @@ - (void)clearRetainedState { NSString* source_path = fake.loader.lastSpecializeURL; ASSERT_NE(source_path, nil); NSURL* source_url = [NSURL fileURLWithPath:source_path]; - EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); fake.loader.restoreMiss = YES; fake.loader.specializeError = bookmark_test_error(); @@ -175,7 +207,8 @@ - (void)clearRetainedState { auto after_failure = saved_bookmark(root.url, key); ASSERT_TRUE(after_failure.ok()); EXPECT_TRUE([after_failure.get() isEqual:original]); - EXPECT_TRUE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + EXPECT_TRUE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); fake.loader.specializeError = nil; ASSERT_EQ(initialize_and_destroy(dict, &data, root.url.path), Error::Ok); EXPECT_EQ(fake.loader->specializations.load(), 3); @@ -186,12 +219,14 @@ - (void)clearRetainedState { NSData* saved = saved_result.get(); ASSERT_NE(saved, nil); EXPECT_FALSE([saved isEqual:original]); - EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); ASSERT_EQ( clear_bookmark_assets(root.url.path, key, fake.loader), Error::Ok); EXPECT_NE(fake.loader.models[original], nil); EXPECT_EQ(fake.loader.models[saved], nil); - EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); fake.loader.restoreMiss = NO; fake.loader.emptyBookmark = YES; EXPECT_NE(initialize_and_destroy(dict, &data, root.url.path), Error::Ok); @@ -205,9 +240,11 @@ - (void)clearRetainedState { } } -TEST_F(CoreAIEvictionTest, BusyAndUnknownEvictionRetainBookmarkUntilConfirmedAbsent) { +TEST_F( + CoreAIEvictionTest, + BusyAndUnknownEvictionRetainBookmarkUntilConfirmedAbsent) { @autoreleasepool { - BookmarkDirectory root; + TestDirectory root; ASSERT_NE(root.url, nil); BookmarkLoaderScope fake; ASSERT_NE(fake.loader, nil); @@ -225,20 +262,21 @@ - (void)clearRetainedState { StorageFaultScope fault(StorageOperation::Remove); ASSERT_EQ( initialize_and_destroy(manifest_dict(), &data, root.url.path), - Error::AccessFailed); + Error::Ok); EXPECT_EQ(fault.hits, 1); } auto saved_result = saved_bookmark(root.url, key); ASSERT_TRUE(saved_result.ok()); NSData* saved = saved_result.get(); ASSERT_NE(saved, nil); - NSURL* source_url = - [[[root.url URLByAppendingPathComponent:@"staging"] URLByAppendingPathComponent:key] - URLByAppendingPathComponent:selected->path.lastPathComponent]; + NSURL* source_url = [[[root.url URLByAppendingPathComponent:@"staging"] + URLByAppendingPathComponent:key] + URLByAppendingPathComponent:selected->path.lastPathComponent]; EXPECT_TRUE( [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); for (bool probe_error : {false, true}) { - SCOPED_TRACE(probe_error ? "unknown restore failure" : "restore hit/busy"); + SCOPED_TRACE( + probe_error ? "unknown restore failure" : "restore hit/busy"); fake.loader.evictError = bookmark_test_error(EBUSY); fake.loader.restoreError = probe_error ? bookmark_test_error() : nil; EXPECT_NE( @@ -266,9 +304,12 @@ - (void)clearRetainedState { auto absent = saved_bookmark(root.url, key); ASSERT_TRUE(absent.ok()); EXPECT_EQ(absent.get(), nil); - EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); fake.loader.evictError = nil; - ASSERT_EQ(initialize_and_destroy(manifest_dict(), &data, root.url.path), Error::Ok); + ASSERT_EQ( + initialize_and_destroy(manifest_dict(), &data, root.url.path), + Error::Ok); { StorageFaultScope fault(StorageOperation::Remove); EXPECT_NE( @@ -283,14 +324,17 @@ - (void)clearRetainedState { auto removed = saved_bookmark(root.url, key); ASSERT_TRUE(removed.ok()); EXPECT_EQ(removed.get(), nil); - EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); EXPECT_EQ(data.requests.load(), data.releases.load()); } } -TEST_F(CoreAIEvictionTest, PreparedPinsBlockEvictionWithoutRetainedCompletionLeak) { +TEST_F( + CoreAIEvictionTest, + PreparedPinsBlockEvictionWithoutRetainedCompletionLeak) { @autoreleasepool { - BookmarkDirectory root; + TestDirectory root; ASSERT_NE(root.url, nil); TestData data; auto selected = bookmark_manifest(); @@ -304,8 +348,8 @@ - (void)clearRetainedState { // failures. bridge.state().bookmark_loader = loader; @autoreleasepool { - auto pinned = - acquire_bookmark_model(selected.get(), &data, root.url.path, @"macOS", @"arch_b", loader); + auto pinned = acquire_bookmark_model( + selected.get(), &data, root.url.path, @"macOS", @"arch_b", loader); ASSERT_TRUE(pinned.ok()); EXPECT_NE(loader.temporary, nil); EXPECT_NE(clear_bookmark_assets(root.url.path, key, loader), Error::Ok); @@ -320,10 +364,11 @@ - (void)clearRetainedState { ASSERT_TRUE(removed.ok()); EXPECT_EQ(removed.get(), nil); EXPECT_EQ(loader->evictions.load(), 2); - NSURL* source_url = - [[[root.url URLByAppendingPathComponent:@"staging"] URLByAppendingPathComponent:key] - URLByAppendingPathComponent:selected->path.lastPathComponent]; - EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:source_url.path]); + NSURL* source_url = [[[root.url URLByAppendingPathComponent:@"staging"] + URLByAppendingPathComponent:key] + URLByAppendingPathComponent:selected->path.lastPathComponent]; + EXPECT_FALSE( + [NSFileManager.defaultManager fileExistsAtPath:source_url.path]); [loader clearRetainedState]; EXPECT_EQ(data.requests.load(), data.releases.load()); } @@ -331,7 +376,7 @@ - (void)clearRetainedState { TEST_F(CoreAIEvictionTest, LoadContendsOnActualKeyLockUntilEvictionFinishes) { @autoreleasepool { - BookmarkDirectory root; + TestDirectory root; ASSERT_NE(root.url, nil); BookmarkLoaderScope fake; ASSERT_NE(fake.loader, nil); @@ -341,7 +386,9 @@ - (void)clearRetainedState { auto key_result = test_bookmark_key(selected.get()); ASSERT_TRUE(key_result.ok()); NSString* key = key_result.get(); - ASSERT_EQ(initialize_and_destroy(manifest_dict(), &data, root.url.path), Error::Ok); + ASSERT_EQ( + initialize_and_destroy(manifest_dict(), &data, root.url.path), + Error::Ok); dispatch_semaphore_t deleting = dispatch_semaphore_create(0); dispatch_semaphore_t proceed = dispatch_semaphore_create(0); dispatch_semaphore_t loading_observed = dispatch_semaphore_create(0); @@ -362,8 +409,13 @@ - (void)clearRetainedState { }; auto load = [&] { @autoreleasepool { - auto acquired = acquire_bookmark_model(selected.get(), &data, root.url.path, @"macOS", - @"arch_b", fake.loader); + auto acquired = acquire_bookmark_model( + selected.get(), + &data, + root.url.path, + @"macOS", + @"arch_b", + fake.loader); load_ok = acquired.ok() && acquired.get() != nil; } dispatch_semaphore_signal(loading_observed); @@ -371,20 +423,23 @@ - (void)clearRetainedState { }; ReleaseAndJoin cleanup{proceed, eviction_done, loading_done, bridge}; const int eviction_start = pthread_create( - &cleanup.eviction, nullptr, + &cleanup.eviction, + nullptr, [](void* context) -> void* { (*static_cast(context))(); return nullptr; }, &evict); cleanup.eviction_started = eviction_start == 0; - const long entered = - cleanup.eviction_started ? dispatch_semaphore_wait(deleting, deadline()) : -1; + const long entered = cleanup.eviction_started + ? dispatch_semaphore_wait(deleting, deadline()) + : -1; int loading_start = -1; long premature_completion = -1; if (entered == 0) { loading_start = pthread_create( - &cleanup.loading, nullptr, + &cleanup.loading, + nullptr, [](void* context) -> void* { (*static_cast(context))(); return nullptr; @@ -393,7 +448,8 @@ - (void)clearRetainedState { cleanup.loading_started = loading_start == 0; if (cleanup.loading_started) { premature_completion = dispatch_semaphore_wait( - loading_observed, dispatch_time(DISPATCH_TIME_NOW, 50 * NSEC_PER_MSEC)); + loading_observed, + dispatch_time(DISPATCH_TIME_NOW, 50 * NSEC_PER_MSEC)); } } cleanup.join(); @@ -415,4 +471,4 @@ - (void)clearRetainedState { EXPECT_EQ(data.requests.load(), data.releases.load()); } } -} // namespace +} // namespace diff --git a/backends/apple/coreai/runtime/test/coreai_storage_test.mm b/backends/apple/coreai/runtime/test/coreai_storage_test.mm index 1afa11b8f8b..5487f8ed864 100644 --- a/backends/apple/coreai/runtime/test/coreai_storage_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_storage_test.mm @@ -798,6 +798,231 @@ Error populate() { EXPECT_EQ(errno, ENOENT); } } + +TEST_F(CoreAIStagingRemovalTest, MaintenanceRemovesKeyedNonDirectoryEntries) { + TestDirectory outside; + ASSERT_NE(outside.url, nil); + NSURL* sentinel = [outside.url URLByAppendingPathComponent:@"sentinel"]; + NSData* expected = [@"untouched" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([expected writeToURL:sentinel atomically:YES]); + NSDictionary* outside_before = nil; + ASSERT_TRUE(asset_tree_snapshot(outside.url, outside_before)); + ASSERT_EQ(mkdirat(root.get(), "staging", 0700), 0); + ASSERT_EQ( + symlinkat( + outside.url.fileSystemRepresentation, + root.get(), + entry().fileSystemRepresentation), + 0); + EXPECT_EQ( + remove_storage_staging( + root.get(), key, true, StorageRemovalMode::Maintenance), + Error::Ok); + struct stat info; + EXPECT_EQ( + fstatat( + root.get(), + entry().fileSystemRepresentation, + &info, + AT_SYMLINK_NOFOLLOW), + -1); + EXPECT_EQ(errno, ENOENT); + ASSERT_EQ(mkfifoat(root.get(), entry().fileSystemRepresentation, 0600), 0); + EXPECT_EQ( + remove_storage_staging( + root.get(), key, true, StorageRemovalMode::Maintenance), + Error::Ok); + EXPECT_EQ( + fstatat( + root.get(), + entry().fileSystemRepresentation, + &info, + AT_SYMLINK_NOFOLLOW), + -1); + EXPECT_EQ(errno, ENOENT); + ASSERT_EQ(write_storage_file(root.get(), entry(), "stale", 5), Error::Ok); + EXPECT_EQ( + remove_storage_staging( + root.get(), key, true, StorageRemovalMode::Maintenance), + Error::Ok); + EXPECT_EQ( + fstatat( + root.get(), + entry().fileSystemRepresentation, + &info, + AT_SYMLINK_NOFOLLOW), + -1); + EXPECT_EQ(errno, ENOENT); + EXPECT_TRUE(snapshot_matches(outside.url, outside_before)); + EXPECT_TRUE([[NSData dataWithContentsOfURL:sentinel] isEqual:expected]); +} + +TEST_F(CoreAIStagingRemovalTest, MaintenanceRemovesNestedSpecialEntries) { + TestDirectory outside; + ASSERT_NE(outside.url, nil); + NSURL* sentinel = [outside.url URLByAppendingPathComponent:@"sentinel"]; + NSData* expected = [@"untouched" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([expected writeToURL:sentinel atomically:YES]); + ASSERT_EQ(populate(), Error::Ok); + NSString* bundle = [entry() stringByAppendingPathComponent:@"model.aimodel"]; + ASSERT_EQ( + symlinkat( + sentinel.fileSystemRepresentation, + root.get(), + [bundle stringByAppendingPathComponent:@"link"] + .fileSystemRepresentation), + 0); + ASSERT_EQ( + mkfifoat( + root.get(), + [bundle stringByAppendingPathComponent:@"pipe"] + .fileSystemRepresentation, + 0600), + 0); + ASSERT_EQ( + write_storage_file( + root.get(), + [bundle stringByAppendingPathComponent:@"nested/extra"], + "extra", + 5), + Error::Ok); + NSDictionary* outside_before = nil; + ASSERT_TRUE(asset_tree_snapshot(outside.url, outside_before)); + EXPECT_EQ( + remove_storage_staging( + root.get(), key, true, StorageRemovalMode::Maintenance), + Error::Ok); + struct stat info; + EXPECT_EQ( + fstatat( + root.get(), + entry().fileSystemRepresentation, + &info, + AT_SYMLINK_NOFOLLOW), + -1); + EXPECT_EQ(errno, ENOENT); + EXPECT_TRUE(snapshot_matches(outside.url, outside_before)); + EXPECT_TRUE([[NSData dataWithContentsOfURL:sentinel] isEqual:expected]); +} + +TEST_F(CoreAIStagingRemovalTest, MaintenanceEntryHandlesBookmarkLikeEntries) { + TestDirectory outside; + ASSERT_NE(outside.url, nil); + NSURL* sentinel = [outside.url URLByAppendingPathComponent:@"sentinel"]; + NSData* expected = [@"untouched" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([expected writeToURL:sentinel atomically:YES]); + NSDictionary* outside_before = nil; + ASSERT_TRUE(asset_tree_snapshot(outside.url, outside_before)); + NSString* bookmark_like = [key stringByAppendingString:@".bookmark"]; + ASSERT_EQ( + write_storage_file(root.get(), bookmark_like, "bookmark", 8), Error::Ok); + EXPECT_EQ(remove_storage_entry(root.get(), bookmark_like), Error::Ok); + ASSERT_EQ( + symlinkat( + sentinel.fileSystemRepresentation, + root.get(), + bookmark_like.fileSystemRepresentation), + 0); + EXPECT_EQ(remove_storage_entry(root.get(), bookmark_like), Error::Ok); + ASSERT_EQ( + mkfifoat(root.get(), bookmark_like.fileSystemRepresentation, 0600), 0); + EXPECT_EQ(remove_storage_entry(root.get(), bookmark_like), Error::Ok); + ASSERT_EQ( + mkdirat(root.get(), bookmark_like.fileSystemRepresentation, 0700), 0); + ASSERT_EQ( + symlinkat( + sentinel.fileSystemRepresentation, + root.get(), + [bookmark_like stringByAppendingPathComponent:@"link"] + .fileSystemRepresentation), + 0); + ASSERT_EQ( + mkfifoat( + root.get(), + [bookmark_like stringByAppendingPathComponent:@"pipe"] + .fileSystemRepresentation, + 0600), + 0); + ASSERT_EQ( + write_storage_file( + root.get(), + [bookmark_like stringByAppendingPathComponent:@"nested/weights"], + "weights", + 7), + Error::Ok); + EXPECT_EQ(remove_storage_entry(root.get(), bookmark_like, false), Error::Ok); + struct stat info; + EXPECT_EQ( + fstatat( + root.get(), + bookmark_like.fileSystemRepresentation, + &info, + AT_SYMLINK_NOFOLLOW), + 0); + EXPECT_TRUE(S_ISDIR(info.st_mode)); + EXPECT_EQ(remove_storage_entry(root.get(), bookmark_like), Error::Ok); + EXPECT_EQ( + fstatat( + root.get(), + bookmark_like.fileSystemRepresentation, + &info, + AT_SYMLINK_NOFOLLOW), + -1); + EXPECT_EQ(errno, ENOENT); + EXPECT_EQ(remove_storage_entry(root.get(), bookmark_like), Error::Ok); + EXPECT_TRUE(snapshot_matches(outside.url, outside_before)); + EXPECT_TRUE([[NSData dataWithContentsOfURL:sentinel] isEqual:expected]); +} + +TEST_F( + CoreAIStagingRemovalTest, + RejectsInvalidEntryNamesWithoutTouchingStorage) { + ASSERT_EQ(write_storage_file(root.get(), @"kept", "kept", 4), Error::Ok); + NSDictionary* before = nil; + ASSERT_TRUE(asset_tree_snapshot(storage.url, before)); + for (NSString* invalid in @[ @"", @".", @"..", @"a/b", @"/absolute" ]) { + EXPECT_EQ( + remove_storage_entry(root.get(), invalid), Error::InvalidArgument); + } + EXPECT_EQ(remove_storage_entry(root.get(), nil), Error::InvalidArgument); + EXPECT_TRUE(snapshot_matches(storage.url, before)); +} + +TEST_F(CoreAIStagingRemovalTest, MaintenanceKeepsKeyValidation) { + ASSERT_EQ(populate(), Error::Ok); + NSDictionary* before = nil; + ASSERT_TRUE(asset_tree_snapshot(storage.url, before)); + for (NSString* invalid in + @[ @"", @"../escape", @"/absolute", @"short", key.uppercaseString ]) { + EXPECT_EQ( + remove_storage_staging( + root.get(), invalid, true, StorageRemovalMode::Maintenance), + Error::InvalidArgument); + } + EXPECT_EQ( + remove_storage_staging( + root.get(), nil, true, StorageRemovalMode::Maintenance), + Error::InvalidArgument); + EXPECT_TRUE(snapshot_matches(storage.url, before)); +} + +TEST_F(CoreAIStagingRemovalTest, MaintenanceRejectsLinkedStagingParent) { + TestDirectory outside; + ASSERT_NE(outside.url, nil); + ASSERT_EQ( + symlinkat(outside.url.fileSystemRepresentation, root.get(), "staging"), + 0); + NSDictionary* before = nil; + NSDictionary* outside_before = nil; + ASSERT_TRUE(asset_tree_snapshot(storage.url, before)); + ASSERT_TRUE(asset_tree_snapshot(outside.url, outside_before)); + EXPECT_EQ( + remove_storage_staging( + root.get(), key, true, StorageRemovalMode::Maintenance), + Error::AccessFailed); + EXPECT_TRUE(snapshot_matches(storage.url, before)); + EXPECT_TRUE(snapshot_matches(outside.url, outside_before)); +} } // namespace TEST(CoreAIStorageTest, RetriesInterruptedPublicationOperations) { diff --git a/tools/cmake/preset/default.cmake b/tools/cmake/preset/default.cmake index 7b017cf36b9..cf3a335dd84 100644 --- a/tools/cmake/preset/default.cmake +++ b/tools/cmake/preset/default.cmake @@ -440,6 +440,10 @@ check_required_options_on( EXECUTORCH_BUILD_EXTENSION_LLM ) +check_required_options_on( + IF_ON EXECUTORCH_BUILD_COREAI REQUIRES EXECUTORCH_BUILD_EXTENSION_DATA_LOADER +) + check_required_options_on( IF_ON EXECUTORCH_BUILD_EXTENSION_MODULE REQUIRES EXECUTORCH_BUILD_EXTENSION_DATA_LOADER EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR From d24b4320103b9c77bc85ce6a3f5522caaaf04919 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Tue, 27 Oct 2026 17:23:26 -0700 Subject: [PATCH 4/5] [CoreAI] Trace cache clearing and verify bounded lock reuse Log eviction request, confirmed SDK absence, and local cleanup completion. Add clear-aware stripe reuse and per-key isolation regressions, and cover durable-save and eviction failure log boundaries. Preserve the previously validated final implementation after placing storage and acquisition changes in their owning PR layers. Test Plan: Exact final CoreAI tree matches the saved validated snapshot: 256 native tests, all five CTest targets, 124 Python tests, and playground rebuild passed. Lower bookmark and acquisition layers were also built and tested independently during placement. --- backends/apple/coreai/README.md | 39 ++++---- .../coreai/runtime/coreai_load_coordinator.mm | 10 +- .../runtime/test/coreai_acquisition_test.mm | 99 +++++++++++++++++++ .../runtime/test/coreai_bookmark_test.mm | 48 ++++++++- .../coreai/runtime/test/coreai_cache_test.mm | 4 +- 5 files changed, 177 insertions(+), 23 deletions(-) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index a6afb391bb9..4510e97cf12 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -209,7 +209,7 @@ threads and processes. The first byte of the validated cache key selects are created on demand, so new model versions do not grow the pool beyond 256. The mapping is part of the on-disk format and is not configurable. Different keys on the same stripe serialize; different stripes remain independent. -Bookmarks and staging remain separate for every key. Runtime +Bookmarks, attribution, and staging remain separate for every key. Runtime operations hold only one stripe at a time. Lock ownership is maintained by the OS, and acquiring a lock does not flush its @@ -223,7 +223,25 @@ root and upgrading them together. Old and new locking schemes must not run concurrently against the same root. Legacy per-key lock files are ignored and left untouched; they can be removed separately during a fully quiescent offline cleanup. The 256-file bound applies to the new stripe pool, not pre-existing -legacy files. Bookmarks and staged bundles need no migration. +legacy files. Bookmarks, metadata, and staged bundles need no migration. + +### Lifecycle logging + +Build with `EXECUTORCH_ENABLE_LOGGING=ON` and `EXECUTORCH_LOG_LEVEL=Info` to +observe cache lookup, bookmark restoration hits, missing-bookmark versus +SDK-confirmed misses, specialization start/success, durable bookmark saves, and +cache-clearing milestones. Each message carries the cache key so events from +multiple partitions can be correlated. These messages do not include bookmark +contents, model data, or per-inference logs. + +A cache hit means that the SDK successfully restored a model from its bookmark. +Specialization start means a request to the SDK's persistent default cache; the +SDK does not report whether that request compiled afresh or reused an internal +artifact. The durable-save message confirms local bookmark publication, +including its final directory flush, not independent verification of SDK-private +artifact storage. During clearing, SDK entry absence and successful local +cleanup are reported separately. Failed operations do not emit their success +messages, and clearing an absent entry remains quiet. ### Clearing cached assets @@ -252,7 +270,7 @@ not load models or materialize assets. Actual SDK eviction still requires a usable bookmark and an available loader. These synchronous calls require callers to unload affected models and prevent -concurrent loads, inference and maintenance across processes. Per-key locks are +concurrent loads, inference and maintenance across processes. Stripe locks are not a whole-root barrier. Clearing first deletes the recorded SDK entry, then staging, then its bookmark. Busy or unknown SDK errors preserve the files; SDK-confirmed absence permits cleanup. Filesystem errors may leave partial @@ -367,21 +385,6 @@ Ninja builds use one architecture per build directory. Current runtime support covers arm64 macOS and iOS device builds. x86_64 builds are blocked by Swift `Float16` availability, and the tested iOS simulator SDKs do not contain Core AI. -### Lifecycle logging - -Build with `EXECUTORCH_ENABLE_LOGGING=ON` and `EXECUTORCH_LOG_LEVEL=Info` to -observe cache lookup, bookmark restoration hits, missing-bookmark versus -SDK-confirmed misses, specialization start/success, and durable bookmark saves. Each message carries the cache key so events from -multiple partitions can be correlated. These messages do not include bookmark -contents, model data, or per-inference logs. - -A cache hit means that the SDK successfully restored a model from its bookmark. -Specialization start means a request to the SDK's persistent default cache; the -SDK does not report whether that request compiled afresh or reused an internal -artifact. The durable-save message confirms local bookmark publication, -including its final directory flush, not independent verification of SDK-private -artifact storage. Failed operations do not emit their success messages. - Enable `EXECUTORCH_BUILD_EXTENSION_DATA_LOADER` for the public PTE-file maintenance API and consumers such as `Module`. Link the CMake target `coreaidelegate`, not just its archive filename. Its diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.mm b/backends/apple/coreai/runtime/coreai_load_coordinator.mm index 7dfde306521..2951eb15dfa 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.mm +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.mm @@ -227,6 +227,7 @@ Error clear_bookmark_assets( return locked.error(); if (locked.get() == nullptr) return Error::Ok; + ET_LOG(Info, "Core AI cache clear started: key=%s", key.UTF8String); auto bookmark = read_bookmark(*locked.get(), /*maintenance=*/true); if (!bookmark.ok() && bookmark.error() != Error::InvalidExternalData) { return bookmark.error(); @@ -252,14 +253,21 @@ Error clear_bookmark_assets( ET_CHECK_OR_RETURN_ERROR( loader != nil, Internal, "Cannot create Core AI model loader"); } + ET_LOG(Info, "Core AI SDK eviction requested: key=%s", key.UTF8String); ET_CHECK_OK_OR_RETURN_ERROR(evict_one(loader, bytes)); + ET_LOG( + Info, + "Core AI SDK cache entry absence confirmed: key=%s", + key.UTF8String); } ET_CHECK_OK_OR_RETURN_ERROR(remove_bookmark_staging( *locked.get(), true, StorageRemovalMode::Maintenance)); if (!bookmark.ok() || bytes != nil) { ET_CHECK_OK_OR_RETURN_ERROR(remove_bookmark(*locked.get())); } - return remove_bookmark_identity(*locked.get()); + ET_CHECK_OK_OR_RETURN_ERROR(remove_bookmark_identity(*locked.get())); + ET_LOG(Info, "Core AI local cache entry cleared: key=%s", key.UTF8String); + return Error::Ok; } } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm index 99735b4ce1f..f5bf032b46f 100644 --- a/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm @@ -889,4 +889,103 @@ void expect_order( EXPECT_EQ(logs.find("Core AI cache bookmark saved"), std::string::npos); } +TEST_F(CoreAILifecycleLogTest, FinalBookmarkFlushFailureDoesNotLogDurableSave) { + BookmarkLoaderScope fake; + { + auto locked = lock_bookmark(root.url.path, key); + ASSERT_TRUE(locked.ok()); + ASSERT_EQ( + ensure_bookmark_identity(*locked.get(), manifest, @"macOS", @"arch_b"), + Error::Ok); + } + StorageFaultScope fault(StorageOperation::FullSync); + fault.match_index = 5; + const auto logs = + capture([&] { return acquire(fake.loader, &data); }, Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + EXPECT_NE(logs.find("Core AI specialization succeeded:"), std::string::npos); + EXPECT_EQ(logs.find("Core AI cache bookmark saved"), std::string::npos); + EXPECT_TRUE([saved() isEqual:fake.loader.lastAcquired.bookmark]); +} + +TEST_F(CoreAILifecycleLogTest, ClearLogsConfirmedAbsenceBeforeLocalCompletion) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + const auto logs = capture( + [&] { return clear_bookmark_assets(root.url.path, key, fake.loader); }); + expect_order( + logs, + { + "Core AI cache clear started:", + "Core AI SDK eviction requested:", + "Core AI SDK cache entry absence confirmed:", + "Core AI local cache entry cleared:", + }); + EXPECT_EQ(saved(), nil); +} + +TEST_F(CoreAILifecycleLogTest, FailedEvictionWithConfirmedMissLogsAbsence) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + fake.loader.evictError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:ENOENT + userInfo:nil]; + fake.loader.restoreMiss = YES; + const auto logs = capture( + [&] { return clear_bookmark_assets(root.url.path, key, fake.loader); }); + expect_order( + logs, + { + "Core AI SDK eviction requested:", + "Core AI SDK cache entry absence confirmed:", + "Core AI local cache entry cleared:", + }); +} + +TEST_F(CoreAILifecycleLogTest, EvictionErrorDoesNotLogAbsenceOrCompletion) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + fake.loader.evictError = [NSError errorWithDomain:NSPOSIXErrorDomain + code:EBUSY + userInfo:nil]; + const auto logs = capture( + [&] { return clear_bookmark_assets(root.url.path, key, fake.loader); }, + Error::Internal); + EXPECT_NE(logs.find("Core AI SDK eviction requested:"), std::string::npos); + EXPECT_EQ( + logs.find("Core AI SDK cache entry absence confirmed:"), + std::string::npos); + EXPECT_EQ(logs.find("Core AI local cache entry cleared:"), std::string::npos); + EXPECT_NE(saved(), nil); +} + +TEST_F( + CoreAILifecycleLogTest, + IdentityRemovalFailureDoesNotLogLocalCompletion) { + BookmarkLoaderScope fake; + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + StorageFaultScope fault(StorageOperation::Remove); + fault.match_index = 2; + const auto logs = capture( + [&] { return clear_bookmark_assets(root.url.path, key, fake.loader); }, + Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + EXPECT_NE( + logs.find("Core AI SDK cache entry absence confirmed:"), + std::string::npos); + EXPECT_EQ(logs.find("Core AI local cache entry cleared:"), std::string::npos); + auto locked = lock_bookmark(root.url.path, key, false); + ASSERT_TRUE(locked.ok()); + ASSERT_NE(locked.get(), nullptr); + auto identity = read_bookmark_identity(*locked.get()); + ASSERT_TRUE(identity.ok()); + EXPECT_NE(identity.get(), nil); +} + +TEST_F(CoreAILifecycleLogTest, AbsentEntryClearRemainsQuiet) { + const auto logs = + capture([&] { return clear_bookmark_assets(root.url.path, key, nil); }); + EXPECT_TRUE(logs.empty()) << logs; +} + } // namespace diff --git a/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm b/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm index 9e2801e0ae6..2cebbc77d61 100644 --- a/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm @@ -16,6 +16,7 @@ #include "coreai_bookmark_fixture.h" #include "coreai_fault_scope.h" #include "coreai_filesystem_fixture.h" +#include "coreai_load_coordinator.h" #include "coreai_manifest_fixture.h" #include "coreai_storage.h" @@ -961,7 +962,7 @@ FileDescriptor file(open( isEqual:test_bookmark_lock_url(root.url, changed_key(key, true))]); } -TEST_F(CoreAIBookmarkPathTest, LockPoolIsBoundedAcrossManyKeys) { +TEST_F(CoreAIBookmarkPathTest, LockPoolIsBoundedAcrossManyKeysAndClears) { std::array inodes{}; NSURL* locks = [root.url URLByAppendingPathComponent:@"locks"]; for (size_t generation = 0; generation < 4; ++generation) { @@ -980,7 +981,13 @@ FileDescriptor file(open( inodes[stripe] = info.st_ino; else EXPECT_EQ(info.st_ino, inodes[stripe]); + auto staging = prepare_bookmark_staging(*locked.get()); + ASSERT_TRUE(staging.ok()); + NSString* stage = [staging.get() stringByAppendingPathComponent:entry]; + ASSERT_EQ(mkdir(stage.fileSystemRepresentation, 0700), 0); locked.get().reset(); + ASSERT_EQ(clear_bookmark_assets(root.url.path, entry, nil), Error::Ok); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:stage]); auto reused = lock_bookmark(root.url.path, entry); ASSERT_TRUE(reused.ok()); ASSERT_EQ(lstat(file.fileSystemRepresentation, &info), 0); @@ -997,6 +1004,9 @@ FileDescriptor file(open( EXPECT_EQ(name.length, 17u); } } + auto inventory = inventory_bookmark_keys(root.url.path); + ASSERT_TRUE(inventory.ok()); + EXPECT_EQ(inventory.get().count, 0u); } TEST_F(CoreAIBookmarkPathTest, CollidingKeysSerializeAcrossThreads) { @@ -1038,6 +1048,41 @@ FileDescriptor file(open( EXPECT_EQ(names.count, 1u); } +TEST_F(CoreAIBookmarkPathTest, ClearingCollidingKeyPreservesOtherKeyData) { + NSString* collision = changed_key(key, true); + NSData* bytes = [@"other bookmark" dataUsingEncoding:NSUTF8StringEncoding]; + for (NSString* entry in @[ key, collision ]) { + auto locked = lock_bookmark(root.url.path, entry); + ASSERT_TRUE(locked.ok()); + auto staging = prepare_bookmark_staging(*locked.get()); + ASSERT_TRUE(staging.ok()); + ASSERT_EQ( + mkdir( + [staging.get() stringByAppendingPathComponent:entry] + .fileSystemRepresentation, + 0700), + 0); + if ([entry isEqual:collision]) + ASSERT_EQ(write_bookmark(*locked.get(), bytes), Error::Ok); + } + struct stat before {}; + ASSERT_EQ(lstat(lock_url().fileSystemRepresentation, &before), 0); + ASSERT_EQ(clear_bookmark_assets(root.url.path, key, nil), Error::Ok); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:[[root.url URLByAppendingPathComponent:@"staging"] + URLByAppendingPathComponent:key] + .path]); + EXPECT_TRUE([NSFileManager.defaultManager + fileExistsAtPath:[[root.url URLByAppendingPathComponent:@"staging"] + URLByAppendingPathComponent:collision] + .path]); + auto saved = saved_bookmark(root.url, collision); + ASSERT_TRUE(saved.ok()); + EXPECT_TRUE([saved.get() isEqual:bytes]); + struct stat after {}; + ASSERT_EQ(lstat(lock_url().fileSystemRepresentation, &after), 0); + EXPECT_EQ(before.st_ino, after.st_ino); +} TEST_F(CoreAIBookmarkPathTest, RejectsSymlinkDirectoryAndFifoStripeFiles) { NSData* bytes = [@"untouched target" dataUsingEncoding:NSUTF8StringEncoding]; @@ -1076,6 +1121,7 @@ FileDescriptor file(open( auto locked = lock_bookmark(root.url.path, key); ASSERT_TRUE(locked.ok()); locked.get().reset(); + ASSERT_EQ(clear_bookmark_assets(root.url.path, key, nil), Error::Ok); struct stat after {}; ASSERT_EQ(lstat(legacy.fileSystemRepresentation, &after), 0); EXPECT_EQ(before.st_ino, after.st_ino); diff --git a/backends/apple/coreai/runtime/test/coreai_cache_test.mm b/backends/apple/coreai/runtime/test/coreai_cache_test.mm index c9e55500c03..dc323122c0a 100644 --- a/backends/apple/coreai/runtime/test/coreai_cache_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_cache_test.mm @@ -148,9 +148,7 @@ void cache_entry_removed(NSURL* root, const CacheEntry& entry) { EXPECT_FALSE(cache_exists(cache_identity_url(root, entry.key))); EXPECT_FALSE(cache_exists(test_bookmark_url(root, entry.key))); EXPECT_FALSE(cache_exists(entry.stage)); - EXPECT_TRUE(cache_exists([[root URLByAppendingPathComponent:@"locks"] - URLByAppendingPathComponent:[entry.key - stringByAppendingString:@".lock"]])); + EXPECT_TRUE(cache_exists(test_bookmark_lock_url(root, entry.key))); } NSMutableDictionary* cache_manifest(char identity) { From 86f50e75a34297cbfa24614db0b4e5e79dc3f508 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 28 Oct 2026 17:37:53 -0700 Subject: [PATCH 5/5] [CoreAI] Unify cache identity and specialization ownership Keep loading and deletion in agreement by giving each model variant a self-describing identity derived from the actual specialization options. This avoids policy drift and lets PTE clearing discover every related device and specialization variant without attribution sidecars. Permit SDK reclamation under storage pressure without invalidating cached models when staged source is removed. Test Plan: All five native CTest targets passed, including independent warm bookmarks for different specialization policies, all-variant PTE clearing, staging-only cleanup, and real-SDK restoration after staged-source deletion. The policy regression confirms storage-pressure-only purging. The complete CoreAI Python suite passed (124 tests). Lint, iOS syntax checks, and whitespace checks passed; the standalone stress runner rebuilt successfully. --- backends/apple/coreai/CMakeLists.txt | 2 + backends/apple/coreai/README.md | 115 ++-- .../apple/coreai/runtime/ETCoreAIBridge.h | 6 +- .../apple/coreai/runtime/ETCoreAIModel.swift | 65 +- backends/apple/coreai/runtime/coreai_assets.h | 3 +- .../apple/coreai/runtime/coreai_assets.mm | 10 +- .../apple/coreai/runtime/coreai_bookmarks.h | 37 +- .../apple/coreai/runtime/coreai_bookmarks.mm | 212 +------ backends/apple/coreai/runtime/coreai_cache.h | 7 +- backends/apple/coreai/runtime/coreai_cache.mm | 66 +- .../apple/coreai/runtime/coreai_cache_key.h | 35 ++ .../apple/coreai/runtime/coreai_cache_key.mm | 154 +++++ .../coreai/runtime/coreai_load_coordinator.h | 2 +- .../coreai/runtime/coreai_load_coordinator.mm | 11 +- .../apple/coreai/runtime/coreai_storage.mm | 24 +- .../runtime/test/SwiftBridgeTests.swift | 80 ++- .../runtime/test/coreai_acquisition_test.mm | 145 ++--- .../runtime/test/coreai_bookmark_fixture.h | 5 + .../runtime/test/coreai_bookmark_fixture.mm | 36 +- .../runtime/test/coreai_bookmark_test.mm | 593 ++++++------------ .../coreai/runtime/test/coreai_cache_test.mm | 140 ++--- .../runtime/test/coreai_delegate_test.mm | 10 +- .../coreai/runtime/test/coreai_fake_loader.h | 5 +- .../coreai/runtime/test/coreai_fake_loader.mm | 21 +- .../runtime/test/coreai_key_format_test.mm | 81 +++ .../runtime/test/coreai_source_fixture.mm | 23 +- .../coreai/runtime/test/coreai_source_test.mm | 55 +- .../runtime/test/coreai_storage_test.mm | 28 +- 28 files changed, 974 insertions(+), 997 deletions(-) create mode 100644 backends/apple/coreai/runtime/coreai_cache_key.h create mode 100644 backends/apple/coreai/runtime/coreai_cache_key.mm create mode 100644 backends/apple/coreai/runtime/test/coreai_key_format_test.mm diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index 481790ff8a4..2c1ba4f1f38 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/backends/apple/coreai/CMakeLists.txt @@ -32,6 +32,7 @@ set(_coreai_runtime_sources runtime/coreai_assets.mm runtime/coreai_storage.mm runtime/coreai_bookmarks.mm + runtime/coreai_cache_key.mm runtime/coreai_load_coordinator.mm runtime/coreai_pte.cpp runtime/coreai_cache.mm @@ -60,6 +61,7 @@ if(EXECUTORCH_BUILD_TESTS AND CMAKE_SYSTEM_NAME STREQUAL "Darwin") runtime/test/coreai_filesystem_fixture.mm runtime/test/coreai_source_test.mm runtime/test/coreai_bookmark_test.mm + runtime/test/coreai_key_format_test.mm runtime/test/coreai_source_fixture.mm runtime/test/coreai_bookmark_fixture.mm runtime/test/coreai_fake_loader.mm diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index 4510e97cf12..f05df35527d 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -98,7 +98,7 @@ are reported separately from an absent architecture entry. Both formats can require device specialization. Every load first attempts SDK bookmark restoration. A missing bookmark or SDK-confirmed cache miss falls back -to source materialization and persistent specialization, then binds the requested +to source materialization and cached specialization, then binds the requested function and ordered I/O on that same acquired model. ### Asset storage @@ -171,24 +171,49 @@ validating source assets only when recovery requires them. There is no synthetic zero-input inference at initialization. SDK specialization and function resource loading can still be substantial. -The runtime always uses `AIModelCache.default` with persistent retention and -stores an opaque bookmark for direct restoration. `coreai_assets_dir` does not -choose the SDK artifact directory. The backend does not inspect SDK-private -files, alter their backup flags, or use an App Group cache. - -Each partition derives a key from its selected export digest/bundle, source/AOT -mode, platform, Core AI device architecture, default SDK cache namespace, -versioned default options and persistent policy. Function bindings and delivery -location are excluded. No model bytes are hashed at runtime. Each key determines -a raw bookmark in `bookmarks`, a bundle directory in `staging`, and one of -256 shared coordination files in `locks`. There is no PTE-wide bookmark list or -serialized lifecycle record. +The runtime always uses `AIModelCache.default` with storage-pressure-only purging +and stores an opaque bookmark for direct restoration. Removing staged source does +not trigger SDK purging, but the SDK may reclaim artifacts when storage is low. +A later bookmark miss rematerializes source and specializes again, so packaged +named data must remain available. This policy applies when specializing; restoring +an existing bookmark does not update that entry's retention policy. +`coreai_assets_dir` does not choose the SDK artifact directory. The backend does +not inspect SDK-private files, alter their backup flags, or use an App Group cache. + +Each partition uses the shared key builder to produce +`--`. The model component identifies +the format version, selected export digest/bundle, and source/AOT mode. The device +component contains filename-safe platform and Core AI architecture names, +separated by an underscore, not an individual machine identifier. The +specialization component hashes the loader's stable specialization descriptor. +Both hashes are lowercase SHA-256 values. Function bindings and delivery location +are excluded, and no model bytes are hashed at runtime. + +The Swift loader owns an immutable snapshot of the actual SDK specialization +options. Its descriptor is derived from that snapshot, and specialization uses +those same options. Options equal to SDK `.default` produce exactly the UTF-8 +string `default`. Other snapshots encode a versioned JSON descriptor with sorted +allowed compute units, the preferred compute unit, and frequent-reshape behavior. +The coordinator reads the descriptor once before lookup; it cannot independently +select an options label. Future configurable SDK fields must be represented in +the descriptor schema and tested before use. No public specialization option is +exposed yet. SDK cache selection and the storage-pressure-only purge policy +remain fixed and are separate from specialization identity. + +The cache-key module owns hashing and filename interpretation, bookmark storage +owns locked I/O and bookmark/staging inventory, and the coordinator owns SDK +acquisition and eviction ordering. PTE clearing selects model prefixes rather +than reconstructing keys from today's loader options. + +Each composite key determines a raw bookmark in `bookmarks`, a bundle directory +in `staging`, and one of 256 shared coordination files in `locks`. No metadata +sidecars or PTE-wide bookmark lists are written. A healthy bookmark hit does not consult NDS and attempts to clean leftover staging. On a missing bookmark or confirmed SDK miss, the coordinator prepares the source, specializes it and atomically writes the returned bookmark bytes. After durable publication, it attempts to remove the entire keyed staging -directory while holding the key's stripe lock. Persistent SDK caching supports +directory while holding the key's stripe lock. The SDK cache policy supports inference and bookmark restoration without that extracted copy. Keep the packaged named data available: a later SDK cache miss requires source materialization again. There is no secondary source-URL cache lookup. @@ -204,12 +229,11 @@ The backend does not track historical or unrecorded SDK entries. Raw bookmarks have an 8 MiB backend allocation limit, not an asserted SDK format maximum, and are published atomically because losing one orphans an SDK entry. A fixed pool of at most 256 disk locks coordinates loading and eviction across -threads and processes. The first byte of the validated cache key selects -`stripe-v1-00.lock` through `stripe-v1-ff.lock` within the lock directory. Files -are created on demand, so new model versions do not grow the pool beyond 256. -The mapping is part of the on-disk format and is not configurable. Different -keys on the same stripe serialize; different stripes remain independent. -Bookmarks, attribution, and staging remain separate for every key. Runtime +threads and processes. The first byte of a SHA-256 digest of the validated full +cache key selects its stripe. Files are created on demand, so new model versions +do not grow the pool beyond 256. The mapping is part of the on-disk format and is +not configurable. Different keys on the same stripe serialize; different stripes +remain independent. Bookmarks and staging remain separate for every key. Runtime operations hold only one stripe at a time. Lock ownership is maintained by the OS, and acquiring a lock does not flush its @@ -218,12 +242,12 @@ clearing or after process exit, so waiters and future openers keep using the same inode. Removing the assets root externally requires all loads, sessions and maintenance to stop. -Upgrading from per-key lock files requires stopping all processes using the -root and upgrading them together. Old and new locking schemes must not run -concurrently against the same root. Legacy per-key lock files are ignored and -left untouched; they can be removed separately during a fully quiescent offline -cleanup. The 256-file bound applies to the new stripe pool, not pre-existing -legacy files. Bookmarks, metadata, and staged bundles need no migration. +Changing cache-key formats requires stopping all processes using the root and +upgrading them together. Different key formats and lock mappings must not run +concurrently against the same root. Older single-hash bookmarks, metadata, and +staged bundles are not migrated or discovered by these maintenance APIs. Retire +old entries using their version's maintenance API while quiescent. Pre-existing +per-key lock files are ignored and are outside the 256-file stripe-pool bound. ### Lifecycle logging @@ -235,7 +259,7 @@ multiple partitions can be correlated. These messages do not include bookmark contents, model data, or per-inference logs. A cache hit means that the SDK successfully restored a model from its bookmark. -Specialization start means a request to the SDK's persistent default cache; the +Specialization start means a request to the SDK's default cache; the SDK does not report whether that request compiled afresh or reused an internal artifact. The durable-save message confirms local bookmark publication, including its final directory flush, not independent verification of SDK-private @@ -260,14 +284,16 @@ when the directory is omitted or null. A supplied path must be absolute; invalid paths fail without falling back to the default. If loading used a custom root, pass that same root when clearing; the PTE does not store runtime overrides. -Root-wide clearing visits keyed bookmarks, identity records, and staged -directories. PTE clearing finds every Core AI delegate in every method and clears -all recorded variants of its bundles, including other platforms and SDK -architectures. Manifests must be structurally valid before any entry is cleared; -their minimum OS and runtime compatibility do not restrict maintenance. -Discovery requires neither SDK availability nor an architecture query, and does -not load models or materialize assets. Actual SDK eviction still requires a -usable bookmark and an available loader. +Root-wide clearing visits keyed bookmarks and staged directories. PTE clearing +collects model hashes for every declared bundle in every Core AI delegate across +all methods. It inventories the cache filenames once and matches the complete +model component of each valid composite key, finding every recorded device and +specialization variant without knowing which options were previously used. +Manifests must be structurally valid before any entry is cleared; their minimum +OS and runtime compatibility do not restrict maintenance. Discovery requires +neither metadata reads, SDK availability, nor an architecture query, and does not +load models or materialize assets. Actual SDK eviction still requires a usable +bookmark and an available loader. These synchronous calls require callers to unload affected models and prevent concurrent loads, inference and maintenance across processes. Stripe locks are @@ -292,25 +318,15 @@ that these APIs cannot target. No history of removed entries is retained. Copied or reused exported artifacts can share keys; clearing them also removes the other copy's warm-cache benefit. Subsequent loads can specialize again. -Identity records are small JSON sidecars beside bookmarks. They retain bundle -identity and cache context, and are checked against their hashed key. Loading -publishes the record before staging or SDK acquisition, and clearing removes it -last. Existing keys and opaque bookmark bytes are unchanged. Older AOT entries -can be found directly from the PTE's architecture list. Older source entries -without attribution cannot always be identified: PTE clearing removes known -matches, logs unresolved keys, and reports incomplete cleanup rather than -removing unrelated entries. A one-time root-wide `clear_cache` handles these -older entries. Missing or corrupt attribution has the same conservative behavior. - ### Persistent SDK Rejection SDK errors are not classified as transient or permanent. Busy errors and all unclassified eviction/restore errors retain the local bookmark for retry; only successful eviction or a confirmed SDK miss permits normal removal. If retries keep failing, first stop every process using the -assets root. Back up or quarantine only the affected key's bookmark, identity -sidecar, and keyed staging outside that root, leaving the stable lock files in -place. A later load can attempt cold acquisition and create new tracking data. +assets root. Back up or quarantine only the affected key's bookmark and keyed +staging outside that root, leaving the stable lock files in place. A later load +can attempt cold acquisition and publish a new bookmark. This manual recovery does not confirm SDK eviction or guarantee disk-space reclamation: keep the backup because it may contain the only remaining eviction handle. Do not perform manual recovery while loading, inference, or maintenance @@ -343,7 +359,8 @@ have a 255-byte limit. The private Swift module `CoreAIBridge` (`runtime/ETCoreAIModel.swift`) wraps the `CoreAI` framework behind the Objective-C protocols in `runtime/ETCoreAIBridge.h`. It restores models from bookmarks, specializes source -bundles with the persistent default SDK cache and binds named functions with +bundles with the default SDK cache and storage-pressure-only purging, and binds +named functions with ordered inputs and outputs. Asynchronous SDK calls run in detached tasks and report through completion blocks; function binding completes synchronously. diff --git a/backends/apple/coreai/runtime/ETCoreAIBridge.h b/backends/apple/coreai/runtime/ETCoreAIBridge.h index df2b4ba09cf..73426ac00c6 100644 --- a/backends/apple/coreai/runtime/ETCoreAIBridge.h +++ b/backends/apple/coreai/runtime/ETCoreAIBridge.h @@ -111,9 +111,13 @@ typedef void (^ETCoreAIAcquisitionCompletion)( NSError* _Nullable error); @protocol ETCoreAIModelLoading +// Stable for this loader's lifetime; derived from its actual specialization +// options. +@property(nonatomic, readonly, copy) NSString* specializationIdentifier; - (void)restoreModelFromBookmark:(NSData*)bookmark completion:(ETCoreAIAcquisitionCompletion)completion; -// Uses only the SDK default cache and persistent policy. Never returns Miss. +// Uses the SDK default cache with storage-pressure-only purging. Never returns +// Miss. - (void)specializeModelAtURL:(NSURL*)url completion:(ETCoreAIAcquisitionCompletion)completion; // Nil error means SDK-confirmed deletion. No busy/absent classification is diff --git a/backends/apple/coreai/runtime/ETCoreAIModel.swift b/backends/apple/coreai/runtime/ETCoreAIModel.swift index 3f04b9414d6..1e7dc71fffb 100644 --- a/backends/apple/coreai/runtime/ETCoreAIModel.swift +++ b/backends/apple/coreai/runtime/ETCoreAIModel.swift @@ -404,9 +404,71 @@ private final class CoreAISwiftPreparedModel: NSObject, ETCoreAIPreparedModel, S } } +@available(macOS 27.0, iOS 27.0, *) +struct CoreAISpecialization: Sendable { + let options: SpecializationOptions + let identifier: String + + static let `default` = Self() + + private init() { + options = .default + identifier = "default" + } + + // New SDK option fields require a schema review and versioned identity changes. + private struct Descriptor: Encodable { + let version: Int + let allowedComputeUnits: [String] + let preferredComputeUnit: String? + let expectFrequentReshapes: Bool + } + + init(options: SpecializationOptions) throws { + if options == .default { + self = Self.default + return + } + let descriptor = Descriptor( + version: 1, + allowedComputeUnits: try options.allowedComputeUnitKinds.map(Self.name).sorted(), + preferredComputeUnit: try options.preferredComputeUnitKind.map(Self.name), + expectFrequentReshapes: options.expectFrequentReshapes) + let encoder = JSONEncoder() + encoder.outputFormatting = [.sortedKeys] + identifier = String(decoding: try encoder.encode(descriptor), as: UTF8.self) + self.options = options + } + + private static func name(_ kind: ComputeUnitKind) throws -> String { + switch kind { + case .cpu: return "cpu" + case .gpu: return "gpu" + case .neuralEngine: return "neuralEngine" + @unknown default: + throw bridgeError(.unsupported, "Unknown specialization compute unit") + } + } +} + @available(macOS 27.0, iOS 27.0, *) @objc(ETCoreAISwiftModelLoader) public final class ETCoreAISwiftModelLoader: NSObject, ETCoreAIModelLoading { + static let cachePolicy = AIModelCache.Policy(purgeConditions: [.storagePressure]) + private let specialization: CoreAISpecialization + + public override init() { + specialization = .default + super.init() + } + + init(specialization: CoreAISpecialization) { + self.specialization = specialization + super.init() + } + + public var specializationIdentifier: String { specialization.identifier } + @objc public static var deviceArchitectureName: String { AIModel.deviceArchitectureName } public func restoreModel( @@ -418,12 +480,13 @@ public final class ETCoreAISwiftModelLoader: NSObject, ETCoreAIModelLoading { } public func specializeModel(at url: URL, completion: @escaping ETCoreAIAcquisitionCompletion) { + let options = specialization.options acquireCoreAIModel(context: "Core AI specialization failed", completion: completion) { guard url.isFileURL else { throw bridgeError(.invalidArgument, "A file URL is required") } let model = try await AIModel.specialize( - contentsOf: url, options: .default, cache: .default, cachePolicy: .persistent) + contentsOf: url, options: options, cache: .default, cachePolicy: Self.cachePolicy) return CoreAISwiftPreparedModel(model: model) } } diff --git a/backends/apple/coreai/runtime/coreai_assets.h b/backends/apple/coreai/runtime/coreai_assets.h index 302fcfd729e..a3fc78f8754 100644 --- a/backends/apple/coreai/runtime/coreai_assets.h +++ b/backends/apple/coreai/runtime/coreai_assets.h @@ -35,8 +35,7 @@ runtime::Result select_assets( NSString* device_architecture, NSString* platform); // Caller holds the key's disk lock and has prepared source_root. -// The key is 64 lowercase hex characters identifying the selected export digest -// and bundle name, optionally including platform or SDK context. +// The key is a composite cache key identifying the model's cached variant. // Bundles are published at source_root/key/bundle. Reuse validates structure, // not contents, and fully flushes source_root before returning. When // replace_corrupt is set, damaged regular-file trees may be removed and diff --git a/backends/apple/coreai/runtime/coreai_assets.mm b/backends/apple/coreai/runtime/coreai_assets.mm index 1a66c48792b..9afd3c36a35 100644 --- a/backends/apple/coreai/runtime/coreai_assets.mm +++ b/backends/apple/coreai/runtime/coreai_assets.mm @@ -7,6 +7,7 @@ */ #import "coreai_assets.h" +#import "coreai_cache_key.h" #include "coreai_file.h" #include @@ -452,14 +453,9 @@ FileDescriptor entry(retry_eintr([&] { InvalidProgram, "Core AI bundle must be selected before preparing its source"); ET_CHECK_OR_RETURN_ERROR( - nonempty_string(key) && key.length == 64 && - [key rangeOfCharacterFromSet: - [[NSCharacterSet - characterSetWithCharactersInString:@"0123456789abcdef"] - invertedSet]] - .location == NSNotFound, + is_cache_key(key), InvalidArgument, - "Core AI staging key must be a 64-character lowercase hex digest"); + "Core AI staging key must be a composite cache key"); auto prepared = prepare_storage_root(source_root, false); if (!prepared.ok()) { return prepared.error(); diff --git a/backends/apple/coreai/runtime/coreai_bookmarks.h b/backends/apple/coreai/runtime/coreai_bookmarks.h index 3fdda5d1467..e33182280e7 100644 --- a/backends/apple/coreai/runtime/coreai_bookmarks.h +++ b/backends/apple/coreai/runtime/coreai_bookmarks.h @@ -9,28 +9,17 @@ #include #include -#import "coreai_assets.h" #include "coreai_storage.h" namespace executorch::backends::coreai { -runtime::Result bookmark_key( - const Manifest& selected, - NSString* platform, - NSString* architecture); -runtime::Result bookmark_identity( - const Manifest& manifest, - NSString* bundle); -// The first key byte selects a stripe in a fixed, versioned lock namespace. -constexpr size_t kBookmarkLockStripeCount = 256; -runtime::Result bookmark_lock_name(NSString* key); // With create=false, a missing root is returned unchanged for no-op // maintenance. runtime::Result resolve_bookmark_root( NSString* path, bool create = true); // Missing roots produce an empty inventory. -runtime::Result*> inventory_bookmark_keys(NSString* root); +runtime::Result*> inventory_cache_keys(NSString* root); class BookmarkLock final { public: @@ -43,21 +32,13 @@ class BookmarkLock final { friend runtime::Result read_bookmark( const BookmarkLock& lock, bool maintenance); - friend runtime::Result read_bookmark_identity( - const BookmarkLock& lock); - friend runtime::Error ensure_bookmark_identity( - const BookmarkLock& lock, - const Manifest& selected, - NSString* platform, - NSString* architecture); - friend runtime::Error remove_bookmark_identity(const BookmarkLock& lock); friend runtime::Error write_bookmark(const BookmarkLock& lock, NSData* data); friend runtime::Error remove_bookmark(const BookmarkLock& lock); friend runtime::Error remove_bookmark_staging( const BookmarkLock& lock, bool remove, StorageRemovalMode mode); - friend runtime::Result prepare_bookmark_staging( + friend runtime::Result prepare_bookmark_staging_root( const BookmarkLock& lock); BookmarkLock( int root, @@ -85,14 +66,6 @@ lock_bookmark(NSString* root, NSString* key, bool create = true); runtime::Result read_bookmark( const BookmarkLock& lock, bool maintenance = false); -// Missing metadata is nil; malformed or misattributed records fail closed. -runtime::Result read_bookmark_identity(const BookmarkLock& lock); -runtime::Error ensure_bookmark_identity( - const BookmarkLock& lock, - const Manifest& selected, - NSString* platform, - NSString* architecture); -runtime::Error remove_bookmark_identity(const BookmarkLock& lock); // A staging-only lock has no bookmark directory; writes return InvalidArgument. runtime::Error write_bookmark(const BookmarkLock& lock, NSData* data); runtime::Error remove_bookmark(const BookmarkLock& lock); @@ -102,7 +75,9 @@ runtime::Error remove_bookmark_staging( const BookmarkLock& lock, bool remove = true, StorageRemovalMode mode = StorageRemovalMode::Strict); -// Prepared only for cold loads; source preparation appends key/bundle. -runtime::Result prepare_bookmark_staging(const BookmarkLock& lock); +// Returns the staging parent; the caller appends key/bundle. Prepared only +// for cold loads. +runtime::Result prepare_bookmark_staging_root( + const BookmarkLock& lock); } // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_bookmarks.mm b/backends/apple/coreai/runtime/coreai_bookmarks.mm index 16d153d6926..59985593603 100644 --- a/backends/apple/coreai/runtime/coreai_bookmarks.mm +++ b/backends/apple/coreai/runtime/coreai_bookmarks.mm @@ -6,9 +6,9 @@ * LICENSE file in the root directory of this source tree. */ #import "coreai_bookmarks.h" +#import "coreai_cache_key.h" #import "coreai_storage.h" -#include #include #include #include @@ -22,115 +22,10 @@ // An allocation bound, not an SDK bookmark format limit. constexpr size_t kBookmarkByteLimit = 8 * 1024 * 1024; -bool text(id value) { - return [value isKindOfClass:NSString.class] && [value length] > 0 && - [value length] <= 4096 && - [value rangeOfString:@"\0"].location == NSNotFound; -} -bool component(id value) { - return text(value) && ![value containsString:@"/"] && - ![value containsString:@"\\"] && ![value isEqual:@"."] && - ![value isEqual:@".."]; -} -bool digest(id value) { - return text(value) && [value length] == 64 && - [value rangeOfCharacterFromSet:[[NSCharacterSet - characterSetWithCharactersInString: - @"0123456789abcdef"] invertedSet]] - .location == NSNotFound; -} NSString* bookmark_name(NSString* key) { return [key stringByAppendingString:@".bookmark"]; } -constexpr size_t kIdentityByteLimit = 128 * 1024; -NSString* identity_name(NSString* key) { - return [key stringByAppendingString:@".metadata"]; -} - -Result*> identity_fields( - const Manifest& selected, - NSString* bundle) { - ET_CHECK_OR_RETURN_ERROR( - component(bundle), - DelegateInvalidCompatibility, - "Invalid Core AI bundle name"); - NSString* exported_digest = selected.bundle_digests[bundle]; - ET_CHECK_OR_RETURN_ERROR( - digest(exported_digest), - DelegateInvalidCompatibility, - "Missing Core AI bundle digest"); - return @[ - @"1", - exported_digest, - selected.aot_compiled ? @"aot" : @"source", - bundle - ]; -} - -Result*> key_fields( - const Manifest& selected, - NSString* platform, - NSString* architecture) { - auto identity = identity_fields(selected, selected.path.lastPathComponent); - if (!identity.ok()) - return identity.error(); - ET_CHECK_OR_RETURN_ERROR( - text(platform) && text(architecture), - DelegateInvalidCompatibility, - "Missing Core AI cache context"); - return [identity.get() arrayByAddingObjectsFromArray:@[ - platform, - architecture, - @"AIModelCache.default", - @"default-options-v1", - @"persistent" - ]]; -} - -Result hash_fields(NSArray* fields) { - NSData* encoded = [NSJSONSerialization dataWithJSONObject:fields - options:0 - error:nil]; - ET_CHECK_OR_RETURN_ERROR( - encoded != nil, InvalidExternalData, "Cannot encode bookmark key"); - unsigned char bytes[CC_SHA256_DIGEST_LENGTH]; - CC_SHA256(encoded.bytes, static_cast(encoded.length), bytes); - NSMutableString* key = [NSMutableString stringWithCapacity:64]; - for (unsigned char byte : bytes) - [key appendFormat:@"%02x", byte]; - return [key copy]; -} - -} // namespace - -Result bookmark_identity( - const Manifest& manifest, - NSString* bundle) { - auto fields = identity_fields(manifest, bundle); - if (!fields.ok()) - return fields.error(); - return hash_fields(fields.get()); -} - -Result bookmark_key( - const Manifest& selected, - NSString* platform, - NSString* architecture) { - auto fields = key_fields(selected, platform, architecture); - if (!fields.ok()) - return fields.error(); - return hash_fields(fields.get()); -} - -Result bookmark_lock_name(NSString* key) { - ET_CHECK_OR_RETURN_ERROR( - digest(key), InvalidArgument, "Invalid Core AI bookmark key"); - return [NSString - stringWithFormat:@"stripe-v1-%@.lock", [key substringToIndex:2]]; -} - -namespace { // Returns -1 when the root does not exist. Result open_root(NSString* root) { const int fd = retry_eintr([&] { @@ -164,9 +59,8 @@ bool digest(id value) { FileDescriptor child(opened.get()); if (child.get() < 0) continue; - NSArray* entries = [name isEqual:@"bookmarks"] - ? @[ bookmark_name(key), identity_name(key) ] - : @[ key ]; + NSArray* entries = + [name isEqual:@"bookmarks"] ? @[ bookmark_name(key) ] : @[ key ]; for (NSString* entry in entries) { struct stat info; const int result = retry_eintr([&] { @@ -188,7 +82,9 @@ bool digest(id value) { Result resolve_bookmark_root(NSString* path, bool create) { ET_CHECK_OR_RETURN_ERROR( - text(path) && [path hasPrefix:@"/"], + [path isKindOfClass:NSString.class] && path.length <= 4096 && + [path hasPrefix:@"/"] && + [path rangeOfString:@"\0"].location == NSNotFound, InvalidArgument, "Core AI assets require an absolute root path"); if (create) { @@ -206,7 +102,7 @@ bool digest(id value) { return path; } -Result*> inventory_bookmark_keys(NSString* root) { +Result*> inventory_cache_keys(NSString* root) { auto resolved = resolve_bookmark_root(root, false); if (!resolved.ok()) return resolved.error(); @@ -230,13 +126,11 @@ bool digest(id value) { for (NSString* entry in children.get()) { NSString* key = entry; if ([name isEqual:@"bookmarks"]) { - NSString* suffix = - [entry hasSuffix:@".bookmark"] ? @".bookmark" : @".metadata"; - if (![entry hasSuffix:suffix]) + if (![entry hasSuffix:@".bookmark"]) continue; - key = [entry substringToIndex:entry.length - suffix.length]; + key = [entry substringToIndex:entry.length - @".bookmark".length]; } - if (digest(key)) + if (is_cache_key(key)) [keys addObject:key]; } } @@ -387,90 +281,6 @@ FileDescriptor file(retry_eintr([&] { maintenance); } -Result read_bookmark_identity(const BookmarkLock& lock) { - auto bytes = read_record( - lock.bookmarks_.get(), - identity_name(lock.key_), - kIdentityByteLimit, - true); - if (!bytes.ok()) - return bytes.error(); - if (bytes.get() == nil) - return static_cast(nil); - id parsed = [NSJSONSerialization JSONObjectWithData:bytes.get() - options:0 - error:nil]; - if (![parsed isKindOfClass:NSArray.class] || [parsed count] != 9) - return Error::InvalidExternalData; - NSArray* fields = parsed; - for (id field in fields) { - if (!text(field)) - return Error::InvalidExternalData; - } - if (![fields[2] isEqual:@"source"] && ![fields[2] isEqual:@"aot"]) - return Error::InvalidExternalData; - Manifest selected; - selected.path = fields[3]; - selected.bundle_digests = @{fields[3] : fields[1]}; - selected.aot_compiled = [fields[2] isEqual:@"aot"]; - auto expected = key_fields(selected, fields[4], fields[5]); - if (!expected.ok() || ![expected.get() isEqual:fields]) - return Error::InvalidExternalData; - auto key = hash_fields(fields); - if (!key.ok()) - return key.error(); - if (![key.get() isEqual:lock.key_]) - return Error::InvalidExternalData; - return bookmark_identity(selected, fields[3]); -} - -Error ensure_bookmark_identity( - const BookmarkLock& lock, - const Manifest& selected, - NSString* platform, - NSString* architecture) { - ET_CHECK_OR_RETURN_ERROR( - lock.bookmarks_.get() >= 0, - InvalidArgument, - "Cannot attribute a staging-only lock"); - auto fields = key_fields(selected, platform, architecture); - if (!fields.ok()) - return fields.error(); - auto key = hash_fields(fields.get()); - if (!key.ok()) - return key.error(); - ET_CHECK_OR_RETURN_ERROR( - [key.get() isEqual:lock.key_], - InvalidArgument, - "Core AI cache identity does not match key"); - NSData* encoded = [NSJSONSerialization dataWithJSONObject:fields.get() - options:0 - error:nil]; - ET_CHECK_OR_RETURN_ERROR( - encoded != nil && encoded.length <= kIdentityByteLimit, - InvalidExternalData, - "Invalid Core AI cache identity"); - auto existing = read_record( - lock.bookmarks_.get(), - identity_name(lock.key_), - kIdentityByteLimit, - true); - if (!existing.ok() && existing.error() != Error::InvalidExternalData) - return existing.error(); - if (existing.ok() && [existing.get() isEqual:encoded]) { - // A prior publication may have failed its post-rename flush. - return sync_storage_directory(lock.bookmarks_.get(), true); - } - return publish_storage_data( - lock.bookmarks_.get(), identity_name(lock.key_), encoded); -} - -Error remove_bookmark_identity(const BookmarkLock& lock) { - if (lock.bookmarks_.get() < 0) - return Error::Ok; - return remove_storage_entry(lock.bookmarks_.get(), identity_name(lock.key_)); -} - Error write_bookmark(const BookmarkLock& lock, NSData* data) { ET_CHECK_OR_RETURN_ERROR( [data isKindOfClass:NSData.class] && data.length > 0 && @@ -498,7 +308,7 @@ Error remove_bookmark_staging( return remove_storage_staging(lock.root_.get(), lock.key_, remove, mode); } -Result prepare_bookmark_staging(const BookmarkLock& lock) { +Result prepare_bookmark_staging_root(const BookmarkLock& lock) { auto staging = open_child(lock.root_.get(), @"staging", true); if (!staging.ok()) return staging.error(); diff --git a/backends/apple/coreai/runtime/coreai_cache.h b/backends/apple/coreai/runtime/coreai_cache.h index b6ce1c06258..6904b399bde 100644 --- a/backends/apple/coreai/runtime/coreai_cache.h +++ b/backends/apple/coreai/runtime/coreai_cache.h @@ -35,10 +35,9 @@ namespace executorch::backends::coreai { * Clears all recorded cache variants referenced by every CoreAI delegate in * the PTE after validating its manifests. Discovery does not require host * compatibility or SDK availability. Actual SDK eviction requires a loader. - * Missing/corrupt attribution can leave legacy source entries unidentified: - * known matches are cleared, then an error reports incomplete cleanup. A - * root-wide clear can recover these entries. Does not initialize delegates or - * materialize Core AI assets. + * Matching model prefixes include all recorded device and specialization + * variants without metadata sidecars. Does not initialize delegates or + * materialize Core AI assets. Older single-hash entries are not discovered. * * A null directory selects the same Caches default as backend initialization. * Supplied directories must be valid explicit roots; they never fall back to diff --git a/backends/apple/coreai/runtime/coreai_cache.mm b/backends/apple/coreai/runtime/coreai_cache.mm index f20155dd843..59fef6ad9e7 100644 --- a/backends/apple/coreai/runtime/coreai_cache.mm +++ b/backends/apple/coreai/runtime/coreai_cache.mm @@ -8,6 +8,8 @@ #include "coreai_cache.h" #import "coreai_assets.h" +#import "coreai_bookmarks.h" +#import "coreai_cache_key.h" #import "coreai_load_coordinator.h" #include "coreai_pte.h" #import "coreai_storage.h" @@ -56,9 +58,7 @@ Error clear_pte(runtime::DataLoader& loader, NSString* root) { auto pte = inspect_coreai_pte(loader); if (!pte.ok()) return pte.error(); - NSMutableSet* identities = [NSMutableSet set]; - NSMutableSet* legacy_aot_keys = [NSMutableSet set]; - bool has_source = false; + NSMutableSet* models = [NSMutableSet set]; for (const auto& buffer : pte->processed) { auto data = buffer.data_safe(); if (!data.ok()) @@ -74,65 +74,23 @@ Error clear_pte(runtime::DataLoader& loader, NSString* root) { if (!manifest.ok()) return manifest.error(); for (NSString* bundle in manifest->bundle_digests) { - auto identity = bookmark_identity(manifest.get(), bundle); - if (!identity.ok()) - return identity.error(); - [identities addObject:identity.get()]; - } - if (manifest->aot_compiled) { - for (NSString* architecture in manifest->archs) { - auto selected = - select_assets(manifest.get(), architecture, manifest->platform); - if (!selected.ok()) - return selected.error(); - auto key = - bookmark_key(selected.get(), manifest->platform, architecture); - if (!key.ok()) - return key.error(); - [legacy_aot_keys addObject:key.get()]; - } - } else { - has_source = true; + auto model = cache_model_key(manifest.get(), bundle); + if (!model.ok()) + return model.error(); + [models addObject:model.get()]; } } - if (identities.count == 0) + if (models.count == 0) return Error::Ok; - auto inventory = inventory_bookmark_keys(root); + auto inventory = inventory_cache_keys(root); if (!inventory.ok()) return inventory.error(); NSMutableArray* keys = [NSMutableArray array]; - Error lookup_error = Error::Ok; for (NSString* key in inventory.get()) { - if ([legacy_aot_keys containsObject:key]) { - [keys addObject:key]; - continue; - } - auto locked = lock_bookmark(root, key, false); - if (!locked.ok()) { - if (lookup_error == Error::Ok) - lookup_error = locked.error(); - continue; - } - if (locked.get() == nullptr) - continue; - auto identity = read_bookmark_identity(*locked.get()); - if (!identity.ok() || identity.get() == nil) { - if (has_source) { - ET_LOG( - Info, - "Core AI cache key %s has missing or invalid attribution; PTE cleanup is incomplete", - key.UTF8String); - if (lookup_error == Error::Ok) - lookup_error = - identity.ok() ? Error::InvalidExternalData : identity.error(); - } - continue; - } - if ([identities containsObject:identity.get()]) + if ([models containsObject:cache_key_model_prefix(key)]) [keys addObject:key]; } - const auto cleared = clear_keys(root, keys); - return lookup_error != Error::Ok ? lookup_error : cleared; + return clear_keys(root, keys); } } // namespace @@ -141,7 +99,7 @@ Error clear_cache(const char* coreai_assets_dir) { auto root = cache_root(coreai_assets_dir); if (!root.ok()) return root.error(); - auto keys = inventory_bookmark_keys(root.get()); + auto keys = inventory_cache_keys(root.get()); if (!keys.ok()) return keys.error(); return clear_keys(root.get(), keys.get()); diff --git a/backends/apple/coreai/runtime/coreai_cache_key.h b/backends/apple/coreai/runtime/coreai_cache_key.h new file mode 100644 index 00000000000..1c27a825ba1 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache_key.h @@ -0,0 +1,35 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +#pragma once + +#import +#include +#include + +namespace executorch::backends::coreai { + +struct Manifest; + +bool is_cache_key(NSString* key); +NSString* cache_key_model_prefix(NSString* key); + +// Model identity excludes device and specialization so clearing finds all +// variants. +runtime::Result cache_model_key( + const Manifest& manifest, + NSString* bundle); +runtime::Result cache_variant_key( + const Manifest& selected, + NSString* platform, + NSString* architecture, + NSString* specialization); + +constexpr size_t kBookmarkLockStripeCount = 256; +runtime::Result bookmark_lock_name(NSString* key); + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_cache_key.mm b/backends/apple/coreai/runtime/coreai_cache_key.mm new file mode 100644 index 00000000000..e5b66d37102 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache_key.mm @@ -0,0 +1,154 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +#import "coreai_cache_key.h" +#include +#include +#import "coreai_assets.h" + +namespace executorch::backends::coreai { +namespace { +using runtime::Error; +using runtime::Result; + +constexpr NSUInteger kHashLength = 2 * CC_SHA256_DIGEST_LENGTH; +constexpr NSUInteger kFixedLength = 2 * (kHashLength + 1); +constexpr NSUInteger kMaxLength = NAME_MAX - sizeof(".bookmark") + 1; + +bool hex_lower(unichar c) { + return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f'); +} + +bool text(id value) { + return [value isKindOfClass:NSString.class] && [value length] > 0 && + [value length] <= 4096 && + [value rangeOfString:@"\0"].location == NSNotFound; +} + +bool valid_device(NSString* device) { + if (![device isKindOfClass:NSString.class] || + device.length > kMaxLength - kFixedLength) { + return false; + } + const NSUInteger separator = [device rangeOfString:@"_"].location; + if (separator == NSNotFound || separator == 0 || + separator + 1 == device.length) + return false; + for (NSUInteger i = 0; i < device.length; ++i) { + const unichar c = [device characterAtIndex:i]; + if (i == separator || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || + (c >= '0' && c <= '9')) + continue; + if (i < separator || (c != '_' && c != '-')) + return false; + } + return true; +} + +Result hash_data(NSData* data) { + ET_CHECK_OR_RETURN_ERROR( + data != nil && data.length <= UINT_MAX, + InvalidExternalData, + "Cannot encode Core AI cache key"); + unsigned char bytes[CC_SHA256_DIGEST_LENGTH]; + CC_SHA256(data.bytes, static_cast(data.length), bytes); + NSMutableString* result = [NSMutableString stringWithCapacity:kHashLength]; + for (unsigned char byte : bytes) + [result appendFormat:@"%02x", byte]; + return [result copy]; +} +} // namespace + +bool is_cache_key(NSString* key) { + if (![key isKindOfClass:NSString.class] || key.length <= kFixedLength || + key.length > kMaxLength) + return false; + // Architectures can contain hyphens; locate the hashes from both ends. + const NSUInteger options_start = key.length - kHashLength; + if ([key characterAtIndex:kHashLength] != '-' || + [key characterAtIndex:options_start - 1] != '-') + return false; + for (NSUInteger i = 0; i < kHashLength; ++i) { + if (!hex_lower([key characterAtIndex:i]) || + !hex_lower([key characterAtIndex:options_start + i])) + return false; + } + return valid_device( + [key substringWithRange:NSMakeRange( + kHashLength + 1, key.length - kFixedLength)]); +} + +NSString* cache_key_model_prefix(NSString* key) { + return is_cache_key(key) ? [key substringToIndex:kHashLength] : nil; +} + +Result cache_model_key(const Manifest& manifest, NSString* bundle) { + ET_CHECK_OR_RETURN_ERROR( + text(bundle) && ![bundle containsString:@"/"] && + ![bundle containsString:@"\\"] && ![bundle isEqual:@"."] && + ![bundle isEqual:@".."], + DelegateInvalidCompatibility, + "Invalid Core AI bundle name"); + NSString* digest = manifest.bundle_digests[bundle]; + ET_CHECK_OR_RETURN_ERROR( + text(digest) && digest.length == kHashLength, + DelegateInvalidCompatibility, + "Missing Core AI bundle digest"); + for (NSUInteger i = 0; i < digest.length; ++i) { + ET_CHECK_OR_RETURN_ERROR( + hex_lower([digest characterAtIndex:i]), + DelegateInvalidCompatibility, + "Invalid Core AI bundle digest"); + } + NSArray* fields = + @[ @"1", digest, manifest.aot_compiled ? @"aot" : @"source", bundle ]; + return hash_data([NSJSONSerialization dataWithJSONObject:fields + options:0 + error:nil]); +} + +Result cache_variant_key( + const Manifest& selected, + NSString* platform, + NSString* architecture, + NSString* specialization) { + auto model = cache_model_key(selected, selected.path.lastPathComponent); + if (!model.ok()) + return model.error(); + ET_CHECK_OR_RETURN_ERROR( + text(platform) && text(architecture) && ![platform containsString:@"_"], + DelegateInvalidCompatibility, + "Invalid Core AI cache context"); + NSString* device = + [NSString stringWithFormat:@"%@_%@", platform, architecture]; + ET_CHECK_OR_RETURN_ERROR( + valid_device(device), + DelegateInvalidCompatibility, + "Invalid Core AI cache device"); + ET_CHECK_OR_RETURN_ERROR( + text(specialization), + DelegateInvalidCompatibility, + "Missing Core AI cache specialization"); + auto options = + hash_data([specialization dataUsingEncoding:NSUTF8StringEncoding]); + if (!options.ok()) + return options.error(); + return [NSString + stringWithFormat:@"%@-%@-%@", model.get(), device, options.get()]; +} + +Result bookmark_lock_name(NSString* key) { + ET_CHECK_OR_RETURN_ERROR( + is_cache_key(key), InvalidArgument, "Invalid Core AI bookmark key"); + auto hash = hash_data([key dataUsingEncoding:NSUTF8StringEncoding]); + if (!hash.ok()) + return hash.error(); + return [NSString + stringWithFormat:@"stripe-v1-%@.lock", [hash.get() substringToIndex:2]]; +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.h b/backends/apple/coreai/runtime/coreai_load_coordinator.h index cfd44efdf00..e5125ece2e8 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.h +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.h @@ -8,7 +8,7 @@ #pragma once #import "ETCoreAIBridge.h" -#import "coreai_bookmarks.h" +#import "coreai_assets.h" namespace executorch::backends::coreai { diff --git a/backends/apple/coreai/runtime/coreai_load_coordinator.mm b/backends/apple/coreai/runtime/coreai_load_coordinator.mm index 2951eb15dfa..55992b1a03d 100644 --- a/backends/apple/coreai/runtime/coreai_load_coordinator.mm +++ b/backends/apple/coreai/runtime/coreai_load_coordinator.mm @@ -6,6 +6,8 @@ * LICENSE file in the root directory of this source tree. */ #import "coreai_load_coordinator.h" +#import "coreai_bookmarks.h" +#import "coreai_cache_key.h" #import "coreai_storage.h" #include @@ -149,7 +151,9 @@ Error evict_one(id loader, NSData* bookmark) { id loader) { ET_CHECK_OR_RETURN_ERROR( loader != nil, InvalidArgument, "Missing Core AI model loader"); - auto key = bookmark_key(selected, platform, architecture); + NSString* specialization = loader.specializationIdentifier; + auto key = + cache_variant_key(selected, platform, architecture, specialization); if (!key.ok()) return key.error(); NSString* cache_key = key.get(); @@ -157,8 +161,6 @@ Error evict_one(id loader, NSData* bookmark) { auto locked = lock_bookmark(root, cache_key); if (!locked.ok()) return locked.error(); - ET_CHECK_OK_OR_RETURN_ERROR(ensure_bookmark_identity( - *locked.get(), selected, platform, architecture)); auto bookmark = read_bookmark(*locked.get()); if (!bookmark.ok()) return bookmark.error(); @@ -182,7 +184,7 @@ Error evict_one(id loader, NSData* bookmark) { ET_LOG( Info, "Core AI cache miss: no bookmark, key=%s", key.get().UTF8String); } - auto staging = prepare_bookmark_staging(*locked.get()); + auto staging = prepare_bookmark_staging_root(*locked.get()); if (!staging.ok()) return staging.error(); auto source = prepare_source_bundle( @@ -265,7 +267,6 @@ Error clear_bookmark_assets( if (!bookmark.ok() || bytes != nil) { ET_CHECK_OK_OR_RETURN_ERROR(remove_bookmark(*locked.get())); } - ET_CHECK_OK_OR_RETURN_ERROR(remove_bookmark_identity(*locked.get())); ET_LOG(Info, "Core AI local cache entry cleared: key=%s", key.UTF8String); return Error::Ok; } diff --git a/backends/apple/coreai/runtime/coreai_storage.mm b/backends/apple/coreai/runtime/coreai_storage.mm index 2c246e0623f..15157b3a2f9 100644 --- a/backends/apple/coreai/runtime/coreai_storage.mm +++ b/backends/apple/coreai/runtime/coreai_storage.mm @@ -7,6 +7,7 @@ */ #import "coreai_storage.h" +#import "coreai_cache_key.h" #include "coreai_file.h" #include @@ -262,7 +263,9 @@ FileDescriptor child(retry_eintr([&] { Result open_storage_directory(int parent_fd, NSString* name, bool create) { ET_CHECK_OR_RETURN_ERROR( - storage_component(name), InvalidArgument, "Invalid storage directory name"); + storage_component(name), + InvalidArgument, + "Invalid storage directory name"); const auto open_directory = [&] { return retry_eintr([&] { return openat( @@ -278,8 +281,9 @@ FileDescriptor child(retry_eintr([&] { storage_fault(StorageOperation::BeforeCreateDirectory) == 0, AccessFailed, "Cannot begin Core AI directory creation"); - const int made = retry_eintr( - [&] { return mkdirat(parent_fd, name.fileSystemRepresentation, 0777); }); + const int made = retry_eintr([&] { + return mkdirat(parent_fd, name.fileSystemRepresentation, 0777); + }); ET_CHECK_OR_RETURN_ERROR( made == 0 || errno == EEXIST, AccessFailed, @@ -374,16 +378,6 @@ Error sync_storage_directory(int fd, bool full_sync) { return Error::Ok; } -namespace { -bool storage_key(NSString* key) { - return [key isKindOfClass:NSString.class] && key.length == 64 && - [key rangeOfCharacterFromSet:[[NSCharacterSet - characterSetWithCharactersInString: - @"0123456789abcdef"] invertedSet]] - .location == NSNotFound; -} -} // namespace - Error remove_storage_entry(int parent_fd, NSString* name, bool remove) { ET_CHECK_OR_RETURN_ERROR( storage_component(name), InvalidArgument, "Invalid Core AI entry name"); @@ -442,7 +436,7 @@ Error remove_storage_staging( bool remove, StorageRemovalMode mode) { ET_CHECK_OR_RETURN_ERROR( - storage_key(key), InvalidArgument, "Invalid Core AI staging key"); + is_cache_key(key), InvalidArgument, "Invalid Core AI staging key"); FileDescriptor staging(retry_eintr([&] { return openat( root_fd, "staging", O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC); @@ -462,7 +456,7 @@ Error remove_storage_source( bool remove, StorageRemovalMode mode) { ET_CHECK_OR_RETURN_ERROR( - storage_key(key), InvalidArgument, "Invalid Core AI staging key"); + is_cache_key(key), InvalidArgument, "Invalid Core AI staging key"); struct stat info; const int inspected = retry_eintr([&] { return fstatat( diff --git a/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift b/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift index d9ec9c42d37..da3dce548a5 100644 --- a/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift +++ b/backends/apple/coreai/runtime/test/SwiftBridgeTests.swift @@ -104,6 +104,71 @@ final class SwiftAcquisitionTests: XCTestCase { } } +@available(macOS 27.0, iOS 27.0, *) +final class SwiftSpecializationTests: XCTestCase { + func testCachePolicyAllowsOnlyStoragePressure() { + XCTAssertEqual(ETCoreAISwiftModelLoader.cachePolicy.purgeConditions, [.storagePressure]) + } + func testDefaultPreservesIdentifier() throws { + let snapshot = try CoreAISpecialization(options: .default) + XCTAssertEqual(snapshot.options, .default) + XCTAssertEqual(snapshot.identifier, "default") + XCTAssertEqual(ETCoreAISwiftModelLoader().specializationIdentifier, snapshot.identifier) + } + + func testAllOptionFieldsParticipateInIdentity() throws { + let choices: [SpecializationOptions] = [ + .default, .cpuOnly, + SpecializationOptions(preferredComputeUnitKind: .cpu), + SpecializationOptions(preferredComputeUnitKind: .gpu), + SpecializationOptions(preferredComputeUnitKind: .neuralEngine), + ] + let options = + choices + + choices.map { original in + var changed = original + changed.expectFrequentReshapes.toggle() + return changed + } + let snapshots = try options.map { try CoreAISpecialization(options: $0) } + for first in snapshots { + XCTAssertEqual(first.identifier, try CoreAISpecialization(options: first.options).identifier) + XCTAssertEqual( + ETCoreAISwiftModelLoader(specialization: first).specializationIdentifier, first.identifier) + for second in snapshots { + XCTAssertEqual(first.options == second.options, first.identifier == second.identifier) + } + } + } + + func testSnapshotDoesNotFollowSubsequentOptionMutation() throws { + var options = SpecializationOptions.cpuOnly + let original = options + let snapshot = try CoreAISpecialization(options: options) + let loader = ETCoreAISwiftModelLoader(specialization: snapshot) + options.expectFrequentReshapes.toggle() + XCTAssertEqual(snapshot.options, original) + XCTAssertEqual(loader.specializationIdentifier, snapshot.identifier) + XCTAssertNotEqual( + loader.specializationIdentifier, try CoreAISpecialization(options: options).identifier) + } + + func testDescriptorIsVersionedAndUsesStableComputeUnitNames() throws { + var options = SpecializationOptions(preferredComputeUnitKind: .gpu) + options.expectFrequentReshapes = !SpecializationOptions.default.expectFrequentReshapes + let snapshot = try CoreAISpecialization(options: options) + let json = try XCTUnwrap( + JSONSerialization.jsonObject(with: Data(snapshot.identifier.utf8)) as? [String: Any]) + XCTAssertEqual(json["version"] as? Int, 1) + XCTAssertEqual(json["preferredComputeUnit"] as? String, "gpu") + XCTAssertEqual(json["expectFrequentReshapes"] as? Bool, options.expectFrequentReshapes) + let allowed = try XCTUnwrap(json["allowedComputeUnits"] as? [String]) + XCTAssertEqual(allowed, allowed.sorted()) + XCTAssertTrue(Set(allowed).isSubset(of: ["cpu", "gpu", "neuralEngine"])) + XCTAssertEqual(allowed.count, options.allowedComputeUnitKinds.count) + } +} + @available(macOS 27.0, iOS 27.0, *) final class SwiftSessionTests: XCTestCase { private struct Manifest: Decodable { @@ -134,7 +199,7 @@ final class SwiftSessionTests: XCTestCase { } } - private let loader = ETCoreAISwiftModelLoader() + private var loader = ETCoreAISwiftModelLoader() private var manifest: Manifest! private var fixtureDirectory: URL! private var model: (any ETCoreAIPreparedModel)? @@ -156,6 +221,15 @@ final class SwiftSessionTests: XCTestCase { XCTAssertEqual(manifest.outputNames.count, 2) } + func testSpecializationUsesTheSnapshotOptions() throws { + let snapshot = try CoreAISpecialization(options: .cpuOnly) + loader = ETCoreAISwiftModelLoader(specialization: snapshot) + try prepareModel() + XCTAssertNotNil( + try AIModelCache.default.model( + for: fixtureDirectory.appendingPathComponent("model.aimodel"), options: snapshot.options)) + } + private func prepareModel() throws { guard model == nil else { return } let completed = expectation(description: "specialized session fixture") @@ -386,7 +460,9 @@ final class SwiftBridgeTests: XCTestCase { ([2, 1, 3], [3, 19, 1]), ([], []), ([2, 0, 3], [3, 3, 1]), ] { let count = shape.reduce(1, *) - let expected = (0..<(count * info.byteWidth)).map { UInt8(truncatingIfNeeded: $0 * 73 + 255) } + let expected = (0..<(count * info.byteWidth)).map { + UInt8(truncatingIfNeeded: $0 * 73 + 255) + } var array = NDArray(shape: shape, scalarType: info.sdkType, strides: strides) do { var view = array.mutableRawView() diff --git a/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm index f5bf032b46f..ff3218c3d93 100644 --- a/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_acquisition_test.mm @@ -99,87 +99,71 @@ ino_t lock_inode() { } }; -TEST_F( - CoreAIAcquisitionTest, - AttributionPublicationFailurePrecedesSdkAndSourceAccess) { - BookmarkLoaderScope fake; - { - StorageFaultScope fault(StorageOperation::Rename); - EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); - EXPECT_EQ(fault.hits, 1); +TEST_F(CoreAIAcquisitionTest, LoaderPoliciesUseIndependentWarmBookmarks) { + BookmarkLoaderScope standard; + BookmarkLoaderScope alternate( + @"{\"allowedComputeUnits\":[\"cpu\"],\"expectFrequentReshapes\":false,\"version\":1}"); + ASSERT_EQ(acquire(standard.loader, &data), Error::Ok); + NSData* original = saved(); + ASSERT_EQ(acquire(alternate.loader, &data), Error::Ok); + NSData* other = alternate.loader.lastAcquired.bookmark; + ASSERT_FALSE([original isEqual:other]); + EXPECT_EQ(standard.loader->restores.load(), 0); + EXPECT_EQ(alternate.loader->restores.load(), 0); + ASSERT_EQ(acquire(standard.loader, nullptr), Error::Ok); + EXPECT_TRUE([standard.loader.lastAcquired.bookmark isEqual:original]); + ASSERT_EQ(acquire(alternate.loader, nullptr), Error::Ok); + EXPECT_TRUE([alternate.loader.lastAcquired.bookmark isEqual:other]); + EXPECT_EQ(standard.loader->specializations.load(), 1); + EXPECT_EQ(alternate.loader->specializations.load(), 1); + EXPECT_EQ(standard.loader->restores.load(), 1); + EXPECT_EQ(alternate.loader->restores.load(), 1); + EXPECT_EQ(standard.loader->specializationQueries.load(), 2); + EXPECT_EQ(alternate.loader->specializationQueries.load(), 2); + EXPECT_TRUE([saved() isEqual:original]); + EXPECT_EQ(data.attempts.load(), 2); +} + +TEST_F(CoreAIAcquisitionTest, InvalidPolicyFailsBeforeFilesystemOrSdkActivity) { + for (NSString* identifier in @[ + @"", + [@"x" stringByPaddingToLength:4097 withString:@"x" startingAtIndex:0] + ]) { + BookmarkLoaderScope fake(identifier); + NSDictionary* before = nil; + ASSERT_TRUE(asset_tree_snapshot(root.url, before)); + EXPECT_EQ(acquire(fake.loader, &data), Error::DelegateInvalidCompatibility); + EXPECT_TRUE(snapshot_matches(root.url, before)); + EXPECT_EQ(fake.loader->specializationQueries.load(), 1); + EXPECT_EQ(fake.loader->specializations.load(), 0); + EXPECT_EQ(fake.loader->restores.load(), 0); + EXPECT_EQ(data.attempts.load(), 0); } - EXPECT_EQ(fake.loader->specializations.load(), 0); - EXPECT_EQ(fake.loader->restores.load(), 0); - EXPECT_EQ(data.attempts.load(), 0); - EXPECT_EQ(saved(), nil); - EXPECT_FALSE(staging_key_exists()); - EXPECT_EQ(acquire(fake.loader, &data), Error::Ok); - auto locked = lock_bookmark(root.url.path, key, false); - ASSERT_TRUE(locked.ok()); - ASSERT_NE(locked.get(), nullptr); - auto identity = read_bookmark_identity(*locked.get()); - ASSERT_TRUE(identity.ok()); - EXPECT_NE(identity.get(), nil); } -TEST_F( - CoreAIAcquisitionTest, - AttributionFlushFailureMustBeRetriedBeforeSdkAccess) { +TEST_F(CoreAIAcquisitionTest, ClearDoesNotQueryCurrentSpecialization) { BookmarkLoaderScope fake; - { - StorageFaultScope fault(StorageOperation::FullSync); - fault.match_index = 2; - EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); - EXPECT_EQ(fault.hits, 1); - } - { - auto locked = lock_bookmark(root.url.path, key, false); - ASSERT_TRUE(locked.ok()); - ASSERT_NE(locked.get(), nullptr); - auto identity = read_bookmark_identity(*locked.get()); - ASSERT_TRUE(identity.ok()); - EXPECT_NE(identity.get(), nil); - } - { - StorageFaultScope fault(StorageOperation::FullSync); - EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); - EXPECT_EQ(fault.hits, 1); - } - EXPECT_EQ(fake.loader->specializations.load(), 0); - EXPECT_EQ(fake.loader->restores.load(), 0); - EXPECT_EQ(data.attempts.load(), 0); + ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); + const int queries = fake.loader->specializationQueries.load(); + ASSERT_EQ(clear_bookmark_assets(root.url.path, key, fake.loader), Error::Ok); + EXPECT_EQ(fake.loader->specializationQueries.load(), queries); + EXPECT_EQ(fake.loader->evictions.load(), 1); EXPECT_EQ(saved(), nil); - EXPECT_FALSE(staging_key_exists()); - EXPECT_EQ(acquire(fake.loader, &data), Error::Ok); } -TEST_F(CoreAIAcquisitionTest, WarmLegacyEntryGetsAttributionWithoutRestaging) { +TEST_F(CoreAIAcquisitionTest, WarmHitDoesNotWriteOrFlushCacheFiles) { BookmarkLoaderScope fake; ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); NSData* original = saved(); - { - auto locked = lock_bookmark(root.url.path, key, false); - ASSERT_TRUE(locked.ok()); - ASSERT_NE(locked.get(), nullptr); - ASSERT_EQ(remove_bookmark_identity(*locked.get()), Error::Ok); - } - { - StorageFaultScope fault(StorageOperation::Write); - EXPECT_EQ(acquire(fake.loader, nullptr), Error::AccessFailed); - EXPECT_EQ(fault.hits, 1); + for (auto operation : {StorageOperation::Write, StorageOperation::FullSync}) { + StorageFaultScope fault(operation); + EXPECT_EQ(acquire(fake.loader, nullptr), Error::Ok); + EXPECT_EQ(fault.hits, 0); } EXPECT_TRUE([saved() isEqual:original]); - EXPECT_EQ(fake.loader->restores.load(), 0); - ASSERT_EQ(acquire(fake.loader, nullptr), Error::Ok); - EXPECT_EQ(fake.loader->restores.load(), 1); + EXPECT_EQ(fake.loader->restores.load(), 2); EXPECT_EQ(fake.loader->specializations.load(), 1); EXPECT_EQ(data.attempts.load(), 1); - auto locked = lock_bookmark(root.url.path, key, false); - ASSERT_TRUE(locked.ok()); - ASSERT_NE(locked.get(), nullptr); - auto identity = read_bookmark_identity(*locked.get()); - ASSERT_TRUE(identity.ok()); - EXPECT_NE(identity.get(), nil); } TEST_F(CoreAIAcquisitionTest, ColdSpecializesStagedSourceAndPersistsBookmark) { @@ -451,16 +435,13 @@ ino_t lock_inode() { { auto locked = lock_bookmark(root.url.path, key); ASSERT_TRUE(locked.ok()); - ASSERT_EQ( - ensure_bookmark_identity(*locked.get(), manifest, @"macOS", @"arch_b"), - Error::Ok); } __block std::vector ops; Error error = Error::Internal; { - // Attribution reuse flushes once, then source and bookmark twice each. + // Source and bookmark publication each flush twice. StorageFaultScope fault(StorageOperation::FullSync); - fault.match_index = 5; + fault.match_index = 4; fault.observe = ^(StorageOperation operation) { ops.push_back(operation); }; @@ -681,12 +662,9 @@ ino_t lock_inode() { { auto locked = lock_bookmark(root.url.path, key); ASSERT_TRUE(locked.ok()); - ASSERT_EQ( - ensure_bookmark_identity(*locked.get(), manifest, @"macOS", @"arch_b"), - Error::Ok); } { - // With attribution present, the bookmark follows source publication. + // Bookmark publication follows source publication. StorageFaultScope fault(StorageOperation::Rename); fault.match_index = 2; EXPECT_EQ(acquire(fake.loader, &data), Error::AccessFailed); @@ -894,12 +872,9 @@ void expect_order( { auto locked = lock_bookmark(root.url.path, key); ASSERT_TRUE(locked.ok()); - ASSERT_EQ( - ensure_bookmark_identity(*locked.get(), manifest, @"macOS", @"arch_b"), - Error::Ok); } StorageFaultScope fault(StorageOperation::FullSync); - fault.match_index = 5; + fault.match_index = 4; const auto logs = capture([&] { return acquire(fake.loader, &data); }, Error::AccessFailed); EXPECT_EQ(fault.hits, 1); @@ -961,11 +936,11 @@ void expect_order( TEST_F( CoreAILifecycleLogTest, - IdentityRemovalFailureDoesNotLogLocalCompletion) { + BookmarkRemovalFailureDoesNotLogLocalCompletion) { BookmarkLoaderScope fake; ASSERT_EQ(acquire(fake.loader, &data), Error::Ok); StorageFaultScope fault(StorageOperation::Remove); - fault.match_index = 2; + fault.match_index = 1; const auto logs = capture( [&] { return clear_bookmark_assets(root.url.path, key, fake.loader); }, Error::AccessFailed); @@ -977,9 +952,9 @@ void expect_order( auto locked = lock_bookmark(root.url.path, key, false); ASSERT_TRUE(locked.ok()); ASSERT_NE(locked.get(), nullptr); - auto identity = read_bookmark_identity(*locked.get()); - ASSERT_TRUE(identity.ok()); - EXPECT_NE(identity.get(), nil); + auto bookmark = read_bookmark(*locked.get()); + ASSERT_TRUE(bookmark.ok()); + EXPECT_NE(bookmark.get(), nil); } TEST_F(CoreAILifecycleLogTest, AbsentEntryClearRemainsQuiet) { diff --git a/backends/apple/coreai/runtime/test/coreai_bookmark_fixture.h b/backends/apple/coreai/runtime/test/coreai_bookmark_fixture.h index 252ed85b08d..2c4ea12a1c3 100644 --- a/backends/apple/coreai/runtime/test/coreai_bookmark_fixture.h +++ b/backends/apple/coreai/runtime/test/coreai_bookmark_fixture.h @@ -12,12 +12,17 @@ #include #include #include +#include "coreai_assets.h" #include "coreai_bookmarks.h" +#include "coreai_cache_key.h" namespace executorch::backends::coreai::testing { runtime::Result bookmark_manifest(bool aot = false); runtime::Result test_bookmark_key(const Manifest& manifest); +runtime::Result test_bookmark_key_with_specialization( + const Manifest& manifest, + NSString* specialization); NSURL* test_bookmark_url(NSURL* root, NSString* key); NSURL* test_bookmark_lock_url(NSURL* root, NSString* key); runtime::Result saved_bookmark(NSURL* root, NSString* key); diff --git a/backends/apple/coreai/runtime/test/coreai_bookmark_fixture.mm b/backends/apple/coreai/runtime/test/coreai_bookmark_fixture.mm index 694a47670e9..16c2fa6f920 100644 --- a/backends/apple/coreai/runtime/test/coreai_bookmark_fixture.mm +++ b/backends/apple/coreai/runtime/test/coreai_bookmark_fixture.mm @@ -75,7 +75,13 @@ bool send_signal_byte(int fd, char byte) { } Result test_bookmark_key(const Manifest& manifest) { - return bookmark_key(manifest, @"macOS", @"arch_b"); + return cache_variant_key(manifest, @"macOS", @"arch_b", @"default"); +} + +Result test_bookmark_key_with_specialization( + const Manifest& manifest, + NSString* specialization) { + return cache_variant_key(manifest, @"macOS", @"arch_b", specialization); } NSURL* test_bookmark_url(NSURL* root, NSString* key) { @@ -144,14 +150,28 @@ int bookmark_child_mode(int argc, char** argv) { auto key = test_bookmark_key(manifest.get()); if (!key.ok()) return 7; - NSMutableString* selected_key = [key.get() mutableCopy]; + NSString* selected_key = key.get(); if (other || collision) { - const NSUInteger index = other ? 0 : 63; - [selected_key - replaceCharactersInRange:NSMakeRange(index, 1) - withString:[selected_key characterAtIndex:index] == '0' - ? @"1" - : @"0"]; + auto original = bookmark_lock_name(selected_key); + if (!original.ok()) + return 7; + selected_key = nil; + for (unsigned i = 0; i < 65536; ++i) { + NSString* candidate = [key.get() + stringByReplacingCharactersInRange:NSMakeRange(0, 64) + withString:[NSString + stringWithFormat:@"%064x", i]]; + auto stripe = bookmark_lock_name(candidate); + if (!stripe.ok()) + return 7; + if ([stripe.get() isEqual:original.get()] == collision && + ![candidate isEqual:key.get()]) { + selected_key = candidate; + break; + } + } + if (selected_key == nil) + return 7; } StorageFaultScope scope(StorageOperation::BeforeBookmarkLock); scope.armed = false; diff --git a/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm b/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm index 2cebbc77d61..45c719cba5b 100644 --- a/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm @@ -6,12 +6,14 @@ * LICENSE file in the root directory of this source tree. */ +#include #include #include #include #include #include #include +#include #include #include "coreai_bookmark_fixture.h" #include "coreai_fault_scope.h" @@ -27,14 +29,66 @@ namespace { constexpr size_t kMiB = 1024 * 1024; constexpr auto kWaitTimeout = std::chrono::seconds(5); +constexpr NSUInteger kModelHashLength = 64; + +NSString* sha256_hex(NSString* input) { + NSData* data = [input dataUsingEncoding:NSUTF8StringEncoding]; + unsigned char bytes[CC_SHA256_DIGEST_LENGTH]; + CC_SHA256(data.bytes, static_cast(data.length), bytes); + NSMutableString* hex = + [NSMutableString stringWithCapacity:2 * CC_SHA256_DIGEST_LENGTH]; + for (unsigned char byte : bytes) + [hex appendFormat:@"%02x", byte]; + return [hex copy]; +} + +// Builds a valid composite key without a manifest. Model/options must be 64 +// lowercase hex chars; device must be platform_arch. +NSString* synthetic_key(NSString* model, NSString* device, NSString* options) { + return [NSString stringWithFormat:@"%@-%@-%@", model, device, options]; +} + +NSString* padded_model(unsigned value) { + return [NSString stringWithFormat:@"%064x", value]; +} + +NSString* fixed_options(char digit) { + return [@"" stringByPaddingToLength:64 + withString:[NSString stringWithFormat:@"%c", digit] + startingAtIndex:0]; +} + +// Derives a valid composite variant by rewriting the model prefix until its +// lock stripe matches (or differs from) the original. Mirrors the child-mode +// search so thread/process contention tests stay deterministic under the +// hash-derived stripe scheme. +NSString* variant_with_stripe(NSString* key, bool same_stripe) { + auto original = bookmark_lock_name(key); + if (!original.ok()) + return nil; + for (unsigned i = 0; i < 65536; ++i) { + NSString* candidate = + [key stringByReplacingCharactersInRange:NSMakeRange(0, kModelHashLength) + withString:padded_model(i)]; + if ([candidate isEqual:key]) + continue; + if (!is_cache_key(candidate)) + continue; + auto stripe = bookmark_lock_name(candidate); + if (!stripe.ok()) + continue; + if ([stripe.get() isEqual:original.get()] == same_stripe) + return candidate; + } + return nil; +} -NSString* changed_key(NSString* key, bool same_stripe) { - NSMutableString* changed = [key mutableCopy]; - const NSUInteger index = same_stripe ? 63 : 0; - [changed replaceCharactersInRange:NSMakeRange(index, 1) - withString:[key characterAtIndex:index] == '0' ? @"1" - : @"0"]; - return changed; +NSString* colliding_key(NSString* key) { + return variant_with_stripe(key, true); +} + +NSString* different_key(NSString* key) { + return variant_with_stripe(key, false); } struct BookmarkWaiter { @@ -145,13 +199,12 @@ FileDescriptor fd(open( return ::testing::AssertionSuccess(); } -NSURL* identity_url(NSURL* root, NSString* key) { - return [[root URLByAppendingPathComponent:@"bookmarks"] - URLByAppendingPathComponent:[key stringByAppendingString:@".metadata"]]; -} - TEST_F(CoreAIBookmarkIdentityTest, IgnoresFunctionBindingsAndDeliveryLocation) { - EXPECT_EQ(key.length, 64u); + ASSERT_TRUE(is_cache_key(key)); + // Composite key: 64-char model hash, readable device, 64-char options hash. + EXPECT_TRUE([key containsString:@"-macOS_arch_b-"]); + EXPECT_TRUE([[key substringFromIndex:key.length - 64] + isEqual:sha256_hex(@"default")]); manifest.function = @"another_function"; manifest.inputs = @[ @"different_input" ]; manifest.outputs = @[ @"different_output" ]; @@ -164,23 +217,64 @@ FileDescriptor fd(open( TEST_F( CoreAIBookmarkIdentityTest, KeySeparatesPlatformArchitectureWeightsAndAotContext) { - auto platform = bookmark_key(manifest, @"iOS", @"arch_b"); - auto architecture = bookmark_key(manifest, @"macOS", @"arch_a"); + auto platform = cache_variant_key(manifest, @"iOS", @"arch_b", @"default"); + auto architecture = + cache_variant_key(manifest, @"macOS", @"arch_a", @"default"); + auto specialization = + test_bookmark_key_with_specialization(manifest, @"other"); Manifest reweighted = manifest; reweighted.bundle_digests = @{@"model.aimodel" : fixture_identity('2')}; auto weights = test_bookmark_key(reweighted); auto aot_manifest = bookmark_manifest(true); ASSERT_TRUE(aot_manifest.ok()); auto aot = test_bookmark_key(aot_manifest.get()); - for (auto* changed : {&platform, &architecture, &weights, &aot}) { + for (auto* changed : + {&platform, &architecture, &specialization, &weights, &aot}) { ASSERT_TRUE(changed->ok()); + EXPECT_TRUE(is_cache_key(changed->get())); EXPECT_FALSE([changed->get() isEqual:key]); } } +TEST_F( + CoreAIBookmarkIdentityTest, + FullKeySharesModelPrefixAcrossDeviceAndOptions) { + auto platform = cache_variant_key(manifest, @"iOS", @"arch_b", @"default"); + auto architecture = + cache_variant_key(manifest, @"macOS", @"arch_a", @"default"); + auto specialization = + test_bookmark_key_with_specialization(manifest, @"other"); + ASSERT_TRUE(platform.ok()); + ASSERT_TRUE(architecture.ok()); + ASSERT_TRUE(specialization.ok()); + NSString* prefix = cache_key_model_prefix(key); + ASSERT_NE(prefix, nil); + EXPECT_EQ(prefix.length, kModelHashLength); + // Same model, so the 64-char prefix is stable while device/options vary. + EXPECT_TRUE([cache_key_model_prefix(platform.get()) isEqual:prefix]); + EXPECT_TRUE([cache_key_model_prefix(architecture.get()) isEqual:prefix]); + EXPECT_TRUE([cache_key_model_prefix(specialization.get()) isEqual:prefix]); + EXPECT_TRUE([platform.get() containsString:@"-iOS_arch_b-"]); + EXPECT_TRUE([architecture.get() containsString:@"-macOS_arch_a-"]); + EXPECT_TRUE([[platform.get() substringFromIndex:platform.get().length - 64] + isEqual:[key substringFromIndex:key.length - 64]]); + EXPECT_TRUE([[specialization.get() + substringFromIndex:specialization.get().length - 64] + isEqual:sha256_hex(@"other")]); +} + TEST_F(CoreAIBookmarkIdentityTest, RejectsMissingArchitecture) { - EXPECT_FALSE(bookmark_key(manifest, @"macOS", nil).ok()); - EXPECT_FALSE(bookmark_key(manifest, @"macOS", @"").ok()); + EXPECT_FALSE(cache_variant_key(manifest, @"macOS", nil, @"default").ok()); + EXPECT_FALSE(cache_variant_key(manifest, @"macOS", @"", @"default").ok()); +} + +TEST_F(CoreAIBookmarkIdentityTest, RejectsMissingPlatformAndSpecialization) { + EXPECT_FALSE(cache_variant_key(manifest, nil, @"arch_b", @"default").ok()); + EXPECT_FALSE(cache_variant_key(manifest, @"", @"arch_b", @"default").ok()); + EXPECT_FALSE(cache_variant_key(manifest, @"macOS", @"arch_b", nil).ok()); + EXPECT_FALSE(cache_variant_key(manifest, @"macOS", @"arch_b", @"").ok()); + // All three inputs are required for the composite key. + EXPECT_FALSE(cache_variant_key(manifest, nil, nil, nil).ok()); } TEST_F(CoreAIBookmarkIdentityTest, RejectsUnselectedBundle) { @@ -277,13 +371,13 @@ FileDescriptor file(open( TEST_F(CoreAIBookmarkRawTest, PreparesSharedStagingParentOnDemand) { NSString* expected = [root.url URLByAppendingPathComponent:@"staging"].path; EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:expected]); - auto staging = prepare_bookmark_staging(*locked); + auto staging = prepare_bookmark_staging_root(*locked); ASSERT_TRUE(staging.ok()); EXPECT_TRUE([staging.get() isEqual:expected]); struct stat info {}; ASSERT_EQ(lstat(staging.get().fileSystemRepresentation, &info), 0); EXPECT_TRUE(S_ISDIR(info.st_mode)); - auto repeated = prepare_bookmark_staging(*locked); + auto repeated = prepare_bookmark_staging_root(*locked); ASSERT_TRUE(repeated.ok()); EXPECT_TRUE([repeated.get() isEqual:expected]); EXPECT_FALSE([NSFileManager.defaultManager @@ -295,7 +389,7 @@ FileDescriptor file(open( ASSERT_EQ(mkdir(reference.fileSystemRepresentation, 0777), 0); struct stat expected; ASSERT_EQ(stat(reference.fileSystemRepresentation, &expected), 0); - ASSERT_TRUE(prepare_bookmark_staging(*locked).ok()); + ASSERT_TRUE(prepare_bookmark_staging_root(*locked).ok()); for (NSString* name in @[ @"locks", @"bookmarks", @"staging" ]) { NSURL* directory = [root.url URLByAppendingPathComponent:name]; struct stat actual; @@ -305,7 +399,7 @@ FileDescriptor file(open( } TEST_F(CoreAIBookmarkRawTest, ReusesDirectoriesWithoutCreation) { - ASSERT_TRUE(prepare_bookmark_staging(*locked).ok()); + ASSERT_TRUE(prepare_bookmark_staging_root(*locked).ok()); ASSERT_EQ( write_bookmark( *locked, [@"bookmark" dataUsingEncoding:NSUTF8StringEncoding]), @@ -316,19 +410,20 @@ FileDescriptor file(open( auto reused = lock_bookmark(root.url.path, key, create); ASSERT_TRUE(reused.ok()); ASSERT_NE(reused.get(), nullptr); - EXPECT_TRUE(prepare_bookmark_staging(*reused.get()).ok()); + EXPECT_TRUE(prepare_bookmark_staging_root(*reused.get()).ok()); } EXPECT_EQ(fault.matches, 0); } TEST_F(CoreAIBookmarkRawTest, RecreatesMissingStagingParent) { - auto staging = prepare_bookmark_staging(*locked); + auto staging = prepare_bookmark_staging_root(*locked); ASSERT_TRUE(staging.ok()); ASSERT_EQ(rmdir(staging.get().fileSystemRepresentation), 0); StorageFaultScope fault(StorageOperation::BeforeCreateDirectory); - EXPECT_EQ(prepare_bookmark_staging(*locked).error(), Error::AccessFailed); + EXPECT_EQ( + prepare_bookmark_staging_root(*locked).error(), Error::AccessFailed); EXPECT_EQ(fault.hits, 1); - auto recreated = prepare_bookmark_staging(*locked); + auto recreated = prepare_bookmark_staging_root(*locked); ASSERT_TRUE(recreated.ok()); EXPECT_TRUE([recreated.get() isEqual:staging.get()]); } @@ -344,7 +439,7 @@ FileDescriptor file(open( EXPECT_EQ(mkdir(staging.fileSystemRepresentation, 0700), 0); } }; - EXPECT_TRUE(prepare_bookmark_staging(*locked).ok()); + EXPECT_TRUE(prepare_bookmark_staging_root(*locked).ok()); EXPECT_EQ(creations, 1); } @@ -357,7 +452,8 @@ FileDescriptor file(open( target.url.fileSystemRepresentation, staging.fileSystemRepresentation), 0); - EXPECT_EQ(prepare_bookmark_staging(*locked).error(), Error::AccessFailed); + EXPECT_EQ( + prepare_bookmark_staging_root(*locked).error(), Error::AccessFailed); NSArray* contents = [NSFileManager.defaultManager contentsOfDirectoryAtPath:target.url.path error:nil]; @@ -369,7 +465,8 @@ FileDescriptor file(open( NSURL* staging = [root.url URLByAppendingPathComponent:@"staging"]; NSData* bytes = [@"not a directory" dataUsingEncoding:NSUTF8StringEncoding]; ASSERT_TRUE([bytes writeToURL:staging atomically:NO]); - EXPECT_EQ(prepare_bookmark_staging(*locked).error(), Error::AccessFailed); + EXPECT_EQ( + prepare_bookmark_staging_root(*locked).error(), Error::AccessFailed); EXPECT_TRUE(([[NSData dataWithContentsOfURL:staging] isEqual:bytes])); } @@ -392,6 +489,9 @@ FileDescriptor file(open( ASSERT_TRUE(allowed.ok()); ASSERT_NE(allowed.get(), nullptr); EXPECT_FALSE(lock_bookmark(root.url.path, @"../escape").ok()); + // Bare 64-char digests are not composite cache keys. + EXPECT_FALSE(lock_bookmark(root.url.path, fixture_identity('a')).ok()); + EXPECT_FALSE(lock_bookmark(root.url.path, @"").ok()); } TEST_F(CoreAIBookmarkPathTest, ExcludesRootWithoutCreatingStaging) { @@ -425,7 +525,9 @@ FileDescriptor file(open( } TEST_F(CoreAIBookmarkPathTest, NoncreatingLookupIgnoresOtherEntries) { - NSString* other_key = changed_key(key, false); + NSString* other_key = different_key(key); + ASSERT_NE(other_key, nil); + ASSERT_FALSE([other_key isEqual:key]); auto other = lock_bookmark(root.url.path, other_key); ASSERT_TRUE(other.ok()); ASSERT_NE(other.get(), nullptr); @@ -504,8 +606,12 @@ FileDescriptor file(open( EXPECT_EQ( waiter.result.wait_for(std::chrono::milliseconds(100)), std::future_status::timeout); - auto independent = - start_bookmark_waiter(root.url.path, changed_key(key, false)); + NSString* independent_key = different_key(key); + ASSERT_NE(independent_key, nil); + // Different stripes hash to different lock files. + EXPECT_FALSE([test_bookmark_lock_url(root.url, independent_key) + isEqual:test_bookmark_lock_url(root.url, key)]); + auto independent = start_bookmark_waiter(root.url.path, independent_key); ASSERT_EQ( independent.result.wait_for(kWaitTimeout), std::future_status::ready); auto other = independent.result.get(); @@ -581,395 +687,106 @@ FileDescriptor file(open( EXPECT_EQ(before.st_ino, after.st_ino); } -TEST_F(CoreAIBookmarkIdentityTest, IdentityExcludesContextAndFunction) { +TEST_F(CoreAIBookmarkIdentityTest, ModelKeyExcludesContextAndFunction) { NSString* bundle = manifest.path.lastPathComponent; - auto identity = bookmark_identity(manifest, bundle); + auto identity = cache_model_key(manifest, bundle); ASSERT_TRUE(identity.ok()); - EXPECT_EQ(identity.get().length, 64u); + EXPECT_EQ(identity.get().length, kModelHashLength); manifest.function = @"another_function"; manifest.inputs = @[ @"different_input" ]; manifest.outputs = @[ @"different_output" ]; manifest.hash = @"another_delivery_location"; manifest.minimum_version = (NSOperatingSystemVersion){27, 0, 0}; - auto unchanged = bookmark_identity(manifest, bundle); + auto unchanged = cache_model_key(manifest, bundle); ASSERT_TRUE(unchanged.ok()); EXPECT_TRUE([unchanged.get() isEqual:identity.get()]); - auto context = bookmark_key(manifest, @"iOS", @"arch_a"); + auto context = cache_variant_key(manifest, @"iOS", @"arch_a", @"default"); ASSERT_TRUE(context.ok()); EXPECT_FALSE([context.get() isEqual:key]); } -TEST_F(CoreAIBookmarkIdentityTest, IdentitySeparatesDigestBundleAndKind) { +TEST_F(CoreAIBookmarkIdentityTest, ModelKeySeparatesDigestBundleAndKind) { NSString* bundle = manifest.path.lastPathComponent; - auto identity = bookmark_identity(manifest, bundle); + auto identity = cache_model_key(manifest, bundle); ASSERT_TRUE(identity.ok()); Manifest reweighted = manifest; reweighted.bundle_digests = @{bundle : fixture_identity('2')}; - auto weights = bookmark_identity(reweighted, bundle); + auto weights = cache_model_key(reweighted, bundle); ASSERT_TRUE(weights.ok()); EXPECT_FALSE([weights.get() isEqual:identity.get()]); Manifest renamed = manifest; renamed.bundle_digests = @{@"other.aimodel" : manifest.bundle_digests[bundle]}; - auto bundle_name = bookmark_identity(renamed, @"other.aimodel"); + auto bundle_name = cache_model_key(renamed, @"other.aimodel"); ASSERT_TRUE(bundle_name.ok()); EXPECT_FALSE([bundle_name.get() isEqual:identity.get()]); Manifest compiled = manifest; compiled.aot_compiled = true; - auto kind = bookmark_identity(compiled, bundle); + auto kind = cache_model_key(compiled, bundle); ASSERT_TRUE(kind.ok()); EXPECT_FALSE([kind.get() isEqual:identity.get()]); } TEST_F( CoreAIBookmarkIdentityTest, - IdentityRejectsUnselectedAndEscapingBundles) { - EXPECT_FALSE(bookmark_identity(manifest, @"other.aimodel").ok()); + ModelKeyRejectsUnselectedAndEscapingBundles) { + EXPECT_FALSE(cache_model_key(manifest, @"other.aimodel").ok()); Manifest escaped = manifest; escaped.bundle_digests = @{@"../escape" : fixture_identity('1')}; - EXPECT_FALSE(bookmark_identity(escaped, @"../escape").ok()); -} - -TEST_F(CoreAIBookmarkRawTest, EnsurePublishesFullKeyFieldsReadableAsIdentity) { - ASSERT_EQ( - ensure_bookmark_identity(*locked, manifest, @"macOS", @"arch_b"), - Error::Ok); - auto identity = read_bookmark_identity(*locked); - ASSERT_TRUE(identity.ok()); - auto expected = bookmark_identity(manifest, manifest.path.lastPathComponent); - ASSERT_TRUE(expected.ok()); - EXPECT_TRUE([identity.get() isEqual:expected.get()]); - NSData* bytes = [NSData dataWithContentsOfURL:identity_url(root.url, key)]; - ASSERT_NE(bytes, nil); - id parsed = [NSJSONSerialization JSONObjectWithData:bytes - options:0 - error:nil]; - ASSERT_TRUE([parsed isKindOfClass:NSArray.class]); - NSArray* fields = parsed; - ASSERT_EQ(fields.count, 9u); - EXPECT_TRUE([fields[0] isEqual:@"1"]); - EXPECT_TRUE([fields[1] - isEqual:manifest.bundle_digests[manifest.path.lastPathComponent]]); - EXPECT_TRUE([fields[2] isEqual:@"source"]); - EXPECT_TRUE([fields[3] isEqual:manifest.path.lastPathComponent]); - EXPECT_TRUE([fields[4] isEqual:@"macOS"]); - EXPECT_TRUE([fields[5] isEqual:@"arch_b"]); - for (id field in fields) { - EXPECT_TRUE([field isKindOfClass:NSString.class] && [field length] > 0); - } - EXPECT_TRUE(no_temporary_bookmarks(root.url)); -} - -TEST_F(CoreAIBookmarkRawTest, EnsureReusesMatchingRecordWithoutRepublishing) { - ASSERT_EQ( - ensure_bookmark_identity(*locked, manifest, @"macOS", @"arch_b"), - Error::Ok); - NSData* published = - [NSData dataWithContentsOfURL:identity_url(root.url, key)]; - ASSERT_NE(published, nil); - for (StorageOperation operation : - {StorageOperation::Write, StorageOperation::Rename}) { - SCOPED_TRACE(static_cast(operation)); - StorageFaultScope fault(operation); - EXPECT_EQ( - ensure_bookmark_identity(*locked, manifest, @"macOS", @"arch_b"), - Error::Ok); - EXPECT_EQ(fault.hits, 0); - } - EXPECT_TRUE([[NSData dataWithContentsOfURL:identity_url(root.url, key)] - isEqual:published]); - EXPECT_TRUE(no_temporary_bookmarks(root.url)); -} - -TEST_F(CoreAIBookmarkRawTest, MissingIdentityReadsAsNil) { - auto identity = read_bookmark_identity(*locked); - ASSERT_TRUE(identity.ok()); - EXPECT_EQ(identity.get(), nil); -} - -TEST_F(CoreAIBookmarkRawTest, EnsureRepairsCorruptAndMisattributedRecords) { - NSURL* metadata = identity_url(root.url, key); - NSData* garbage = [@"not metadata" dataUsingEncoding:NSUTF8StringEncoding]; - ASSERT_TRUE([garbage writeToURL:metadata atomically:NO]); - ASSERT_EQ( - ensure_bookmark_identity(*locked, manifest, @"macOS", @"arch_b"), - Error::Ok); - auto repaired = read_bookmark_identity(*locked); - ASSERT_TRUE(repaired.ok()); - auto expected = bookmark_identity(manifest, manifest.path.lastPathComponent); - ASSERT_TRUE(expected.ok()); - EXPECT_TRUE([repaired.get() isEqual:expected.get()]); - Manifest reweighted = manifest; - reweighted.bundle_digests = - @{manifest.path.lastPathComponent : fixture_identity('2')}; - auto foreign_key = test_bookmark_key(reweighted); - ASSERT_TRUE(foreign_key.ok()); - auto foreign = lock_bookmark(root.url.path, foreign_key.get()); - ASSERT_TRUE(foreign.ok()); - ASSERT_EQ( - ensure_bookmark_identity(*foreign.get(), reweighted, @"macOS", @"arch_b"), - Error::Ok); - NSData* foreign_bytes = - [NSData dataWithContentsOfURL:identity_url(root.url, foreign_key.get())]; - ASSERT_NE(foreign_bytes, nil); - ASSERT_TRUE([foreign_bytes writeToURL:metadata atomically:NO]); - ASSERT_EQ( - ensure_bookmark_identity(*locked, manifest, @"macOS", @"arch_b"), - Error::Ok); - auto restored = read_bookmark_identity(*locked); - ASSERT_TRUE(restored.ok()); - EXPECT_TRUE([restored.get() isEqual:expected.get()]); -} - -TEST_F(CoreAIBookmarkRawTest, RejectsMalformedIdentityRecords) { - NSString* bundle = manifest.path.lastPathComponent; - NSString* digest = manifest.bundle_digests[bundle]; - NSArray* base = @[ - @"1", - digest, - @"source", - bundle, - @"macOS", - @"arch_b", - @"AIModelCache.default", - @"default-options-v1", - @"persistent" - ]; - NSMutableArray* payloads = [NSMutableArray array]; - [payloads addObject:[@"not json" dataUsingEncoding:NSUTF8StringEncoding]]; - for (id document in @[ - @{@"version" : @"1"}, - @"metadata", - [base subarrayWithRange:NSMakeRange(0, 8)], - [base arrayByAddingObject:@"extra"] - ]) { - [payloads addObject:[NSJSONSerialization - dataWithJSONObject:document - options:NSJSONWritingFragmentsAllowed - error:nil]]; - } - NSMutableArray* numeric = [base mutableCopy]; - numeric[1] = @1; - NSMutableArray* empty = [base mutableCopy]; - empty[4] = @""; - NSMutableArray* kind = [base mutableCopy]; - kind[2] = @"compiled"; - NSMutableArray* undigestable = [base mutableCopy]; - undigestable[1] = [@"" stringByPaddingToLength:64 - withString:@"z" - startingAtIndex:0]; - NSMutableArray* nested = [base mutableCopy]; - nested[3] = @"a/b"; - NSMutableArray* retargeted = [base mutableCopy]; - retargeted[8] = @"other"; - for (NSArray* variant in - @[ numeric, empty, kind, undigestable, nested, retargeted ]) { - [payloads addObject:[NSJSONSerialization dataWithJSONObject:variant - options:0 - error:nil]]; - } - NSURL* metadata = identity_url(root.url, key); - for (NSData* payload in payloads) { - SCOPED_TRACE(payload.length); - ASSERT_NE(payload, nil); - ASSERT_TRUE([payload writeToURL:metadata atomically:NO]); - EXPECT_EQ( - read_bookmark_identity(*locked).error(), Error::InvalidExternalData); - EXPECT_TRUE([[NSData dataWithContentsOfURL:metadata] isEqual:payload]); - } -} - -TEST_F(CoreAIBookmarkRawTest, RejectsOversizedIdentityRecord) { - NSURL* metadata = identity_url(root.url, key); - FileDescriptor file(open( - metadata.fileSystemRepresentation, - O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, - 0600)); - ASSERT_GE(file.get(), 0); - ASSERT_EQ(ftruncate(file.get(), 128 * 1024 + 1), 0); - ASSERT_EQ(close(file.release()), 0); - EXPECT_EQ( - read_bookmark_identity(*locked).error(), Error::InvalidExternalData); - struct stat info {}; - ASSERT_EQ(lstat(metadata.fileSystemRepresentation, &info), 0); - EXPECT_EQ(info.st_size, static_cast(128 * 1024 + 1)); -} - -TEST_F(CoreAIBookmarkRawTest, RejectsLinkedIdentityRecordWithoutReadingTarget) { - NSURL* metadata = identity_url(root.url, key); - NSURL* target = [root.url URLByAppendingPathComponent:@"target"]; - NSData* bytes = [@"target bytes" dataUsingEncoding:NSUTF8StringEncoding]; - ASSERT_TRUE([bytes writeToURL:target atomically:NO]); - ASSERT_EQ( - symlink( - target.fileSystemRepresentation, metadata.fileSystemRepresentation), - 0); - EXPECT_EQ( - read_bookmark_identity(*locked).error(), Error::InvalidExternalData); - EXPECT_TRUE(([[NSData dataWithContentsOfURL:target] isEqual:bytes])); - struct stat info {}; - ASSERT_EQ(lstat(metadata.fileSystemRepresentation, &info), 0); - EXPECT_TRUE(S_ISLNK(info.st_mode)); -} - -TEST_F(CoreAIBookmarkRawTest, RejectsDirectoryAndFifoIdentityRecords) { - NSURL* metadata = identity_url(root.url, key); - ASSERT_EQ(mkdir(metadata.fileSystemRepresentation, 0700), 0); - EXPECT_EQ( - read_bookmark_identity(*locked).error(), Error::InvalidExternalData); - ASSERT_EQ(rmdir(metadata.fileSystemRepresentation), 0); - ASSERT_EQ(mkfifo(metadata.fileSystemRepresentation, 0600), 0); - EXPECT_EQ( - read_bookmark_identity(*locked).error(), Error::InvalidExternalData); -} - -TEST_F(CoreAIBookmarkRawTest, RemoveIdentityPreservesBookmarkAndStaging) { - NSData* bytes = [@"bookmark" dataUsingEncoding:NSUTF8StringEncoding]; - ASSERT_EQ(write_bookmark(*locked, bytes), Error::Ok); - ASSERT_EQ( - ensure_bookmark_identity(*locked, manifest, @"macOS", @"arch_b"), - Error::Ok); - auto staging = prepare_bookmark_staging(*locked); - ASSERT_TRUE(staging.ok()); - ASSERT_EQ(remove_bookmark_identity(*locked), Error::Ok); - EXPECT_FALSE([NSFileManager.defaultManager - fileExistsAtPath:identity_url(root.url, key).path]); - EXPECT_TRUE(bookmark_bytes_equal(*locked, bytes)); - EXPECT_TRUE([NSFileManager.defaultManager fileExistsAtPath:staging.get()]); - auto reread = read_bookmark_identity(*locked); - ASSERT_TRUE(reread.ok()); - EXPECT_EQ(reread.get(), nil); - EXPECT_EQ(remove_bookmark_identity(*locked), Error::Ok); + EXPECT_FALSE(cache_model_key(escaped, @"../escape").ok()); } -TEST_F(CoreAIBookmarkPathTest, EnsureRejectsSelectedOrContextOutsideHeldKey) { - auto locked = lock_bookmark(root.url.path, key); - ASSERT_TRUE(locked.ok()); - Manifest reweighted = manifest; - reweighted.bundle_digests = - @{manifest.path.lastPathComponent : fixture_identity('2')}; - EXPECT_EQ( - ensure_bookmark_identity(*locked.get(), reweighted, @"macOS", @"arch_b"), - Error::InvalidArgument); - EXPECT_EQ( - ensure_bookmark_identity(*locked.get(), manifest, @"iOS", @"arch_b"), - Error::InvalidArgument); - EXPECT_EQ( - ensure_bookmark_identity(*locked.get(), manifest, @"macOS", @"arch_a"), - Error::InvalidArgument); - EXPECT_FALSE([NSFileManager.defaultManager - fileExistsAtPath:identity_url(root.url, key).path]); -} - -TEST_F(CoreAIBookmarkPathTest, EnsureRejectsStagingOnlyLock) { - NSURL* staged = [[root.url URLByAppendingPathComponent:@"staging"] - URLByAppendingPathComponent:key]; - ASSERT_TRUE([NSFileManager.defaultManager createDirectoryAtURL:staged - withIntermediateDirectories:YES - attributes:nil - error:nil]); - auto existing = lock_bookmark(root.url.path, key, false); - ASSERT_TRUE(existing.ok()); - ASSERT_NE(existing.get(), nullptr); - EXPECT_EQ( - ensure_bookmark_identity(*existing.get(), manifest, @"macOS", @"arch_b"), - Error::InvalidArgument); - EXPECT_FALSE([NSFileManager.defaultManager - fileExistsAtPath:identity_url(root.url, key).path]); -} - -TEST_F(CoreAIBookmarkPathTest, RemoveIdentityThroughStagingOnlyLockSucceeds) { - NSURL* staged = [[root.url URLByAppendingPathComponent:@"staging"] - URLByAppendingPathComponent:key]; - ASSERT_TRUE([NSFileManager.defaultManager createDirectoryAtURL:staged - withIntermediateDirectories:YES - attributes:nil - error:nil]); - auto existing = lock_bookmark(root.url.path, key, false); - ASSERT_TRUE(existing.ok()); - ASSERT_NE(existing.get(), nullptr); - EXPECT_EQ(remove_bookmark_identity(*existing.get()), Error::Ok); -} - -TEST_F(CoreAIBookmarkPathTest, RejectsMetadataWrittenForAnotherKey) { - Manifest reweighted = manifest; - reweighted.bundle_digests = - @{manifest.path.lastPathComponent : fixture_identity('2')}; - auto foreign_key = test_bookmark_key(reweighted); - ASSERT_TRUE(foreign_key.ok()); - auto foreign = lock_bookmark(root.url.path, foreign_key.get()); - ASSERT_TRUE(foreign.ok()); - ASSERT_EQ( - ensure_bookmark_identity(*foreign.get(), reweighted, @"macOS", @"arch_b"), - Error::Ok); - NSData* foreign_bytes = - [NSData dataWithContentsOfURL:identity_url(root.url, foreign_key.get())]; - ASSERT_NE(foreign_bytes, nil); - auto locked = lock_bookmark(root.url.path, key); - ASSERT_TRUE(locked.ok()); - ASSERT_TRUE([foreign_bytes writeToURL:identity_url(root.url, key) - atomically:NO]); - EXPECT_EQ( - read_bookmark_identity(*locked.get()).error(), - Error::InvalidExternalData); - EXPECT_TRUE([[NSData dataWithContentsOfURL:identity_url(root.url, key)] - isEqual:foreign_bytes]); -} - -TEST_F(CoreAIBookmarkPathTest, MetadataOnlyEntriesListLockAndRemove) { - auto locked = lock_bookmark(root.url.path, key); - ASSERT_TRUE(locked.ok()); - ASSERT_EQ( - ensure_bookmark_identity(*locked.get(), manifest, @"macOS", @"arch_b"), - Error::Ok); - locked.get().reset(); - auto inventory = inventory_bookmark_keys(root.url.path); - ASSERT_TRUE(inventory.ok()); - EXPECT_TRUE([inventory.get() containsObject:key]); - auto existing = lock_bookmark(root.url.path, key, false); - ASSERT_TRUE(existing.ok()); - ASSERT_NE(existing.get(), nullptr); - auto identity = read_bookmark_identity(*existing.get()); - ASSERT_TRUE(identity.ok()); - auto expected = bookmark_identity(manifest, manifest.path.lastPathComponent); - ASSERT_TRUE(expected.ok()); - EXPECT_TRUE([identity.get() isEqual:expected.get()]); - EXPECT_TRUE(bookmark_bytes_equal(*existing.get(), nil)); - ASSERT_EQ(remove_bookmark_identity(*existing.get()), Error::Ok); - EXPECT_FALSE([NSFileManager.defaultManager - fileExistsAtPath:identity_url(root.url, key).path]); - auto relisted = inventory_bookmark_keys(root.url.path); - ASSERT_TRUE(relisted.ok()); - EXPECT_FALSE([relisted.get() containsObject:key]); - auto absent = lock_bookmark(root.url.path, key, false); - ASSERT_TRUE(absent.ok()); - EXPECT_EQ(absent.get(), nullptr); -} - -TEST_F(CoreAIBookmarkPathTest, LockNamesUseOnlyValidatedFirstByte) { +TEST_F(CoreAIBookmarkPathTest, LockNamesDeriveFromFullKeyHash) { EXPECT_EQ(bookmark_lock_name(nil).error(), Error::InvalidArgument); for (NSString* invalid in - @[ @"", @"../escape", @"00", [fixture_identity('a') uppercaseString] ]) { + @[ @"", @"../escape", @"00", fixture_identity('a') ]) { EXPECT_EQ(bookmark_lock_name(invalid).error(), Error::InvalidArgument); + EXPECT_FALSE(is_cache_key(invalid)); } - auto first = bookmark_lock_name(fixture_identity('0')); - auto last = bookmark_lock_name(fixture_identity('f')); + // Stripe names hash the full composite key, not the first byte. + auto first = bookmark_lock_name(key); ASSERT_TRUE(first.ok()); - ASSERT_TRUE(last.ok()); - EXPECT_TRUE([first.get() isEqual:@"stripe-v1-00.lock"]); - EXPECT_TRUE([last.get() isEqual:@"stripe-v1-ff.lock"]); + EXPECT_TRUE([first.get() hasPrefix:@"stripe-v1-"]); + EXPECT_EQ(first.get().length, 17u); + NSString* collision = colliding_key(key); + ASSERT_NE(collision, nil); + ASSERT_FALSE([collision isEqual:key]); + ASSERT_TRUE(is_cache_key(collision)); + auto collided = bookmark_lock_name(collision); + ASSERT_TRUE(collided.ok()); + EXPECT_TRUE([collided.get() isEqual:first.get()]); EXPECT_TRUE([test_bookmark_lock_url(root.url, key) - isEqual:test_bookmark_lock_url(root.url, changed_key(key, true))]); + isEqual:test_bookmark_lock_url(root.url, collision)]); + NSString* independent = different_key(key); + ASSERT_NE(independent, nil); + auto separated = bookmark_lock_name(independent); + ASSERT_TRUE(separated.ok()); + EXPECT_FALSE([separated.get() isEqual:first.get()]); + // Synthetic composite keys stay valid and hash stably. + NSString* synthetic = + synthetic_key(fixed_options('a'), @"macOS_arch_b", fixed_options('b')); + ASSERT_TRUE(is_cache_key(synthetic)); + EXPECT_TRUE([cache_key_model_prefix(synthetic) isEqual:fixed_options('a')]); + auto hashed = bookmark_lock_name(synthetic); + ASSERT_TRUE(hashed.ok()); + EXPECT_TRUE([hashed.get() hasPrefix:@"stripe-v1-"]); } TEST_F(CoreAIBookmarkPathTest, LockPoolIsBoundedAcrossManyKeysAndClears) { - std::array inodes{}; + std::map inodes; NSURL* locks = [root.url URLByAppendingPathComponent:@"locks"]; + NSString* options = fixed_options('b'); for (size_t generation = 0; generation < 4; ++generation) { - for (size_t stripe = 0; stripe < kBookmarkLockStripeCount; ++stripe) { + for (size_t index = 0; index < kBookmarkLockStripeCount; ++index) { @autoreleasepool { - NSString* entry = - [NSString stringWithFormat:@"%02zx%062zx", stripe, generation]; + // Distinct composite keys: model prefix varies, device/options stay + // readable. Stripes derive from the full-key hash. + NSString* entry = synthetic_key( + padded_model(static_cast( + generation * kBookmarkLockStripeCount + index)), + @"macOS_arch_b", + options); + ASSERT_TRUE(is_cache_key(entry)); auto locked = lock_bookmark(root.url.path, entry); ASSERT_TRUE(locked.ok()); struct stat info {}; @@ -977,11 +794,13 @@ FileDescriptor file(open( ASSERT_EQ(lstat(file.fileSystemRepresentation, &info), 0); ASSERT_TRUE(S_ISREG(info.st_mode)); EXPECT_EQ(info.st_size, 0); - if (generation == 0) - inodes[stripe] = info.st_ino; + std::string name = file.lastPathComponent.UTF8String; + auto known = inodes.find(name); + if (known == inodes.end()) + inodes[name] = info.st_ino; else - EXPECT_EQ(info.st_ino, inodes[stripe]); - auto staging = prepare_bookmark_staging(*locked.get()); + EXPECT_EQ(info.st_ino, known->second); + auto staging = prepare_bookmark_staging_root(*locked.get()); ASSERT_TRUE(staging.ok()); NSString* stage = [staging.get() stringByAppendingPathComponent:entry]; ASSERT_EQ(mkdir(stage.fileSystemRepresentation, 0700), 0); @@ -991,20 +810,21 @@ FileDescriptor file(open( auto reused = lock_bookmark(root.url.path, entry); ASSERT_TRUE(reused.ok()); ASSERT_EQ(lstat(file.fileSystemRepresentation, &info), 0); - EXPECT_EQ(info.st_ino, inodes[stripe]); + EXPECT_EQ(info.st_ino, inodes[name]); } } NSArray* children = [NSFileManager.defaultManager contentsOfDirectoryAtPath:locks.path error:nil]; ASSERT_NE(children, nil); - ASSERT_EQ(children.count, kBookmarkLockStripeCount); + // Hash-derived stripes stay within the bounded pool. + EXPECT_LE(children.count, kBookmarkLockStripeCount); for (NSString* name in children) { EXPECT_TRUE([name hasPrefix:@"stripe-v1-"]); EXPECT_EQ(name.length, 17u); } } - auto inventory = inventory_bookmark_keys(root.url.path); + auto inventory = inventory_cache_keys(root.url.path); ASSERT_TRUE(inventory.ok()); EXPECT_EQ(inventory.get().count, 0u); } @@ -1012,7 +832,8 @@ FileDescriptor file(open( TEST_F(CoreAIBookmarkPathTest, CollidingKeysSerializeAcrossThreads) { auto holder = lock_bookmark(root.url.path, key); ASSERT_TRUE(holder.ok()); - NSString* collision = changed_key(key, true); + NSString* collision = colliding_key(key); + ASSERT_NE(collision, nil); ASSERT_FALSE([collision isEqual:key]); auto waiter = start_bookmark_waiter(root.url.path, collision); ASSERT_EQ( @@ -1032,8 +853,9 @@ FileDescriptor file(open( NoncreatingCollisionMissDoesNotWaitOrCreateFiles) { auto holder = lock_bookmark(root.url.path, key); ASSERT_TRUE(holder.ok()); - auto waiter = - start_bookmark_waiter(root.url.path, changed_key(key, true), false); + NSString* collision = colliding_key(key); + ASSERT_NE(collision, nil); + auto waiter = start_bookmark_waiter(root.url.path, collision, false); EXPECT_EQ(waiter.result.wait_for(kWaitTimeout), std::future_status::ready); holder.get().reset(); ASSERT_EQ(waiter.result.wait_for(kWaitTimeout), std::future_status::ready); @@ -1049,12 +871,13 @@ FileDescriptor file(open( } TEST_F(CoreAIBookmarkPathTest, ClearingCollidingKeyPreservesOtherKeyData) { - NSString* collision = changed_key(key, true); + NSString* collision = colliding_key(key); + ASSERT_NE(collision, nil); NSData* bytes = [@"other bookmark" dataUsingEncoding:NSUTF8StringEncoding]; for (NSString* entry in @[ key, collision ]) { auto locked = lock_bookmark(root.url.path, entry); ASSERT_TRUE(locked.ok()); - auto staging = prepare_bookmark_staging(*locked.get()); + auto staging = prepare_bookmark_staging_root(*locked.get()); ASSERT_TRUE(staging.ok()); ASSERT_EQ( mkdir( diff --git a/backends/apple/coreai/runtime/test/coreai_cache_test.mm b/backends/apple/coreai/runtime/test/coreai_cache_test.mm index dc323122c0a..beb162769da 100644 --- a/backends/apple/coreai/runtime/test/coreai_cache_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_cache_test.mm @@ -68,7 +68,7 @@ AssertionResult seed_cache( bool stage = true, NSString* architecture = @"arch_b", NSString* platform = @"macOS", - bool metadata = true) { + NSString* specialization = @"default") { if (root == nil || loader == nil) { return AssertionFailure() << "Cache seed requires a root and fake loader"; } @@ -80,7 +80,8 @@ AssertionResult seed_cache( if (!selected.ok()) { return AssertionFailure() << "Cannot select cache seed assets"; } - auto key = bookmark_key(selected.get(), platform, architecture); + auto key = + cache_variant_key(selected.get(), platform, architecture, specialization); if (!key.ok()) { return AssertionFailure() << "Cannot compute cache seed key"; } @@ -93,13 +94,8 @@ AssertionResult seed_cache( if (!lock.ok()) { return AssertionFailure() << "Cannot lock cache seed bookmark"; } - if (metadata && - ensure_bookmark_identity( - *lock.get(), selected.get(), platform, architecture) != Error::Ok) { - return AssertionFailure() << "Cannot attribute cache seed"; - } if (stage) { - auto staging = prepare_bookmark_staging(*lock.get()); + auto staging = prepare_bookmark_staging_root(*lock.get()); if (!staging.ok()) { return AssertionFailure() << "Cannot prepare cache seed staging"; } @@ -139,13 +135,7 @@ bool cache_exists(NSURL* url) { return [NSFileManager.defaultManager fileExistsAtPath:url.path]; } -NSURL* cache_identity_url(NSURL* root, NSString* key) { - return [[test_bookmark_url(root, key) URLByDeletingPathExtension] - URLByAppendingPathExtension:@"metadata"]; -} - void cache_entry_removed(NSURL* root, const CacheEntry& entry) { - EXPECT_FALSE(cache_exists(cache_identity_url(root, entry.key))); EXPECT_FALSE(cache_exists(test_bookmark_url(root, entry.key))); EXPECT_FALSE(cache_exists(entry.stage)); EXPECT_TRUE(cache_exists(test_bookmark_lock_url(root, entry.key))); @@ -384,7 +374,7 @@ FBSyntheticPTE pte( } } -TEST_F(CoreAICacheTest, ClearsLegacyAotWithoutAttribution) { +TEST_F(CoreAICacheTest, ClearsAotBundles) { @autoreleasepool { TestDirectory root; BookmarkLoaderScope fake; @@ -399,7 +389,7 @@ FBSyntheticPTE pte( false, @"arch_a", @"macOS", - false)); + @"default")); ASSERT_TRUE(seed_cache( root.url, dict, @@ -409,7 +399,7 @@ FBSyntheticPTE pte( false, @"arch_b", @"macOS", - false)); + @"default")); FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", encode(dict)}}}}); EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); EXPECT_EQ(fake.loader->evictions.load(), 2); @@ -418,60 +408,56 @@ FBSyntheticPTE pte( } } -TEST_F( - CoreAICacheTest, - UnattributedEntriesReportIncompleteButClearKnownMatches) { - for (bool corrupt : {false, true}) { - SCOPED_TRACE(corrupt); - @autoreleasepool { - TestDirectory root; - BookmarkLoaderScope fake; - CacheEntry unknown, known; - ASSERT_TRUE(seed_cache( - root.url, - manifest_dict(), - fake.loader, - unknown, - true, - false, - @"old_arch", - @"macOS", - corrupt)); - ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, known)); - if (corrupt) { - ASSERT_TRUE([[@"{}" dataUsingEncoding:NSUTF8StringEncoding] - writeToURL:cache_identity_url(root.url, unknown.key) - options:0 - error:nil]); - } - FBSyntheticPTE pte( - {{"forward", {{"CoreAIBackend", encode(manifest_dict())}}}}); - EXPECT_EQ( - clear_cache_for_pte(pte, root.url.path.UTF8String), - Error::InvalidExternalData); - EXPECT_EQ(fake.loader->evictions.load(), 1); - cache_entry_removed(root.url, known); - EXPECT_TRUE(cache_exists(test_bookmark_url(root.url, unknown.key))); - EXPECT_NE(fake.loader.models[unknown.bookmark], nil); - EXPECT_EQ(clear_cache(root.url.path.UTF8String), Error::Ok); - cache_entry_removed(root.url, unknown); - } - } -} - -TEST_F(CoreAICacheTest, ClearsAttributionOnlyAndStagingEntriesWithoutSdk) { +TEST_F(CoreAICacheTest, ClearsOptionAndDeviceVariantsPreservingUnrelatedModel) { @autoreleasepool { TestDirectory root; BookmarkLoaderScope fake; - CacheEntry metadata_only, staged; + CacheEntry base, specialized, device, unrelated; + ASSERT_TRUE(seed_cache(root.url, manifest_dict(), fake.loader, base)); ASSERT_TRUE(seed_cache( root.url, manifest_dict(), fake.loader, - metadata_only, - false, - false, + specialized, + true, + true, + @"arch_b", + @"macOS", + @"tuned")); + ASSERT_TRUE(seed_cache( + root.url, + manifest_dict(), + fake.loader, + device, + true, + true, @"other_arch")); + ASSERT_TRUE( + seed_cache(root.url, cache_manifest('3'), fake.loader, unrelated)); + FBSyntheticPTE pte( + {{"forward", {{"CoreAIBackend", encode(manifest_dict())}}}}); + EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); + drain_callbacks(); + EXPECT_EQ(fake.loader->evictions.load(), 3); + EXPECT_EQ(fake.loader->specializationQueries.load(), 0); + cache_entry_removed(root.url, base); + cache_entry_removed(root.url, specialized); + cache_entry_removed(root.url, device); + EXPECT_TRUE(cache_exists(test_bookmark_url(root.url, unrelated.key))); + EXPECT_TRUE(cache_exists(unrelated.stage)); + EXPECT_NE(fake.loader.models[unrelated.bookmark], nil); + EXPECT_EQ(fake.loader.models[base.bookmark], nil); + EXPECT_EQ(fake.loader.models[specialized.bookmark], nil); + EXPECT_EQ(fake.loader.models[device.bookmark], nil); + EXPECT_TRUE(pte.check_released()); + } +} + +TEST_F(CoreAICacheTest, ClearsStagingOnlyEntriesWithoutSdk) { + @autoreleasepool { + TestDirectory root; + BookmarkLoaderScope fake; + CacheEntry staged; ASSERT_TRUE(seed_cache( root.url, manifest_dict(), fake.loader, staged, false, true)); bridge.state().available = false; @@ -479,7 +465,6 @@ FBSyntheticPTE pte( {{"forward", {{"CoreAIBackend", encode(manifest_dict())}}}}); EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); EXPECT_EQ(fake.loader->evictions.load(), 0); - cache_entry_removed(root.url, metadata_only); cache_entry_removed(root.url, staged); } } @@ -819,33 +804,6 @@ FBSyntheticPTE pte( } } -TEST_F(CoreAICacheTest, AttributionRemovalFailureLeavesDiscoverableRetry) { - @autoreleasepool { - TestDirectory root; - BookmarkLoaderScope fake; - CacheEntry entry; - ASSERT_TRUE( - seed_cache(root.url, manifest_dict(), fake.loader, entry, true, false)); - FBSyntheticPTE pte( - {{"forward", {{"CoreAIBackend", encode(manifest_dict())}}}}); - { - StorageFaultScope fault(StorageOperation::Remove); - fault.match_index = 2; - EXPECT_EQ( - clear_cache_for_pte(pte, root.url.path.UTF8String), - Error::AccessFailed); - EXPECT_EQ(fault.hits, 1); - } - EXPECT_FALSE(cache_exists(test_bookmark_url(root.url, entry.key))); - EXPECT_TRUE(cache_exists(cache_identity_url(root.url, entry.key))); - EXPECT_EQ(fake.loader->evictions.load(), 1); - bridge.state().available = false; - EXPECT_EQ(clear_cache_for_pte(pte, root.url.path.UTF8String), Error::Ok); - EXPECT_EQ(fake.loader->evictions.load(), 1); - cache_entry_removed(root.url, entry); - } -} - TEST_F(CoreAICacheTest, IndependentKeysContinueInSortedOrderAfterFailures) { @autoreleasepool { TestDirectory root; diff --git a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm index 76c8e7b047a..178459db506 100644 --- a/backends/apple/coreai/runtime/test/coreai_delegate_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_delegate_test.mm @@ -145,7 +145,7 @@ DelegateFixture fixture( const int queries = bridge.state().architecture_queries.load(); ASSERT_EQ(fixture.init(), Error::Ok); fixture.reset(); - auto keys = inventory_bookmark_keys(fixture.storage.url.path); + auto keys = inventory_cache_keys(fixture.storage.url.path); ASSERT_TRUE(keys.ok()); ASSERT_EQ(keys.get().count, 1); auto bookmark = saved_bookmark(fixture.storage.url, keys.get()[0]); @@ -162,7 +162,7 @@ FBSyntheticPTE pte( EXPECT_EQ(fake.loader->specializations.load(), 1); EXPECT_EQ(fake.loader->evictions.load(), 1); EXPECT_EQ(fake.loader.models[bookmark.get()], nil); - auto remaining = inventory_bookmark_keys(fixture.storage.url.path); + auto remaining = inventory_cache_keys(fixture.storage.url.path); ASSERT_TRUE(remaining.ok()); EXPECT_EQ(remaining.get().count, 0); EXPECT_TRUE(pte.check_released()); @@ -181,7 +181,7 @@ FBSyntheticPTE pte( DefaultAssetsRootScope default_root(fixture.storage.url.path); ASSERT_EQ(fixture.init_with_options({}), Error::Ok); fixture.reset(); - auto keys = inventory_bookmark_keys(fixture.storage.url.path); + auto keys = inventory_cache_keys(fixture.storage.url.path); ASSERT_TRUE(keys.ok()); ASSERT_EQ(keys.get().count, 1); FBSyntheticPTE pte({{"forward", {{"CoreAIBackend", fixture.json}}}}); @@ -206,7 +206,7 @@ FBSyntheticPTE pte( } EXPECT_EQ(fake.loader->evictions.load(), 1); EXPECT_EQ(fake.loader.models.count, 0); - auto remaining = inventory_bookmark_keys(fixture.storage.url.path); + auto remaining = inventory_cache_keys(fixture.storage.url.path); ASSERT_TRUE(remaining.ok()); EXPECT_EQ(remaining.get().count, 0); EXPECT_TRUE(pte.check_released()); @@ -431,7 +431,7 @@ void expect_preflight_rejection( auto selected = select_assets(parsed.get(), arch, @"macOS"); ASSERT_TRUE(selected.ok()); ASSERT_EQ(fixture.init(), Error::Ok); - auto key = bookmark_key(selected.get(), @"macOS", arch); + auto key = cache_variant_key(selected.get(), @"macOS", arch, @"default"); ASSERT_TRUE(key.ok()); NSString* expected = [fixture.storage.url.path stringByAppendingPathComponent: diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.h b/backends/apple/coreai/runtime/test/coreai_fake_loader.h index faefbf15f94..b7d1c6f96d3 100644 --- a/backends/apple/coreai/runtime/test/coreai_fake_loader.h +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.h @@ -88,7 +88,10 @@ class ScopedFakeBridgeState final { std::atomic specializations; std::atomic bindings; std::atomic evictions; + std::atomic specializationQueries; } +@property(nonatomic, readonly, copy) NSString* specializationIdentifier; +- (instancetype)initWithSpecializationIdentifier:(NSString*)identifier; @property(nonatomic, strong) NSMutableDictionary* models; @property(nonatomic, strong) NSError* restoreError; @@ -111,7 +114,7 @@ namespace executorch::backends::coreai::testing { struct BookmarkLoaderScope final { BookmarkFakeLoader* loader = nil; - BookmarkLoaderScope(); + explicit BookmarkLoaderScope(NSString* specialization = @"default"); ~BookmarkLoaderScope(); BookmarkLoaderScope(const BookmarkLoaderScope&) = delete; BookmarkLoaderScope& operator=(const BookmarkLoaderScope&) = delete; diff --git a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm index d53b6542748..84cb02e46ce 100644 --- a/backends/apple/coreai/runtime/test/coreai_fake_loader.mm +++ b/backends/apple/coreai/runtime/test/coreai_fake_loader.mm @@ -216,6 +216,9 @@ @interface FakeLoader () @end @implementation FakeLoader +- (NSString*)specializationIdentifier { + return @"default"; +} - (instancetype)init { if (current_context == nil) { return nil; @@ -318,14 +321,25 @@ - (void)loadFunctionNamed:(NSString*)functionName } @end -@implementation BookmarkFakeLoader +@implementation BookmarkFakeLoader { + NSString* _specializationIdentifier; +} +- (NSString*)specializationIdentifier { + ++specializationQueries; + return _specializationIdentifier; +} - (instancetype)init { + return [self initWithSpecializationIdentifier:@"default"]; +} +- (instancetype)initWithSpecializationIdentifier:(NSString*)identifier { if (current_context == nil) { return nil; } if ((self = [super init])) { self.context = current_context; self.models = [NSMutableDictionary dictionary]; + _specializationIdentifier = [identifier copy]; + specializationQueries = 0; restores = 0; specializations = 0; bindings = 0; @@ -454,13 +468,14 @@ - (void)clearRetainedState { return ::wait_for_callbacks(context_, deadline); } -BookmarkLoaderScope::BookmarkLoaderScope() +BookmarkLoaderScope::BookmarkLoaderScope(NSString* specialization) : context_(current_context), previous_(context_ == nil ? nil : context_->state.bookmark_loader) { EXPECT_NE(context_, nil) << "BookmarkLoaderScope requires ScopedFakeBridgeState"; if (context_ != nil) { - loader = [[BookmarkFakeLoader alloc] init]; + loader = [[BookmarkFakeLoader alloc] + initWithSpecializationIdentifier:specialization]; context_->state.bookmark_loader = loader; } } diff --git a/backends/apple/coreai/runtime/test/coreai_key_format_test.mm b/backends/apple/coreai/runtime/test/coreai_key_format_test.mm new file mode 100644 index 00000000000..475a08345dd --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_key_format_test.mm @@ -0,0 +1,81 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +#import "../coreai_cache_key.h" + +#include +#include + +namespace executorch::backends::coreai::testing { +namespace { +NSString* key_for_device(NSString* device) { + NSString* model = [@"" stringByPaddingToLength:64 + withString:@"a" + startingAtIndex:0]; + NSString* options = [@"" stringByPaddingToLength:64 + withString:@"b" + startingAtIndex:0]; + return [NSString stringWithFormat:@"%@-%@-%@", model, device, options]; +} + +TEST(CoreAIKeyFormatTest, AcceptsReadableDeviceComponents) { + for (NSString* device in + @[ @"macOS_h17p", @"iOS_arch_b", @"visionOS_h17p-high" ]) { + NSString* key = key_for_device(device); + EXPECT_TRUE(is_cache_key(key)); + EXPECT_TRUE( + [cache_key_model_prefix(key) isEqual:[key substringToIndex:64]]); + } +} + +TEST(CoreAIKeyFormatTest, RejectsUnsafeOrIncompleteDevices) { + for (NSString* device in @[ + @"", + @"iOS", + @"iOS_", + @"_h17p", + @"iOS_../h17p", + @"iOS_h17p\\extra", + @"iOS_h17p extra", + @"iOS_h\u00e9", + @"iOS_h17p.bookmark" + ]) { + NSString* key = key_for_device(device); + EXPECT_FALSE(is_cache_key(key)); + EXPECT_EQ(cache_key_model_prefix(key), nil); + } +} + +TEST(CoreAIKeyFormatTest, LeavesRoomForBookmarkExtension) { + constexpr NSUInteger max_device_length = NAME_MAX - 130 - 9; + NSString* device = [@"iOS_" stringByPaddingToLength:max_device_length + withString:@"a" + startingAtIndex:0]; + NSString* key = key_for_device(device); + ASSERT_EQ(key.length + 9, NAME_MAX); + EXPECT_TRUE(is_cache_key(key)); + EXPECT_FALSE( + is_cache_key(key_for_device([device stringByAppendingString:@"a"]))); +} + +TEST(CoreAIKeyFormatTest, RequiresExactHashBoundaries) { + NSString* key = key_for_device(@"iOS_h17p"); + EXPECT_FALSE(is_cache_key(nil)); + EXPECT_FALSE(is_cache_key(@"")); + NSString* legacy = [key substringToIndex:64]; + EXPECT_FALSE(is_cache_key(legacy)); + EXPECT_EQ(cache_key_model_prefix(legacy), nil); + EXPECT_FALSE(is_cache_key([key substringFromIndex:1])); + EXPECT_FALSE(is_cache_key([key stringByAppendingString:@"0"])); + EXPECT_FALSE(is_cache_key(key.uppercaseString)); + EXPECT_FALSE(is_cache_key([key stringByAppendingString:@".bookmark"])); + EXPECT_FALSE( + is_cache_key([key stringByReplacingCharactersInRange:NSMakeRange(64, 1) + withString:@"_"])); +} +} // namespace +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_source_fixture.mm b/backends/apple/coreai/runtime/test/coreai_source_fixture.mm index afcceab8d2e..45e38dcbbb7 100644 --- a/backends/apple/coreai/runtime/test/coreai_source_fixture.mm +++ b/backends/apple/coreai/runtime/test/coreai_source_fixture.mm @@ -13,6 +13,7 @@ #include #include #include +#include "coreai_cache_key.h" #include "coreai_file.h" #include "coreai_storage.h" @@ -82,11 +83,15 @@ void aot_data(TestData& data, NSString* arch) { const Manifest& manifest, const runtime::NamedDataMap* data, NSString* source_root) { - NSString* key = manifest.path == nil - ? nil - : manifest.bundle_digests[manifest.path.lastPathComponent]; + auto key_result = + cache_variant_key(manifest, @"macOS", @"arch_b", @"default"); + if (!key_result.ok()) { + ADD_FAILURE() << "Cannot compute source fixture key"; + return key_result.error(); + } + NSString* key = key_result.get(); FileDescriptor lock(-1); - if (key != nil) { + { FileDescriptor parent(open( locks_.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_NOFOLLOW | O_CLOEXEC)); @@ -95,9 +100,13 @@ FileDescriptor parent(open( << strerror(errno); return Error::AccessFailed; } - NSString* name = [key stringByAppendingString:@".lock"]; - lock.reset( - open_storage_shared_file(parent.get(), name.fileSystemRepresentation)); + auto name_result = bookmark_lock_name(key); + if (!name_result.ok()) { + ADD_FAILURE() << "Cannot compute source fixture lock name"; + return name_result.error(); + } + lock.reset(open_storage_shared_file( + parent.get(), name_result.get().fileSystemRepresentation)); if (lock.get() < 0 || retry_eintr([&] { return flock(lock.get(), LOCK_EX); }) != 0) { ADD_FAILURE() << "Cannot acquire source fixture lock: " diff --git a/backends/apple/coreai/runtime/test/coreai_source_test.mm b/backends/apple/coreai/runtime/test/coreai_source_test.mm index a95b3e07c2e..cdf9430a5ff 100644 --- a/backends/apple/coreai/runtime/test/coreai_source_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_source_test.mm @@ -10,6 +10,7 @@ #include #include #include "coreai_bookmarks.h" +#include "coreai_cache_key.h" #include "coreai_fault_scope.h" #include "coreai_source_fixture.h" @@ -18,6 +19,12 @@ using runtime::Result; namespace { +NSString* source_key(const Manifest& manifest) { + auto key = cache_variant_key(manifest, @"macOS", @"arch_b", @"default"); + EXPECT_TRUE(key.ok()); + return key.ok() ? key.get() : nil; +} + Result source_manifest(bool aot = false) { auto parsed = parse_manifest(encode(aot ? aot_manifest_dict() : manifest_dict())); @@ -447,10 +454,13 @@ ASSERT_TRUE(manifest.ok()); auto root = resolve_bookmark_root(directory.url.path); ASSERT_TRUE(root.ok()); - NSString* key = fixture_identity('9'); + auto variant = cache_variant_key( + manifest.get(), @"macOS", @"arch_b", @"explicit-test-options"); + ASSERT_TRUE(variant.ok()); + NSString* key = variant.get(); auto lock = lock_bookmark(root.get(), key); ASSERT_TRUE(lock.ok()); - auto staging = prepare_bookmark_staging(*lock.get()); + auto staging = prepare_bookmark_staging_root(*lock.get()); ASSERT_TRUE(staging.ok()); auto first = executorch::backends::coreai::prepare_source_bundle( manifest.get(), &data, staging.get(), key); @@ -621,11 +631,10 @@ EXPECT_EQ(data.releases, 0); EXPECT_EQ(staging_directories(storage.url).count, 0u); EXPECT_FALSE([NSFileManager.defaultManager - fileExistsAtPath: - [storage.url - URLByAppendingPathComponent: - manifest->bundle_digests[manifest->path.lastPathComponent]] - .path]); + fileExistsAtPath:[storage.url + URLByAppendingPathComponent:source_key( + manifest.get())] + .path]); } } @@ -661,11 +670,10 @@ EXPECT_EQ(data.releases, 1); EXPECT_EQ(staging_directories(storage.url).count, 0u); EXPECT_FALSE([NSFileManager.defaultManager - fileExistsAtPath: - [storage.url - URLByAppendingPathComponent: - manifest->bundle_digests[manifest->path.lastPathComponent]] - .path]); + fileExistsAtPath:[storage.url + URLByAppendingPathComponent:source_key( + manifest.get())] + .path]); } } @@ -676,9 +684,8 @@ aot_data(data, @"arch_b"); auto manifest = source_manifest(true); ASSERT_TRUE(manifest.ok()); - NSURL* final = [storage.url - URLByAppendingPathComponent:manifest->bundle_digests - [manifest->path.lastPathComponent]]; + NSURL* final = + [storage.url URLByAppendingPathComponent:source_key(manifest.get())]; NSURL* root = storage.url; const NSUInteger file_count = manifest->files.count; __block int synced_files = 0; @@ -727,9 +734,8 @@ TestData data; auto manifest = source_manifest(); ASSERT_TRUE(manifest.ok()); - NSURL* final = [storage.url - URLByAppendingPathComponent:manifest->bundle_digests - [manifest->path.lastPathComponent]]; + NSURL* final = + [storage.url URLByAppendingPathComponent:source_key(manifest.get())]; { StorageFaultScope fault(operation); EXPECT_EQ( @@ -786,9 +792,8 @@ TestData data; auto manifest = source_manifest(); ASSERT_TRUE(manifest.ok()); - NSURL* final = [storage.url - URLByAppendingPathComponent:manifest->bundle_digests - [manifest->path.lastPathComponent]]; + NSURL* final = + [storage.url URLByAppendingPathComponent:source_key(manifest.get())]; { StorageFaultScope fault(StorageOperation::FullSync); fault.match_index = 2; @@ -845,9 +850,7 @@ NSDictionary* snapshot = nil; ASSERT_TRUE(asset_tree_snapshot(winner_entry, snapshot)); NSURL* root = storage.url; - NSURL* final = - [root URLByAppendingPathComponent:manifest->bundle_digests - [manifest->path.lastPathComponent]]; + NSURL* final = [root URLByAppendingPathComponent:source_key(manifest.get())]; __block int publications = 0; __block bool callback_ok = true; { @@ -892,9 +895,7 @@ NSDictionary* snapshot = nil; ASSERT_TRUE(asset_tree_snapshot(winner_entry, snapshot)); NSURL* root = storage.url; - NSURL* final = - [root URLByAppendingPathComponent:manifest->bundle_digests - [manifest->path.lastPathComponent]]; + NSURL* final = [root URLByAppendingPathComponent:source_key(manifest.get())]; __block int publications = 0; __block bool callback_ok = true; { diff --git a/backends/apple/coreai/runtime/test/coreai_storage_test.mm b/backends/apple/coreai/runtime/test/coreai_storage_test.mm index 5487f8ed864..46303ac82c1 100644 --- a/backends/apple/coreai/runtime/test/coreai_storage_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_storage_test.mm @@ -58,7 +58,8 @@ TestDirectory storage; ASSERT_NE(storage.url, nil); FileDescriptor root(open( - storage.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); ASSERT_GE(root.get(), 0); auto missing = open_storage_directory(root.get(), @"child", false); ASSERT_TRUE(missing.ok()); @@ -95,7 +96,8 @@ FileDescriptor root(open( TestDirectory storage; ASSERT_NE(storage.url, nil); FileDescriptor root(open( - storage.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); ASSERT_GE(root.get(), 0); NSURL* child = [storage.url URLByAppendingPathComponent:@"child"]; StorageFaultScope fault(StorageOperation::BeforeCreateDirectory); @@ -115,7 +117,8 @@ FileDescriptor root(open( TestDirectory storage; ASSERT_NE(storage.url, nil); FileDescriptor root(open( - storage.url.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + storage.url.fileSystemRepresentation, + O_RDONLY | O_DIRECTORY | O_CLOEXEC)); ASSERT_GE(root.get(), 0); EXPECT_EQ( open_storage_directory(root.get(), nil, false).error(), @@ -159,12 +162,13 @@ FileDescriptor root(open( ASSERT_EQ(stat(reference.fileSystemRepresentation, &expected), 0); NSURL* root = [storage.url URLByAppendingPathComponent:@"cache"]; ASSERT_TRUE(prepare_storage_root(root.path, true).ok()); - FileDescriptor directory(open( - root.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + FileDescriptor directory( + open(root.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); ASSERT_GE(directory.get(), 0); ASSERT_EQ( write_storage_file(directory.get(), @"nested/file", "x", 1), Error::Ok); - for (NSURL* entry in @[ root, [root URLByAppendingPathComponent:@"nested"] ]) { + for (NSURL* entry in + @[ root, [root URLByAppendingPathComponent:@"nested"] ]) { struct stat actual; ASSERT_EQ(stat(entry.fileSystemRepresentation, &actual), 0); EXPECT_EQ(actual.st_mode & 0777, expected.st_mode & 0777); @@ -627,9 +631,10 @@ FileDescriptor root(open( protected: TestDirectory storage; FileDescriptor root{-1}; - NSString* key = [@"a" stringByPaddingToLength:64 - withString:@"a" - startingAtIndex:0]; + NSString* hash = [@"a" stringByPaddingToLength:64 + withString:@"a" + startingAtIndex:0]; + NSString* key = [NSString stringWithFormat:@"%@-macOS_arch_b-%@", hash, hash]; void SetUp() override { ASSERT_NE(storage.url, nil); @@ -684,9 +689,8 @@ Error populate() { .fileSystemRepresentation, 0700), 0); - NSString* other = [@"b" stringByPaddingToLength:64 - withString:@"b" - startingAtIndex:0]; + NSString* other = [key stringByReplacingCharactersInRange:NSMakeRange(0, 1) + withString:@"b"]; NSString* other_path = [[@"staging" stringByAppendingPathComponent:other] stringByAppendingPathComponent:@"keep"]; ASSERT_EQ(write_storage_file(root.get(), other_path, "other", 5), Error::Ok);