验证中...
8月18日(周六)成都源创会火热报名中,四位一线行业大牛与你面对面,探讨区块链技术热潮下的冷思考。
语言: C/C++
分类: 内核开发
最后更新于 2018-03-20 14:44
cpu_lock_ex.h
原始数据 复制代码
#pragma once
namespace ddk
{
class cpu_swticher
{
public:
cpu_swticher()
{
bset = false;
previous_affinity = {};
}
~cpu_swticher()
{
if (bset)
{
KeRevertToUserGroupAffinityThread(&previous_affinity);
}
}
private:
GROUP_AFFINITY previous_affinity;
bool bset;
public:
void unswitch()
{
if (bset)
KeRevertToUserGroupAffinityThread(&previous_affinity);
bset = false;
}
bool switch_cpu(ULONG Index)
{
if (bset)
{
return false;
}
PROCESSOR_NUMBER processor_number = {};
auto status =
KeGetProcessorNumberFromIndex(Index, &processor_number);
if (!NT_SUCCESS(status)) {
return false;
}
GROUP_AFFINITY affinity = {};
affinity.Group = processor_number.Group;
affinity.Mask = 1ull << processor_number.Number;
KeSetSystemGroupAffinityThread(&affinity, &previous_affinity);
bset = true;
return true;
}
};
class cpu_lock_ex
{
public:
cpu_lock_ex()
{
g_ExclpNumberOfLockedProcessors = 0;
g_ExclpReleaseAllProcessors = 0;
g_CpuNumber = 0;
}
~cpu_lock_ex()
{
while (!InterlockedCompareExchange(&g_ExclpReleaseAllProcessors, 1, 1))
{
KeStallExecutionProcessor(10);
}
}
private:
cpu_swticher cpu1;
void _lock()
{
InterlockedIncrement(&g_ExclpNumberOfLockedProcessors);
LOG_DEBUG("Core: %d is Waiting... \r\n", KeGetCurrentProcessorNumberEx(nullptr));
while (!InterlockedCompareExchange(&g_ExclpReleaseAllProcessors, 1, 1))
{
//_mm_pause();
//YieldProcessor();
KeStallExecutionProcessor(10);
}
InterlockedDecrement(&g_ExclpNumberOfLockedProcessors);
}
public:
void lock()
{
NT_ASSERT(InterlockedAdd(&g_ExclpNumberOfLockedProcessors, 0) == 0);
InterlockedAnd(&g_ExclpReleaseAllProcessors, 0);
g_CpuNumber = KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS);
cpu1.switch_cpu(0);
for (auto processor_index = 1ul;
processor_index < g_CpuNumber;
processor_index++) {
auto thread = ddk::thread(
std::bind(&ddk::cpu_lock_ex::doLock, this, std::placeholders::_1),
processor_index);
thread.detach();
}
const auto needToBeLocked = g_CpuNumber - 1;
while (_InterlockedCompareExchange(&g_ExclpNumberOfLockedProcessors,
needToBeLocked, needToBeLocked) !=
static_cast<LONG>(needToBeLocked))
{
KeStallExecutionProcessor(10);
}
}
void unlock()
{
InterlockedIncrement(&g_ExclpReleaseAllProcessors);
// Wait until all other processors were unlocked.
while (InterlockedCompareExchange(&g_ExclpNumberOfLockedProcessors, 0, 0))
{
KeStallExecutionProcessor(10);
}
}
void doLock(ULONG cpuIndex)
{
cpu_swticher cpu2;
cpu2.switch_cpu(cpuIndex);
KIRQL oldirql;
KeRaiseIrql(DISPATCH_LEVEL, &oldirql);
_lock();
KeLowerIrql(oldirql);
cpu2.unswitch();
}
private:
LONG g_ExclpReleaseAllProcessors;
LONG g_ExclpNumberOfLockedProcessors;
ULONG g_CpuNumber;
};
};

评论列表( 0 )

你可以在登录后,发表评论