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