eGLock/c_src/eNifLock/eNifLock.cc

#include "erl_nif.h"
#include <atomic>
using namespace std;

const int LockSize = 2097152;
const int HashSalt = 786234121;
atomic<uint64_t> LockSlot[LockSize];

ERL_NIF_TERM atomTrue;
ERL_NIF_TERM atomFalse;
ERL_NIF_TERM atomUndefined;

typedef struct KeyNode_r
{
    int KeyIx;
    struct KeyNode_r *next;
} KeyNode;

bool isNotLocked(KeyNode *LockedHead, int KeyIx)
{
    KeyNode *temp = LockedHead;
    while (temp != NULL)
    {
        if (temp->KeyIx == KeyIx)
            return false;
        temp = temp->next;
    }
    return true;
}

int nifLoad(ErlNifEnv *env, void **priv_data, ERL_NIF_TERM)
{
    atomTrue = enif_make_atom(env, "true");
    atomFalse = enif_make_atom(env, "false");
    atomUndefined = enif_make_atom(env, "undefined");
    return 0;
}

bool lockOne(ErlNifEnv *env, ErlNifPid *ThePid, int KeyIx, uint64_t Val)
{
    uint64_t Expected = 0;
    if (LockSlot[KeyIx].compare_exchange_strong(Expected, Val))
    {
        return true;
    }
    else
    {
        ThePid->pid = (ERL_NIF_TERM)Expected;
        if (enif_is_process_alive(env, ThePid))
        {
            return false;
        }
        else
        {
            if (LockSlot[KeyIx].compare_exchange_strong(Expected, Val))
            {
                return true;
            }
            else
            {
                return false;
            }
        }
    }
}

ERL_NIF_TERM tryLock(ErlNifEnv *env, int, const ERL_NIF_TERM argv[])
{
    if (enif_is_list(env, argv[0]))
    {
        ERL_NIF_TERM allList = argv[0];
        ERL_NIF_TERM head;
        ErlNifPid ThePid;
        enif_self(env, &ThePid);
        uint64_t Val = (uint64_t)(ThePid.pid);
        int KeyIx;
        KeyNode *LockedHead = NULL;
        while (enif_get_list_cell(env, allList, &head, &allList))
        {
            KeyIx = enif_hash(ERL_NIF_INTERNAL_HASH, head, HashSalt) % LockSize;
            KeyNode OneKeyNode = {KeyIx, LockedHead};
            if (isNotLocked(LockedHead, KeyIx))
            {
                if (lockOne(env, &ThePid, KeyIx, Val))
                {
                    LockedHead = &OneKeyNode;
                }
                else
                {
                    uint64_t RExpected;
                    KeyNode *temp = LockedHead;
                    while (temp != NULL)
                    {
                        RExpected = Val;
                        LockSlot[temp->KeyIx].compare_exchange_strong(RExpected, 0);
                        temp = temp->next;
                    }
                }
            }
        }
        return atomTrue;
    }
    else
    {
        int KeyIx;
        KeyIx = enif_hash(ERL_NIF_INTERNAL_HASH, argv[0], HashSalt) % LockSize;
        ErlNifPid ThePid;
        enif_self(env, &ThePid);
        uint64_t Val = (uint64_t)(ThePid.pid);

        if (lockOne(env, &ThePid, KeyIx, Val))
        {
            return atomTrue;
        }
        else
        {
            return atomFalse;
        }
    }
}

ERL_NIF_TERM releaseLock(ErlNifEnv *env, int, const ERL_NIF_TERM argv[])
{
    if (enif_is_list(env, argv[0]))
    {
        ERL_NIF_TERM allList = argv[0];
        ERL_NIF_TERM head;
        ErlNifPid ThePid;
        enif_self(env, &ThePid);
        uint64_t Expected = (uint64_t)(ThePid.pid);
        uint64_t RExpected;
        int KeyIx;
        int isAllOk = 1;

        while (enif_get_list_cell(env, allList, &head, &allList))
        {
            KeyIx = enif_hash(ERL_NIF_INTERNAL_HASH, head, HashSalt) % LockSize;
            RExpected = Expected;
            if (!LockSlot[KeyIx].compare_exchange_strong(RExpected, 0))
            {
                isAllOk = 0;
            }
        }
        return isAllOk > 0 ? atomTrue : atomFalse;
    }
    else
    {
        int KeyIx;
        KeyIx = enif_hash(ERL_NIF_INTERNAL_HASH, argv[0], HashSalt) % LockSize;
        ErlNifPid ThePid;
        enif_self(env, &ThePid);
        uint64_t Expected = (uint64_t)(ThePid.pid);

        if (LockSlot[KeyIx].compare_exchange_strong(Expected, 0))
        {
            return atomTrue;
        }
        else
        {
            return atomFalse;
        }
    }
}

ERL_NIF_TERM getLockPid(ErlNifEnv *env, int, const ERL_NIF_TERM argv[])
{
    int KeyIx;
    KeyIx = enif_hash(ERL_NIF_INTERNAL_HASH, argv[0], HashSalt) % LockSize;
    ErlNifPid ThePid;

    uint64_t Var = LockSlot[KeyIx].load();
    if (Var > 0)
    {
        ThePid.pid = (ERL_NIF_TERM)Var;
        return enif_make_pid(env, &ThePid);
    }
    else
    {
        return atomUndefined;
    }
}

static ErlNifFunc nifFuncs[] = {
    {"tryLock", 1, tryLock},
    {"releaseLock", 1, releaseLock},
    {"getLockPid", 1, getLockPid}};

ERL_NIF_INIT(eNifLock, nifFuncs, &nifLoad, NULL, NULL, NULL)