DriverPlatform-V3.1/CcosLock/CcosLock.cpp

#include "CcosLock.h"
#include <pthread.h>
#include <unistd.h>
#include <map>
#include <iostream>
#include <stdexcept>

// 进程锁实现
class ProcessLock {
    pthread_mutex_t m_Mutex = PTHREAD_MUTEX_INITIALIZER;
    pthread_cond_t m_Cond = PTHREAD_COND_INITIALIZER;
    std::atomic<bool> m_Locked{ false };
    
public:
    DWORD Lock(DWORD timeout) {
        if (pthread_mutex_lock(&m_Mutex) != 0) {
            throw std::runtime_error("Failed to lock process mutex");
        }
        
        if (!m_Locked) {
            m_Locked = true;
            pthread_mutex_unlock(&m_Mutex);
            return WAIT_OBJECT_0;
        }
        
        if (timeout == 0) {
            pthread_mutex_unlock(&m_Mutex);
            return WAIT_TIMEOUT;
        }
        
        struct timespec ts;
        if (timeout != INFINITE) {
            clock_gettime(CLOCK_REALTIME, &ts);
            ts.tv_sec += timeout / 1000;
            ts.tv_nsec += (timeout % 1000) * 1000000;
            if (ts.tv_nsec >= 1000000000) {
                ts.tv_sec++;
                ts.tv_nsec -= 1000000000;
            }
        }
        
        int ret = 0;
        while (m_Locked && ret == 0) {
            ret = (timeout == INFINITE) 
                ? pthread_cond_wait(&m_Cond, &m_Mutex)
                : pthread_cond_timedwait(&m_Cond, &m_Mutex, &ts);
        }

        if (ret == ETIMEDOUT) {  
            pthread_mutex_unlock(&m_Mutex);
            return WAIT_TIMEOUT;
        }
        
        m_Locked = true;
        pthread_mutex_unlock(&m_Mutex);
        return WAIT_OBJECT_0;
    }

    void Unlock() {
        if (pthread_mutex_lock(&m_Mutex) != 0) {
            throw std::runtime_error("Failed to lock process mutex");
        }
        m_Locked = false;
        pthread_cond_signal(&m_Cond);
        pthread_mutex_unlock(&m_Mutex);
    }
};

ProcessLock& GetProcessLock() {
    static ProcessLock instance;
    return instance;
}

pthread_mutex_t& GetMapMutex() {
    static pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;
    return mutex;
}

std::map<UINT64, DWORD>& GetLockedMap() {
    static std::map<UINT64, DWORD> instance;
    return instance;
}

std::map<UINT64, CcosLock*>& GetObjectMap() {
    static std::map<UINT64, CcosLock*> instance;
    return instance;
}

// 注册/注销锁对象
void RegisterLock(CcosLock* p) {
    pthread_mutex_lock(&GetMapMutex());
     GetObjectMap()[reinterpret_cast<UINT64>(p)] = p;
    pthread_mutex_unlock(&GetMapMutex());
}

void UnregisterLock(CcosLock* p) {
    pthread_mutex_lock(&GetMapMutex());
    UINT64 key = reinterpret_cast<UINT64>(p);
    GetObjectMap().erase(key);
    GetLockedMap().erase(key);
    pthread_mutex_unlock(&GetMapMutex());
}

// CcosLock 成员函数实现
CcosLock::CcosLock() {
    if (pthread_mutex_init(&m_Mutex, nullptr) != 0) {
        throw std::runtime_error("Failed to initialize mutex");
    }
    if (pthread_cond_init(&m_Cond, nullptr) != 0) {
        pthread_mutex_destroy(&m_Mutex);
        throw std::runtime_error("Failed to initialize condition variable");
    }
    RegisterLock(this);
}

CcosLock::~CcosLock() {
    m_IsDestroyed = true;
    UnregisterLock(this);
    // 等待所有等待线程唤醒后再销毁
    pthread_mutex_lock(&m_Mutex);
    pthread_cond_broadcast(&m_Cond);
    pthread_mutex_unlock(&m_Mutex);

    pthread_mutex_destroy(&m_Mutex);
    pthread_cond_destroy(&m_Cond);
}

bool CcosLock::CcosInternal_EnterCriticalSection(DWORD TryCount) {
    const DWORD tid = GetCurrentThreadId();
    DWORD expected = 0;
    
    // 尝试获取锁
    if (m_InterLock.compare_exchange_strong(expected, tid,
        std::memory_order_acq_rel,
        std::memory_order_relaxed)) {
        m_RefCount = 1;
        pthread_mutex_lock(&GetMapMutex());
        GetLockedMap()[reinterpret_cast<UINT64>(this)] = tid;
        pthread_mutex_unlock(&GetMapMutex());
        return true;
    }
    
    // 检查是否已持有锁
    if (expected == tid) {
        m_RefCount++;
        return true;
    }
    
    // 尝试多次获取锁
    while (TryCount-- > 0) {
        expected = 0;
        if (m_InterLock.compare_exchange_strong(expected, tid,
            std::memory_order_acq_rel,
            std::memory_order_relaxed)) {
            m_RefCount = 1;
            pthread_mutex_lock(&GetMapMutex());
            GetLockedMap()[reinterpret_cast<UINT64>(this)] = tid;
            pthread_mutex_unlock(&GetMapMutex());
            return true;
        }
        usleep(1000); // 1ms 延迟
    }
    return false;
}

