From a11fb37fceb2c3f1bb473c134ae892c58cf3466f Mon Sep 17 00:00:00 2001
From: SisMaker <1713699517@qq.com>
Date: Tue, 17 Dec 2024 17:30:26 +0800
Subject: [PATCH] =?UTF-8?q?ft:=20=E6=B7=BB=E5=8A=A0transaction=20=E4=BA=8B?=
 =?UTF-8?q?=E5=8A=A1=E6=89=A7=E8=A1=8C?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/eGLock.erl | 92 +++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 72 insertions(+), 20 deletions(-)

diff --git a/src/eGLock.erl b/src/eGLock.erl
index 05ce52a..ecca3a6 100644
--- a/src/eGLock.erl
+++ b/src/eGLock.erl
@@ -10,6 +10,9 @@
 %% 数组数量
 -define(eGLockSize, 2097152).
 
+%% 没有ets 表的key
+-define(undefTab, undefTab).
+
 -export([
 	tryLock/1
 	, tryLock/2
@@ -19,13 +22,11 @@
 	, lockApply/3
 ]).
 
-
 -export([
 	lockGet/1
 	, lockGet/2
 	, transaction/2
 	, transaction/3
-
 ]).
 
 -spec tryLock(KeyOrKeys :: term() | [term()]) -> true | lockTimeout.
@@ -162,6 +163,21 @@ lockGet(KeyOrKeys) ->
 	lockGet(KeyOrKeys, ?LockTimeOut).
 
 -spec lockGet(KeyOrKeys :: term() | [term()], TimeOut :: integer() | infinity) -> ok.
+lockGet({?undefTab, Key = GetKey}, TimeOut) ->
+	KeyIx = erlang:phash2(GetKey, ?eGLockSize),
+	case doTryLock(KeyIx, TimeOut) of
+		true ->
+			try
+				#{{?undefTab, GetKey} => undefined}
+			catch C:R:S ->
+				{error, {lockGetError, {?undefTab, Key}, {C, R, S}}}
+			after
+				eNifLock:releaseLock(KeyIx),
+				ok
+			end;
+		lockTimeout ->
+			{error, lockTimeout}
+	end;
 lockGet({EtsTab, Key} = GetKey, TimeOut) ->
 	KeyIx = erlang:phash2(GetKey, ?eGLockSize),
 	case doTryLock(KeyIx, TimeOut) of
@@ -177,6 +193,21 @@ lockGet({EtsTab, Key} = GetKey, TimeOut) ->
 		lockTimeout ->
 			{error, lockTimeout}
 	end;
+lockGet({?undefTab, Key = GetKey, DefValue}, TimeOut) ->
+	KeyIx = erlang:phash2(GetKey, ?eGLockSize),
+	case doTryLock(KeyIx, TimeOut) of
+		true ->
+			try
+				#{{?undefTab, GetKey} => getDefValue(DefValue)}
+			catch C:R:S ->
+				{error, {lockGetError, {?undefTab, Key, DefValue}, {C, R, S}}}
+			after
+				eNifLock:releaseLock(KeyIx),
+				ok
+			end;
+		lockTimeout ->
+			{error, lockTimeout}
+	end;
 lockGet({EtsTab, Key, DefValue}, TimeOut) ->
 	GetKey = {EtsTab, Key},
 	KeyIx = erlang:phash2(GetKey, ?eGLockSize),
@@ -194,9 +225,7 @@ lockGet({EtsTab, Key, DefValue}, TimeOut) ->
 			{error, lockTimeout}
 	end;
 lockGet(EtsTabKeys, TimeOut) ->
