内核同步的一段代码

看到一段很有意思的代码。DPC/Dispatch Level的同步操作,除了当前CPU,其他的都插入一个DPC,利用两个全局变量同步。


PKDPC GainExclusivity();
NTSTATUS ReleaseExclusivity(PVOID pKDpc);
VOID RaiseCPUIrqlAndWait(PKDPC Dpc, PVOID DeferredContext, PVOID SysArg1, PVOID SysArg2);
DWORD g_count, g_now;

PKDPC 
GainExclusivity()
{
	NTSTATUS us;
	ULONG u_currentCPU;
	CCHAR i;
	PKDPC pKDpc, temp;

	// Check level
	if (KeGetCurrentIrql() != DISPATCH_LEVEL) return NULL;

	// Interlocked zero the globals
	InterlockedAnd(&g_count, 0);
	InterlockedAnd(&g_now, 0);

	// Setup nonpaged space for DPC functions
	temp = (PKDPC)ExAllocatePool(NonPagedPool, KeNumberProcessors * sizeof(KDPC));
	if (temp == NULL) return NULL;
	
	u_currentCPU = KeGetCurrentProcessorNumber();
	pKDpc = temp;

	for (i = 0; i < KeNumberProcessors; i++, temp++)
	{
		// The DPC must not run on the current CPU
		if(i == u_currentCPU) continue;
		
		KeInitializeDpc(temp, RaiseCPUIrqlAndWait, NULL);
		KeSetTargetProcessorDpc(temp, i);
		KeInsertQueueDpc(temp, NULL, NULL);
	}

	// Wait
	while(InterlockedCompareExchange(&g_now, KeNumberProcessors-1, KeNumberProcessors-1) != KeNumberProcessors-1)
	{
		__asm nop;
	}
	
	// Return
	return pKDpc;
}


NTSTATUS 
ReleaseExclusivity(PVOID pKDpc)
{
	// Inc CPU counter
	InterlockedIncrement(&g_count);

	// Wait for...
	while(InterlockedCompareExchange(&g_now, 0, 0))
	{
		__asm nop;
	}

	// Free mem
	if (pKDpc != NULL)
	{
		ExFreePool(pKDpc);
		pKDpc = NULL;
	}
	
	return STATUS_SUCCESS;
}

VOID 
RaiseCPUIrqlAndWait(PKDPC Dpc, PVOID DeferredContext, PVOID SysArg1, PVOID SysArg2)
{
	InterlockedIncrement(&g_now);
	while (!InterlockedCompareExchange(&g_count, 1, 1))
	{
		__asm nop;
	}
	InterlockedDecrement(&g_now);
}

void DoSomeThing()
{
	// Get exclusivity
	KIRQL CurrentIrql, OldIrql;
	PKDPC pKDpc;
	
	// Get into dispatch mode
	CurrentIrql = KeGetCurrentIrql();
	OldIrql = CurrentIrql;
	if (CurrentIrql < DISPATCH_LEVEL) KeRaiseIrql(DISPATCH_LEVEL, &OldIrql);

	// Gain exclusivity on multicore/multiCPU machines
	pKDpc = GainExclusivity();

	// Release exclusivity and leave DISPATCH_LEVEL
	ReleaseExclusivity(pKDpc);
	KeLowerIrql(OldIrql);
}

其实这个实现跟SPIN_LOCK自旋锁实现类似,都是让CPU空跑。很轻量,但是略浪费。