-
Notifications
You must be signed in to change notification settings - Fork 11
/
ConcurrentCallbacks.h
309 lines (264 loc) · 9.21 KB
/
ConcurrentCallbacks.h
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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
// Concurrent callback support
#pragma once
#include <list>
/*
Minimalist Windows thread pool API object to run simple callbacks concurrently to take advantage
of multiple CPU cores; spreading out work for performance gain.
Instance a "new ConcurrentCallbackGroup" to initialize needed resources (local thread pool and support data),
then "delete ConcurrentCallbackGroup" the object when done to kill callback threads et al.
Intended to be used one of two ways:
1) Queue up all the callbacks first via "Add()" then fire them up all at once using "Start()".
2) Or one at the time add callbacks via "Add()" and either do optional "start = TRUE" argument, or
call "Start()" imeadiatly after.
In either case, queue/add all the callbacks jobs first and start them up before using either the
"Wait()" or "Poll()" methods for completion status since these rely on a simple counting mechanism internally.
The WORKER_CALLBACK callbacks should return FALSE for success or TRUE on error.
Pass opaque data to the callbacks via the "PVOID lParm" argument.
It's up to the callback do any atomic operations to be thread safe if external shared resource data are used.
Version: 1.1.0 9/17/2022
Author: Kevin Weatherman
*/
class ConcurrentCallbackGroup
{
public:
// Callbacks return TRUE to indicate error, else FALSE to signal success
typedef BOOL (WINAPI* WORKER_CALLBACK)(__in PVOID lParm);
// Add a new worker callback to the queue
// Returns HRESULT ERROR_SUCCESS on success else error code on failure
HRESULT Add(__in ConcurrentCallbackGroup::WORKER_CALLBACK callback, __in PVOID lParm = NULL, BOOL start = FALSE)
{
workers_lock();
m_workers.push_back({ NULL, FALSE, this, callback, lParm });
WORKER *wp = &m_workers.back();
workers_unlock();
wp->work = CreateThreadpoolWork(CallbackWrapper, wp, &m_callbackEnv);
HRESULT hr = (wp->work != NULL) ? ERROR_SUCCESS : HRESULT_FROM_WIN32(GetLastError());
// Optionally start up the callback(s) now
if ((m_started || start) && (hr == ERROR_SUCCESS))
{
Start();
// If not enabled already via the previous Start(), do it now
if (m_started && (InterlockedCompareExchange((PLONG) &wp->enabled, TRUE, FALSE) == FALSE))
SubmitThreadpoolWork(wp->work);
}
return hr;
}
// Start up the parallel callback queue
// Returns ERROR_SUCCESS on success
// On failure our object should be deleted for cleanup
void Start()
{
if (!m_started)
{
m_started = TRUE;
// Set the pool size max to the CPU physical core count if the user didn't override it
// The thread pool API doesn't allow us to force or suggest physical cores.
// SMT threads won't help the scan performance of complex Yara rules, the compute of actual physical cores will.
if (m_maxThreads == 0)
m_maxThreads = GetPhysicalCoreCount();
SetThreadpoolThreadMinimum(m_pool, 0);
SetThreadpoolThreadMaximum(m_pool, m_maxThreads);
// Fire up any queued callbacks..
workers_lock();
{
for (WORKER &w: m_workers)
{
if (InterlockedCompareExchange((PLONG) &w.enabled, TRUE, FALSE) == FALSE)
SubmitThreadpoolWork(w.work);
}
}
workers_unlock();
}
}
// Wait for started queue callbacks to complete
// Returns ERROR_SUCCESS if completed, else returns E_FAIL on errors
HRESULT Wait(__out long &errorCount)
{
if (m_started)
{
#ifdef _DEBUG
// Should be enabled here
workers_lock();
{
for (WORKER &w: m_workers)
{
_ASSERT(InterlockedCompareExchange((PLONG) &w.enabled, TRUE, FALSE) == TRUE);
}
}
workers_unlock();
#endif
// Loop until all callbacks have returned or until one errored out
long queueSize = MAXINT;
while (!m_workerErrors)
{
workers_lock();
queueSize = (long) m_workers.size();
workers_unlock();
if (m_completed >= queueSize)
break;
// Give up our caller thread for a while..
Sleep(50);
};
errorCount = m_workerErrors;
return ((m_workerErrors == 0) && (m_completed >= queueSize)) ? ERROR_SUCCESS : E_FAIL;
}
else
{
errorCount = 0;
return E_FAIL;
}
}
// Check if the callback queue has completed
// Returns ERROR_SUCCESS if completed or E_PENDING if still pending; E_FAIL on error
HRESULT Poll(__out long &errorCount)
{
if (m_started)
{
errorCount = m_workerErrors;
if (errorCount == 0)
{
workers_lock();
long queueSize = (long) m_workers.size();
workers_unlock();
return ((m_completed >= queueSize)) ? ERROR_SUCCESS : E_PENDING;
}
}
else
errorCount = 0;
return E_FAIL;
}
// Abort the running queue for cases where the user requests it, when or app is closing, or when our DLL is unloading.
// Will block until all active callback threads are done.
// Same effect as simply destructing our object
void Abort()
{
Cleanup();
}
// Return max thread count. Valid after Start() called
UINT32 MaxTheads()
{
if (m_started)
return m_maxThreads;
else
return 0;
}
// Construct object:
// initResult = HRESULT initialize result. ERROR_SUCCESS on success, else FAILED status
// maxThreadCount = Optionally limit the max pool threads to this count, default '0' to use max one thread per physical core.
ConcurrentCallbackGroup(__out HRESULT &initResult, __in_opt UINT32 maxThreadCount = 0) :
m_completed(0), m_workerErrors(0), m_started(FALSE), m_pool(NULL), m_cleanupGroup(NULL)
{
initResult = E_FAIL;
m_maxThreads = maxThreadCount;
ZeroMemory(&m_callbackEnv, sizeof(m_callbackEnv));
// Lock for m_workers list access from potentially multiple threads
InitializeCriticalSectionAndSpinCount(&m_workersLock, 20);
// Create our private thread pool
m_pool = CreateThreadpool(NULL);
if (!m_pool)
{
initResult = HRESULT_FROM_WIN32(GetLastError());
goto exit;
}
// Setup pool for worker callbacks
InitializeThreadpoolEnvironment(&m_callbackEnv);
SetThreadpoolCallbackPool(&m_callbackEnv, m_pool);
//SetThreadpoolCallbackRunsLong(&m_callbackEnv);
SetThreadpoolCallbackPriority(&m_callbackEnv, TP_CALLBACK_PRIORITY_HIGH);
// Create a group for automated cleanup later
m_cleanupGroup = CreateThreadpoolCleanupGroup();
if (!m_cleanupGroup)
{
initResult = HRESULT_FROM_WIN32(GetLastError());
goto exit;
}
SetThreadpoolCallbackCleanupGroup(&m_callbackEnv, m_cleanupGroup, NULL);
initResult = ERROR_SUCCESS;
exit:;
if (initResult != ERROR_SUCCESS)
Cleanup();
}
// Note: Queued callbacks are canceled, but started callbacks will block until they are completed
~ConcurrentCallbackGroup() { Cleanup(); }
// Get CPU physical core count
static UINT32 GetPhysicalCoreCount()
{
UINT32 cores = 0;
// 1st pass get buffer size
PBYTE buffer = NULL;
DWORD length = 0;
if (!GetLogicalProcessorInformationEx(RelationProcessorCore, (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX) buffer, &length))
{
if (GetLastError() == ERROR_INSUFFICIENT_BUFFER)
buffer = new BYTE[length];
if (buffer)
{
// 2nd pass read CPU info
if (GetLogicalProcessorInformationEx(RelationProcessorCore, (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX) buffer, &length))
{
// Physical core count is the count of 'RelationProcessorCore' entries
cores = (length / ((PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX) buffer)->Size);
}
delete buffer;
}
}
return max(1, cores);
}
private:
struct WORKER
{
PTP_WORK work;
BOOL enabled;
ConcurrentCallbackGroup *self;
WORKER_CALLBACK callback;
PVOID lParm;
};
std::list<WORKER> m_workers;
CRITICAL_SECTION m_workersLock;
inline void workers_lock() { EnterCriticalSection(&m_workersLock); }
inline void workers_unlock() { LeaveCriticalSection(&m_workersLock); }
long m_completed, m_workerErrors;
UINT32 m_maxThreads; // On init optional max thread count override, after Start() the max threads in use
BOOL m_started;
PTP_POOL m_pool;
PTP_CLEANUP_GROUP m_cleanupGroup;
TP_CALLBACK_ENVIRON m_callbackEnv;
// Local worker callback wrapper
// Needs to be thread safe within our context
static void CALLBACK CallbackWrapper(__in PTP_CALLBACK_INSTANCE pci, __inout PVOID lParm, __in PTP_WORK work)
{
WORKER &wrk = *((WORKER*) lParm);
_ASSERT(wrk.work == work);
// Call worker callback
BOOL result = TRUE;
__try
{
result = wrk.callback(wrk.lParm);
}
__except (ReportException("ConcurrentCallbackGroup::CallbackWrapper", GetExceptionInformation(), TRUE)){ result = TRUE; }
if (result)
{
// return TRUE from the user callback indicates error status
_InterlockedIncrement(&wrk.self->m_workerErrors);
}
_InterlockedIncrement(&wrk.self->m_completed);
}
void Cleanup()
{
if (m_cleanupGroup)
{
// Cancels any pending callbacks and blocks for any running until they return, then releases
// all of our worker objects for us.
CloseThreadpoolCleanupGroupMembers(m_cleanupGroup, TRUE, NULL);
CloseThreadpoolCleanupGroup(m_cleanupGroup);
m_cleanupGroup = NULL;
}
if (m_pool)
{
CloseThreadpool(m_pool);
m_pool = NULL;
}
m_workers.clear();
DeleteCriticalSection(&m_workersLock);
}
};