void CcosLock::CcosInternal_LeaveCriticalSection() {
    const DWORD tid = GetCurrentThreadId();
    const DWORD current_tid = m_InterLock.load(std::memory_order_acquire);

    if (current_tid == tid) {
        if (--m_RefCount == 0) {
            // 清除锁状态前获取全局映射锁
            pthread_mutex_lock(&GetMapMutex());
            GetLockedMap().erase(reinterpret_cast<UINT64>(this));
            pthread_mutex_unlock(&GetMapMutex());

            m_InterLock.store(0, std::memory_order_release);

            // 通知等待的线程
            pthread_mutex_lock(&m_Mutex);
            pthread_cond_broadcast(&m_Cond);
            pthread_mutex_unlock(&m_Mutex);
        }
    }
    else {
        std::cerr << "Thread " << tid << " attempted to unlock lock owned by " << current_tid << std::endl;
    }
}

DWORD CcosLock::Thread_Lock(DWORD timeout) {
    if (IsDestroyed()) return WAIT_TIMEOUT;

    if (CcosInternal_EnterCriticalSection(1)) {
        return WAIT_OBJECT_0;
    }

    const DWORD start = GetTickCount();
    while (!IsDestroyed()) {
        const DWORD elapsed = GetTickCount() - start;
        if (timeout != INFINITE && elapsed >= timeout) {
            return WAIT_TIMEOUT;
        }
        
        const DWORD remaining = (timeout == INFINITE) ? 100 : (timeout - elapsed);
        const DWORD waitTime = std::min<DWORD>(remaining, 100);
        const DWORD ret = Thread_WaitUnlockNotify(waitTime);
        
        if (ret == WAIT_OBJECT_0 && CcosInternal_EnterCriticalSection(1)) {
            return WAIT_OBJECT_0;
        }
    }
    return WAIT_TIMEOUT;
}

void CcosLock::Thread_UnLock() {
    if (!IsDestroyed()) {
        CcosInternal_LeaveCriticalSection();
    }
}

bool CcosLock::Thread_Clear(DWORD tid) {
    if (m_InterLock == tid) {
        m_RefCount = 0;
        pthread_mutex_lock(&GetMapMutex());
        GetLockedMap().erase(reinterpret_cast<UINT64>(this));
        pthread_mutex_unlock(&GetMapMutex());
        m_InterLock.store(0, std::memory_order_release);
        pthread_cond_broadcast(&m_Cond);
        return true;
    }
    return false;
}

DWORD CcosLock::Thread_WaitUnlockNotify(DWORD timeout) {
    if (IsDestroyed()) return WAIT_TIMEOUT;

    pthread_mutex_lock(&m_Mutex);
    
    struct timespec ts;
    if (timeout != INFINITE) {
        clock_gettime(CLOCK_REALTIME, &ts);
        ts.tv_sec += timeout / 1000;
        ts.tv_nsec += (timeout % 1000) * 1000000;
        if (ts.tv_nsec >= 1000000000) {
            ts.tv_sec++;
            ts.tv_nsec -= 1000000000;
        }
    }

    int ret = 0;
    while (m_InterLock != 0 && ret == 0 && !IsDestroyed()) {
        ret = (timeout == INFINITE)
            ? pthread_cond_wait(&m_Cond, &m_Mutex)
            : pthread_cond_timedwait(&m_Cond, &m_Mutex, &ts);
    }
    
    pthread_mutex_unlock(&m_Mutex);
   
    if (IsDestroyed()) return WAIT_TIMEOUT;
    return (m_InterLock == 0) ? WAIT_OBJECT_0 : WAIT_TIMEOUT;
}

int CcosLock::FindProcessThreads(DWORD tid) {
    return IsThreadAlive(tid) ? 1 : 0;
}

bool CcosLock::Try_Clear_DeadThread() {
    return (m_InterLock != 0 && !FindProcessThreads(m_InterLock))
        ? Thread_Clear(m_InterLock) : false;
}

// 全局清理函数
void CleanupForThread(DWORD tid) {
    pthread_mutex_lock(&GetMapMutex());
    
    auto it = GetLockedMap().begin();
    while (it != GetLockedMap().end()) {
        if (it->second == tid) {
            if ( GetObjectMap().find(it->first) !=  GetObjectMap().end()) {
                 GetObjectMap()[it->first]->Thread_Clear(tid);
            }
            it = GetLockedMap().erase(it);
        } else {
            ++it;
        }
    }
    
    pthread_mutex_unlock(&GetMapMutex());
}

// C接口函数实现
bool Ccos_ThreadLock(ThreadLock *pLock) {
    if (!pLock) return false;
    const DWORD tid = GetCurrentThreadId();
    DWORD expected = 0;

    if (pLock->thread_id.compare_exchange_strong(expected, tid,
        std::memory_order_acq_rel,
        std::memory_order_relaxed)) {
        pLock->ref_count = 1;
        return true;
    }

    if (expected == tid) {
        pLock->ref_count++;
        return true;
    }
    
    return false;
}

void Ccos_ThreadUnLock(ThreadLock* pLock) {
    if (!pLock) return;

    const DWORD tid = GetCurrentThreadId();
    if (pLock->thread_id == tid) {
        if (--pLock->ref_count == 0) {
            pLock->thread_id.store(0, std::memory_order_release);
        }
    }
}

void Ccos_ThreadClearForTid(ThreadLock* pLock, DWORD tid) {
    if (!pLock) return;

    if (pLock->thread_id == tid) {
        pLock->ref_count = 0;
        pLock->thread_id.store(0, std::memory_order_release);
    }
}

DWORD Proc_Lock(DWORD timeout) {
    return GetProcessLock().Lock(timeout);
}

void Proc_UnLock() {
    GetProcessLock().Unlock();
}

DWORD Thread_GetUniqTick() {
    static std::atomic<DWORD> counter(1);
    return counter.fetch_add(1, std::memory_order_relaxed);
}