diff --git a/libraries/libstratosphere/include/stratosphere/fs/fs_memory_management.hpp b/libraries/libstratosphere/include/stratosphere/fs/fs_memory_management.hpp index b7cbf5c92..c59068cb8 100644 --- a/libraries/libstratosphere/include/stratosphere/fs/fs_memory_management.hpp +++ b/libraries/libstratosphere/include/stratosphere/fs/fs_memory_management.hpp @@ -28,6 +28,94 @@ namespace ams::fs { void *Allocate(size_t size); void Deallocate(void *ptr, size_t size); + void LockAllocatorMutex(); + void UnlockAllocatorMutex(); + + void *AllocateUnsafe(size_t size); + void DeallocateUnsafe(void *ptr, size_t size); + + class AllocatorImpl { + public: + static ALWAYS_INLINE void *Allocate(size_t size) { return ::ams::fs::impl::Allocate(size); } + static ALWAYS_INLINE void *AllocateUnsafe(size_t size) { return ::ams::fs::impl::AllocateUnsafe(size); } + + static ALWAYS_INLINE void Deallocate(void *ptr, size_t size) { return ::ams::fs::impl::Deallocate(ptr, size); } + static ALWAYS_INLINE void DeallocateUnsafe(void *ptr, size_t size) { return ::ams::fs::impl::DeallocateUnsafe(ptr, size); } + + static ALWAYS_INLINE void LockAllocatorMutex() { return ::ams::fs::impl::LockAllocatorMutex(); } + static ALWAYS_INLINE void UnlockAllocatorMutex() { return ::ams::fs::impl::UnlockAllocatorMutex(); } + }; + + template + class AllocatorTemplate : public std::allocator { + public: + template + struct rebind { + using other = AllocatorTemplate; + }; + private: + bool m_allocation_failed; + private: + static ALWAYS_INLINE T *AllocateImpl(::std::size_t n) { + if constexpr (AllocateWhileLocked) { + auto * const p = Impl::AllocateUnsafe(sizeof(T) * n); + Impl::UnlockAllocatorMutex(); + return static_cast(p); + } else { + return static_cast(Impl::Allocate(sizeof(T) * n)); + } + } + public: + AllocatorTemplate() : m_allocation_failed(false) { /* ... */ } + + template + AllocatorTemplate(const AllocatorTemplate &rhs) : m_allocation_failed(rhs.IsAllocationFailed()) { /* ... */ } + + bool IsAllocationFailed() const { return m_allocation_failed; } + + [[nodiscard]] T *allocate(::std::size_t n) { + auto * const p = AllocateImpl(n); + if (AMS_UNLIKELY(p == nullptr) && n) { + m_allocation_failed = true; + } + return p; + } + + void deallocate(T *p, ::std::size_t n) { + Impl::Deallocate(p, sizeof(T) * n); + } + }; + + template + using AllocatorTemplateForAllocateShared = AllocatorTemplate; + + template class AllocatorTemplateT, typename Impl, typename... Args> + std::shared_ptr AllocateSharedImpl(Args &&... args) { + /* Try to allocate. */ + { + /* Acquire exclusive access to the allocator. */ + Impl::LockAllocatorMutex(); + + /* Check that we can allocate memory (using overestimate of 0x80 + sizeof(T)). */ + if (auto * const p = Impl::Allocate(0x80 + sizeof(T)); AMS_LIKELY(p != nullptr)) { + /* Free the memory we allocated. */ + Impl::Deallocate(p, 0x80 + sizeof(T)); + + /* Get allocator type. */ + using AllocatorType = AllocatorTemplateT; + + /* Allocate the shared pointer. */ + return std::allocate_shared(AllocatorType{}, std::forward(args)...); + } else { + /* We can't allocate. */ + Impl::UnlockAllocatorMutex(); + } + } + + /* We failed. */ + return nullptr; + } + class Deleter { private: size_t m_size; @@ -59,4 +147,9 @@ namespace ams::fs { } + template + std::shared_ptr AllocateShared(Args &&... args) { + return ::ams::fs::impl::AllocateSharedImpl(std::forward(args)...); + } + } diff --git a/libraries/libstratosphere/source/fs/fs_memory_management.cpp b/libraries/libstratosphere/source/fs/fs_memory_management.cpp index 70a71f86f..d73b91850 100644 --- a/libraries/libstratosphere/source/fs/fs_memory_management.cpp +++ b/libraries/libstratosphere/source/fs/fs_memory_management.cpp @@ -19,7 +19,7 @@ namespace ams::fs { namespace { - bool g_used_default_allocator; + constinit bool g_used_default_allocator = false; void *DefaultAllocate(size_t size) { g_used_default_allocator = true; @@ -31,7 +31,7 @@ namespace ams::fs { ams::Free(ptr); } - constinit os::SdkMutex g_lock; + constinit os::SdkMutex g_mutex; constinit AllocateFunction g_allocate_func = DefaultAllocate; constinit DeallocateFunction g_deallocate_func = DefaultDeallocate; @@ -59,26 +59,61 @@ namespace ams::fs { namespace impl { - void *Allocate(size_t size) { - void *ptr; - { - std::scoped_lock lk(g_lock); - ptr = g_allocate_func(size); - if (!util::IsAligned(reinterpret_cast(ptr), RequiredAlignment)) { - R_ABORT_UNLESS(fs::ResultAllocatorAlignmentViolation()); - } + void LockAllocatorMutex() { + g_mutex.Lock(); + } + + void UnlockAllocatorMutex() { + g_mutex.Unlock(); + } + + void *AllocateUnsafe(size_t size) { + /* Check pre-conditions. */ + AMS_ASSERT(g_mutex.IsLockedByCurrentThread()); + + /* Allocate. */ + void * const ptr = g_allocate_func(size); + + /* Check alignment. */ + if (AMS_UNLIKELY(!util::IsAligned(reinterpret_cast(ptr), RequiredAlignment))) { + R_ABORT_UNLESS(fs::ResultAllocatorAlignmentViolation()); } + + /* Return allocated pointer. */ return ptr; } - void Deallocate(void *ptr, size_t size) { - if (ptr == nullptr) { - return; - } - std::scoped_lock lk(g_lock); + void DeallocateUnsafe(void *ptr, size_t size) { + /* Check pre-conditions. */ + AMS_ASSERT(g_mutex.IsLockedByCurrentThread()); + + /* Deallocate the pointer. */ g_deallocate_func(ptr, size); } + void *Allocate(size_t size) { + /* Check pre-conditions. */ + AMS_ASSERT(g_allocate_func != nullptr); + + /* Lock the allocator. */ + std::scoped_lock lk(g_mutex); + + return AllocateUnsafe(size); + } + + void Deallocate(void *ptr, size_t size) { + /* Check pre-conditions. */ + AMS_ASSERT(g_deallocate_func != nullptr); + + /* If the pointer is non-null, deallocate it. */ + if (ptr != nullptr) { + /* Lock the allocator. */ + std::scoped_lock lk(g_mutex); + + DeallocateUnsafe(ptr, size); + } + } + } } \ No newline at end of file diff --git a/libraries/libstratosphere/source/fssystem/fssystem_nca_file_system_driver.cpp b/libraries/libstratosphere/source/fssystem/fssystem_nca_file_system_driver.cpp index 188446052..d4581f3dd 100644 --- a/libraries/libstratosphere/source/fssystem/fssystem_nca_file_system_driver.cpp +++ b/libraries/libstratosphere/source/fssystem/fssystem_nca_file_system_driver.cpp @@ -352,7 +352,7 @@ namespace ams::fssystem { R_UNLESS(storage_size > 0, fs::ResultInvalidNcaHeader()); /* Allocate a substorage. */ - *out = AllocateShared>(m_reader->GetBodyStorage(), storage_offset, storage_size, m_reader); + *out = fssystem::AllocateShared>(m_reader->GetBodyStorage(), storage_offset, storage_size, m_reader); R_UNLESS(*out != nullptr, fs::ResultAllocationFailureInAllocateShared()); return ResultSuccess();