Skip to content

Commit 36d19dc

Browse files
ColdwingsColdwings
committed
FlexLockfree queues
1 parent 0747f45 commit 36d19dc

2 files changed

Lines changed: 129 additions & 70 deletions

File tree

‎common/lockfree_queue.h‎

Lines changed: 108 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,11 @@ limitations under the License.
1919
#include <atomic>
2020
#include <cstddef>
2121
#include <cstdint>
22+
#include <cstdlib>
2223
#include <cstring>
2324
#include <memory>
2425
#include <thread>
26+
#include <type_traits>
2527
#include <utility>
2628
#ifndef __aarch64__
2729
#include <immintrin.h>
@@ -31,38 +33,6 @@ limitations under the License.
3133
#include <photon/common/utility.h>
3234
#include <photon/thread/thread.h>
3335

34-
#define size_t uint64_t
35-
36-
template <size_t x>
37-
struct Capacity_2expN {
38-
constexpr static size_t capacity = Capacity_2expN<(x >> 1)>::capacity << 1;
39-
constexpr static size_t mask = capacity - 1;
40-
constexpr static size_t shift = Capacity_2expN<(x >> 1)>::shift + 1;
41-
constexpr static size_t lshift = Capacity_2expN<(x >> 1)>::lshift - 1;
42-
43-
static_assert(shift + lshift == sizeof(size_t) * 8, "...");
44-
};
45-
46-
template <size_t x>
47-
constexpr size_t Capacity_2expN<x>::capacity;
48-
49-
template <size_t x>
50-
constexpr size_t Capacity_2expN<x>::mask;
51-
52-
template <>
53-
struct Capacity_2expN<0> {
54-
constexpr static size_t capacity = 2;
55-
constexpr static size_t mask = 1;
56-
constexpr static size_t shift = 1;
57-
constexpr static size_t lshift = 8 * sizeof(size_t) - shift;
58-
};
59-
60-
template <>
61-
struct Capacity_2expN<1> : public Capacity_2expN<0> {};
62-
63-
template <>
64-
struct Capacity_2expN<2> : public Capacity_2expN<0> {};
65-
6636
struct PauseBase {};
6737

