[waterflow] fix: clean flow locks

Author: loveTsong

framework/waterflow/java/waterflow-core/src/main/java/modelengine/fit/waterflow/domain/context/repo/flowlock/FlowLocks.java

@@ -9,11 +9,7 @@
 import modelengine.fit.waterflow.domain.common.Constants;
 import modelengine.fitframework.util.StringUtils;
 
-import java.util.Map;
-import java.util.Optional;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
 
 /**
  * 流程实例的锁接口
@@ -22,10 +18,6 @@
  * @since 1.0
  */
 public interface FlowLocks {
-    /**
-     * 本地锁全局静态对象
-     */
-    Map<String, Lock> locks = new ConcurrentHashMap<>();
 
     /**
      * 节点分布式锁key前缀
@@ -38,9 +30,7 @@ public interface FlowLocks {
      * @param key 获取本地锁的key值，一般是流程版本的streamID
      * @return {@link Lock} 锁对象
      */
-    default Lock getLocalLock(String key) {
-        return Optional.ofNullable(locks.putIfAbsent(key, new ReentrantLock())).orElseGet(() -> locks.get(key));
-    }
+    Lock getLocalLock(String key);
 
     /**
      * 获取分布式锁
@@ -50,16 +40,6 @@ default Lock getLocalLock(String key) {
      */
     Lock getDistributeLock(String key);
 
-    /**
-     * 删除本地锁
-     * TODO xiangyu 删除流程定义的时候需要删除该定义的本地锁资源
-     *
-     * @param key 删除本地锁的key值，一般是流程版本的streamID
-     */
-    default void removeLocalLock(String key) {
-        locks.remove(key);
-    }
-
     /**
      * 获取节点分布式锁key值
      * 获取分布式锁的key值，一般是prefix-streamID-nodeID-type

framework/waterflow/java/waterflow-core/src/main/java/modelengine/fit/waterflow/domain/context/repo/flowlock/FlowLocksMemo.java

@@ -6,7 +6,13 @@
 
 package modelengine.fit.waterflow.domain.context.repo.flowlock;
 
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 
 /**
  * 流程锁，内存版本的实现
@@ -15,6 +21,18 @@
  * @since 1.0
  */
 public class FlowLocksMemo implements FlowLocks {
+    private final Map<String, MemLockWrapper> locks = new ConcurrentHashMap<>();
+
+    @Override
+    public Lock getLocalLock(String key) {
+        return locks.compute(key, (__, value) -> {
+            if (value == null) {
+                return new MemLockWrapper(key, new ReentrantLock(), this);
+            }
+            return value;
+        });
+    }
+
     /**
      * 获取分布式锁
      * 获取分布式锁的key值，一般是prefix-streamID-nodeID-suffixes
@@ -27,4 +45,69 @@ public class FlowLocksMemo implements FlowLocks {
     public Lock getDistributeLock(String key) {
         return getLocalLock(key);
     }
+
+    private void tryCleanLocalLock(String key) {
+        this.locks.compute(key, (__, value) -> {
+            if (value == null) {
+                return null;
+            }
+            if (value.getRefCount() == 0) {
+                return null;
+            }
+            return value;
+        });
+    }
+
+    private static class MemLockWrapper implements Lock {
+        private final String key;
+        private final AtomicInteger refCount = new AtomicInteger(1);
+        private final ReentrantLock target;
+        private final FlowLocksMemo locksMemo;
+
+        private MemLockWrapper(String key, ReentrantLock target, FlowLocksMemo locksMemo) {
+            this.key = key;
+            this.target = target;
+            this.locksMemo = locksMemo;
+        }
+
+        @Override
+        public void lock() {
+            this.target.lock();
+        }
+
+        @Override
+        public void lockInterruptibly() throws InterruptedException {
+            this.target.lockInterruptibly();
+        }
+
+        @Override
+        public boolean tryLock() {
+            return this.target.tryLock();
+        }
+
+        @Override
+        public boolean tryLock(long time, TimeUnit unit) throws InterruptedException {
+            return this.target.tryLock(time, unit);
+        }
+
+        @Override
+        public void unlock() {
+            this.target.unlock();
+            this.refCount.decrementAndGet();
+            this.locksMemo.tryCleanLocalLock(this.key);
+        }
+
+        @Override
+        public Condition newCondition() {
+            return this.target.newCondition();
+        }
+
+        private void addRef() {
+            this.refCount.incrementAndGet();
+        }
+
+        private int getRefCount() {
+            return this.refCount.get();
+        }
+    }
 }