1 /*
2  * Copyright (c) 2023 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 UTIL_SLAB_HPP
17 #define UTIL_SLAB_HPP
18 
19 #include <new>
20 #include <vector>
21 #include <mutex>
22 #ifdef FFRT_BBOX_ENABLE
23 #include <unordered_set>
24 #endif
25 #include <sys/mman.h>
26 #include "dfx/log/ffrt_log_api.h"
27 #include "spmc_queue.h"
28 #include "sync/sync.h"
29 
30 namespace ffrt {
31 const std::size_t BatchAllocSize = 32 * 1024;
32 #ifdef FFRT_BBOX_ENABLE
33 constexpr uint32_t ALLOCATOR_DESTRUCT_TIMESOUT = 1000;
34 #endif
35 
36 #ifndef FFRT_ALLOCATOR_MMAP_SIZE
37 #define FFRT_ALLOCATOR_MMAP_SIZE (8 * 1024 * 1024)
38 #endif
39 
40 template <typename T, size_t MmapSz = BatchAllocSize>
41 class SimpleAllocator {
42 public:
43     SimpleAllocator(const SimpleAllocator&) = delete;
44     SimpleAllocator(SimpleAllocator&&) = delete;
45     SimpleAllocator& operator=(const SimpleAllocator&) = delete;
46     SimpleAllocator& operator=(SimpleAllocator&&) = delete;
47     fast_mutex lock;
48 
Instance(std::size_t size = sizeof(T))49     static SimpleAllocator<T>* Instance(std::size_t size = sizeof(T))
50     {
51         static SimpleAllocator<T> ins(size);
52         return &ins;
53     }
54 
55     // NOTE: call constructor after AllocMem
AllocMem()56     static T* AllocMem()
57     {
58         return Instance()->Alloc();
59     }
60 
61     // NOTE: call destructor before FreeMem
FreeMem(T* t)62     static void FreeMem(T* t)
63     {
64         t->~T();
65         // unlock()内部lck记录锁的状态为非持有状态，析构时访问状态变量为非持有状态，则不访问实际持有的mutex
66         // return之前的lck析构不产生UAF问题，因为return之前随着root析构，锁的内存被释放
67         Instance()->free(t);
68     }
69 
70     // only used for BBOX
getUnfreedMem()71     static std::vector<void *> getUnfreedMem()
72     {
73         return Instance()->getUnfreed();
74     }
75 
LockMem()76     static void LockMem()
77     {
78         return Instance()->SimpleAllocatorLock();
79     }
80 
UnlockMem()81     static void UnlockMem()
82     {
83         return Instance()->SimpleAllocatorUnLock();
84     }
85 private:
86     SpmcQueue primaryCache;
87 #ifdef FFRT_BBOX_ENABLE
88     std::unordered_set<T*> secondaryCache;
89 #endif
90     std::size_t TSize;
91     T* basePtr = nullptr;
92     std::size_t count = 0;
93 
getUnfreed()94     std::vector<void *> getUnfreed()
95     {
96         std::vector<void *> ret;
97 #ifdef FFRT_BBOX_ENABLE
98         ret.reserve(MmapSz / TSize + secondaryCache.size());
99         char* p = reinterpret_cast<char*>(basePtr);
100         for (std::size_t i = 0; i + TSize <= MmapSz; i += TSize) {
101             if (basePtr != nullptr &&
102                 primaryCache.FindElement(reinterpret_cast<void *>(p + i)) == false) {
103                 ret.push_back(reinterpret_cast<void *>(p + i));
104             }
105         }
106         for (auto ite = secondaryCache.cbegin(); ite != secondaryCache.cend(); ite++) {
107             ret.push_back(reinterpret_cast<void *>(*ite));
108         }
109 #endif
110         return ret;
111     }
112 
SimpleAllocatorLock()113     void SimpleAllocatorLock()
114     {
115         lock.lock();
116     }
117 
SimpleAllocatorUnLock()118     void SimpleAllocatorUnLock()
119     {
120         lock.unlock();
121     }
122 
init()123     void init()
124     {
125         lock.lock();
126         if (basePtr) {
127             lock.unlock();
128             return;
129         }
130 
131         char* p = reinterpret_cast<char*>(operator new(MmapSz));
132         count = MmapSz / TSize;
133         primaryCache.Init(count);
134         for (std::size_t i = 0; i + TSize <= MmapSz; i += TSize) {
135             primaryCache.PushTail(p + i);
136         }
137         basePtr = reinterpret_cast<T*>(p);
138         lock.unlock();
139     }
140 
Alloc()141     T* Alloc()
142     {
143         // 未初始化
144         if (primaryCache.GetCapacity() == 0) {
145             init();
146         }
147         // 一级cache已耗尽
148         if (primaryCache.GetLength()) {
149             T* t = reinterpret_cast<T*>(::operator new(TSize));
150 #ifdef FFRT_BBOX_ENABLE
151             lock.lock();
152             secondaryCache.insert(t);
153             lock.unlock();
154 #endif
155             return t;
156         }
157         return static_cast<T*>(primaryCache.PopHead());
158     }
159 
free(T* t)160     void free(T* t)
161     {
162         if (basePtr && basePtr <= t &&
163             static_cast<size_t>(reinterpret_cast<uintptr_t>(t)) <
164             static_cast<size_t>(reinterpret_cast<uintptr_t>(basePtr)) + MmapSz) {
165             primaryCache.PushTail(t);
166             return;
167         }
168 
169 #ifdef FFRT_BBOX_ENABLE
170         lock.lock();
171         secondaryCache.erase(t);
172         lock.unlock();
173 #endif
174         ::operator delete(t);
175     }
176 
SimpleAllocator(std::size_t size = sizeof(T))177     SimpleAllocator(std::size_t size = sizeof(T)) : TSize(size)
178     {
179     }
~SimpleAllocator()180     ~SimpleAllocator()
181     {
182         std::unique_lock<decltype(lock)> lck(lock);
183         if (basePtr == nullptr) {
184             return;
185         }
186 #ifdef FFRT_BBOX_ENABLE
187         uint32_t try_cnt = ALLOCATOR_DESTRUCT_TIMESOUT;
188         std::size_t reserved = MmapSz / TSize;
189         while (try_cnt > 0) {
190             if (primaryCache.GetLength() == reserved && secondaryCache.size() == 0) {
191                 break;
192             }
193             lck.unlock();
194             usleep(1000);
195             try_cnt--;
196             lck.lock();
197         }
198         if (try_cnt == 0) {
199             FFRT_LOGE("clear allocator failed");
200         }
201         for (auto ite = secondaryCache.cbegin(); ite != secondaryCache.cend(); ite++) {
202             ::operator delete(*ite);
203         }
204 #endif
205         ::operator delete(basePtr);
206         FFRT_LOGI("destruct SimpleAllocator");
207     }
208 };
209 
210 template <typename T, std::size_t MmapSz = FFRT_ALLOCATOR_MMAP_SIZE>
211 class QSimpleAllocator {
212     std::size_t TSize;
213     std::size_t curAllocated;
214     std::size_t maxAllocated;
215     std::mutex lock;
216     std::vector<T*> cache;
217     uint32_t flags = MAP_ANONYMOUS | MAP_PRIVATE;
218 
expand()219     bool expand()
220     {
221         const int prot = PROT_READ | PROT_WRITE;
222         char* p = reinterpret_cast<char*>(mmap(nullptr, MmapSz, prot, flags, -1, 0));
223         if (p == (char*)MAP_FAILED) {
224             if ((flags & MAP_HUGETLB) != 0) {
225                 flags = MAP_ANONYMOUS | MAP_PRIVATE;
226                 p = reinterpret_cast<char*>(mmap(nullptr, MmapSz, prot, flags, -1, 0));
227             }
228             if (p == (char*)MAP_FAILED) {
229                 perror("mmap");
230                 return false;
231             }
232         }
233         for (std::size_t i = 0; i + TSize <= MmapSz; i += TSize) {
234             cache.push_back(reinterpret_cast<T*>(p + i));
235         }
236         return true;
237     }
238 
Alloc()239     T* Alloc()
240     {
241         T* p = nullptr;
242         lock.lock();
243         if (cache.empty()) {
244             if (!expand()) {
245                 lock.unlock();
246                 return nullptr;
247             }
248         }
249         p = cache.back();
250         ++curAllocated;
251         maxAllocated = std::max(curAllocated, maxAllocated);
252         cache.pop_back();
253         lock.unlock();
254         return p;
255     }
256 
free(T* p)257     void free(T* p)
258     {
259         lock.lock();
260         --curAllocated;
261         cache.push_back(p);
262         lock.unlock();
263     }
264 
release()265     void release()
266     {
267         T* p = nullptr;
268         lock.lock();
269         FFRT_LOGD("coroutine release with waterline %d, cur occupied %d, cached size %d",
270             maxAllocated, curAllocated, cache.size());
271         size_t reservedCnt = maxAllocated - curAllocated + 1; // reserve additional one for robustness
272         maxAllocated = curAllocated;
273         while (cache.size() > reservedCnt) {
274             p = cache.back();
275             cache.pop_back();
276             int ret = munmap(p, TSize);
277             if (ret != 0) {
278                 FFRT_LOGE("munmap failed with errno: %d", errno);
279             }
280         }
281         lock.unlock();
282     }
283 
QSimpleAllocator()284     QSimpleAllocator()
285     {
286     }
287 
288 public:
QSimpleAllocator(std::size_t size = sizeof(T))289     explicit QSimpleAllocator(std::size_t size = sizeof(T)) : curAllocated(0), maxAllocated(0)
290     {
291         std::size_t p_size = static_cast<std::size_t>(getpagesize());
292         // manually align the size to the page size
293         TSize = (size - 1 + p_size) & -p_size;
294         if (MmapSz % TSize != 0) {
295             FFRT_LOGE("MmapSz is not divisible by TSize which may cause memory leak!");
296         }
297     }
298     QSimpleAllocator(QSimpleAllocator const&) = delete;
299     void operator=(QSimpleAllocator const&) = delete;
300 
Instance(std::size_t size)301     static QSimpleAllocator<T, MmapSz>* Instance(std::size_t size)
302     {
303         static QSimpleAllocator<T, MmapSz> ins(size);
304         return &ins;
305     }
306 
AllocMem(std::size_t size = sizeof(T))307     static T* AllocMem(std::size_t size = sizeof(T))
308     {
309         return Instance(size)->Alloc();
310     }
311 
FreeMem(T* p, std::size_t size = sizeof(T))312     static void FreeMem(T* p, std::size_t size = sizeof(T))
313     {
314         Instance(size)->free(p);
315     }
316 
releaseMem(std::size_t size = sizeof(T))317     static void releaseMem(std::size_t size = sizeof(T))
318     {
319         Instance(size)->release();
320     }
321 };
322 } // namespace ffrt
323 #endif /* UTIL_SLAB_H */
324