1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
|
//
// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2025
//
// Distributed under the Boost Software License, Version 1.0. (See accompanying
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
//
#pragma once
#include "td/utils/common.h"
#include "td/utils/misc.h"
#include <array>
#include <atomic>
namespace td {
template <class T, size_t N = 256>
class StealingQueue {
public:
static_assert(N > 0 && (N & (N - 1)) == 0, "");
// tries to put a value
// returns if succeeded
// only owner is allowed to do this
template <class F>
void local_push(T value, F &&overflow_f) {
while (true) {
auto tail = tail_.load(std::memory_order_relaxed);
auto head = head_.load(); // TODO: memory order
if (static_cast<size_t>(tail - head) < N) {
buf_[tail & MASK].store(value, std::memory_order_relaxed);
tail_.store(tail + 1, std::memory_order_release);
return;
}
// queue is full
// TODO: batch insert into global queue?
auto n = N / 2 + 1;
auto new_head = head + n;
if (!head_.compare_exchange_strong(head, new_head)) {
continue;
}
for (size_t i = 0; i < n; i++) {
overflow_f(buf_[(i + head) & MASK].load(std::memory_order_relaxed));
}
overflow_f(value);
return;
}
}
// tries to pop a value
// returns if succeeded
// only owner is allowed to do this
bool local_pop(T &value) {
auto tail = tail_.load(std::memory_order_relaxed);
auto head = head_.load();
if (head == tail) {
return false;
}
value = buf_[head & MASK].load(std::memory_order_relaxed);
return head_.compare_exchange_strong(head, head + 1);
}
bool steal(T &value, StealingQueue<T, N> &other) {
while (true) {
auto tail = tail_.load(std::memory_order_relaxed);
auto head = head_.load(); // TODO: memory order
auto other_head = other.head_.load();
auto other_tail = other.tail_.load(std::memory_order_acquire);
if (other_tail < other_head) {
continue;
}
auto n = narrow_cast<size_t>(other_tail - other_head);
if (n > N) {
continue;
}
n -= n / 2;
n = td::min(n, static_cast<size_t>(head + N - tail));
if (n == 0) {
return false;
}
for (size_t i = 0; i < n; i++) {
buf_[(i + tail) & MASK].store(other.buf_[(i + other_head) & MASK].load(std::memory_order_relaxed),
std::memory_order_relaxed);
}
if (!other.head_.compare_exchange_strong(other_head, other_head + n)) {
continue;
}
n--;
value = buf_[(tail + n) & MASK].load(std::memory_order_relaxed);
tail_.store(tail + n, std::memory_order_release);
return true;
}
}
StealingQueue() {
for (auto &x : buf_) {
// workaround for https://gcc.gnu.org/bugzilla/show_bug.cgi?id=64658
#if TD_GCC && GCC_VERSION <= 40902
x = T();
#else
std::atomic_init(&x, T());
#endif
}
std::atomic_thread_fence(std::memory_order_seq_cst);
}
private:
std::atomic<int64> head_{0};
std::atomic<int64> tail_{0};
static constexpr size_t MASK{N - 1};
std::array<std::atomic<T>, N> buf_;
};
} // namespace td
|
