// DiosLock.cpp : 定义 DLL 应用程序的导出函数。
//

#include "stdafx.h"
#include <map>
#include <vector>
#include <assert.h>
#include <tlhelp32.h>

#include "DiosLock.h"

using namespace std;

#define MIN_CHECK_TIMEPERIOD (INFINITE)


class _ProcLock {
	CRITICAL_SECTION m_Section;
	HANDLE m_ProcHandle;
public:
	_ProcLock()
	{
		InitializeCriticalSectionAndSpinCount(&m_Section, 4000);//ms recommended number is 4000

		m_ProcHandle = CreateEvent(0, 0, 0, 0);
	};
	~_ProcLock()
	{
		DeleteCriticalSection(&m_Section);
		CloseHandle(m_ProcHandle);
	};

	DWORD SetSpinWaitCount(DWORD waitcount)
	{
		return SetCriticalSectionSpinCount(&m_Section, waitcount);
	}

	DWORD Proc_Lock(DWORD timeout)
	{
		if (TryEnterCriticalSection(&m_Section) == FALSE)
		{
			if (timeout == 0)
			{
				return WAIT_TIMEOUT;
			}

			//failed get
			while (timeout > 0)
			{
				DWORD FirstHit = GetTickCount();
				DWORD ret = WaitForSingleObject(m_ProcHandle, timeout);
				if (ret == WAIT_OBJECT_0)
				{
					//try it again
					if (TryEnterCriticalSection(&m_Section))
					{
						return WAIT_OBJECT_0;
					}
					//failed get again
					DWORD Usedtime = GetTickCount() - FirstHit;
					if (Usedtime < timeout)
					{
						timeout -= Usedtime;
					}
					else
					{
						return WAIT_TIMEOUT;
					}

				}
				else
				{
					return ret;
				}
			}
		}

		ResetEvent(m_ProcHandle);

		return WAIT_OBJECT_0;
	}

	void Proc_UnLock()
	{
		LeaveCriticalSection(&m_Section);
		SetEvent(m_ProcHandle);
	}

};

_ProcLock g_m_ProcLock;

DWORD Proc_Lock(DWORD timeout)
{
	return g_m_ProcLock.Proc_Lock(timeout);
}

void Proc_UnLock()
{
	return g_m_ProcLock.Proc_UnLock();
}


//TID ->vector<DiosLocks>
volatile UINT64 g_csDiosLock = 0;//线程之间进行互锁

//【Locked DiosLock】 maps 【ThreadID】
//only one Thread Holds the Lock
map<UINT64,DWORD> g_InternalLockedMap;
map<UINT64, DiosLock*> g_InternalDiosObjectMap;

void InitDiosLock()
{
	//InitializeCriticalSection(&g_csDiosLock);
}

void DestroyDiosLock()
{
	//DeleteCriticalSection(&g_csDiosLock);
}

void RegistDiosLock(DiosLock* p)
{
	UINT64 val = (UINT64)p;
	Dios_ThreadLock(&g_csDiosLock);
	g_InternalDiosObjectMap[val] = p;
	Dios_ThreadUnLock(&g_csDiosLock);

}

void UnRegistDiosLock(DiosLock* p)
{
	UINT64 val = (UINT64)p;
	Dios_ThreadLock(&g_csDiosLock);
	//clean up object map
	map<UINT64, DiosLock*>::iterator iter = g_InternalDiosObjectMap.find(val);
	if (iter != g_InternalDiosObjectMap.end())
	{
		g_InternalDiosObjectMap.erase(iter);
	}

	//clean up locked map
	map<UINT64, DWORD>::iterator iter_locked = g_InternalLockedMap.find(val);
	if (iter_locked != g_InternalLockedMap.end())
	{
		g_InternalLockedMap.erase(iter_locked);
	}
	Dios_ThreadUnLock(&g_csDiosLock);

}

