From 8268b755ef147ce0cef88fcd8b47cc1821bb20c9 Mon Sep 17 00:00:00 2001
From: Jermy Li <lizhangmei@baidu.com>
Date: Tue, 31 Dec 2019 15:06:33 +0800
Subject: [PATCH] add generic parameter for RowLock (#42)

Change-Id: I309b89846ca15d58ea39095d406a0b0b5de510e7
---
 .../java/com/baidu/hugegraph/concurrent/LockGroup.java   | 9 ++++++---
 .../baidu/hugegraph/unit/concurrent/LockGroupTest.java   | 6 +++---
 .../com/baidu/hugegraph/unit/concurrent/RowLockTest.java | 6 ++----
 3 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/src/main/java/com/baidu/hugegraph/concurrent/LockGroup.java b/src/main/java/com/baidu/hugegraph/concurrent/LockGroup.java
index a3f26ef9f..7d29c98ed 100644
--- a/src/main/java/com/baidu/hugegraph/concurrent/LockGroup.java
+++ b/src/main/java/com/baidu/hugegraph/concurrent/LockGroup.java
@@ -71,11 +71,14 @@ public KeyLock keyLock(String lockName, int size) {
         return (KeyLock) this.locksMap.get(lockName);
     }
 
-    public RowLock rowLock(String lockName) {
+    public <K extends Comparable<K>> RowLock<K> rowLock(String lockName) {
         if (!this.locksMap.containsKey(lockName)) {
-            this.locksMap.putIfAbsent(lockName, new RowLock());
+            this.locksMap.putIfAbsent(lockName, new RowLock<>());
         }
-        return (RowLock) this.locksMap.get(lockName);
+        Object value = this.locksMap.get(lockName);
+        @SuppressWarnings("unchecked")
+        RowLock<K> lock = (RowLock<K>) value;
+        return lock;
     }
 
     public String name() {
diff --git a/src/test/java/com/baidu/hugegraph/unit/concurrent/LockGroupTest.java b/src/test/java/com/baidu/hugegraph/unit/concurrent/LockGroupTest.java
index ed7c7425b..7a5230402 100644
--- a/src/test/java/com/baidu/hugegraph/unit/concurrent/LockGroupTest.java
+++ b/src/test/java/com/baidu/hugegraph/unit/concurrent/LockGroupTest.java
@@ -28,8 +28,8 @@
 
 import com.baidu.hugegraph.concurrent.AtomicLock;
 import com.baidu.hugegraph.concurrent.KeyLock;
-import com.baidu.hugegraph.concurrent.RowLock;
 import com.baidu.hugegraph.concurrent.LockGroup;
+import com.baidu.hugegraph.concurrent.RowLock;
 import com.baidu.hugegraph.testutil.Assert;
 
 public class LockGroupTest {
@@ -80,9 +80,9 @@ public void testKeyLockWithSize() {
 
     @Test
     public void testRowLock() {
-        RowLock lock = this.group.rowLock("lock");
+        RowLock<?> lock = this.group.rowLock("lock");
         Assert.assertNotNull(lock);
-        RowLock lock1 = this.group.rowLock("lock");
+        RowLock<?> lock1 = this.group.rowLock("lock");
         Assert.assertSame(lock, lock1);
     }
 
diff --git a/src/test/java/com/baidu/hugegraph/unit/concurrent/RowLockTest.java b/src/test/java/com/baidu/hugegraph/unit/concurrent/RowLockTest.java
index d7dd15326..bd8419bb0 100644
--- a/src/test/java/com/baidu/hugegraph/unit/concurrent/RowLockTest.java
+++ b/src/test/java/com/baidu/hugegraph/unit/concurrent/RowLockTest.java
@@ -62,10 +62,9 @@ public void testRowLockMultiRows() {
         lock.unlockAll(ImmutableSet.of(1, 2, 3));
     }
 
-    @SuppressWarnings("unchecked")
     @Test
     public void testRowLockWithMultiThreads() {
-        RowLock lock = new RowLock();
+        RowLock<Integer> lock = new RowLock<>();
         Set<String> names = new HashSet<>(THREADS_NUM);
         List<Integer> keys = new ArrayList<>(5);
         Random random = new Random();
@@ -84,10 +83,9 @@ public void testRowLockWithMultiThreads() {
         Assert.assertEquals(THREADS_NUM, names.size());
     }
 
-    @SuppressWarnings("unchecked")
     @Test
     public void testRowLockWithMultiThreadsWithRandomKey() {
-        RowLock lock = new RowLock();
+        RowLock<Integer> lock = new RowLock<>();
         Set<String> names = new HashSet<>(THREADS_NUM);
 
         Assert.assertEquals(0, names.size());