6838
struct CPUPause : PauseBase {
@@ -111,15 +81,32 @@ class LockfreeRingQueueBase {
11181
#endif
11282

11383
constexpr static size_t CACHELINE_SIZE = 64;
84+
constexpr static size_t SLOTS_NUM =
85+
N ? 1UL << (8 * sizeof(size_t) - __builtin_clzll(N - 1)) : 0;
11486

115-
constexpr static size_t capacity = Capacity_2expN<N>::capacity;
116-
constexpr static size_t mask = Capacity_2expN<N>::mask;
117-
constexpr static size_t shift = Capacity_2expN<N>::shift;
118-
constexpr static size_t lshift = Capacity_2expN<N>::lshift;
87+
const size_t capacity;
88+
const size_t mask;
89+
const size_t shift;
90+
const size_t lshift;
11991

12092
alignas(CACHELINE_SIZE) std::atomic<size_t> tail{0};
12193
alignas(CACHELINE_SIZE) std::atomic<size_t> head{0};
12294

95+
// For only flexible queue
96+
explicit LockfreeRingQueueBase(size_t c)
97+
: capacity(c > 1 ? 1UL << (8 * sizeof(size_t) - __builtin_clzll(c - 1))
98+
: 2),
99+
mask(capacity - 1),
100+
shift(__builtin_ctzll(capacity)),
101+
lshift(8 * sizeof(size_t) - shift) {}
102+
103+
// For only deterministic queue
104+
LockfreeRingQueueBase()
105+
: capacity(N > 1 ? 1UL << (8 * sizeof(size_t) - __builtin_clzll(N - 1))
106+
: 2),
107+
mask(capacity - 1),
108+
shift(__builtin_ctzll(capacity)),
109+
lshift(8 * sizeof(size_t) - shift) {}
123110
bool empty() {
124111
return check_empty(head.load(std::memory_order_relaxed),
125112
tail.load(std::memory_order_relaxed));
@@ -156,6 +143,45 @@ class LockfreeRingQueueBase {
156143
size_t turn(size_t x) const { return x >> shift; }
157144
};
158145

146+
template <typename T, typename Q>
147+
class FlexQueue : public Q {
148+
protected:
149+
constexpr static size_t CACHELINE_SIZE = 64;
150+
FlexQueue(size_t c) : Q(c) {}
151+
~FlexQueue() {}
152+
153+
public:
154+
static FlexQueue* create(size_t c) {
155+
size_t space = required_space(c);
156+
void* ptr = nullptr;
157+
posix_memalign(&ptr, CACHELINE_SIZE, required_space(c));
158+
if (!ptr) return nullptr;
159+
memset(ptr, 0, space);
160+
return new (ptr) FlexQueue(c);
161+
}
162+
163+
static size_t required_space(size_t c) {
164+
size_t capacity =
165+
c > 1 ? 1UL << (8 * sizeof(size_t) - __builtin_clzll(c - 1)) : 2;
166+
return sizeof(Q) + capacity * sizeof(T);
167+
}
168+
169+
static FlexQueue* init_on(void* ptr, size_t c) {
170+
if (!ptr) return nullptr;
171+
return new (ptr) FlexQueue(c);
172+
}
173+
174+
static void deinit(FlexQueue* q) {
175+
if (!q) return;
176+
q->~FlexQueue();
177+
}
178+
179+
static void destroy(FlexQueue* q) {
180+
deinit(q);
181+
free(q);
182+
}
183+
};
184+
159185
// !!NOTICE: DO NOT USE LockfreeMPMCRingQueue in IPC
160186
// This queue may block if one of processes crashed during push / pop
161187
// Do not use as IPC base. Use it to collect data and send by
@@ -169,12 +195,14 @@ class LockfreeMPMCRingQueue : public LockfreeRingQueueBase<T, N> {
169195
using Base::idx;
170196
using Base::tail;
171197

198+
public:
172199
struct alignas(Base::CACHELINE_SIZE) packedslot {
173200
T data;
174201
std::atomic<MarkType> mark{0};
175202
};
176-
177-
packedslot slots[Base::capacity];
203+
204+
protected:
205+
packedslot slots[Base::SLOTS_NUM];
178206

179207
MarkType this_turn_write(const uint64_t x) const {
180208
return (Base::turn(x) << 1) + 1;
@@ -188,10 +216,14 @@ class LockfreeMPMCRingQueue : public LockfreeRingQueueBase<T, N> {
188216
return Base::turn(x) << 1;
189217
}
190218

219+
explicit LockfreeMPMCRingQueue(size_t c) : Base(c) {}
220+
191221
public:
192222
using Base::empty;
193223
using Base::full;
194224

225+
explicit LockfreeMPMCRingQueue() : Base() {}
226+
195227
bool push(const T& x) {
196228
auto t = tail.load(std::memory_order_acquire);
197229
for (;;) {
@@ -268,6 +300,11 @@ class LockfreeMPMCRingQueue : public LockfreeRingQueueBase<T, N> {
268300
}
269301
};
270302

303+
template <typename T>
304+
using FlexLockfreeMPMCRingQueue =
305+
FlexQueue<typename LockfreeMPMCRingQueue<T, 0>::packedslot,
306+
LockfreeMPMCRingQueue<T, 0>>;
307+
271308
template <typename T, size_t N>
272309
class LockfreeBatchMPMCRingQueue : public LockfreeRingQueueBase<T, N> {
273310
protected:
@@ -283,33 +320,24 @@ class LockfreeBatchMPMCRingQueue : public LockfreeRingQueueBase<T, N> {
283320
alignas(Base::CACHELINE_SIZE) std::atomic<uint64_t> write_head;
284321
alignas(Base::CACHELINE_SIZE) std::atomic<uint64_t> read_tail;
285322

286-
T slots[Base::capacity];
287-
288-
uint64_t this_turn_write(const uint64_t x) const {
289-
return (Base::turn(x) << 1) + 1;
290-
}
291-
292-
uint64_t this_turn_read(const uint64_t x) const {
293-
return (Base::turn(x) << 1) + 2;
294-
}
295-
296-
uint64_t last_turn_read(const uint64_t x) const {
297-
return Base::turn(x) << 1;
298-
}
323+
T slots[Base::SLOTS_NUM];
324+
explicit LockfreeBatchMPMCRingQueue(size_t c) : Base(c) {}
299325

300326
public:
301327
using Base::empty;
302328
using Base::full;
303329

304-
size_t push_batch(const T *x, size_t n) {
330+
explicit LockfreeBatchMPMCRingQueue() : Base() {}
331+
332+
size_t push_batch(const T* x, size_t n) {
305333
size_t rh, wt;
306334
wt = tail.load(std::memory_order_acquire);
307335
for (;;) {
308336
rh = head.load(std::memory_order_acquire);
309337
auto wn = std::min(n, Base::capacity - (wt - rh));
310-
if (wn == 0)
311-
return 0;
312-
if (!tail.compare_exchange_strong(wt, wt + wn, std::memory_order_acq_rel))
338+
if (wn == 0) return 0;
339+
if (!tail.compare_exchange_strong(wt, wt + wn,
340+
std::memory_order_acq_rel))
313341
continue;
314342
auto first_idx = idx(wt);
315343
auto part_length = Base::capacity - first_idx;
@@ -318,28 +346,28 @@ class LockfreeBatchMPMCRingQueue : public LockfreeRingQueueBase<T, N> {
318346
} else {
319347
if (likely(part_length))
320348
memcpy(&slots[first_idx], x, sizeof(T) * (part_length));
321-
memcpy(&slots[0], x + part_length, sizeof(T) * (wn - part_length));
349+
memcpy(&slots[0], x + part_length,
350+
sizeof(T) * (wn - part_length));
322351
}
323352
auto wh = wt;
324-
while (!write_head.compare_exchange_strong(wh, wt + wn, std::memory_order_acq_rel))
353+
while (!write_head.compare_exchange_strong(
354+
wh, wt + wn, std::memory_order_acq_rel))
325355
wh = wt;
326356
return wn;
327357
}
328358
}
329359

330-
bool push(const T &x) {
331-
return push_batch(&x, 1) == 1;
332-
}
360+
bool push(const T& x) { return push_batch(&x, 1) == 1; }
333361

334-
size_t pop_batch(T *x, size_t n) {
362+
size_t pop_batch(T* x, size_t n) {
335363
size_t rt, wh;
336364
rt = read_tail.load(std::memory_order_acquire);
337365
for (;;) {
338366
wh = write_head.load(std::memory_order_acquire);
339367
auto rn = std::min(n, wh - rt);
340-
if (rn == 0)
341-
return 0;
342-
if (!read_tail.compare_exchange_strong(rt, rt + rn, std::memory_order_acq_rel))
368+
if (rn == 0) return 0;
369+
if (!read_tail.compare_exchange_strong(rt, rt + rn,
370+
std::memory_order_acq_rel))
343371
continue;
344372
auto first_idx = idx(rt);
345373
auto part_length = Base::capacity - first_idx;
@@ -348,10 +376,12 @@ class LockfreeBatchMPMCRingQueue : public LockfreeRingQueueBase<T, N> {
348376
} else {
349377
if (likely(part_length))
350378
memcpy(x, &slots[first_idx], sizeof(T) * (part_length));
351-
memcpy(x + part_length, &slots[0], sizeof(T) * (rn - part_length));
379+
memcpy(x + part_length, &slots[0],
380+
sizeof(T) * (rn - part_length));
352381
}
353382
auto rh = rt;
354-
while (!head.compare_exchange_strong(rh, rt + rn, std::memory_order_acq_rel))
383+
while (!head.compare_exchange_strong(rh, rt + rn,
384+
std::memory_order_acq_rel))
355385
rh = rt;
356386
return rn;
357387
}
@@ -417,6 +447,10 @@ class LockfreeBatchMPMCRingQueue : public LockfreeRingQueueBase<T, N> {
417447
}
418448
};
419449

450+
template <typename T>
451+
using FlexLockfreeBatchMPMCRingQueue =
452+
FlexQueue<T, LockfreeBatchMPMCRingQueue<T, 0>>;
453+
420454
template <typename T, size_t N>
421455
class LockfreeSPSCRingQueue : public LockfreeRingQueueBase<T, N> {
422456
protected:
@@ -425,12 +459,16 @@ class LockfreeSPSCRingQueue : public LockfreeRingQueueBase<T, N> {
425459
using Base::idx;
426460
using Base::tail;
427461

428-
T slots[Base::capacity];
462+
T slots[Base::SLOTS_NUM];
463+
464+
explicit LockfreeSPSCRingQueue(size_t c) : Base(c) {}
429465

430466
public:
431467
using Base::empty;
432468
using Base::full;
433469

470+
explicit LockfreeSPSCRingQueue() : Base() {}
471+
434472
bool push(const T& x) {
435473
auto t = tail.load(std::memory_order_acquire);
436474
if (unlikely(Base::check_full(head, t))) return false;
@@ -520,6 +558,9 @@ class LockfreeSPSCRingQueue : public LockfreeRingQueueBase<T, N> {
520558
}
521559
};
522560

561+
template <typename T>
562+
using FlexLockfreeSPSCRingQueue = FlexQueue<T, LockfreeSPSCRingQueue<T, 0>>;
563+
523564
namespace photon {
524565
namespace common {
525566

@@ -599,6 +640,3 @@ class RingChannel : public QueueType {
599640

600641
} // namespace common
601642
} // namespace photon
602-
603-
#undef size_t
604-

‎common/test/test_lockfree.cpp‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,8 +53,20 @@ std::array<int, sender_num> scnt;
5353
std::array<std::atomic<int>, items_num / sender_num> sc, rc;
5454

5555
LockfreeMPMCRingQueue<int, capacity> lqueue;
56+
FlexLockfreeMPMCRingQueue<int> &luqueue = [] {
57+
static auto q = FlexLockfreeMPMCRingQueue<int>::create(capacity);
58+
return std::ref(*q);
59+
}();
5660
LockfreeBatchMPMCRingQueue<int, capacity> lbqueue;
61+
FlexLockfreeBatchMPMCRingQueue<int> &lubqueue = [] {
62+
static auto q = FlexLockfreeBatchMPMCRingQueue<int>::create(capacity);
63+
return std::ref(*q);
64+
}();
5765
LockfreeSPSCRingQueue<int, capacity> cqueue;
66+
FlexLockfreeSPSCRingQueue<int> &cuqueue = [] {
67+
static auto q = FlexLockfreeSPSCRingQueue<int>::create(capacity);
68+
return std::ref(*q);
69+
}();
5870
std::mutex rlock, wlock;
5971

6072
#ifdef TESTING_ENABLE_BOOST
@@ -230,19 +242,28 @@ TEST(lockfree_queue, test_queue) {
230242
test_queue<NoLock, NoLock>("BoostQueue", bqueue);
231243
#endif
232244
test_queue<NoLock, NoLock>("PhotonLockfreeMPMCQueue", lqueue);
245+
test_queue<NoLock, NoLock>("FlexPhotonLockfreeMPMCQueue", luqueue);
233246
test_queue<NoLock, NoLock>("PhotonLockfreeBatchMPMCQueue", lbqueue);
234247
test_queue_batch<NoLock, NoLock>("PhotonLockfreeBatchMPMCQueue+Batch",
235248
lbqueue);
249+
test_queue<NoLock, NoLock>("FlexPhotonLockfreeBatchMPMCQueue", lubqueue);
250+
test_queue_batch<NoLock, NoLock>("FlexPhotonLockfreeBatchMPMCQueue+Batch",
251+
lubqueue);
236252
#ifdef TESTING_ENABLE_BOOST
237253
test_queue<WithLock, WithLock>("BoostSPSCQueue", squeue);
238254
#endif
239255
test_queue<WithLock, WithLock>("PhotonSPSCQueue", cqueue);
240256
test_queue_batch<WithLock, WithLock>("PhotonSPSCQueue+Batch", cqueue);
257+
test_queue<WithLock, WithLock>("FlexPhotonSPSCQueue", cuqueue);
258+
test_queue_batch<WithLock, WithLock>("FlexPhotonSPSCQueue+Batch", cuqueue);
241259
}
242260

243261
int main(int argc, char **arg) {
244262
if (!photon::is_using_default_engine()) return 0;
245263
::testing::InitGoogleTest(&argc, arg);
246264
gflags::ParseCommandLineFlags(&argc, &arg, true);
265+
DEFER(luqueue.destroy(&luqueue));
266+
DEFER(lubqueue.destroy(&lubqueue));
267+
DEFER(cuqueue.destroy(&cuqueue));
247268
return RUN_ALL_TESTS();
248269
}

0 commit comments

Comments
 (0)