void UnMapLockForThread(DiosLock *p)
{
	UINT64 ID = (UINT64)p;
	Dios_ThreadLock(&g_csDiosLock);
	//find locked map
	map<UINT64, DWORD>::iterator iter_locked = g_InternalLockedMap.find(ID);
#ifdef _DEBUG
	if (iter_locked == g_InternalLockedMap.end())
	{
		assert(0);//buggy
	}

	if (iter_locked->second != GetCurrentThreadId())
	{
		assert(0);//buggy
	}
#endif


	if (iter_locked != g_InternalLockedMap.end())
	{
		//incase of exp...
		g_InternalLockedMap.erase(iter_locked);
	}



	Dios_ThreadUnLock(&g_csDiosLock);

}

void MapLockForThread(DiosLock *p)
{
	UINT64 ID = (UINT64)p;
	Dios_ThreadLock(&g_csDiosLock);
	//find locked map
#ifdef _DEBUG
	map<UINT64, DWORD>::iterator iter_locked = g_InternalLockedMap.find(ID);
	if (iter_locked != g_InternalLockedMap.end())
	{
		assert(0);//buggy
	}
#endif
	g_InternalLockedMap[ID] = GetCurrentThreadId();


	Dios_ThreadUnLock(&g_csDiosLock);

}

void CleanupForThread(DWORD Tid)
{
	/*
		try to unlock g_csDiosLock, if it's ThreadId matches.
	*/

	Dios_ThreadClearForTid(&g_csDiosLock, Tid);//try clear g_cs...

	Dios_ThreadLock(&g_csDiosLock);

	//find locked map
	map<UINT64, DWORD>::iterator iter_locked = g_InternalLockedMap.begin();
	while (iter_locked != g_InternalLockedMap.end())
	{
		if (iter_locked->second == Tid)
		{
			//found one 
			DiosLock* pLock = (DiosLock*)(iter_locked->first);


			//check existance of plock
			map<UINT64, DiosLock*>::iterator iter = g_InternalDiosObjectMap.find(iter_locked->first);
			if (iter != g_InternalDiosObjectMap.end())
			{
				//exist
				pLock->Thread_Clear(Tid);
			}
			else
			{
				//the pLock not exit
			}

			iter_locked = g_InternalLockedMap.erase(iter_locked);


			continue;
		}

		++iter_locked;
	}

	//clear the lock for the thread


	Dios_ThreadUnLock(&g_csDiosLock);

}


DiosLock::DiosLock()
{
	m_InterLock = 0;

	m_RefCount = 0;

	m_LockHandle = CreateEvent(0, 1, 1, 0);

	RegistDiosLock(this);
	
	return;
}


DiosLock::~DiosLock()
{

	if ((m_InterLock) > 0)
	{
		//printf("Warning:ThreadId:%d is in Lock-----------\n", (m_InterLock));
	}



	UnRegistDiosLock(this);

	m_InterLock = 0;

	CloseHandle(m_LockHandle);
	return;
}

bool DiosLock::DiosInternal_EnterCriticalSection(DWORD TryCount)
{

	if (TryCount == 0)
	{
		return false;
	}
	
	DWORD Tid = GetCurrentThreadId();

	//first try
	if (InterlockedCompareExchangeAcquire(&m_InterLock, Tid, 0) == 0)
	{
		++m_RefCount;
		MapLockForThread(this);
		return true;
	}

	//check first
	if (InterlockedCompareExchangeAcquire(&m_InterLock, Tid, Tid) == Tid)
	{
		//already locked
		++m_RefCount;
		return true;
	}

	//do loop
	while (TryCount > 1)
	{
		if (InterlockedCompareExchangeAcquire(&m_InterLock, Tid, 0) == 0)
		{
			++m_RefCount;
			MapLockForThread(this);
			return true;
		}

		--TryCount;
	}

	return false;
}