-	KeysMap = #{{EtsTab, Key} => DefValue || OneEtsTabKey <- EtsTabKeys, case OneEtsTabKey of {EtsTab, Key} -> DefValue = undefined, true; {EtsTab, Key, DefValue} -> true end},
-	KeyOrKeys = maps:keys(KeysMap),
-	KeyIxs = getKexIxs(KeyOrKeys, []),
+	{KeyIxs, KeysMap} = getKeyIxAndMaps(EtsTabKeys, [], #{}),
 	case doTryLocks(KeyIxs, TimeOut) of
 		true ->
 			try
@@ -211,6 +240,28 @@ lockGet(EtsTabKeys, TimeOut) ->
 			{error, lockTimeout}
 	end.
 
+getKeyIxAndMaps([], IxAcc, KeysMap) -> {IxAcc, KeysMap};
+getKeyIxAndMaps([Key | Keys], IxAcc, KeysMap) ->
+	case Key of
+		{?undefTab, JustKey} ->
+			GetKey = Key,
+			LDefValue = undefined,
+			KeyIx = erlang:phash2(Key, ?eGLockSize);
+		{?undefTab, JustKey, DefValue} ->
+			GetKey = {?undefTab, JustKey},
+			LDefValue = DefValue,
+			KeyIx = erlang:phash2(Key, ?eGLockSize);
+		{EtsTab, TabKey} ->
+			GetKey = Key,
+			LDefValue = undefined,
+			KeyIx = erlang:phash2(Key, ?eGLockSize);
+		{EtsTab, TabKey, DefValue} ->
+			GetKey = {EtsTab, TabKey},
+			LDefValue = DefValue,
+			KeyIx = erlang:phash2(Key, ?eGLockSize)
+	end,
+	getKeyIxAndMaps(Keys, ?CASE(lists:member(KeyIx, IxAcc), IxAcc, [KeyIx | IxAcc]), KeysMap#{GetKey => LDefValue}).
+
 transactionApply({M, F, Args}, EtsTabValue) ->
 	apply(M, F, [EtsTabValue | Args]);
 transactionApply({Fun, Args}, EtsTabValue) ->
@@ -222,22 +273,22 @@ transaction(EtsTabKeys, MFAOrFun) ->
 
 -spec transaction(KeyOrKeys :: term() | [term()], MFAOrFun :: {M :: atom(), F :: atom(), Args :: list()} | {Fun :: function(), Args :: list()}, TimeOut :: integer() | infinity) -> term().
 transaction(EtsTabKeys, MFAOrFun, TimeOut) ->
-	KeysMap = #{{EtsTab, Key} => DefValue || OneEtsTabKey <- EtsTabKeys, case OneEtsTabKey of {EtsTab, Key} -> DefValue = undefined, true; {EtsTab, Key, DefValue} -> true end},
-	KeyOrKeys = maps:keys(KeysMap),
-	KeyIxs = getKexIxs(KeyOrKeys, []),
+	{KeyIxs, KeysMap} = getKeyIxAndMaps(EtsTabKeys, [], #{}),
 	case doTryLocks(KeyIxs, TimeOut) of
 		true ->
 			try
 				EtsTabValue = #{OneGetKey => getEtsTabValue(OneEtsTab, OneKey, OneDefValue) || {OneEtsTab, OneKey} = OneGetKey := OneDefValue <- KeysMap},
 				case transactionApply(MFAOrFun, EtsTabValue) of
-					{Ret, ChangeEtsTab} ->
+					{ok, Ret, ChangeEtsTab} ->
 						[changeEtsTabValue(OneEtsTab, OneKey, ChangeValue) || {OneEtsTab, OneKey} := ChangeValue <- ChangeEtsTab],
 						Ret;
-					Ret ->
-						Ret
+					{ok, Ret} ->
+						Ret;
+					{error, Err} ->
+						Err
 				end
 			catch
-				throw:Throw-> Throw;
+				throw:Throw -> Throw;
 				C:R:S ->
 					{error, {lockTransactionError, {MFAOrFun, EtsTabKeys}, {C, R, S}}}
 			after
@@ -248,22 +299,23 @@ transaction(EtsTabKeys, MFAOrFun, TimeOut) ->
 			{error, lockTimeout}
 	end.
 
-getEtsTabValue(undefined, _Key, DefValue) ->
-	DefValue;
+getDefValue(undefined) -> undefined;
+getDefValue({DefFun, Args}) when is_function(DefFun) -> erlang:apply(DefFun, Args);
+getDefValue(DefValue) -> DefValue.
+
+getEtsTabValue(?undefTab, _Key, DefValue) ->
+	getDefValue(DefValue);
 getEtsTabValue(Ets, Key, DefValue) ->
 	case ets:lookup(Ets, Key) of
 		[] ->
-			DefValue;
+			getDefValue(DefValue);
 		[OneValue] ->
 			OneValue
 	end.
 
-changeEtsTabValue(undefined, _Key, _Value) ->
+changeEtsTabValue(?undefTab, _Key, _Value) ->
 	ok;
 changeEtsTabValue(Ets, Key, delete) ->
 	ets:delete(Ets, Key);
 changeEtsTabValue(Ets, _Key, Value) ->
-	ets:insert(Ets, Value).
-
-	
-
+	ets:insert(Ets, Value).
\ No newline at end of file