fix lock

star-whale · Oct 23, 2023 · d0baf45 · d0baf45
1 parent 1a716ae
commit d0baf45
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 20 deletions.
diff --git a/server/controller/src/main/java/ai/starwhale/mlops/datastore/DataStore.java b/server/controller/src/main/java/ai/starwhale/mlops/datastore/DataStore.java
@@ -201,7 +201,7 @@ public String update(String tableName,
         this.updateHandle.offer(new Object());
         var table = this.getTable(tableName, true, true);
         //noinspection ConstantConditions
-        table.lock();
+        table.lock(false);
         try {
             var ts = table.update(schema, records);
             synchronized (this.dirtyTables) {
@@ -213,7 +213,7 @@ public String update(String tableName,
             synchronized (updateHandle) {
                 updateHandle.notifyAll();
             }
-            table.unlock();
+            table.unlock(false);
         }
 
     }
@@ -227,7 +227,7 @@ public RecordList query(DataStoreQueryRequest req) {
         if (table == null) {
             return new RecordList(Collections.emptyMap(), Collections.emptyMap(), Collections.emptyList(), null, null);
         }
-        table.lock();
+        table.lock(true);
         try {
             int skipCount = req.getStart();
             if (skipCount < 0) {
@@ -303,7 +303,7 @@ public RecordList query(DataStoreQueryRequest req) {
                     lastKey,
                     lastKeyType);
         } finally {
-            table.unlock();
+            table.unlock(true);
         }
     }
 
@@ -327,7 +327,7 @@ public RecordList scan(DataStoreScanRequest req) {
                         .collect(Collectors.toList());
 
         for (var table : tablesToLock) {
-            table.lock();
+            table.lock(true);
         }
         try {
             class TableMeta {
@@ -495,7 +495,7 @@ record = new HashMap<>();
                     BaseValue.getColumnType(lastKey).name());
         } finally {
             for (var table : tablesToLock) {
-                table.unlock();
+                table.unlock(true);
             }
         }
     }
@@ -625,13 +625,13 @@ private void clearWalLogFiles() {
             if (table == null) {
                 continue;
             }
-            table.lock();
+            table.lock(false);
             try {
                 if (table.getFirstWalLogId() >= 0 && table.getFirstWalLogId() < minWalLogIdToRetain) {
                     minWalLogIdToRetain = table.getFirstWalLogId();
                 }
             } finally {
-                table.unlock();
+                table.unlock(false);
             }
         }
         try {

diff --git a/server/controller/src/main/java/ai/starwhale/mlops/datastore/MemoryTable.java b/server/controller/src/main/java/ai/starwhale/mlops/datastore/MemoryTable.java
@@ -50,9 +50,9 @@ Iterator<RecordResult> scan(
             boolean endInclusive,
             boolean keepNone);
 
-    void lock();
+    void lock(boolean forRead);
 
-    void unlock();
+    void unlock(boolean forRead);
 
     void save() throws IOException;
 

diff --git a/server/controller/src/main/java/ai/starwhale/mlops/datastore/impl/MemoryTableImpl.java b/server/controller/src/main/java/ai/starwhale/mlops/datastore/impl/MemoryTableImpl.java
@@ -64,7 +64,8 @@
 import java.util.TimeZone;
 import java.util.TreeMap;
 import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import lombok.Getter;
@@ -115,7 +116,9 @@ public class MemoryTableImpl implements MemoryTable {
 
     private final TreeMap<BaseValue, List<MemoryRecord>> recordMap = new TreeMap<>();
 
-    private final Lock lock = new ReentrantLock();
+    private final ReadWriteLock lock = new ReentrantReadWriteLock();
+    private final Lock readLock = lock.readLock();
+    private final Lock writeLock = lock.writeLock();
 
     public MemoryTableImpl(String tableName,
             WalManager walManager,
@@ -194,7 +197,7 @@ var record = reader.read();
     public void save() throws IOException {
         String metadata;
         var columnSchema = new HashMap<String, ColumnSchema>();
-        this.lock();
+        this.lock(false);
         var lastRevision = this.lastRevision;
         var firstWalLogId = this.firstWalLogId;
         this.firstWalLogId = -1;
@@ -220,7 +223,7 @@ public void save() throws IOException {
                 columnSchema.put(REVISION_COLUMN_NAME, timestampColumnSchema);
                 columnSchema.put(DELETED_FLAG_COLUMN_NAME, deletedFlagColumnSchema);
             } finally {
-                this.unlock();
+                this.unlock(false);
             }
             var currentSnapshots = this.storageAccessService.list(this.dataPathPrefix).collect(Collectors.toList());
             var path = this.dataPathPrefix + this.dataPathSuffixFormat.format(new Date());
@@ -255,7 +258,7 @@ public Map<String, BaseValue> next() {
                         }
 
                         private void getNext() {
-                            MemoryTableImpl.this.lock();
+                            MemoryTableImpl.this.lock(false);
                             try {
                                 NavigableMap<BaseValue, List<MemoryRecord>> target;
                                 if (this.lastKey == null) {
@@ -284,7 +287,7 @@ private void getNext() {
                                     }
                                 }
                             } finally {
-                                MemoryTableImpl.this.unlock();
+                                MemoryTableImpl.this.unlock(false);
                             }
                         }
                     });
@@ -305,13 +308,21 @@ private void getNext() {
     }
 
     @Override
-    public void lock() {
-        this.lock.lock();
+    public void lock(boolean forRead) {
+        if (forRead) {
+            this.readLock.lock();
+        } else {
+            this.writeLock.lock();
+        }
     }
 
     @Override
-    public void unlock() {
-        this.lock.unlock();
+    public void unlock(boolean forRead) {
+        if (forRead) {
+            this.readLock.unlock();
+        } else {
+            this.writeLock.unlock();
+        }
     }
 
     @Override