Add Lock Data Incr and Append operation support and add related test …

…cases
snower · Apr 24, 2024 · ed49618 · ed49618
1 parent 1d91771
commit ed49618
Show file tree

Hide file tree

Showing 4 changed files with 181 additions and 12 deletions.
diff --git a/client/lock_test.go b/client/lock_test.go
@@ -525,6 +525,70 @@ func TestLock_WithData(t *testing.T) {
 			t.Errorf("Lock Unlock Result LockData Fail %v", result.GetLockData())
 			return
 		}
+
+		lock = client.Lock(testString2Key("TestData3"), 50, 10)
+		lock.SetCount(10)
+		result, err = lock.LockWithData(protocol.NewLockCommandDataIncrData(2))
+		if err != nil {
+			t.Errorf("Lock LockWithData Incr Fail %v", err)
+			return
+		}
+		if result.GetLockData() != nil {
+			t.Errorf("Lock LockWithData Incr Result LockData Fail %v", result.GetLockData())
+			return
+		}
+		ulock1 = client.Lock(testString2Key("TestData3"), 50, 0)
+		ulock1.SetCount(10)
+		result, err = ulock1.LockWithData(protocol.NewLockCommandDataIncrData(-3))
+		if err != nil {
+			t.Errorf("Lock LockWithData1 Incr Expried Fail %v", err)
+			return
+		}
+		if result.GetLockData() == nil || result.GetLockData().GetIncrValue() != 2 {
+			t.Errorf("Lock LockWithData1 Incr Expried Result LockData Fail %v", result.GetLockData())
+			return
+		}
+		result, err = lock.Unlock()
+		if err != nil {
+			t.Errorf("Lock Unlock Incr Fail %v", err)
+			return
+		}
+		if result.GetLockData() == nil || result.GetLockData().GetIncrValue() != -1 {
+			t.Errorf("Lock Unlock Incr Result LockData Fail %v", result.GetLockData())
+			return
+		}
+
+		lock = client.Lock(testString2Key("TestData4"), 50, 10)
+		lock.SetCount(10)
+		result, err = lock.LockWithData(protocol.NewLockCommandDataAppendString("aaa"))
+		if err != nil {
+			t.Errorf("Lock LockWithData Append Fail %v", err)
+			return
+		}
+		if result.GetLockData() != nil {
+			t.Errorf("Lock LockWithData Append Result LockData Fail %v", result.GetLockData())
+			return
+		}
+		ulock1 = client.Lock(testString2Key("TestData4"), 50, 0)
+		ulock1.SetCount(10)
+		result, err = ulock1.LockWithData(protocol.NewLockCommandDataAppendString("bbb"))
+		if err != nil {
+			t.Errorf("Lock LockWithData1 Append Expried Fail %v", err)
+			return
+		}
+		if result.GetLockData() == nil || result.GetLockData().GetStringData() != "aaa" {
+			t.Errorf("Lock LockWithData1 Append Expried Result LockData Fail %v", result.GetLockData())
+			return
+		}
+		result, err = lock.Unlock()
+		if err != nil {
+			t.Errorf("Lock Unlock Append Fail %v", err)
+			return
+		}
+		if result.GetLockData() == nil || result.GetLockData().GetStringData() != "aaabbb" {
+			t.Errorf("Lock Unlock Append Result LockData Fail %v", result.GetLockData())
+			return
+		}
 	})
 }
 

