diff --git a/c_src/eNifLock/eNifLock.cc b/c_src/eNifLock/eNifLock.cc index 2839a41..84527e3 100644 --- a/c_src/eNifLock/eNifLock.cc +++ b/c_src/eNifLock/eNifLock.cc @@ -1,5 +1,6 @@ #include "erl_nif.h" #include +#include using namespace std; const int LockSize = 2097152; @@ -10,20 +11,6 @@ ERL_NIF_TERM atomTrue; ERL_NIF_TERM atomFalse; ERL_NIF_TERM atomUndefined; -typedef struct KeyNode_r{ - int KeyIx; - struct KeyNode_r *next; -} KeyNode; - -bool isNotCurLocked(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"); @@ -50,80 +37,79 @@ bool lockOne(ErlNifEnv *env, ErlNifPid *ThePid, int KeyIx, uint64_t Val){ } 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 (isNotCurLocked(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 atomFalse; - } - } - } + int KeyIx; + enif_get_int(env, argv[0], &KeyIx); + ErlNifPid ThePid; + enif_self(env, &ThePid); + uint64_t Val = (uint64_t)(ThePid.pid); + + if (lockOne(env, &ThePid, KeyIx, Val)){ 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; + return atomFalse; + } +} + +ERL_NIF_TERM tryLocks(ErlNifEnv *env, int, const ERL_NIF_TERM argv[]){ + 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; + int cnt = -1; + + while (enif_get_list_cell(env, allList, &head, &allList)){ + enif_get_int(env, head, &KeyIx); + if(lockOne(env, &ThePid, KeyIx, Val)){ + cnt++; }else{ + uint64_t RExpected; + allList = argv[0]; + for(int i = 0; i <= cnt; i++){ + enif_get_list_cell(env, allList, &head, &allList); + enif_get_int(env, head, &KeyIx); + RExpected = Val; + LockSlot[KeyIx].compare_exchange_strong(RExpected, 0); + } return atomFalse; } } + return atomTrue; } 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; + int KeyIx; + enif_get_int(env, argv[0], &KeyIx); + ErlNifPid ThePid; + enif_self(env, &ThePid); + uint64_t Expected = (uint64_t)(ThePid.pid); + + if (LockSlot[KeyIx].compare_exchange_strong(Expected, 0)){ + return atomTrue; }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; + return atomFalse; + } +} + +ERL_NIF_TERM releaseLocks(ErlNifEnv *env, int, const ERL_NIF_TERM argv[]){ + 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)){ + enif_get_int(env, head, &KeyIx); + RExpected = Expected; + if (!LockSlot[KeyIx].compare_exchange_strong(RExpected, 0)){ + isAllOk = 0; } } + return isAllOk > 0 ? atomTrue : atomFalse; } ERL_NIF_TERM getLockPid(ErlNifEnv *env, int, const ERL_NIF_TERM argv[]){ @@ -142,7 +128,9 @@ ERL_NIF_TERM getLockPid(ErlNifEnv *env, int, const ERL_NIF_TERM argv[]){ static ErlNifFunc nifFuncs[] = { {"tryLock", 1, tryLock}, + {"tryLocks", 1, tryLocks}, {"releaseLock", 1, releaseLock}, + {"releaseLocks", 1, releaseLocks}, {"getLockPid", 1, getLockPid} }; diff --git a/priv/eNifLock.so b/priv/eNifLock.so index 20b89d7..4a31770 100644 Binary files a/priv/eNifLock.so and b/priv/eNifLock.so differ diff --git a/src/eGLock.erl b/src/eGLock.erl index 76768fb..5f4b44e 100644 --- a/src/eGLock.erl +++ b/src/eGLock.erl @@ -6,6 +6,10 @@ -define(LockTimeOut, 5000). %% 超时重试时间单位:Ms -define(ReTryTime, 3). +%% 数组数量 +-define(eGLockSize, 2097152). + +-compile([export_all]). -export([ tryLock/1 @@ -21,45 +25,81 @@ tryLock(KeyOrKeys) -> tryLock(KeyOrKeys, ?LockTimeOut). tryLock(KeyOrKeys, TimeOut) -> - doTryLock(KeyOrKeys, TimeOut). + case is_list(KeyOrKeys) of + true -> + KeyIxs = getKexIxs(KeyOrKeys, []), + doTryLocks(KeyIxs, TimeOut); + _ -> + doTryLock(erlang:phash2(KeyOrKeys, ?eGLockSize), TimeOut) + end. -doTryLock(KeyOrKeys, TimeOut) -> - case eNifLock:tryLock(KeyOrKeys) of +doTryLock(KeyIx, TimeOut) -> + case eNifLock:tryLock(KeyIx) of true -> true; _ -> - loopLock(KeyOrKeys, TimeOut) + loopLock(KeyIx, TimeOut) end. -loopLock(KeyOrKeys, TimeOut) -> +loopLock(KeyIx, TimeOut) -> receive after ?ReTryTime -> LTimeOut = ?CASE(TimeOut == infinity, TimeOut, TimeOut - ?ReTryTime), case LTimeOut >= 0 of true -> - case eNifLock:tryLock(KeyOrKeys) of + case eNifLock:tryLock(KeyIx) of true -> true; _ -> - loopLock(KeyOrKeys, LTimeOut) + loopLock(KeyIx, LTimeOut) end; _ -> ltimeout end end. +doTryLocks(KeyIxs, TimeOut) -> + case eNifLock:tryLocks(KeyIxs) of + true -> + true; + _ -> + loopLocks(KeyIxs, TimeOut) + end. + +loopLocks(KeyIxs, TimeOut) -> + receive + after ?ReTryTime -> + LTimeOut = ?CASE(TimeOut == infinity, TimeOut, TimeOut - ?ReTryTime), + case LTimeOut >= 0 of + true -> + case eNifLock:tryLocks(KeyIxs) of + true -> + true; + _ -> + loopLocks(KeyIxs, LTimeOut) + end; + _ -> + ltimeout + end + end. -spec releaseLock(KeyOrKeys :: term() | [term()]) -> ok. releaseLock(KeyOrKeys) -> - eNifLock:releaseLock(KeyOrKeys). + case is_list(KeyOrKeys) of + true -> + KeyIxs = getKexIxs(KeyOrKeys, []), + eNifLock:releaseLocks(KeyIxs); + _ -> + eNifLock:releaseLock(erlang:phash2(KeyOrKeys, ?eGLockSize)) + end. -spec getLockPid(KeyOrKeys :: term() | [term()]) -> ok. getLockPid(KeyOrKeys) -> case is_list(KeyOrKeys) of true -> - [{OneKey, eNifLock:getLockPid(OneKey)} || OneKey <- KeyOrKeys]; + [{OneKey, eNifLock:getLockPid(erlang:phash2(OneKey, ?eGLockSize))} || OneKey <- KeyOrKeys]; _ -> - {KeyOrKeys, eNifLock:getLockPid(KeyOrKeys)} + {KeyOrKeys, eNifLock:getLockPid(erlang:phash2(KeyOrKeys, ?eGLockSize))} end. -spec lockApply(KeyOrKeys :: term() | [term()], MFAOrFun :: {M :: atom(), F :: atom(), Args :: list()} | {Fun :: function(), Args :: list()}) -> term() | {error, ltimeout} | {error, {lock_apply_error, term()}}. @@ -68,19 +108,48 @@ lockApply(KeyOrKeys, MFAOrFun) -> -spec lockApply(KeyOrKeys :: term() | [term()], MFAOrFun :: {M :: atom(), F :: atom(), Args :: list()} | {Fun :: function(), Args :: list()}, TimeOut :: integer() | infinity) -> term(). lockApply(KeyOrKeys, MFAOrFun, TimeOut) -> - case doTryLock(KeyOrKeys, TimeOut) of + case is_list(KeyOrKeys) of true -> - try doApply(MFAOrFun) - catch C:R:S -> - {error, {lock_apply_error, {C, R, S}}} - after - eNifLock:releaseLock(KeyOrKeys), - ok + KeyIxs = getKexIxs(KeyOrKeys, []), + case doTryLocks(KeyIxs, TimeOut) of + true -> + try doApply(MFAOrFun) + catch C:R:S -> + {error, {lock_apply_error, {C, R, S}}} + after + eNifLock:releaseLocks(KeyIxs), + ok + end; + ltimeout -> + {error, ltimeout} end; - ltimeout -> - {error, ltimeout} + _ -> + KeyIx = erlang:phash2(KeyOrKeys, ?eGLockSize), + case doTryLock(KeyIx, TimeOut) of + true -> + try doApply(MFAOrFun) + catch C:R:S -> + {error, {lock_apply_error, {C, R, S}}} + after + eNifLock:releaseLock(KeyIx), + ok + end; + ltimeout -> + {error, ltimeout} + end end. +getKexIxs([], IxAcc) -> IxAcc; +getKexIxs([Key | Keys], IxAcc) -> + KeyIx = erlang:phash2(Key, ?eGLockSize), + getKexIxs(Keys, ?CASE(lists:member(KeyIx, IxAcc), IxAcc, [KeyIx | IxAcc])). + + +getKexIxs1([], IxAcc) -> IxAcc; +getKexIxs1([Key | Keys], IxAcc) -> + KeyIx = erlang:phash2(Key, ?eGLockSize), + getKexIxs1(Keys, ordsets:add_element(KeyIx, IxAcc)). + doApply({M, F, Args}) -> apply(M, F, Args); doApply({Fun, Args}) -> diff --git a/src/eNifLock.erl b/src/eNifLock.erl index 273d556..0aef094 100644 --- a/src/eNifLock.erl +++ b/src/eNifLock.erl @@ -2,7 +2,9 @@ -export([ tryLock/1 + , tryLocks/1 , releaseLock/1 + , releaseLocks/1 , getLockPid/1 ]). @@ -23,14 +25,22 @@ init() -> end, erlang:load_nif(SoName, 0). --spec tryLock(_KeyOrKeys :: term() | [term()]) -> true | false. -tryLock(_KeyOrKeys) -> +-spec tryLock(KeyIx :: non_neg_integer()) -> true | false. +tryLock(_KeyIx) -> erlang:nif_error({not_loaded, [{module, ?MODULE}, {line, ?LINE}]}). --spec releaseLock(_KeyOrKeys :: term() | [term()]) -> true | false. -releaseLock(_KeyOrKeys) -> +-spec tryLocks(KeyIxs :: [non_neg_integer()]) -> true | false. +tryLocks(_KeyIxs) -> erlang:nif_error({not_loaded, [{module, ?MODULE}, {line, ?LINE}]}). --spec getLockPid(OneKey :: term()) -> pid() | undefined. -getLockPid(_OneKey) -> +-spec releaseLock(KeyIx :: non_neg_integer()) -> true | false. +releaseLock(_KeyIx) -> + erlang:nif_error({not_loaded, [{module, ?MODULE}, {line, ?LINE}]}). + +-spec releaseLocks(KeyIxs :: [non_neg_integer()]) -> true | false. +releaseLocks(_KeyIxs) -> + erlang:nif_error({not_loaded, [{module, ?MODULE}, {line, ?LINE}]}). + +-spec getLockPid(KeyIx :: non_neg_integer()) -> pid() | undefined. +getLockPid(_KeyIx) -> erlang:nif_error({not_loaded, [{module, ?MODULE}, {line1, ?LINE}]}).