1 /* 2 * Copyright (c) 2025 Huawei Device Co., Ltd. 3 * Licensed under the Apache License, Version 2.0 (the "License"); 4 * you may not use this file except in compliance with the License. 5 * You may obtain a copy of the License at 6 * 7 * http://www.apache.org/licenses/LICENSE-2.0 8 * 9 * Unless required by applicable law or agreed to in writing, software 10 * distributed under the License is distributed on an "AS IS" BASIS, 11 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 * See the License for the specific language governing permissions and 13 * limitations under the License. 14 */ 15 16 #ifndef COMMON_COMPONENTS_HEAP_ALLOCATOR_ALLOC_BUFFER_MANAGER 17 #define COMMON_COMPONENTS_HEAP_ALLOCATOR_ALLOC_BUFFER_MANAGER 18 19 #include <atomic> 20 #include <functional> 21 #include <unordered_set> 22 23 #include "common_components/heap/allocator/alloc_buffer.h" 24 #include "common_components/base/atomic_spin_lock.h" 25 #include "common_components/common/page_allocator.h" 26 #include "common_components/common/run_type.h" 27 namespace common { 28 using AllocBufferVisitor = std::function<void(AllocationBuffer& buffer)>; 29 class AllocBufferManager { 30 public: 31 using AllocBuffersSet = std::unordered_set<AllocationBuffer*, std::hash<AllocationBuffer*>, 32 std::equal_to<AllocationBuffer*>, StdContainerAllocator<AllocationBuffer*, ALLOCATOR>>; 33 using HungryBuffers = AllocBuffersSet; AllocBufferManager()34 AllocBufferManager() {} ~AllocBufferManager()35 ~AllocBufferManager() 36 { 37 for (auto* buffer : allocBuffers_) { 38 if (buffer != nullptr) { 39 delete buffer; 40 } 41 } 42 allocBuffers_.clear(); 43 }; 44 RegisterAllocBuffer(AllocationBuffer & buffer)45 void RegisterAllocBuffer(AllocationBuffer& buffer) 46 { 47 allocBufferLock_.Lock(); 48 allocBuffers_.insert(&buffer); 49 allocBufferLock_.Unlock(); 50 } 51 UnregisterAllocBuffer(AllocationBuffer & buffer)52 void UnregisterAllocBuffer(AllocationBuffer& buffer) 53 { 54 allocBufferLock_.Lock(); 55 allocBuffers_.erase(&buffer); 56 allocBufferLock_.Unlock(); 57 } 58 59 template <typename AllocBufferVisitor> VisitAllocBuffers(const AllocBufferVisitor & visitor)60 void VisitAllocBuffers(const AllocBufferVisitor& visitor) 61 { 62 allocBufferLock_.Lock(); 63 for (auto* buffer : allocBuffers_) { 64 visitor(*buffer); 65 } 66 allocBufferLock_.Unlock(); 67 } 68 AddHungryBuffer(AllocationBuffer & buffer)69 void AddHungryBuffer(AllocationBuffer& buffer) 70 { 71 std::lock_guard<std::mutex> lg(hungryBuffersLock_); 72 hungryBuffers_.insert(&buffer); 73 } SwapHungryBuffers(HungryBuffers & getBufferSet)74 void SwapHungryBuffers(HungryBuffers& getBufferSet) 75 { 76 std::lock_guard<std::mutex> lg(hungryBuffersLock_); 77 hungryBuffers_.swap(getBufferSet); 78 } 79 80 private: 81 AllocBuffersSet allocBuffers_; 82 HungryBuffers hungryBuffers_; 83 std::mutex hungryBuffersLock_; 84 AtomicSpinLock allocBufferLock_; 85 }; 86 } // namespace common 87 #endif // ~COMMON_COMPONENTS_HEAP_ALLOCATOR_ALLOC_BUFFER_MANAGER 88