void DiosLock::DiosInternal_LeaveCriticalSection()
{
	DWORD Tid = GetCurrentThreadId();
	if (InterlockedCompareExchangeAcquire(&m_InterLock, Tid, Tid) == Tid)
	{
		//lock owner is right
		if (m_RefCount > 0)
		{
			--m_RefCount;
			if (m_RefCount > 0)
			{
				return;
			}

			//unmap first
			UnMapLockForThread(this);

			//clear second
			if (InterlockedCompareExchangeRelease(&m_InterLock, 0, Tid) == Tid)
			{
				return;
			}

		}
	}

	//ignore wrong unlocks

}


DWORD DiosLock::Thread_Lock(DWORD timeout)
{
	//timeout == 0 优化
	if (timeout == 0)
	{
		if (DiosInternal_EnterCriticalSection(1) == FALSE)
		{
			return WAIT_TIMEOUT;
		}
		m_LockThreadId = GetCurrentThreadId();
		ResetEvent(m_LockHandle);
		return WAIT_OBJECT_0;
	}


	if (DiosInternal_EnterCriticalSection(800) == FALSE)
	{
		while (timeout > 0)
		{
			DWORD FirstHit = GetTickCount();
			//failed get
			DWORD ret = Thread_WaitUnlockNotify(timeout);
			if (ret == WAIT_OBJECT_0)
			{
				//try it again
				if (DiosInternal_EnterCriticalSection(800))
				{
					m_LockThreadId = GetCurrentThreadId();
					ResetEvent(m_LockHandle);
					return WAIT_OBJECT_0;
				}
			}

			//failed get again
			DWORD Usedtime = GetTickCount() - FirstHit;
			if (Usedtime < timeout)
			{
				timeout -= Usedtime;
			}
			else
			{
				return WAIT_TIMEOUT;
			}

		}


	}

	//ResetEvent(m_LockHandle);

	m_LockThreadId = GetCurrentThreadId();
	ResetEvent(m_LockHandle);
	return WAIT_OBJECT_0;
}

void DiosLock::Thread_UnLock()
{
	DiosInternal_LeaveCriticalSection();

	if (m_RefCount == 0)
	{
		m_LockThreadId = 0;
		SetEvent(m_LockHandle);
	}

}

bool DiosLock::Thread_Clear(DWORD Tid)
{
	if (InterlockedCompareExchangeAcquire(&m_InterLock, Tid, Tid) == Tid)
	{
		//lock owner is right
		//if (m_RefCount > 0)
		{
			m_RefCount = 0;
			if (InterlockedCompareExchangeRelease(&m_InterLock, 0, Tid) == Tid)
			{
				SetEvent(m_LockHandle);
				return true;
			}

		}
	}
	return false;
	//ignore wrong unlocks
}


DWORD DiosLock::Thread_WaitUnlockNotify(DWORD timeout)
{

	while (timeout > 0)
	{
		DWORD FirstHit = GetTickCount();

		DWORD WaitPeriod = min(timeout, 1000);

		DWORD ret = WaitForSingleObject(m_LockHandle, WaitPeriod);
		if (ret == WAIT_OBJECT_0)
		{
			return WAIT_OBJECT_0;
		}

		//try clear once
		Try_Clear_DeadThread();

		//failed get again
		DWORD Usedtime = GetTickCount() - FirstHit;
		if (Usedtime < timeout)
		{
			timeout -= Usedtime;
		}
		else
		{
			return WAIT_TIMEOUT;
		}


	}

	return WAIT_TIMEOUT;

}

