diff --git a/backends/apple/coreai/CMakeLists.txt b/backends/apple/coreai/CMakeLists.txt index bb63606a1a3..3d36ecb2b78 100644 --- a/backends/apple/coreai/CMakeLists.txt +++ b/backends/apple/coreai/CMakeLists.txt @@ -11,8 +11,12 @@ endif() enable_language(OBJCXX) set(_coreai_runtime_sources - runtime/coreai_cache_root.mm runtime/coreai_fault_injection.mm - runtime/coreai_manifest.mm runtime/coreai_validation.mm + runtime/coreai_bookmarks.mm + runtime/coreai_cache_key.mm + runtime/coreai_cache_root.mm + runtime/coreai_fault_injection.mm + runtime/coreai_manifest.mm + runtime/coreai_validation.mm runtime/coreai_storage.mm ) @@ -32,9 +36,16 @@ endfunction() if(EXECUTORCH_BUILD_TESTS AND CMAKE_SYSTEM_NAME STREQUAL "Darwin") add_executable( coreai_host_test - runtime/test/coreai_host_test.mm runtime/test/coreai_manifest_test.mm - runtime/test/coreai_cache_root_test.mm runtime/test/coreai_storage_test.mm - runtime/test/coreai_filesystem_fixture.mm ${_coreai_runtime_sources} + runtime/test/coreai_host_test.mm + runtime/test/coreai_manifest_test.mm + runtime/test/coreai_cache_root_test.mm + runtime/test/coreai_storage_test.mm + runtime/test/coreai_filesystem_fixture.mm + runtime/test/coreai_cache_key_test.mm + runtime/test/coreai_bookmark_test.mm + runtime/test/coreai_bookmark_contention_test.mm + runtime/test/coreai_bookmark_fixture.mm + ${_coreai_runtime_sources} ) coreai_configure_objc_target(coreai_host_test) target_compile_definitions(coreai_host_test PRIVATE COREAI_ASSETS_TESTING=1) diff --git a/backends/apple/coreai/README.md b/backends/apple/coreai/README.md index 1a3d96adc7a..f4f90c49708 100644 --- a/backends/apple/coreai/README.md +++ b/backends/apple/coreai/README.md @@ -82,13 +82,38 @@ and enumeration use explicit policies, with a missing child distinguishable from an error. Fault injection is a separate, thread-local test interface and compiles away in production. +Each selected bundle has a model key derived from its export digest, bundle +name, and source/AOT mode. A readable variant key combines that model key with +the platform and Core AI device architecture, plus a hash of a versioned +specialization identity supplied as plain data by the SDK boundary. Function +bindings and delivery location are excluded, and no model bytes are hashed at +runtime. Each variant 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. + +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. A SHA-256 hash of the full validated cache key selects a +stripe modulo the configured pool size. Files are created on demand, so new +model versions do not grow the pool beyond 256. 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 +file or directory. Stripe files are never removed, including during cache +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. + ## Host Tests On macOS, `EXECUTORCH_BUILD_COREAI=ON` with `EXECUTORCH_BUILD_TESTS=ON` registers the standalone `coreai_host_test` executable with CTest. It tests -manifest parsing, AOT asset selection, and filesystem storage using Foundation, -without Swift or CoreAI SDK calls. The host test target is omitted for Apple -cross-builds. +manifest parsing, AOT asset selection, cache identity, bookmark I/O, lock +contention, and filesystem storage using Foundation, without Swift or CoreAI SDK +calls. The host test target is omitted for Apple cross-builds. CI builds and runs these tests on macOS 27. To use the same configuration from the repository root: diff --git a/backends/apple/coreai/runtime/coreai_bookmarks.h b/backends/apple/coreai/runtime/coreai_bookmarks.h new file mode 100644 index 00000000000..23b364b45dc --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_bookmarks.h @@ -0,0 +1,78 @@ +/* + * 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 +#include "coreai_file.h" + +namespace executorch::backends::coreai { + +// With OpenExisting, a missing root is returned unchanged for no-op +// maintenance. +runtime::Result resolve_bookmark_root( + NSString* path, + StorageCreationMode mode); + +enum class CacheEntryPersistence { + BookmarkCapable, + StagingOnly, +}; + +class CacheEntryLock final { + public: + CacheEntryLock(const CacheEntryLock&) = delete; + CacheEntryLock& operator=(const CacheEntryLock&) = delete; + + CacheEntryPersistence persistence() const { + return persistence_; + } + // Missing is nil. Empty, oversized, inaccessible, or nonregular files fail + // closed. + runtime::Result read_bookmark() const; + runtime::Error write_bookmark(NSData* data) const; + // Returns the staging parent. Source preparation appends key and bundle. + runtime::Result prepare_staging_root() const; + + private: + friend runtime::Result> + acquire_cache_entry_lock( + NSString* root, + NSString* key, + StorageCreationMode mode); + + CacheEntryLock( + int root, + int bookmarks, + int lock_file, + NSString* root_path, + NSString* key, + CacheEntryPersistence persistence); + + FileDescriptor root_; + FileDescriptor bookmarks_; + FileDescriptor lock_file_; + NSString* root_path_; + NSString* key_; + CacheEntryPersistence persistence_; +}; + +// Each acquisition opens an independent descriptor on one stable stripe. +// Different keys can share a stripe; callers must not hold multiple entry locks +// at once. Stripe files are never removed by the backend. +runtime::Result> lock_cache_entry( + NSString* root, + NSString* key); +// Missing roots and entries return a null lock without changing storage. +// Existing staging-only entries return a lock with StagingOnly persistence. +runtime::Result> lock_existing_cache_entry( + NSString* root, + NSString* key); + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_bookmarks.mm b/backends/apple/coreai/runtime/coreai_bookmarks.mm new file mode 100644 index 00000000000..8243fb77df2 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_bookmarks.mm @@ -0,0 +1,283 @@ +/* + * 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_bookmarks.h" + +#import "coreai_cache_key.h" +#import "coreai_cache_root.h" +#import "coreai_fault_injection.h" +#import "coreai_storage.h" +#import "coreai_validation.h" + +#include +#include +#include +#include +#include + +namespace executorch::backends::coreai { +namespace { +using runtime::Error; +using runtime::Result; + +constexpr size_t kBookmarkByteLimit = 8 * 1024 * 1024; +NSString* const kBookmarksDirectory = @"bookmarks"; +NSString* const kStagingDirectory = @"staging"; +NSString* const kLocksDirectory = @"locks"; +NSString* const kBookmarkExtension = @".bookmark"; + +NSString* bookmark_name(NSString* key) { + return [key stringByAppendingString:kBookmarkExtension]; +} + +// Returns -1 when the root does not exist. +Result open_root(NSString* root) { + const int fd = retry_eintr([&] { + return open( + root.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC); + }); + if (fd < 0 && errno == ENOENT) { + return -1; + } + ET_CHECK_OR_RETURN_ERROR( + fd >= 0, AccessFailed, "Cannot open Core AI assets root"); + return fd; +} + +// Returns -1 when the child does not exist in OpenExisting mode. +Result open_child(int root_fd, NSString* name, StorageCreationMode mode) { + auto opened = open_storage_directory(root_fd, name, mode); + if (!opened.ok()) { + return opened.error(); + } + FileDescriptor child(opened.get().fd); + if (opened.get().created) { + ET_CHECK_OK_OR_RETURN_ERROR(sync_storage_directory(root_fd)); + } + return child.release(); +} + +Result entry_exists(int root_fd, NSString* key) { + for (NSString* directory_name in + @[ kBookmarksDirectory, kStagingDirectory ]) { + auto opened = + open_child(root_fd, directory_name, StorageCreationMode::OpenExisting); + if (!opened.ok()) { + return opened.error(); + } + FileDescriptor directory(opened.get()); + if (directory.get() < 0) { + continue; + } + NSString* entry = + [directory_name isEqual:kBookmarksDirectory] ? bookmark_name(key) : key; + struct stat info; + if (lstat_at(directory.get(), entry.fileSystemRepresentation, &info) == 0) { + return true; + } + ET_CHECK_OR_RETURN_ERROR( + errno == ENOENT, AccessFailed, "Cannot inspect Core AI cache entry"); + } + return false; +} +} // namespace + +Result> acquire_cache_entry_lock( + NSString* root, + NSString* key, + StorageCreationMode mode) { + auto lock_name = cache_entry_lock_name(key); + if (!lock_name.ok()) { + return lock_name.error(); + } + auto resolved = resolve_bookmark_root(root, mode); + if (!resolved.ok()) { + return resolved.error(); + } + root = [resolved.get() copy]; + auto opened = open_root(root); + if (!opened.ok()) { + return opened.error(); + } + FileDescriptor directory(opened.get()); + const bool create = mode == StorageCreationMode::Create; + if (directory.get() < 0) { + ET_CHECK_OR_RETURN_ERROR( + !create, AccessFailed, "Core AI assets root disappeared"); + return std::unique_ptr(); + } + if (!create) { + auto exists = entry_exists(directory.get(), key); + if (!exists.ok()) { + return exists.error(); + } + if (!exists.get()) { + return std::unique_ptr(); + } + } + auto locks = + open_child(directory.get(), kLocksDirectory, StorageCreationMode::Create); + if (!locks.ok()) { + return locks.error(); + } + FileDescriptor lock_directory(locks.get()); + FileDescriptor lock_file(open_storage_shared_file( + lock_directory.get(), lock_name.get().fileSystemRepresentation)); + struct stat info; + ET_CHECK_OR_RETURN_ERROR( + lock_file.get() >= 0 && fstat(lock_file.get(), &info) == 0 && + S_ISREG(info.st_mode), + AccessFailed, + "Invalid Core AI cache entry lock"); + ET_CHECK_OR_RETURN_ERROR( + fault_point(FaultPoint::BeforeCacheEntryLock) == 0, + AccessFailed, + "Cannot begin Core AI cache entry lock acquisition"); + ET_CHECK_OR_RETURN_ERROR( + retry_eintr([&] { return flock(lock_file.get(), LOCK_EX); }) == 0, + AccessFailed, + "Cannot lock Core AI cache entry"); + if (!create) { + auto exists = entry_exists(directory.get(), key); + if (!exists.ok()) { + return exists.error(); + } + if (!exists.get()) { + return std::unique_ptr(); + } + } + auto bookmarks = open_child(directory.get(), kBookmarksDirectory, mode); + if (!bookmarks.ok()) { + return bookmarks.error(); + } + FileDescriptor bookmark_directory(bookmarks.get()); + const CacheEntryPersistence persistence = bookmark_directory.get() >= 0 + ? CacheEntryPersistence::BookmarkCapable + : CacheEntryPersistence::StagingOnly; + return std::unique_ptr(new CacheEntryLock( + directory.release(), + bookmark_directory.release(), + lock_file.release(), + root, + [key copy], + persistence)); +} + +Result resolve_bookmark_root( + NSString* path, + StorageCreationMode mode) { + ET_CHECK_OR_RETURN_ERROR( + nonempty_string(path) && [path hasPrefix:@"/"], + InvalidArgument, + "Core AI assets require an absolute root path"); + if (mode == StorageCreationMode::Create) { + return prepare_storage_root(path, mode); + } + struct stat info; + const int result = + retry_eintr([&] { return stat(path.fileSystemRepresentation, &info); }); + if (result < 0 && errno == ENOENT) { + return path; + } + ET_CHECK_OR_RETURN_ERROR( + result == 0 && S_ISDIR(info.st_mode), + AccessFailed, + "Invalid Core AI assets root"); + return path; +} + +CacheEntryLock::CacheEntryLock( + int root, + int bookmarks, + int lock_file, + NSString* root_path, + NSString* key, + CacheEntryPersistence persistence) + : root_(root), + bookmarks_(bookmarks), + lock_file_(lock_file), + root_path_(root_path), + key_(key), + persistence_(persistence) {} + +Result CacheEntryLock::read_bookmark() const { + if (persistence_ == CacheEntryPersistence::StagingOnly) { + return static_cast(nil); + } + FileDescriptor file(retry_eintr([&] { + return openat( + bookmarks_.get(), + bookmark_name(key_).fileSystemRepresentation, + O_RDONLY | O_NOFOLLOW | O_NONBLOCK | O_CLOEXEC); + })); + if (file.get() < 0 && errno == ENOENT) { + return static_cast(nil); + } + 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, + InvalidExternalData, + "Invalid Core AI bookmark file"); + NSMutableData* bytes = [NSMutableData dataWithLength:info.st_size]; + ET_CHECK_OR_RETURN_ERROR( + bytes != nil, MemoryAllocationFailed, "Cannot allocate Core AI bookmark"); + size_t offset = 0; + while (offset < bytes.length) { + const ssize_t count = ::read( + file.get(), + static_cast(bytes.mutableBytes) + offset, + bytes.length - offset); + if (count < 0 && errno == EINTR) { + continue; + } + ET_CHECK_OR_RETURN_ERROR( + count > 0, AccessFailed, "Cannot read Core AI bookmark"); + offset += count; + } + return bytes; +} + +Error CacheEntryLock::write_bookmark(NSData* data) const { + ET_CHECK_OR_RETURN_ERROR( + [data isKindOfClass:NSData.class] && data.length > 0 && + data.length <= kBookmarkByteLimit, + InvalidExternalData, + "Invalid Core AI bookmark data"); + ET_CHECK_OR_RETURN_ERROR( + persistence_ == CacheEntryPersistence::BookmarkCapable, + InvalidArgument, + "Cannot write through a staging-only Core AI cache entry lock"); + return publish_storage_data(bookmarks_.get(), bookmark_name(key_), data); +} + +Result CacheEntryLock::prepare_staging_root() const { + auto staging = + open_child(root_.get(), kStagingDirectory, StorageCreationMode::Create); + if (!staging.ok()) { + return staging.error(); + } + FileDescriptor directory(staging.get()); + return [root_path_ stringByAppendingPathComponent:kStagingDirectory]; +} + +Result> lock_cache_entry( + NSString* root, + NSString* key) { + return acquire_cache_entry_lock(root, key, StorageCreationMode::Create); +} + +Result> lock_existing_cache_entry( + NSString* root, + NSString* key) { + return acquire_cache_entry_lock(root, key, StorageCreationMode::OpenExisting); +} + +} // namespace executorch::backends::coreai 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..4ea30a6e1c0 --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache_key.h @@ -0,0 +1,39 @@ +/* + * 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 { + +class Manifest; +class SelectedBundle; + +bool is_cache_key(NSString* key); +NSString* cache_key_model_prefix(NSString* key); + +// Model identity excludes device and specialization so maintenance can find all +// variants of one exported bundle. +runtime::Result cache_model_key( + const Manifest& manifest, + NSString* bundle_name); +runtime::Result cache_model_key(const SelectedBundle& bundle); +// specialization_identity is stable, versioned plain data supplied by the SDK +// boundary. This module does not depend on SDK types. +runtime::Result cache_variant_key( + const SelectedBundle& bundle, + NSString* platform, + NSString* architecture, + NSString* specialization_identity); + +constexpr size_t kCacheEntryLockStripeCount = 256; +runtime::Result cache_entry_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..87b0f9a9f2b --- /dev/null +++ b/backends/apple/coreai/runtime/coreai_cache_key.mm @@ -0,0 +1,190 @@ +/* + * 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_manifest.h" +#import "coreai_validation.h" + +namespace executorch::backends::coreai { +namespace { +using runtime::Result; + +constexpr NSUInteger kHashLength = 2 * CC_SHA256_DIGEST_LENGTH; +constexpr NSUInteger kFixedLength = 2 * (kHashLength + 1); +constexpr NSUInteger kBookmarkExtensionLength = sizeof(".bookmark") - 1; +constexpr NSUInteger kMaxKeyLength = NAME_MAX - kBookmarkExtensionLength; + +static_assert(kCacheEntryLockStripeCount > 0); +static_assert(kCacheEntryLockStripeCount <= 256); + +bool valid_device(NSString* device) { + if (![device isKindOfClass:NSString.class] || + device.length > kMaxKeyLength - 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 encode_json(id value) { + NSError* error = nil; + NSData* data = [NSJSONSerialization dataWithJSONObject:value + options:0 + error:&error]; + ET_CHECK_OR_RETURN_ERROR( + data != nil, + InvalidExternalData, + "Cannot encode Core AI cache identity: %s", + error.localizedDescription.UTF8String); + return data; +} + +Result hash_data(NSData* data) { + ET_CHECK_OR_RETURN_ERROR( + data != nil && data.length <= UINT_MAX, + InvalidExternalData, + "Cannot hash Core AI cache identity"); + 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]; +} + +Result +model_key(NSString* digest, NSString* bundle_name, bool aot_compiled) { + ET_CHECK_OR_RETURN_ERROR( + lowercase_hex(digest, kHashLength) && relative_path(bundle_name) && + ![bundle_name containsString:@"/"], + DelegateInvalidCompatibility, + "Invalid Core AI bundle identity"); + auto encoded = encode_json( + @[ @"1", digest, aot_compiled ? @"aot" : @"source", bundle_name ]); + if (!encoded.ok()) { + return encoded.error(); + } + return hash_data(encoded.get()); +} + +Result utf8_data(NSString* value) { + ET_CHECK_OR_RETURN_ERROR( + nonempty_string(value), + DelegateInvalidCompatibility, + "Missing Core AI cache specialization"); + NSData* data = [value dataUsingEncoding:NSUTF8StringEncoding]; + ET_CHECK_OR_RETURN_ERROR( + data != nil, InvalidExternalData, "Cannot encode Core AI specialization"); + return data; +} +} // namespace + +bool is_cache_key(NSString* key) { + if (![key isKindOfClass:NSString.class] || key.length <= kFixedLength || + key.length > kMaxKeyLength) { + return false; + } + const NSUInteger options_start = key.length - kHashLength; + if ([key characterAtIndex:kHashLength] != '-' || + [key characterAtIndex:options_start - 1] != '-') { + return false; + } + NSString* model = [key substringToIndex:kHashLength]; + NSString* options = [key substringFromIndex:options_start]; + if (!lowercase_hex(model, kHashLength) || + !lowercase_hex(options, kHashLength)) { + 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_name) { + return model_key( + manifest.bundle_digests[bundle_name], bundle_name, manifest.aot_compiled); +} + +Result cache_model_key(const SelectedBundle& bundle) { + return model_key(bundle.digest, bundle.bundle_name, bundle.aot_compiled); +} + +Result cache_variant_key( + const SelectedBundle& bundle, + NSString* platform, + NSString* architecture, + NSString* specialization_identity) { + auto model = cache_model_key(bundle); + if (!model.ok()) { + return model.error(); + } + ET_CHECK_OR_RETURN_ERROR( + nonempty_string(platform) && nonempty_string(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"); + auto specialization = utf8_data(specialization_identity); + if (!specialization.ok()) { + return specialization.error(); + } + auto options = hash_data(specialization.get()); + if (!options.ok()) { + return options.error(); + } + return [NSString + stringWithFormat:@"%@-%@-%@", model.get(), device, options.get()]; +} + +Result cache_entry_lock_name(NSString* key) { + ET_CHECK_OR_RETURN_ERROR( + is_cache_key(key), InvalidArgument, "Invalid Core AI cache key"); + NSData* encoded = [key dataUsingEncoding:NSUTF8StringEncoding]; + auto hash = hash_data(encoded); + if (!hash.ok()) { + return hash.error(); + } + unsigned int first_byte = 0; + const bool parsed = + [[NSScanner scannerWithString:[hash.get() substringToIndex:2]] + scanHexInt:&first_byte]; + ET_CHECK_OR_RETURN_ERROR( + parsed, InvalidExternalData, "Cannot derive Core AI cache lock stripe"); + const size_t stripe = first_byte % kCacheEntryLockStripeCount; + return [NSString stringWithFormat:@"stripe-v1-%02zx.lock", stripe]; +} + +} // namespace executorch::backends::coreai diff --git a/backends/apple/coreai/runtime/coreai_fault_injection.h b/backends/apple/coreai/runtime/coreai_fault_injection.h index cba60c5dfb8..84947db16c8 100644 --- a/backends/apple/coreai/runtime/coreai_fault_injection.h +++ b/backends/apple/coreai/runtime/coreai_fault_injection.h @@ -17,6 +17,7 @@ enum class FaultPoint { FullSync, Rename, BeforeCreateDirectory, + BeforeCacheEntryLock, }; #if defined(COREAI_ASSETS_TESTING) && COREAI_ASSETS_TESTING diff --git a/backends/apple/coreai/runtime/test/coreai_bookmark_contention_test.mm b/backends/apple/coreai/runtime/test/coreai_bookmark_contention_test.mm new file mode 100644 index 00000000000..d45614c2f38 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_bookmark_contention_test.mm @@ -0,0 +1,287 @@ +/* + * 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 +#include "coreai_bookmark_fixture.h" +#include "coreai_fault_scope.h" +#include "coreai_filesystem_fixture.h" + +using namespace executorch::runtime; +using namespace executorch::backends::coreai; +using namespace executorch::backends::coreai::testing; + +namespace { +constexpr auto kWaitTimeout = std::chrono::seconds(5); + +struct CacheEntryWaiter { + std::future attempting; + std::future>> result; +}; + +CacheEntryWaiter start_cache_entry_waiter( + NSString* path, + NSString* key, + bool existing_only = false) { + std::promise attempting; + std::promise>> result; + CacheEntryWaiter waiter{attempting.get_future(), result.get_future()}; + // Detached promise-backed work cannot block test teardown if acquisition + // stalls after a failed assertion. + std::thread([path, + key, + existing_only, + attempting = std::move(attempting), + result = std::move(result)]() mutable { + @autoreleasepool { + StorageFaultScope scope(FaultPoint::BeforeCacheEntryLock); + scope.armed = false; + auto* signal = &attempting; + scope.observe = ^(FaultPoint point) { + if (point == FaultPoint::BeforeCacheEntryLock) { + signal->set_value(); + } + }; + result.set_value( + existing_only ? lock_existing_cache_entry(path, key) + : lock_cache_entry(path, key)); + } + }).detach(); + return waiter; +} + +class CoreAIBookmarkContentionTest : public ::testing::Test { + protected: + TestDirectory root; + std::unique_ptr bundle; + NSString* key = nil; + + void SetUp() override { + ASSERT_NE(root.url, nil); + auto selected = bookmark_bundle(); + ASSERT_TRUE(selected.ok()); + bundle = std::make_unique(selected.get()); + auto result = test_cache_key(*bundle); + ASSERT_TRUE(result.ok()); + key = result.get(); + } + + NSURL* lock_url() { + return test_cache_lock_url(root.url, key); + } +}; + +TEST_F( + CoreAIBookmarkContentionTest, + ThreadContentionSerializesSameKeyButNotDifferentStripes) { + auto holder = lock_cache_entry(root.url.path, key); + ASSERT_TRUE(holder.ok()); + ASSERT_NE(holder.get(), nullptr); + auto waiter = start_cache_entry_waiter(root.url.path, key); + ASSERT_EQ( + waiter.attempting.wait_for(kWaitTimeout), std::future_status::ready); + EXPECT_EQ( + waiter.result.wait_for(std::chrono::milliseconds(100)), + std::future_status::timeout); + + auto independent_key = different_cache_key(*bundle, key, false); + ASSERT_TRUE(independent_key.ok()); + auto independent = + start_cache_entry_waiter(root.url.path, independent_key.get()); + ASSERT_EQ( + independent.result.wait_for(kWaitTimeout), std::future_status::ready); + auto other = independent.result.get(); + ASSERT_TRUE(other.ok()); + ASSERT_NE(other.get(), nullptr); + + holder.get().reset(); + ASSERT_EQ(waiter.result.wait_for(kWaitTimeout), std::future_status::ready); + auto acquired = waiter.result.get(); + ASSERT_TRUE(acquired.ok()); + EXPECT_NE(acquired.get(), nullptr); +} + +TEST_F(CoreAIBookmarkContentionTest, CollidingKeysSerializeAcrossThreads) { + auto collision = different_cache_key(*bundle, key, true); + ASSERT_TRUE(collision.ok()); + auto holder = lock_cache_entry(root.url.path, key); + ASSERT_TRUE(holder.ok()); + auto waiter = start_cache_entry_waiter(root.url.path, collision.get()); + ASSERT_EQ( + waiter.attempting.wait_for(kWaitTimeout), std::future_status::ready); + EXPECT_EQ( + waiter.result.wait_for(std::chrono::milliseconds(100)), + std::future_status::timeout); + holder.get().reset(); + ASSERT_EQ(waiter.result.wait_for(kWaitTimeout), std::future_status::ready); + auto acquired = waiter.result.get(); + ASSERT_TRUE(acquired.ok()); + EXPECT_NE(acquired.get(), nullptr); +} + +TEST_F( + CoreAIBookmarkContentionTest, + MissingCollidingLookupDoesNotWaitOrCreateAnotherLock) { + auto collision = different_cache_key(*bundle, key, true); + ASSERT_TRUE(collision.ok()); + auto holder = lock_cache_entry(root.url.path, key); + ASSERT_TRUE(holder.ok()); + auto waiter = start_cache_entry_waiter(root.url.path, collision.get(), true); + EXPECT_EQ(waiter.result.wait_for(kWaitTimeout), std::future_status::ready); + holder.get().reset(); + auto absent = waiter.result.get(); + ASSERT_TRUE(absent.ok()); + EXPECT_EQ(absent.get(), nullptr); + NSArray* names = [NSFileManager.defaultManager + contentsOfDirectoryAtPath:[root.url URLByAppendingPathComponent:@"locks"] + .path + error:nil]; + ASSERT_NE(names, nil); + EXPECT_EQ(names.count, 1u); +} + +TEST_F(CoreAIBookmarkContentionTest, ExistingWaiterRechecksEvictedEntry) { + auto holder = lock_cache_entry(root.url.path, key); + ASSERT_TRUE(holder.ok()); + NSData* bytes = [@"bookmark" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_EQ(holder.get()->write_bookmark(bytes), Error::Ok); + auto waiter = start_cache_entry_waiter(root.url.path, key, true); + ASSERT_EQ( + waiter.attempting.wait_for(kWaitTimeout), std::future_status::ready); + ASSERT_EQ( + waiter.result.wait_for(std::chrono::milliseconds(100)), + std::future_status::timeout); + NSURL* bookmark = test_bookmark_url(root.url, key); + ASSERT_EQ(unlink(bookmark.fileSystemRepresentation), 0); + holder.get().reset(); + ASSERT_EQ(waiter.result.wait_for(kWaitTimeout), std::future_status::ready); + auto absent = waiter.result.get(); + ASSERT_TRUE(absent.ok()); + EXPECT_EQ(absent.get(), nullptr); +} + +TEST_F( + CoreAIBookmarkContentionTest, + WaiterOpensBookmarkDirectoryAfterAcquiringLock) { + auto holder = lock_cache_entry(root.url.path, key); + ASSERT_TRUE(holder.ok()); + auto waiter = start_cache_entry_waiter(root.url.path, key); + ASSERT_EQ( + waiter.attempting.wait_for(kWaitTimeout), std::future_status::ready); + ASSERT_EQ( + waiter.result.wait_for(std::chrono::milliseconds(100)), + std::future_status::timeout); + NSURL* bookmarks = [root.url URLByAppendingPathComponent:@"bookmarks"]; + ASSERT_EQ(rmdir(bookmarks.fileSystemRepresentation), 0); + ASSERT_EQ(mkdir(bookmarks.fileSystemRepresentation, 0700), 0); + holder.get().reset(); + ASSERT_EQ(waiter.result.wait_for(kWaitTimeout), std::future_status::ready); + auto acquired = waiter.result.get(); + ASSERT_TRUE(acquired.ok()); + ASSERT_NE(acquired.get(), nullptr); + NSData* bytes = [@"bookmark" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_EQ(acquired.get()->write_bookmark(bytes), Error::Ok); + EXPECT_TRUE([[NSData dataWithContentsOfURL:test_bookmark_url(root.url, key)] + isEqual:bytes]); +} + +TEST_F( + CoreAIBookmarkContentionTest, + ProcessContentionAllowsIndependentKeyAndRecoversAfterHolderDeath) { + BookmarkChild holder; + ASSERT_TRUE(holder.spawn(@"--bookmark-hold", root.url.path, @"same")); + ASSERT_TRUE(holder.receive('L')); + NSURL* file = lock_url(); + struct stat before {}; + ASSERT_EQ(stat(file.fileSystemRepresentation, &before), 0); + BookmarkChild waiter; + ASSERT_TRUE(waiter.spawn(@"--bookmark-lock", root.url.path, @"same")); + ASSERT_TRUE(waiter.receive('B')); + ASSERT_TRUE(waiter.expect_blocked()); + BookmarkChild independent; + ASSERT_TRUE(independent.spawn(@"--bookmark-lock", root.url.path, @"other")); + ASSERT_TRUE(independent.receive('L')); + ASSERT_TRUE(independent.expect_exit(0)); + ASSERT_TRUE(holder.kill_and_reap()); + ASSERT_TRUE(waiter.receive('L')); + ASSERT_TRUE(waiter.expect_exit(0)); + struct stat after {}; + ASSERT_EQ(stat(file.fileSystemRepresentation, &after), 0); + EXPECT_EQ(before.st_dev, after.st_dev); + EXPECT_EQ(before.st_ino, after.st_ino); +} + +TEST_F( + CoreAIBookmarkContentionTest, + CollidingProcessWaiterRecoversAfterHolderDeath) { + BookmarkChild holder; + ASSERT_TRUE(holder.spawn(@"--bookmark-hold", root.url.path, @"same")); + ASSERT_TRUE(holder.receive('L')); + struct stat before {}; + ASSERT_EQ(lstat(lock_url().fileSystemRepresentation, &before), 0); + BookmarkChild waiter; + ASSERT_TRUE(waiter.spawn(@"--bookmark-lock", root.url.path, @"collision")); + ASSERT_TRUE(waiter.receive('B')); + ASSERT_TRUE(waiter.expect_blocked()); + ASSERT_TRUE(holder.kill_and_reap()); + ASSERT_TRUE(waiter.receive('L')); + ASSERT_TRUE(waiter.expect_exit(0)); + struct stat after {}; + ASSERT_EQ(lstat(lock_url().fileSystemRepresentation, &after), 0); + EXPECT_EQ(before.st_dev, after.st_dev); + EXPECT_EQ(before.st_ino, after.st_ino); +} + +TEST_F(CoreAIBookmarkContentionTest, LockInventoryIsBoundedAndStable) { + NSMutableDictionary* representatives = + [NSMutableDictionary dictionary]; + for (NSUInteger index = 0; + index < 100000 && representatives.count < kCacheEntryLockStripeCount; + ++index) { + auto candidate = test_cache_key( + *bundle, [NSString stringWithFormat:@"inventory-%lu", index]); + ASSERT_TRUE(candidate.ok()); + auto name = cache_entry_lock_name(candidate.get()); + ASSERT_TRUE(name.ok()); + representatives[name.get()] = candidate.get(); + } + ASSERT_EQ(representatives.count, kCacheEntryLockStripeCount); + + NSMutableDictionary* inodes = + [NSMutableDictionary dictionary]; + for (NSString* name in representatives) { + auto locked = lock_cache_entry(root.url.path, representatives[name]); + ASSERT_TRUE(locked.ok()); + NSURL* file = [[root.url URLByAppendingPathComponent:@"locks"] + URLByAppendingPathComponent:name]; + struct stat info {}; + ASSERT_EQ(lstat(file.fileSystemRepresentation, &info), 0); + ASSERT_TRUE(S_ISREG(info.st_mode)); + EXPECT_EQ(info.st_size, 0); + inodes[name] = @(info.st_ino); + } + NSURL* locks = [root.url URLByAppendingPathComponent:@"locks"]; + NSArray* children = + [NSFileManager.defaultManager contentsOfDirectoryAtPath:locks.path + error:nil]; + ASSERT_EQ(children.count, kCacheEntryLockStripeCount); + + for (NSString* name in representatives) { + auto locked = lock_cache_entry(root.url.path, representatives[name]); + ASSERT_TRUE(locked.ok()); + struct stat info {}; + NSURL* file = [locks URLByAppendingPathComponent:name]; + ASSERT_EQ(lstat(file.fileSystemRepresentation, &info), 0); + EXPECT_EQ(info.st_ino, inodes[name].longLongValue); + } +} + +} // namespace diff --git a/backends/apple/coreai/runtime/test/coreai_bookmark_fixture.h b/backends/apple/coreai/runtime/test/coreai_bookmark_fixture.h new file mode 100644 index 00000000000..90e43b5b617 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_bookmark_fixture.h @@ -0,0 +1,59 @@ +/* + * 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 +#include +#include +#include +#include "coreai_bookmarks.h" +#include "coreai_cache_key.h" +#include "coreai_manifest.h" + +namespace executorch::backends::coreai::testing { + +runtime::Result bookmark_bundle(bool aot = false); +runtime::Result test_cache_key(const SelectedBundle& bundle); +runtime::Result test_cache_key( + const SelectedBundle& bundle, + NSString* specialization_identity); +runtime::Result different_cache_key( + const SelectedBundle& bundle, + NSString* key, + bool same_stripe); +NSURL* test_bookmark_url(NSURL* root, NSString* key); +NSURL* test_cache_lock_url(NSURL* root, NSString* key); +runtime::Result saved_bookmark(NSURL* root, NSString* key); +runtime::Result sized_bookmark(size_t size, uint32_t tag = 1); + +void set_bookmark_test_executable(const char* path); +// Returns -1 for normal test invocations, otherwise a child-mode exit code. +int bookmark_child_mode(int argc, char** argv); + +class BookmarkChild { + public: + BookmarkChild() = default; + ~BookmarkChild(); + BookmarkChild(const BookmarkChild&) = delete; + BookmarkChild& operator=(const BookmarkChild&) = delete; + + ::testing::AssertionResult + spawn(NSString* mode, NSString* root, NSString* argument); + ::testing::AssertionResult receive(char expected); + ::testing::AssertionResult expect_blocked(); + ::testing::AssertionResult expect_exit(int expected); + ::testing::AssertionResult kill_and_reap(); + + private: + pid_t pid_ = -1; + int signal_fd_ = -1; + ::testing::AssertionResult reap(int& status); +}; + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_bookmark_fixture.mm b/backends/apple/coreai/runtime/test/coreai_bookmark_fixture.mm new file mode 100644 index 00000000000..d40351abeb8 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_bookmark_fixture.mm @@ -0,0 +1,404 @@ +/* + * 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_bookmark_fixture.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "coreai_fault_scope.h" +#include "coreai_manifest_fixture.h" + +extern char** environ; + +namespace executorch::backends::coreai::testing { +namespace { +using runtime::Error; +using runtime::Result; +using ::testing::AssertionFailure; +using ::testing::AssertionSuccess; +using Clock = std::chrono::steady_clock; +constexpr auto kChildTimeout = std::chrono::seconds(10); +constexpr auto kContentionObservation = std::chrono::milliseconds(100); +constexpr int kChildSignalFD = 3; +std::string executable; + +int remaining_ms(Clock::time_point deadline) { + const auto remaining = std::chrono::duration_cast( + deadline - Clock::now()); + return remaining.count() <= 0 ? 0 : static_cast(remaining.count()); +} + +bool send_signal_byte(int fd, char byte) { + const auto deadline = Clock::now() + kChildTimeout; + while (const int remaining = remaining_ms(deadline)) { + pollfd entry{fd, POLLOUT, 0}; + const int ready = poll(&entry, 1, remaining); + if (ready < 0 && errno == EINTR) { + continue; + } + if (ready <= 0 || !(entry.revents & POLLOUT)) { + return false; + } + const ssize_t written = write(fd, &byte, 1); + if (written == 1) { + return true; + } + if (written < 0 && (errno == EINTR || errno == EAGAIN)) { + continue; + } + return false; + } + return false; +} +} // namespace + +Result bookmark_bundle(bool aot) { + NSData* encoded = encode(aot ? aot_manifest_dict() : manifest_dict()); + if (encoded == nil) { + return Error::InvalidExternalData; + } + auto parsed = parse_manifest(encoded); + if (!parsed.ok()) { + return parsed.error(); + } + return select_assets(parsed.get(), @"arch_b", @"macOS"); +} + +Result test_cache_key(const SelectedBundle& bundle) { + return test_cache_key(bundle, @"default"); +} + +Result test_cache_key( + const SelectedBundle& bundle, + NSString* specialization_identity) { + return cache_variant_key( + bundle, @"macOS", @"arch_b", specialization_identity); +} + +Result different_cache_key( + const SelectedBundle& bundle, + NSString* key, + bool same_stripe) { + auto original_name = cache_entry_lock_name(key); + if (!original_name.ok()) { + return original_name.error(); + } + for (NSUInteger index = 0; index < 10000; ++index) { + NSString* identity = + [NSString stringWithFormat:@"specialization-%lu", (unsigned long)index]; + auto candidate = test_cache_key(bundle, identity); + if (!candidate.ok()) { + return candidate.error(); + } + if ([candidate.get() isEqual:key]) { + continue; + } + auto candidate_name = cache_entry_lock_name(candidate.get()); + if (!candidate_name.ok()) { + return candidate_name.error(); + } + if ([candidate_name.get() isEqual:original_name.get()] == same_stripe) { + return candidate.get(); + } + } + return Error::NotFound; +} + +NSURL* test_bookmark_url(NSURL* root, NSString* key) { + return [[root URLByAppendingPathComponent:@"bookmarks"] + URLByAppendingPathComponent:[key stringByAppendingString:@".bookmark"]]; +} + +NSURL* test_cache_lock_url(NSURL* root, NSString* key) { + auto name = cache_entry_lock_name(key); + EXPECT_TRUE(name.ok()); + if (!name.ok()) { + return nil; + } + return [[root URLByAppendingPathComponent:@"locks"] + URLByAppendingPathComponent:name.get()]; +} + +Result saved_bookmark(NSURL* root, NSString* key) { + auto locked = lock_cache_entry(root.path, key); + if (!locked.ok()) { + return locked.error(); + } + return locked.get()->read_bookmark(); +} + +Result sized_bookmark(size_t size, uint32_t tag) { + if (size < sizeof(tag)) { + return Error::InvalidArgument; + } + NSMutableData* bytes = [NSMutableData dataWithLength:size]; + if (bytes == nil) { + return Error::MemoryAllocationFailed; + } + std::memcpy(bytes.mutableBytes, &tag, sizeof(tag)); + return [bytes copy]; +} + +void set_bookmark_test_executable(const char* path) { + executable = path == nullptr ? "" : path; +} + +int bookmark_child_mode(int argc, char** argv) { + if (argc < 2 || std::strncmp(argv[1], "--bookmark-", 11) != 0) { + return -1; + } + if (argc != 5) { + return 1; + } + const bool hold = std::strcmp(argv[1], "--bookmark-hold") == 0; + const bool lock = std::strcmp(argv[1], "--bookmark-lock") == 0; + const bool other = std::strcmp(argv[3], "other") == 0; + const bool same = std::strcmp(argv[3], "same") == 0; + const bool collision = std::strcmp(argv[3], "collision") == 0; + if ((!hold && !lock) || (!other && !same && !collision)) { + return 2; + } + char* end = nullptr; + errno = 0; + const long value = std::strtol(argv[4], &end, 10); + if (errno != 0 || end == argv[4] || *end != '\0' || value < 0 || + value > INT_MAX) { + return 3; + } + FileDescriptor signal_fd(static_cast(value)); + if (fcntl(signal_fd.get(), F_GETFD) < 0) { + return 4; + } + signal(SIGPIPE, SIG_IGN); + NSString* root = [NSString stringWithUTF8String:argv[2]]; + if (root == nil) { + return 5; + } + auto bundle = bookmark_bundle(); + if (!bundle.ok()) { + return 6; + } + auto key = test_cache_key(bundle.get()); + if (!key.ok()) { + return 7; + } + NSString* selected_key = key.get(); + if (other || collision) { + auto changed = different_cache_key(bundle.get(), key.get(), collision); + if (!changed.ok()) { + return 8; + } + selected_key = changed.get(); + } + StorageFaultScope scope(FaultPoint::BeforeCacheEntryLock); + scope.armed = false; + const int notify_fd = signal_fd.get(); + scope.observe = ^(FaultPoint point) { + if (lock && (same || collision) && + point == FaultPoint::BeforeCacheEntryLock && + !send_signal_byte(notify_fd, 'B')) { + _exit(12); + } + }; + auto locked = lock_cache_entry(root, selected_key); + if (!locked.ok()) { + return 9; + } + if (!send_signal_byte(signal_fd.get(), 'L')) { + return 10; + } + if (close(signal_fd.release()) != 0) { + return 11; + } + if (hold) { + // The parent owns this PID and always kills and reaps it. + for (;;) { + pause(); + } + } + return 0; +} + +BookmarkChild::~BookmarkChild() { + if (pid_ > 0) { + const auto cleaned = kill_and_reap(); + if (!cleaned) { + ADD_FAILURE() << cleaned.message(); + } + } + if (signal_fd_ >= 0 && close(signal_fd_) != 0) { + ADD_FAILURE() << "Cannot close child signal descriptor: " << errno; + } +} + +::testing::AssertionResult +BookmarkChild::spawn(NSString* mode, NSString* root, NSString* argument) { + if (pid_ > 0 || signal_fd_ >= 0 || executable.empty() || root == nil) { + return AssertionFailure() << "Invalid bookmark child setup"; + } + int descriptors[2]; + if (pipe(descriptors) != 0) { + return AssertionFailure() << "pipe: " << errno; + } + FileDescriptor reader(descriptors[0]); + FileDescriptor writer(descriptors[1]); + for (int fd : descriptors) { + if (fcntl(fd, F_SETFD, FD_CLOEXEC) != 0 || + fcntl(fd, F_SETFL, O_NONBLOCK) != 0) { + return AssertionFailure() << "fcntl pipe: " << errno; + } + } + posix_spawn_file_actions_t actions; + int error = posix_spawn_file_actions_init(&actions); + if (error != 0) { + return AssertionFailure() << "spawn actions: " << error; + } + error = + posix_spawn_file_actions_adddup2(&actions, writer.get(), kChildSignalFD); + std::string descriptor = std::to_string(kChildSignalFD); + char* arguments[] = { + const_cast(executable.c_str()), + const_cast(mode.UTF8String), + const_cast(root.fileSystemRepresentation), + const_cast(argument.UTF8String), + descriptor.data(), + nullptr}; + if (error == 0) { + error = + posix_spawn(&pid_, arguments[0], &actions, nullptr, arguments, environ); + } + const int destroyed = posix_spawn_file_actions_destroy(&actions); + if (error != 0) { + pid_ = -1; + return AssertionFailure() << "posix_spawn: " << error; + } + signal_fd_ = reader.release(); + if (destroyed != 0) { + return AssertionFailure() << "destroy spawn actions: " << destroyed; + } + return AssertionSuccess(); +} + +::testing::AssertionResult BookmarkChild::receive(char expected) { + const auto deadline = Clock::now() + kChildTimeout; + while (const int remaining = remaining_ms(deadline)) { + pollfd entry{signal_fd_, POLLIN, 0}; + const int ready = poll(&entry, 1, remaining); + if (ready < 0 && errno == EINTR) { + continue; + } + if (ready < 0) { + return AssertionFailure() << "poll: " << errno; + } + if (ready == 0) { + break; + } + if (!(entry.revents & POLLIN)) { + return AssertionFailure() + << "Child closed signal pipe before " << expected; + } + char byte = 0; + const ssize_t count = read(signal_fd_, &byte, 1); + if (count < 0 && (errno == EINTR || errno == EAGAIN)) { + continue; + } + if (count != 1 || byte != expected) { + return AssertionFailure() + << "Expected child signal " << expected << ", got " << int(byte); + } + return AssertionSuccess(); + } + return AssertionFailure() + << "Timed out waiting for child signal " << expected; +} + +::testing::AssertionResult BookmarkChild::expect_blocked() { + const auto deadline = Clock::now() + kContentionObservation; + while (const int remaining = remaining_ms(deadline)) { + pollfd entry{signal_fd_, POLLIN, 0}; + const int ready = poll(&entry, 1, remaining); + if (ready < 0 && errno == EINTR) { + continue; + } + if (ready < 0) { + return AssertionFailure() << "poll: " << errno; + } + if (ready > 0) { + return AssertionFailure() << "Child completed while lock was held"; + } + return AssertionSuccess(); + } + return AssertionSuccess(); +} + +::testing::AssertionResult BookmarkChild::reap(int& status) { + if (pid_ <= 0) { + return AssertionFailure() << "No owned child to reap"; + } + const auto deadline = Clock::now() + kChildTimeout; + while (const int remaining = remaining_ms(deadline)) { + const pid_t result = waitpid(pid_, &status, WNOHANG); + if (result == pid_) { + pid_ = -1; + return AssertionSuccess(); + } + if (result < 0 && errno == EINTR) { + continue; + } + if (result < 0) { + return AssertionFailure() << "waitpid: " << errno; + } + poll(nullptr, 0, std::min(remaining, 10)); + } + return AssertionFailure() << "Timed out reaping child " << pid_; +} + +::testing::AssertionResult BookmarkChild::expect_exit(int expected) { + int status = 0; + auto result = reap(status); + if (!result) { + return result; + } + if (!WIFEXITED(status) || WEXITSTATUS(status) != expected) { + return AssertionFailure() + << "Expected child exit " << expected << ", wait status " << status; + } + return AssertionSuccess(); +} + +::testing::AssertionResult BookmarkChild::kill_and_reap() { + if (pid_ <= 0) { + return AssertionFailure() << "No owned child to kill"; + } + const int killed = kill(pid_, SIGKILL); + const int error = errno; + int status = 0; + auto result = reap(status); + if (!result) { + return result; + } + if (killed != 0 || !WIFSIGNALED(status) || WTERMSIG(status) != SIGKILL) { + return AssertionFailure() + << "Expected SIGKILL, kill errno " << (killed == 0 ? 0 : error) + << ", wait status " << status; + } + return AssertionSuccess(); +} + +} // namespace executorch::backends::coreai::testing diff --git a/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm b/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm new file mode 100644 index 00000000000..949c5b4ade8 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_bookmark_test.mm @@ -0,0 +1,449 @@ +/* + * 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 "coreai_bookmark_fixture.h" +#include "coreai_fault_scope.h" +#include "coreai_filesystem_fixture.h" +#include "coreai_storage.h" + +using namespace executorch::runtime; +using namespace executorch::backends::coreai; +using namespace executorch::backends::coreai::testing; + +namespace { +constexpr size_t kMiB = 1024 * 1024; + +class CoreAIBookmarkPathTest : public ::testing::Test { + protected: + TestDirectory root; + std::unique_ptr bundle; + NSString* key = nil; + + void SetUp() override { + ASSERT_NE(root.url, nil); + auto selected = bookmark_bundle(); + ASSERT_TRUE(selected.ok()); + bundle = std::make_unique(selected.get()); + auto result = test_cache_key(*bundle); + ASSERT_TRUE(result.ok()); + key = result.get(); + } + + NSURL* lock_url() { + return test_cache_lock_url(root.url, key); + } +}; + +class CoreAIBookmarkIOTest : public CoreAIBookmarkPathTest { + protected: + std::unique_ptr locked; + NSURL* url = nil; + + void SetUp() override { + CoreAIBookmarkPathTest::SetUp(); + if (HasFatalFailure()) { + return; + } + auto result = lock_cache_entry(root.url.path, key); + ASSERT_TRUE(result.ok()); + locked = std::move(result.get()); + ASSERT_NE(locked, nullptr); + ASSERT_EQ(locked->persistence(), CacheEntryPersistence::BookmarkCapable); + url = test_bookmark_url(root.url, key); + } +}; + +::testing::AssertionResult bookmark_bytes_equal( + const CacheEntryLock& lock, + NSData* expected) { + auto result = lock.read_bookmark(); + if (!result.ok()) { + return ::testing::AssertionFailure() + << "Bookmark read error " << int(result.error()); + } + if (expected == nil ? result.get() != nil + : ![result.get() isEqual:expected]) { + return ::testing::AssertionFailure() << "Bookmark bytes differ"; + } + return ::testing::AssertionSuccess(); +} + +::testing::AssertionResult no_temporary_bookmarks(NSURL* root) { + NSURL* directory = [root URLByAppendingPathComponent:@"bookmarks"]; + FileDescriptor fd(open( + directory.fileSystemRepresentation, O_RDONLY | O_DIRECTORY | O_CLOEXEC)); + if (fd.get() < 0) { + return ::testing::AssertionFailure() + << "open bookmark directory: " << errno; + } + auto children = storage_children(fd.get()); + if (!children.ok()) { + return ::testing::AssertionFailure() + << "Cannot enumerate bookmark directory"; + } + for (NSString* name in children.get()) { + if ([name hasPrefix:@".tmp-"]) { + return ::testing::AssertionFailure() << "Leaked " << name.UTF8String; + } + } + return ::testing::AssertionSuccess(); +} + +TEST_F(CoreAIBookmarkIOTest, MissingBookmarkIsAnEmptySuccessfulRead) { + EXPECT_TRUE(bookmark_bytes_equal(*locked, nil)); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:url.path]); +} + +TEST_F(CoreAIBookmarkIOTest, RoundTripsMaximumBookmarkAndReopens) { + auto bytes = sized_bookmark(8 * kMiB); + ASSERT_TRUE(bytes.ok()); + ASSERT_EQ(locked->write_bookmark(bytes.get()), Error::Ok); + EXPECT_TRUE(bookmark_bytes_equal(*locked, bytes.get())); + EXPECT_TRUE([[NSData dataWithContentsOfURL:url] isEqual:bytes.get()]); + locked.reset(); + auto saved = saved_bookmark(root.url, key); + ASSERT_TRUE(saved.ok()); + EXPECT_TRUE([saved.get() isEqual:bytes.get()]); +} + +TEST_F(CoreAIBookmarkIOTest, InvalidWritesPreserveExistingBookmark) { + auto maximum = sized_bookmark(8 * kMiB); + ASSERT_TRUE(maximum.ok()); + ASSERT_EQ(locked->write_bookmark(maximum.get()), Error::Ok); + auto oversized = sized_bookmark(8 * kMiB + 1); + ASSERT_TRUE(oversized.ok()); + for (NSData* invalid : {[NSData data], oversized.get()}) { + SCOPED_TRACE(invalid.length); + EXPECT_EQ(locked->write_bookmark(invalid), Error::InvalidExternalData); + EXPECT_TRUE([[NSData dataWithContentsOfURL:url] isEqual:maximum.get()]); + EXPECT_TRUE(bookmark_bytes_equal(*locked, maximum.get())); + EXPECT_TRUE(no_temporary_bookmarks(root.url)); + } +} + +TEST_F(CoreAIBookmarkIOTest, RejectsEmptyStoredBookmark) { + FileDescriptor file(open( + url.fileSystemRepresentation, + O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, + 0600)); + ASSERT_GE(file.get(), 0); + EXPECT_EQ(locked->read_bookmark().error(), Error::InvalidExternalData); +} + +TEST_F(CoreAIBookmarkIOTest, RejectsSymlinkWithoutReadingTarget) { + NSURL* target = [root.url URLByAppendingPathComponent:@"target"]; + NSData* bytes = [@"target bytes" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([bytes writeToURL:target atomically:NO]); + ASSERT_EQ( + symlink(target.fileSystemRepresentation, url.fileSystemRepresentation), + 0); + EXPECT_EQ(locked->read_bookmark().error(), Error::AccessFailed); + EXPECT_TRUE([[NSData dataWithContentsOfURL:target] isEqual:bytes]); +} + +TEST_F(CoreAIBookmarkIOTest, RejectsNonregularStoredBookmarks) { + ASSERT_EQ(mkdir(url.fileSystemRepresentation, 0700), 0); + EXPECT_EQ(locked->read_bookmark().error(), Error::InvalidExternalData); + ASSERT_EQ(rmdir(url.fileSystemRepresentation), 0); + ASSERT_EQ(mkfifo(url.fileSystemRepresentation, 0600), 0); + EXPECT_EQ(locked->read_bookmark().error(), Error::InvalidExternalData); +} + +TEST_F(CoreAIBookmarkIOTest, RejectsInaccessibleBookmark) { + if (geteuid() == 0) { + GTEST_SKIP() << "Root can bypass file read permissions"; + } + NSData* bytes = [@"bookmark" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_EQ(locked->write_bookmark(bytes), Error::Ok); + ASSERT_EQ(chmod(url.fileSystemRepresentation, 0000), 0); + EXPECT_EQ(locked->read_bookmark().error(), Error::AccessFailed); + ASSERT_EQ(chmod(url.fileSystemRepresentation, 0600), 0); + EXPECT_TRUE(bookmark_bytes_equal(*locked, bytes)); +} + +TEST_F(CoreAIBookmarkIOTest, RejectsOversizedStoredBookmark) { + FileDescriptor file(open( + url.fileSystemRepresentation, + O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, + 0600)); + ASSERT_GE(file.get(), 0); + ASSERT_EQ(ftruncate(file.get(), 8 * kMiB + 1), 0); + ASSERT_EQ(close(file.release()), 0); + EXPECT_EQ(locked->read_bookmark().error(), Error::InvalidExternalData); + struct stat info {}; + ASSERT_EQ(lstat(url.fileSystemRepresentation, &info), 0); + EXPECT_EQ(info.st_size, static_cast(8 * kMiB + 1)); +} + +TEST_F(CoreAIBookmarkIOTest, ExistingLookupReadsWithoutPreparingStaging) { + NSData* bytes = [@"bookmark" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_EQ(locked->write_bookmark(bytes), Error::Ok); + locked.reset(); + auto existing = lock_existing_cache_entry(root.url.path, key); + ASSERT_TRUE(existing.ok()); + ASSERT_NE(existing.get(), nullptr); + EXPECT_EQ( + existing.get()->persistence(), CacheEntryPersistence::BookmarkCapable); + EXPECT_TRUE(bookmark_bytes_equal(*existing.get(), bytes)); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:[root.url URLByAppendingPathComponent:@"staging"].path]); +} + +TEST_F(CoreAIBookmarkIOTest, PreparesSharedStagingParentOnDemand) { + NSString* expected = [root.url URLByAppendingPathComponent:@"staging"].path; + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:expected]); + auto staging = locked->prepare_staging_root(); + 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 = locked->prepare_staging_root(); + ASSERT_TRUE(repeated.ok()); + EXPECT_TRUE([repeated.get() isEqual:expected]); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:[expected stringByAppendingPathComponent:key]]); +} + +TEST_F(CoreAIBookmarkIOTest, CreatesDirectoriesWithDefaultPermissions) { + NSURL* reference = [root.url URLByAppendingPathComponent:@"reference"]; + ASSERT_EQ(mkdir(reference.fileSystemRepresentation, 0777), 0); + struct stat expected; + ASSERT_EQ(stat(reference.fileSystemRepresentation, &expected), 0); + ASSERT_TRUE(locked->prepare_staging_root().ok()); + for (NSString* name in @[ @"locks", @"bookmarks", @"staging" ]) { + NSURL* directory = [root.url URLByAppendingPathComponent:name]; + struct stat actual; + ASSERT_EQ(stat(directory.fileSystemRepresentation, &actual), 0); + EXPECT_EQ(actual.st_mode & 0777, expected.st_mode & 0777); + } +} + +TEST_F(CoreAIBookmarkIOTest, ReusesDirectoriesWithoutCreation) { + ASSERT_TRUE(locked->prepare_staging_root().ok()); + ASSERT_EQ( + locked->write_bookmark( + [@"bookmark" dataUsingEncoding:NSUTF8StringEncoding]), + Error::Ok); + locked.reset(); + StorageFaultScope fault(FaultPoint::BeforeCreateDirectory); + auto created = lock_cache_entry(root.url.path, key); + ASSERT_TRUE(created.ok()); + ASSERT_NE(created.get(), nullptr); + EXPECT_TRUE(created.get()->prepare_staging_root().ok()); + created.get().reset(); + auto existing = lock_existing_cache_entry(root.url.path, key); + ASSERT_TRUE(existing.ok()); + ASSERT_NE(existing.get(), nullptr); + EXPECT_TRUE(existing.get()->prepare_staging_root().ok()); + EXPECT_EQ(fault.matches, 0); +} + +TEST_F(CoreAIBookmarkIOTest, RecreatesMissingStagingParent) { + auto staging = locked->prepare_staging_root(); + ASSERT_TRUE(staging.ok()); + ASSERT_EQ(rmdir(staging.get().fileSystemRepresentation), 0); + StorageFaultScope fault(FaultPoint::BeforeCreateDirectory); + EXPECT_EQ(locked->prepare_staging_root().error(), Error::AccessFailed); + EXPECT_EQ(fault.hits, 1); + auto recreated = locked->prepare_staging_root(); + ASSERT_TRUE(recreated.ok()); + EXPECT_TRUE([recreated.get() isEqual:staging.get()]); +} + +TEST_F(CoreAIBookmarkIOTest, AcceptsConcurrentStagingCreation) { + NSURL* staging = [root.url URLByAppendingPathComponent:@"staging"]; + StorageFaultScope fault(FaultPoint::BeforeCreateDirectory); + fault.armed = false; + __block int creations = 0; + fault.observe = ^(FaultPoint point) { + if (point == FaultPoint::BeforeCreateDirectory) { + ++creations; + EXPECT_EQ(mkdir(staging.fileSystemRepresentation, 0700), 0); + } + }; + EXPECT_TRUE(locked->prepare_staging_root().ok()); + EXPECT_EQ(creations, 1); +} + +TEST_F(CoreAIBookmarkIOTest, RejectsInvalidStagingParent) { + TestDirectory target; + ASSERT_NE(target.url, nil); + NSURL* staging = [root.url URLByAppendingPathComponent:@"staging"]; + ASSERT_EQ( + symlink( + target.url.fileSystemRepresentation, + staging.fileSystemRepresentation), + 0); + EXPECT_EQ(locked->prepare_staging_root().error(), Error::AccessFailed); + ASSERT_EQ(unlink(staging.fileSystemRepresentation), 0); + NSData* bytes = [@"not a directory" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_TRUE([bytes writeToURL:staging atomically:NO]); + EXPECT_EQ(locked->prepare_staging_root().error(), Error::AccessFailed); + EXPECT_TRUE([[NSData dataWithContentsOfURL:staging] isEqual:bytes]); +} + +TEST_F(CoreAIBookmarkPathTest, AllowsTemporaryRootsAndRejectsInvalidKeys) { + auto allowed = lock_cache_entry(root.url.path, key); + ASSERT_TRUE(allowed.ok()); + ASSERT_NE(allowed.get(), nullptr); + EXPECT_FALSE(lock_cache_entry(root.url.path, @"../escape").ok()); +} + +TEST_F(CoreAIBookmarkPathTest, SnapshotsMutableRootAtAcquisition) { + TestDirectory other; + ASSERT_NE(other.url, nil); + NSMutableString* path = [root.url.path mutableCopy]; + auto locked = lock_cache_entry(path, key); + ASSERT_TRUE(locked.ok()); + ASSERT_NE(locked.get(), nullptr); + [path setString:other.url.path]; + + auto staging = locked.get()->prepare_staging_root(); + ASSERT_TRUE(staging.ok()); + NSString* expected = [root.url URLByAppendingPathComponent:@"staging"].path; + NSString* unrelated = [other.url URLByAppendingPathComponent:@"staging"].path; + EXPECT_TRUE([NSFileManager.defaultManager fileExistsAtPath:expected]); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:unrelated]); + EXPECT_TRUE([staging.get() isEqual:expected]) + << "Returned staging path: " << staging.get().UTF8String + << "; descriptor-pinned staging path: " << expected.UTF8String; +} + +TEST_F(CoreAIBookmarkPathTest, ExcludesRootWithoutCreatingStaging) { + auto locked = lock_cache_entry(root.url.path, key); + ASSERT_TRUE(locked.ok()); + EXPECT_TRUE(backup_excluded(root.url)); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:[root.url URLByAppendingPathComponent:@"staging"].path]); +} + +TEST_F(CoreAIBookmarkPathTest, MissingLookupLeavesRootAndEntryUntouched) { + ASSERT_TRUE(set_backup_excluded(root.url, false)); + NSString* missing = [root.url URLByAppendingPathComponent:@"missing"].path; + for (NSString* path in @[ root.url.path, missing ]) { + SCOPED_TRACE(path.UTF8String); + auto resolved = + resolve_bookmark_root(path, StorageCreationMode::OpenExisting); + ASSERT_TRUE(resolved.ok()); + EXPECT_TRUE([resolved.get() isEqual:path]); + auto absent = lock_existing_cache_entry(path, key); + ASSERT_TRUE(absent.ok()); + EXPECT_EQ(absent.get(), nullptr); + } + NSArray* contents = + [NSFileManager.defaultManager contentsOfDirectoryAtPath:root.url.path + error:nil]; + ASSERT_NE(contents, nil); + EXPECT_EQ(contents.count, 0u); + EXPECT_FALSE(backup_excluded(root.url)); +} + +TEST_F(CoreAIBookmarkPathTest, MissingLookupIgnoresOtherEntries) { + auto other_key = different_cache_key(*bundle, key, false); + ASSERT_TRUE(other_key.ok()); + auto other = lock_cache_entry(root.url.path, other_key.get()); + ASSERT_TRUE(other.ok()); + NSData* bytes = [@"other" dataUsingEncoding:NSUTF8StringEncoding]; + ASSERT_EQ(other.get()->write_bookmark(bytes), Error::Ok); + auto absent = lock_existing_cache_entry(root.url.path, key); + ASSERT_TRUE(absent.ok()); + EXPECT_EQ(absent.get(), nullptr); + EXPECT_FALSE([NSFileManager.defaultManager fileExistsAtPath:lock_url().path]); + EXPECT_TRUE(bookmark_bytes_equal(*other.get(), bytes)); +} + +TEST_F(CoreAIBookmarkPathTest, StagingOnlyLookupHasExplicitPersistence) { + NSURL* staged = [[root.url URLByAppendingPathComponent:@"staging"] + URLByAppendingPathComponent:key]; + ASSERT_TRUE([NSFileManager.defaultManager createDirectoryAtURL:staged + withIntermediateDirectories:YES + attributes:nil + error:nil]); + auto existing = lock_existing_cache_entry(root.url.path, key); + ASSERT_TRUE(existing.ok()); + ASSERT_NE(existing.get(), nullptr); + EXPECT_EQ(existing.get()->persistence(), CacheEntryPersistence::StagingOnly); + EXPECT_TRUE(bookmark_bytes_equal(*existing.get(), nil)); + NSData* bytes = [@"bookmark" dataUsingEncoding:NSUTF8StringEncoding]; + EXPECT_EQ(existing.get()->write_bookmark(bytes), Error::InvalidArgument); + EXPECT_TRUE([NSFileManager.defaultManager fileExistsAtPath:lock_url().path]); + EXPECT_FALSE([NSFileManager.defaultManager + fileExistsAtPath:[root.url URLByAppendingPathComponent:@"bookmarks"] + .path]); +} + +TEST_F(CoreAIBookmarkPathTest, ResolvesAndCreatesRootWithBackupExclusion) { + ASSERT_TRUE(set_backup_excluded(root.url, false)); + NSString* path = [root.url URLByAppendingPathComponent:@"nested/assets"].path; + auto resolved = resolve_bookmark_root(path, StorageCreationMode::Create); + ASSERT_TRUE(resolved.ok()); + EXPECT_TRUE([resolved.get() isEqual:path]); + EXPECT_TRUE(backup_excluded([NSURL fileURLWithPath:resolved.get()])); + EXPECT_FALSE(backup_excluded(root.url)); + auto existing = + resolve_bookmark_root(path, StorageCreationMode::OpenExisting); + ASSERT_TRUE(existing.ok()); + EXPECT_TRUE([existing.get() isEqual:path]); +} + +TEST_F(CoreAIBookmarkPathTest, RejectsInvalidResolvedRoots) { + for (StorageCreationMode mode : + {StorageCreationMode::OpenExisting, StorageCreationMode::Create}) { + for (NSString* path in + @[ @"relative", @"~", @"~/cache", @"~nobody/cache" ]) { + SCOPED_TRACE(path.UTF8String); + EXPECT_EQ( + resolve_bookmark_root(path, mode).error(), Error::InvalidArgument); + } + EXPECT_EQ(resolve_bookmark_root(nil, mode).error(), Error::InvalidArgument); + } + NSURL* file = [root.url URLByAppendingPathComponent:@"file"]; + ASSERT_TRUE([[NSData data] writeToURL:file atomically:NO]); + EXPECT_EQ( + resolve_bookmark_root(file.path, StorageCreationMode::OpenExisting) + .error(), + Error::AccessFailed); + EXPECT_EQ( + resolve_bookmark_root( + [file.path stringByAppendingPathComponent:@"child"], + StorageCreationMode::OpenExisting) + .error(), + Error::AccessFailed); +} + +TEST_F(CoreAIBookmarkPathTest, RejectsInvalidStripeFiles) { + NSData* bytes = [@"untouched target" dataUsingEncoding:NSUTF8StringEncoding]; + NSURL* target = [root.url URLByAppendingPathComponent:@"target"]; + ASSERT_TRUE([bytes writeToURL:target atomically:NO]); + for (int kind = 0; kind < 3; ++kind) { + TestDirectory isolated; + ASSERT_NE(isolated.url, nil); + NSURL* locks = [isolated.url URLByAppendingPathComponent:@"locks"]; + ASSERT_EQ(mkdir(locks.fileSystemRepresentation, 0700), 0); + NSURL* stripe = test_cache_lock_url(isolated.url, key); + if (kind == 0) { + ASSERT_EQ( + symlink( + target.fileSystemRepresentation, stripe.fileSystemRepresentation), + 0); + } else if (kind == 1) { + ASSERT_EQ(mkdir(stripe.fileSystemRepresentation, 0700), 0); + } else { + ASSERT_EQ(mkfifo(stripe.fileSystemRepresentation, 0600), 0); + } + EXPECT_EQ( + lock_cache_entry(isolated.url.path, key).error(), Error::AccessFailed); + EXPECT_TRUE([[NSData dataWithContentsOfURL:target] isEqual:bytes]); + } +} + +} // namespace diff --git a/backends/apple/coreai/runtime/test/coreai_cache_key_test.mm b/backends/apple/coreai/runtime/test/coreai_cache_key_test.mm new file mode 100644 index 00000000000..368e0fc6fa4 --- /dev/null +++ b/backends/apple/coreai/runtime/test/coreai_cache_key_test.mm @@ -0,0 +1,209 @@ +/* + * 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 +#include "coreai_bookmark_fixture.h" +#include "coreai_manifest_fixture.h" + +using namespace executorch::runtime; +using namespace executorch::backends::coreai; +using namespace executorch::backends::coreai::testing; + +namespace { + +Result selected_bundle(NSMutableDictionary* dictionary) { + auto manifest = parse_manifest(encode(dictionary)); + if (!manifest.ok()) { + return manifest.error(); + } + return select_assets(manifest.get(), @"arch_b", @"macOS"); +} + +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(CoreAICacheKeyIdentityTest, IgnoresBindingsAndDeliveryLocation) { + auto baseline = bookmark_bundle(); + ASSERT_TRUE(baseline.ok()); + auto expected = test_cache_key(baseline.get()); + ASSERT_TRUE(expected.ok()); + + auto dictionary = manifest_dict(); + dictionary[@"hash"] = @"cd"; + dictionary[@"path"] = @"cd/model.aimodel"; + dictionary[@"function"] = @"another_function"; + dictionary[@"input_names"] = @[ @"different_input" ]; + dictionary[@"output_names"] = @[ @"different_output" ]; + auto changed = selected_bundle(dictionary); + ASSERT_TRUE(changed.ok()); + auto actual = test_cache_key(changed.get()); + ASSERT_TRUE(actual.ok()); + EXPECT_TRUE([actual.get() isEqual:expected.get()]); +} + +TEST(CoreAICacheKeyIdentityTest, SeparatesExportedBundleIdentity) { + auto baseline = bookmark_bundle(); + ASSERT_TRUE(baseline.ok()); + auto expected = test_cache_key(baseline.get()); + ASSERT_TRUE(expected.ok()); + + auto digest_changed = manifest_dict(); + digest_changed[@"bundle_digests"] = + @{@"model.aimodel" : fixture_identity('2')}; + auto reweighted = selected_bundle(digest_changed); + ASSERT_TRUE(reweighted.ok()); + + auto aot = bookmark_bundle(true); + ASSERT_TRUE(aot.ok()); + for (const SelectedBundle* bundle : {&reweighted.get(), &aot.get()}) { + auto actual = test_cache_key(*bundle); + ASSERT_TRUE(actual.ok()); + EXPECT_FALSE([actual.get() isEqual:expected.get()]); + } +} + +TEST(CoreAICacheKeyIdentityTest, SeparatesBundleNameWithSameDigestAndMode) { + auto dictionary = aot_manifest_dict(); + dictionary[@"bundle_digests"] = @{ + @"model.arch_a.aimodelc" : fixture_identity('b'), + @"model.arch_b.aimodelc" : fixture_identity('b') + }; + auto manifest = parse_manifest(encode(dictionary)); + ASSERT_TRUE(manifest.ok()); + auto first = select_assets(manifest.get(), @"arch_a", @"macOS"); + auto second = select_assets(manifest.get(), @"arch_b", @"macOS"); + ASSERT_TRUE(first.ok()); + ASSERT_TRUE(second.ok()); + auto first_key = cache_model_key(first.get()); + auto second_key = cache_model_key(second.get()); + ASSERT_TRUE(first_key.ok()); + ASSERT_TRUE(second_key.ok()); + EXPECT_FALSE([first_key.get() isEqual:second_key.get()]); +} + +TEST(CoreAICacheKeyIdentityTest, SeparatesDeviceAndSpecializationIdentity) { + auto bundle = bookmark_bundle(); + ASSERT_TRUE(bundle.ok()); + auto baseline = test_cache_key(bundle.get()); + auto platform = + cache_variant_key(bundle.get(), @"iOS", @"arch_b", @"default"); + auto architecture = + cache_variant_key(bundle.get(), @"macOS", @"arch_a", @"default"); + auto specialization = cache_variant_key( + bundle.get(), + @"macOS", + @"arch_b", + @"{\"version\":1,\"compute\":[\"cpu\"]}"); + ASSERT_TRUE(baseline.ok()); + for (auto* changed : {&platform, &architecture, &specialization}) { + ASSERT_TRUE(changed->ok()); + EXPECT_FALSE([changed->get() isEqual:baseline.get()]); + EXPECT_TRUE([cache_key_model_prefix(changed->get()) + isEqual:cache_key_model_prefix(baseline.get())]); + } +} + +TEST(CoreAICacheKeyIdentityTest, RejectsIncompletePlainIdentity) { + auto bundle = bookmark_bundle(); + ASSERT_TRUE(bundle.ok()); + EXPECT_FALSE( + cache_variant_key(bundle.get(), nil, @"arch_b", @"default").ok()); + EXPECT_FALSE(cache_variant_key(bundle.get(), @"macOS", @"", @"default").ok()); + EXPECT_FALSE(cache_variant_key(bundle.get(), @"macOS", @"arch_b", nil).ok()); + EXPECT_FALSE( + cache_variant_key(bundle.get(), @"mac_OS", @"arch_b", @"default").ok()); +} + +TEST(CoreAICacheKeyFormatTest, 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(CoreAICacheKeyFormatTest, 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(CoreAICacheKeyFormatTest, 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(CoreAICacheKeyFormatTest, RequiresExactHashBoundaries) { + NSString* key = key_for_device(@"iOS_h17p"); + EXPECT_FALSE(is_cache_key(nil)); + EXPECT_FALSE(is_cache_key(@"")); + EXPECT_FALSE(is_cache_key([key substringToIndex:64])); + 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:@"_"])); +} + +TEST(CoreAICacheKeyFormatTest, LockStripesStayWithinConfiguredPool) { + auto bundle = bookmark_bundle(); + ASSERT_TRUE(bundle.ok()); + NSMutableSet* names = [NSMutableSet set]; + for (NSUInteger index = 0; index < 4096; ++index) { + auto key = test_cache_key( + bundle.get(), [NSString stringWithFormat:@"identity-%lu", index]); + ASSERT_TRUE(key.ok()); + auto name = cache_entry_lock_name(key.get()); + ASSERT_TRUE(name.ok()); + ASSERT_TRUE([name.get() hasPrefix:@"stripe-v1-"]); + ASSERT_TRUE([name.get() hasSuffix:@".lock"]); + NSString* encoded = + [name.get() substringWithRange:NSMakeRange(10, name.get().length - 15)]; + unsigned int stripe = 0; + ASSERT_TRUE([[NSScanner scannerWithString:encoded] scanHexInt:&stripe]); + EXPECT_LT(stripe, kCacheEntryLockStripeCount); + [names addObject:name.get()]; + } + EXPECT_GT(names.count, 1u); + EXPECT_LE(names.count, kCacheEntryLockStripeCount); +} + +} // namespace diff --git a/backends/apple/coreai/runtime/test/coreai_host_test.mm b/backends/apple/coreai/runtime/test/coreai_host_test.mm index 5c857dfddf4..679321c248b 100644 --- a/backends/apple/coreai/runtime/test/coreai_host_test.mm +++ b/backends/apple/coreai/runtime/test/coreai_host_test.mm @@ -9,10 +9,17 @@ #import #include #include +#include "coreai_bookmark_fixture.h" int main(int argc, char** argv) { executorch::runtime::runtime_init(); @autoreleasepool { + executorch::backends::coreai::testing::set_bookmark_test_executable( + argv[0]); + const int child_result = + executorch::backends::coreai::testing::bookmark_child_mode(argc, argv); + if (child_result >= 0) + return child_result; ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); }