tdlight-team/tdlight
tdlight-team/tdlight
1
1
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add thread::set_affinity_mask/get_affinity_mask.

Browse Source
This commit is contained in:
levlam 2022-09-14 14:38:14 +03:00
parent 5f98548e02
commit 0eddd8d405
4 changed files with 70 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tdutils/td/utils/port/detail/ThreadPthread.cpp
View File

@ -94,6 +94,14 @@ int ThreadPthread::do_pthread_create(pthread_t *thread, const pthread_attr_t *at
return pthread_create(thread, attr, start_routine, arg);
}
Status ThreadPthread::set_affinity_mask(id thread_id, uint64 mask) {
return Status::Error("Unsupported");
}
uint64 ThreadPthread::get_affinity_mask(id thread_id) {
return 0;
}
namespace this_thread_pthread {
ThreadPthread::id get_id() {
return pthread_self();

5
tdutils/td/utils/port/detail/ThreadPthread.h
View File

@ -17,6 +17,7 @@
#include "td/utils/port/detail/ThreadIdGuard.h"
#include "td/utils/port/thread_local.h"
#include "td/utils/Slice.h"
#include "td/utils/Status.h"
#include <tuple>
#include <type_traits>
@ -66,6 +67,10 @@ class ThreadPthread {
static void send_real_time_signal(id thread_id, int real_time_signal_number);
static Status set_affinity_mask(id thread_id, uint64 mask);
static uint64 get_affinity_mask(id thread_id);
private:
MovableValue<bool> is_inited_;
pthread_t thread_;

30
tdutils/td/utils/port/detail/ThreadStl.h
View File

@ -15,6 +15,7 @@
#include "td/utils/port/detail/ThreadIdGuard.h"
#include "td/utils/port/thread_local.h"
#include "td/utils/Slice.h"
#include "td/utils/Status.h"
#include <thread>
#include <tuple>
@ -75,6 +76,35 @@ class ThreadStl {
// not supported
}
static Status set_affinity_mask(id thread_id, uint64 mask) {
#if TD_WINDOWS
if (static_cast<DWORD_PTR>(mask) != mask) {
return Status::Error("Invalid thread affinity mask specified");
}
if (SetThreadAffinityMask(thread_id, static_cast<DWORD_PTR>(mask))) {
return Status::OK();
}
return OS_ERROR("Failed to set thread affinity mask");
#else
return Status::Error("Unsupported");
#endif
}
static uint64 get_affinity_mask(id thread_id) {
#if TD_WINDOWS
DWORD_PTR process_mask = 0;
DWORD_PTR system_mask = 0;
if (GetProcessAffinityMask(GetCurrentProcess(), &process_mask, &system_mask)) {
auto result = SetThreadAffinityMask(thread_id, process_mask);
if (result != 0 && result != process_mask) {
SetThreadAffinityMask(thread_id, result);
}
return result;
}
#endif
return 0;
}
private:
std::thread thread_;

27
tdutils/test/port.cpp
View File

@ -284,3 +284,30 @@ TEST(Port, EventFdAndSignals) {
}
#endif
#endif
#if !TD_THREAD_UNSUPPORTED
TEST(Port, ThreadAffinityMask) {
auto thread_id = td::this_thread::get_id();
auto old_mask = td::thread::get_affinity_mask(thread_id);
LOG(INFO) << "Initial thread affinity mask: " << old_mask;
for (size_t i = 0; i < 64; i++) {
auto mask = td::thread::get_affinity_mask(thread_id);
LOG(INFO) << mask;
auto result = td::thread::set_affinity_mask(thread_id, static_cast<td::uint64>(1) << i);
LOG(INFO) << i << ": " << result;
mask = td::thread::get_affinity_mask(thread_id);
LOG(INFO) << mask;
if (i <= 1) {
td::thread thread([] {
auto mask = td::thread::get_affinity_mask(td::this_thread::get_id());
LOG(INFO) << "New thread affinity mask: " << mask;
});
}
}
auto result = td::thread::set_affinity_mask(thread_id, old_mask);
LOG(INFO) << result;
old_mask = td::thread::get_affinity_mask(thread_id);
LOG(INFO) << old_mask;
}
#endif