INT DiosLock::FindProcessThreads(DWORD dwOwnerTID)
{
	INT ret = -1;
	DWORD dwOwnerPID = GetCurrentProcessId();
	HANDLE hThreadSnap = INVALID_HANDLE_VALUE;
	THREADENTRY32 te32;


	// Take a snapshot of all running threads  
	hThreadSnap = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
	if (hThreadSnap == INVALID_HANDLE_VALUE)
		return (ret);

	// Fill in the size of the structure before using it. 
	te32.dwSize = sizeof(THREADENTRY32);

	// Retrieve information about the first thread,
	// and exit if unsuccessful
	if (!Thread32First(hThreadSnap, &te32))
	{
		CloseHandle(hThreadSnap);    // Must clean up the
		//   snapshot object!
		return (ret);
	}

	// Now walk the thread list of the system,
	// and display information about each thread
	// associated with the specified process

	ret = 0;//set not found flag first

	do
	{

		if (te32.th32OwnerProcessID == dwOwnerPID && te32.th32ThreadID == dwOwnerTID)
		{
			//found it
			ret = 1;
			break;
		}
	} while (Thread32Next(hThreadSnap, &te32));

	//  Don't forget to clean up the snapshot object.
	CloseHandle(hThreadSnap);


	return (ret);
}



bool DiosLock::Try_Clear_DeadThread()
{
	DWORD Tid = m_InterLock;//just get it
	if (Tid == 0)
	{
		return true;
	}

	int FoundT = FindProcessThreads(Tid);
	if (FoundT == 0)
	{
		//目标不存在
		return Thread_Clear(Tid);
	}
	else if (FoundT > 0)
	{
		return false;
	}

	return false;



}

#pragma warning( push )
#pragma warning( disable : 4114 )


DIOSLOCK_C_API bool Dios_ThreadLock(volatile UINT64 volatile *pLock)
{

	DWORD *pLockData = (DWORD*)pLock;
	DWORD *pRefCount = (DWORD*)&pLockData[1];
	DWORD Tid = GetCurrentThreadId();

	//first try
	if (InterlockedCompareExchangeAcquire(pLockData, Tid, 0) == 0)
	{
		++(*pRefCount);
		return true;
	}

	//check first
	if (InterlockedCompareExchangeAcquire(pLockData, Tid, Tid) == Tid)
	{
		//already locked
		++(*pRefCount);
		return true;
	}

	//do loop
	while (1)
	{
		if (InterlockedCompareExchangeAcquire(pLockData, Tid, 0) == 0)
		{
			++(*pRefCount);
			return true;
		}

		//check exist tid


	}

	return false;
}

DIOSLOCK_C_API void Dios_ThreadUnLock(volatile UINT64 volatile *pLock)
{
	DWORD *pLockData = (DWORD*)pLock;
	DWORD *pRefCount = (DWORD*)&pLockData[1];
	DWORD Tid = GetCurrentThreadId();
	if (InterlockedCompareExchangeAcquire(pLockData, Tid, Tid) == Tid)
	{
		//lock owner is right
		if ((*pRefCount) > 0)
		{
			--(*pRefCount);
			if ((*pRefCount) > 0)
			{
				return;
			}

			if (InterlockedCompareExchangeRelease(pLockData, 0, Tid) == Tid)
			{
				return;
			}

		}
	}

	//ignore wrong unlocks

}

DIOSLOCK_C_API void Dios_ThreadClearForTid(volatile UINT64 volatile *pLock,DWORD Tid)
{
	DWORD *pLockData = (DWORD*)pLock;
	DWORD *pRefCount = (DWORD*)&pLockData[1];
	if (InterlockedCompareExchangeAcquire(pLockData, Tid, Tid) == Tid)
	{
		(*pRefCount) = 0;

		if (InterlockedCompareExchangeRelease(pLockData, 0, Tid) == Tid)
		{
			return;
		}

	}

	//ignore wrong unlocks

}


#pragma warning( pop )



DIOSLOCK_C_API DWORD Thread_GetUniqTick()
{
	static volatile UINT64 g_IdxLock = 0;

	static volatile DWORD oldTick = 1;
	DWORD tick = 0;

	if (Dios_ThreadLock(&g_IdxLock))
	{
		if (oldTick == 0)
		{
			oldTick = 1;
		}


		tick = oldTick++;


		Dios_ThreadUnLock(&g_IdxLock);

	}


	return tick;

}