ft: 完善的版本

c_src/eNifLock/eNifLock.cc

@@ -1,5 +1,6 @@
 #include "erl_nif.h"
 #include <atomic>
+#include <set>
 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}
 };
 

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}) ->

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}]}).