diff --git a/protocol/command.go b/protocol/command.go
@@ -95,8 +95,10 @@ const (
 )
 
 const (
-	LOCK_DATA_COMMAND_TYPE_SET   = 0
-	LOCK_DATA_COMMAND_TYPE_UNSET = 1
+	LOCK_DATA_COMMAND_TYPE_SET    = 0
+	LOCK_DATA_COMMAND_TYPE_UNSET  = 1
+	LOCK_DATA_COMMAND_TYPE_INCR   = 2
+	LOCK_DATA_COMMAND_TYPE_APPEND = 3
 )
 
 var ERROR_MSG []string = []string{
@@ -371,7 +373,21 @@ func NewLockCommandDataSetString(data string) *LockCommandData {
 }
 
 func NewLockCommandDataUnsetData() *LockCommandData {
-	return NewLockCommandDataFromBytes([]byte{}, LOCK_DATA_COMMAND_TYPE_UNSET, 0)
+	return &LockCommandData{[]byte{2, 0, 0, 0, LOCK_DATA_COMMAND_TYPE_UNSET, 0}, LOCK_DATA_COMMAND_TYPE_UNSET, 0}
+}
+
+func NewLockCommandDataIncrData(incrValue int64) *LockCommandData {
+	return &LockCommandData{[]byte{10, 0, 0, 0, LOCK_DATA_COMMAND_TYPE_INCR, 0,
+		byte(incrValue), byte(incrValue >> 8), byte(incrValue >> 16), byte(incrValue >> 24), byte(incrValue >> 32), byte(incrValue >> 40), byte(incrValue >> 48), byte(incrValue >> 56)},
+		LOCK_DATA_COMMAND_TYPE_INCR, 0}
+}
+
+func NewLockCommandDataAppendData(data []byte) *LockCommandData {
+	return NewLockCommandDataFromBytes(data, LOCK_DATA_COMMAND_TYPE_APPEND, 0)
+}
+
+func NewLockCommandDataAppendString(data string) *LockCommandData {
+	return NewLockCommandDataFromString(data, LOCK_DATA_COMMAND_TYPE_APPEND, 0)
 }
 
 func (self *LockCommandData) GetBytesData() []byte {
@@ -388,6 +404,24 @@ func (self *LockCommandData) GetStringData() string {
 	return string(self.Data[6:])
 }
 
+func (self *LockCommandData) GetIncrValue() int64 {
+	if self.Data == nil || self.CommandType == LOCK_DATA_COMMAND_TYPE_UNSET {
+		return 0
+	}
+	value := int64(0)
+	for i := 0; i < 8; i++ {
+		if i+6 >= len(self.Data) {
+			break
+		}
+		if i > 0 {
+			value |= int64(self.Data[i+6]) << (i * 8)
+		} else {
+			value |= int64(self.Data[i+6])
+		}
+	}
+	return value
+}
+
 type LockCommand struct {
 	Command
 	Flag        uint8
@@ -528,6 +562,24 @@ func (self *LockResultCommandData) GetStringData() string {
 	return string(self.Data[6:])
 }
 
+func (self *LockResultCommandData) GetIncrValue() int64 {
+	if self.Data == nil || self.CommandType == LOCK_DATA_COMMAND_TYPE_UNSET {
+		return 0
+	}
+	value := int64(0)
+	for i := 0; i < 8; i++ {
+		if i+6 >= len(self.Data) {
+			break
+		}
+		if i > 0 {
+			value |= int64(self.Data[i+6]) << (i * 8)
+		} else {
+			value |= int64(self.Data[i+6])
+		}
+	}
+	return value
+}
+
 type LockResultCommand struct {
 	ResultCommand
 	Flag    uint8

diff --git a/run-tests b/run-tests
@@ -74,7 +74,8 @@ please run 'go run benchmarks/benchmark/benchmark.go' on your changes before com
     exit 1
 fi
 
-echo "Benchmark Multi Key Single Lock Info:
+echo "
+Benchmark Multi Key Single Lock Info:
 $BENCH_OUTPUTS
 "
 
@@ -92,7 +93,8 @@ please run 'go run benchmarks/benchmark/benchmark.go' on your changes before com
     exit 1
 fi
 
-echo "Benchmark Single Key Multi Lock Info:
+echo "
+Benchmark Single Key Multi Lock Info:
 $BENCH_OUTPUTS
 "
 
@@ -110,7 +112,8 @@ please run 'go run benchmarks/benchmark/benchmark.go' on your changes before com
     exit 1
 fi
 
-echo "Benchmark Multi Key Multi Lock Require Ack Info:
+echo "
+Benchmark Multi Key Multi Lock Require Ack Info:
 $BENCH_OUTPUTS
 "
 
@@ -128,7 +131,8 @@ please run 'go run benchmarks/benchmark/benchmark.go' on your changes before com
     exit 1
 fi
 
-echo "Benchmark Multi Key Multi Lock With Data Info:
+echo "
+Benchmark Multi Key Multi Lock With Data Info:
 $BENCH_OUTPUTS
 "
 
@@ -146,7 +150,8 @@ please run 'go run benchmarks/benchmark/benchmark.go' on your changes before com
     exit 1
 fi
 
-echo "Benchmark Multi Key Multi Require Ack And Lock With Data Info:
+echo "
+Benchmark Multi Key Multi Require Ack And Lock With Data Info:
 $BENCH_OUTPUTS
 "
 
@@ -164,7 +169,8 @@ please run 'go run benchmarks/benchmark/benchmark.go' on your changes before com
     exit 1
 fi
 
-echo "Benchmark Parallel Lock Info:
+echo "
+Benchmark Parallel Lock Info:
 $BENCH_OUTPUTS
 "
 

diff --git a/server/lock.go b/server/lock.go
@@ -362,7 +362,7 @@ func (self *LockManager) ProcessLockData(command *protocol.LockCommand) {
 				}
 			}
 		}
-		self.currentData = NewLockData(lockCommandData.Data)
+		self.currentData = NewLockData(lockCommandData.Data, protocol.LOCK_DATA_COMMAND_TYPE_SET)
 	case protocol.LOCK_DATA_COMMAND_TYPE_UNSET:
 		if command.CommandType == protocol.COMMAND_LOCK && command.ExpriedFlag&0x4440 == 0 && command.Expried == 0 {
 			if self.currentData != nil && self.currentData.commandType == protocol.LOCK_DATA_COMMAND_TYPE_UNSET {
@@ -371,6 +371,35 @@ func (self *LockManager) ProcessLockData(command *protocol.LockCommand) {
 			}
 		}
 		self.currentData = NewLockDataUnsetData()
+	case protocol.LOCK_DATA_COMMAND_TYPE_INCR:
+		incrValue := command.GetLockData().GetIncrValue()
+		if self.currentData != nil && self.currentData.GetData() != nil {
+			incrValue += self.currentData.GetIncrValue()
+			if len(command.Data.Data) == 14 {
+				data := command.Data.Data
+				data[4], data[5] = protocol.LOCK_DATA_COMMAND_TYPE_SET, 0
+				data[6], data[7], data[8], data[9], data[10], data[11], data[12], data[13] = byte(incrValue), byte(incrValue>>8), byte(incrValue>>16), byte(incrValue>>24), byte(incrValue>>32), byte(incrValue>>40), byte(incrValue>>48), byte(incrValue>>56)
+				self.currentData = NewLockData(data, protocol.LOCK_DATA_COMMAND_TYPE_INCR)
+				command.Data = nil
+				return
+			}
+		}
+		self.currentData = NewLockData([]byte{10, 0, 0, 0, protocol.LOCK_DATA_COMMAND_TYPE_SET, 0,
+			byte(incrValue), byte(incrValue >> 8), byte(incrValue >> 16), byte(incrValue >> 24), byte(incrValue >> 32), byte(incrValue >> 40), byte(incrValue >> 48), byte(incrValue >> 56)},
+			protocol.LOCK_DATA_COMMAND_TYPE_INCR)
+	case protocol.LOCK_DATA_COMMAND_TYPE_APPEND:
+		if self.currentData == nil || self.currentData.GetData() == nil {
+			lockCommandData.Data[4] = protocol.LOCK_DATA_COMMAND_TYPE_SET
+			self.currentData = NewLockData(lockCommandData.Data, protocol.LOCK_DATA_COMMAND_TYPE_APPEND)
+		} else {
+			dataLen := len(lockCommandData.Data) + len(self.currentData.data) - 10
+			data := make([]byte, dataLen+4)
+			data[0], data[1], data[2], data[3] = byte(dataLen), byte(dataLen>>8), byte(dataLen>>16), byte(dataLen>>24)
+			data[4], data[5] = protocol.LOCK_DATA_COMMAND_TYPE_SET, 0
+			copy(data[6:], self.currentData.data[6:])
+			copy(data[len(self.currentData.data):], lockCommandData.Data[6:])
+			self.currentData = NewLockData(data, protocol.LOCK_DATA_COMMAND_TYPE_APPEND)
+		}
 	}
 	command.Data = nil
 }
@@ -414,8 +443,8 @@ type LockData struct {
 	isAof       bool
 }
 
-func NewLockData(data []byte) *LockData {
-	return &LockData{data, nil, nil, data[4], false}
+func NewLockData(data []byte, commandType uint8) *LockData {
+	return &LockData{data, nil, nil, commandType, false}
 }
 
 func NewLockDataUnsetData() *LockData {
@@ -430,6 +459,24 @@ func (self *LockData) GetData() []byte {
 	return nil
 }
 
+func (self *LockData) GetIncrValue() int64 {
+	if self.data == nil || self.commandType == protocol.LOCK_DATA_COMMAND_TYPE_UNSET {
+		return 0
+	}
+	value := int64(0)
+	for i := 0; i < 8; i++ {
+		if i+6 >= len(self.data) {
+			break
+		}
+		if i > 0 {
+			value |= int64(self.data[i+6]) << (i * 8)
+		} else {
+			value |= int64(self.data[i+6])
+		}
+	}
+	return value
+}
+
 func (self *LockData) Equal(lockData []byte) bool {
 	return len(self.data) == len(lockData) && string(self.data) == string(lockData)
 }