Seregon/ShadPKG

A tool for deriving PKG packet encryption keys for ps4 written in c++

C++/47.3 KB/No license
common/bounded_threadsafe_queue.h
raw
ShadPKG / common / bounded_threadsafe_queue.h
1// SPDX-FileCopyrightText: Copyright 2023 yuzu Emulator Project
2// SPDX-License-Identifier: GPL-2.0-or-later
3 
4#pragma once
5 
6#include <atomic>
7#include <condition_variable>
8#include <cstddef>
9#include <memory>
10#include <mutex>
11#include "common/polyfill_thread.h"
12 
13namespace Common {
14 
15namespace detail {
16constexpr std::size_t DefaultCapacity = 0x1000;
17} // namespace detail
18 
19template <typename T, std::size_t Capacity = detail::DefaultCapacity>
20class SPSCQueue {
21 static_assert((Capacity & (Capacity - 1)) == 0, "Capacity must be a power of two.");
22 
23public:
24 template <typename... Args>
25 bool TryEmplace(Args&&... args) {
26 return Emplace<PushMode::Try>(std::forward<Args>(args)...);
27 }
28 
29 template <typename... Args>
30 void EmplaceWait(Args&&... args) {
31 Emplace<PushMode::Wait>(std::forward<Args>(args)...);
32 }
33 
34 bool TryPop(T& t) {
35 return Pop<PopMode::Try>(t);
36 }
37 
38 void PopWait(T& t) {
39 Pop<PopMode::Wait>(t);
40 }
41 
42 void PopWait(T& t, std::stop_token stop_token) {
43 Pop<PopMode::WaitWithStopToken>(t, stop_token);
44 }
45 
46 T PopWait() {
47 T t;
48 Pop<PopMode::Wait>(t);
49 return t;
50 }
51 
52 T PopWait(std::stop_token stop_token) {
53 T t;
54 Pop<PopMode::WaitWithStopToken>(t, stop_token);
55 return t;
56 }
57 
58private:
59 enum class PushMode {
60 Try,
61 Wait,
62 Count,
63 };
64 
65 enum class PopMode {
66 Try,
67 Wait,
68 WaitWithStopToken,
69 Count,
70 };
71 
72 template <PushMode Mode, typename... Args>
73 bool Emplace(Args&&... args) {
74 const std::size_t write_index = m_write_index.load(std::memory_order::relaxed);
75 
76 if constexpr (Mode == PushMode::Try) {
77 // Check if we have free slots to write to.
78 if ((write_index - m_read_index.load(std::memory_order::acquire)) == Capacity) {
79 return false;
80 }
81 } else if constexpr (Mode == PushMode::Wait) {
82 // Wait until we have free slots to write to.
83 std::unique_lock lock{producer_cv_mutex};
84 producer_cv.wait(lock, [this, write_index] {
85 return (write_index - m_read_index.load(std::memory_order::acquire)) < Capacity;
86 });
87 } else {
88 static_assert(Mode < PushMode::Count, "Invalid PushMode.");
89 }
90 
91 // Determine the position to write to.
92 const std::size_t pos = write_index % Capacity;
93 
94 // Emplace into the queue.
95 new (std::addressof(m_data[pos])) T(std::forward<Args>(args)...);
96 
97 // Increment the write index.
98 ++m_write_index;
99 
100 // Notify the consumer that we have pushed into the queue.
101 std::scoped_lock lock{consumer_cv_mutex};
102 consumer_cv.notify_one();
103 
104 return true;
105 }
106 
107 template <PopMode Mode>
108 bool Pop(T& t, [[maybe_unused]] std::stop_token stop_token = {}) {
109 const std::size_t read_index = m_read_index.load(std::memory_order::relaxed);
110 
111 if constexpr (Mode == PopMode::Try) {
112 // Check if the queue is empty.
113 if (read_index == m_write_index.load(std::memory_order::acquire)) {
114 return false;
115 }
116 } else if constexpr (Mode == PopMode::Wait) {
117 // Wait until the queue is not empty.
118 std::unique_lock lock{consumer_cv_mutex};
119 consumer_cv.wait(lock, [this, read_index] {
120 return read_index != m_write_index.load(std::memory_order::acquire);
121 });
122 } else if constexpr (Mode == PopMode::WaitWithStopToken) {
123 // Wait until the queue is not empty.
124 std::unique_lock lock{consumer_cv_mutex};
125 Common::CondvarWait(consumer_cv, lock, stop_token, [this, read_index] {
126 return read_index != m_write_index.load(std::memory_order::acquire);
127 });
128 if (stop_token.stop_requested()) {
129 return false;
130 }
131 } else {
132 static_assert(Mode < PopMode::Count, "Invalid PopMode.");
133 }
134 
135 // Determine the position to read from.
136 const std::size_t pos = read_index % Capacity;
137 
138 // Pop the data off the queue, moving it.
139 t = std::move(m_data[pos]);
140 
141 // Increment the read index.
142 ++m_read_index;
143 
144 // Notify the producer that we have popped off the queue.
145 std::scoped_lock lock{producer_cv_mutex};
146 producer_cv.notify_one();
147 
148 return true;
149 }
150 
151 alignas(128) std::atomic_size_t m_read_index{0};
152 alignas(128) std::atomic_size_t m_write_index{0};
153 
154 std::array<T, Capacity> m_data;
155 
156 std::condition_variable_any producer_cv;
157 std::mutex producer_cv_mutex;
158 std::condition_variable_any consumer_cv;
159 std::mutex consumer_cv_mutex;
160};
161 
162template <typename T, std::size_t Capacity = detail::DefaultCapacity>
163class MPSCQueue {
164public:
165 template <typename... Args>
166 bool TryEmplace(Args&&... args) {
167 std::scoped_lock lock{write_mutex};
168 return spsc_queue.TryEmplace(std::forward<Args>(args)...);
169 }
170 
171 template <typename... Args>
172 void EmplaceWait(Args&&... args) {
173 std::scoped_lock lock{write_mutex};
174 spsc_queue.EmplaceWait(std::forward<Args>(args)...);
175 }
176 
177 bool TryPop(T& t) {
178 return spsc_queue.TryPop(t);
179 }
180 
181 void PopWait(T& t) {
182 spsc_queue.PopWait(t);
183 }
184 
185 void PopWait(T& t, std::stop_token stop_token) {
186 spsc_queue.PopWait(t, stop_token);
187 }
188 
189 T PopWait() {
190 return spsc_queue.PopWait();
191 }
192 
193 T PopWait(std::stop_token stop_token) {
194 return spsc_queue.PopWait(stop_token);
195 }
196 
197private:
198 SPSCQueue<T, Capacity> spsc_queue;
199 std::mutex write_mutex;
200};
201 
202template <typename T, std::size_t Capacity = detail::DefaultCapacity>
203class MPMCQueue {
204public:
205 template <typename... Args>
206 bool TryEmplace(Args&&... args) {
207 std::scoped_lock lock{write_mutex};
208 return spsc_queue.TryEmplace(std::forward<Args>(args)...);
209 }
210 
211 template <typename... Args>
212 void EmplaceWait(Args&&... args) {
213 std::scoped_lock lock{write_mutex};
214 spsc_queue.EmplaceWait(std::forward<Args>(args)...);
215 }
216 
217 bool TryPop(T& t) {
218 std::scoped_lock lock{read_mutex};
219 return spsc_queue.TryPop(t);
220 }
221 
222 void PopWait(T& t) {
223 std::scoped_lock lock{read_mutex};
224 spsc_queue.PopWait(t);
225 }
226 
227 void PopWait(T& t, std::stop_token stop_token) {
228 std::scoped_lock lock{read_mutex};
229 spsc_queue.PopWait(t, stop_token);
230 }
231 
232 T PopWait() {
233 std::scoped_lock lock{read_mutex};
234 return spsc_queue.PopWait();
235 }
236 
237 T PopWait(std::stop_token stop_token) {
238 std::scoped_lock lock{read_mutex};
239 return spsc_queue.PopWait(stop_token);
240 }
241 
242private:
243 SPSCQueue<T, Capacity> spsc_queue;
244 std::mutex write_mutex;
245 std::mutex read_mutex;
246};
247 
248} // namespace Common
249