把MatchWindow放在FlowSessionCache统一管理

Author: relat-ivity

framework/waterflow/java/waterflow-core/src/main/java/modelengine/fit/waterflow/domain/context/MatchWindow.java

@@ -6,12 +6,13 @@
 
 package modelengine.fit.waterflow.domain.context;
 
+import modelengine.fit.waterflow.domain.context.repo.flowsession.FlowSessionRepo;
+
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 
 /**
@@ -21,8 +22,6 @@
  * @since 1.0
  */
 public class MatchWindow extends Window {
-    private static final Map<String, MatchWindow> all = new ConcurrentHashMap<>();
-
     private final Set<MatchWindow> arms = new HashSet<>();
 
     /**
@@ -41,24 +40,26 @@ public MatchWindow(Window source, UUID id, Object data) {
     /**
      * 创建一个MatchWindow
      *
+     * @param flowId 流程ID
      * @param source 源窗口
      * @param id 窗口ID
      * @param data 窗口数据
      * @return 返回创建的MatchWindow对象
      */
-    public static synchronized MatchWindow from(Window source, UUID id, Object data) {
-        // Use composite key: sessionId + UUID to prevent cross-session pollution
-        String cacheKey = source.getSession().getId() + ":" + id.toString();
-        MatchWindow window = all.get(cacheKey);
+    public static synchronized MatchWindow from(String flowId, Window source, UUID id, Object data) {
+        // 从 FlowSessionRepo 获取缓存
+        Map<UUID, MatchWindow> cache = FlowSessionRepo.getMatchWindowCache(flowId, source.getSession());
+
+        MatchWindow window = cache.get(id);
         if (window == null) {
             window = new MatchWindow(source, id, data);
             FlowSession session = new FlowSession(source.getSession());
             session.setWindow(window);
-            all.put(cacheKey, window);
+            cache.put(id, window);
         }
         WindowToken token = window.createToken();
         token.beginConsume();
-        List<MatchWindow> arms = all.values().stream().filter(t -> t.from == source).collect(Collectors.toList());
+        List<MatchWindow> arms = cache.values().stream().filter(t -> t.from == source).collect(Collectors.toList());
         for (MatchWindow a : arms) {
             a.setArms(arms);
         }
@@ -82,23 +83,4 @@ public void complete() {
     public boolean fulfilled() {
         return this.from.isComplete() && this.from.isOngoing();
     }
-
-    @Override
-    public void tryFinish() {
-        super.tryFinish();
-        // 当 MatchWindow 完全完成时，从缓存中移除以防止内存泄漏
-        if (this.isDone()) {
-            cleanup();
-        }
-    }
-
-    /**
-     * 从缓存中移除当前 MatchWindow
-     */
-    private void cleanup() {
-        if (this.getSession() != null) {
-            String cacheKey = this.getSession().getId() + ":" + this.key().toString();
-            all.remove(cacheKey);
-        }
-    }
 }

framework/waterflow/java/waterflow-core/src/main/java/modelengine/fit/waterflow/domain/context/repo/flowsession/FlowSessionRepo.java

@@ -8,6 +8,7 @@
 
 import modelengine.fit.waterflow.domain.context.FlatMapSourceWindow;
 import modelengine.fit.waterflow.domain.context.FlowSession;
+import modelengine.fit.waterflow.domain.context.MatchWindow;
 import modelengine.fit.waterflow.domain.context.Window;
 import modelengine.fit.waterflow.domain.context.repo.flowcontext.FlowContextRepo;
 import modelengine.fitframework.inspection.Validation;
@@ -90,6 +91,19 @@ public static FlatMapSourceWindow getFlatMapSource(String flowId, Window window,
                 .getFlatMapSourceWindow(window, repo);
     }
 
+    /**
+     * 获取 MatchWindow 缓存 Map，用于存储和检索 MatchWindow 实例
+     *
+     * @param flowId The unique identifier of the flow.
+     * @param session The current session context.
+     * @return MatchWindow 缓存 Map
+     */
+    public static Map<UUID, MatchWindow> getMatchWindowCache(String flowId, FlowSession session) {
+        Validation.notNull(flowId, "Flow id cannot be null.");
+        Validation.notNull(session, "Session cannot be null.");
+        return getFlowSessionCache(flowId, session).getMatchWindowCache();
+    }
+
     /**
      * Releases all resources associated with a specific flow session.
      *
@@ -137,6 +151,12 @@ private static class FlowSessionCache {
          */
         private final Map<UUID, FlatMapSourceWindow> flatMapSourceWindows = new ConcurrentHashMap<>();
 
+        /**
+         * 记录流程中条件匹配节点产生的窗口信息，用于将同一批数据汇聚。
+         * 其中索引为 match window 的唯一标识。
+         */
+        private final Map<UUID, MatchWindow> matchWindows = new ConcurrentHashMap<>();
+
         private final Map<String, Integer> accOrders = new ConcurrentHashMap<>();
 
         private FlowSession getNextToSession(FlowSession session) {
@@ -165,6 +185,10 @@ private FlatMapSourceWindow getFlatMapSourceWindow(Window window, FlowContextRep
             });
         }
 
+        private Map<UUID, MatchWindow> getMatchWindowCache() {
+            return this.matchWindows;
+        }
+
         private int getNextAccOrder(String nodeId) {
             return this.accOrders.compute(nodeId, (key, value) -> {
                 if (value == null) {

framework/waterflow/java/waterflow-core/src/main/java/modelengine/fit/waterflow/domain/states/MatchHappen.java

@@ -50,8 +50,8 @@ public MatchHappen<O, D, I, F> match(Operators.Whether<I> whether,
             Operators.BranchProcessor<O, D, I, F> processor) {
         UUID id = UUID.randomUUID();
         State<I, D, I, F> branchStart = new State<>(this.node.publisher()
-                .just(input -> input.setSession(
-                        MatchWindow.from(input.getWindow(), id, input.getData()).getSession()), whether)
+                .just(input -> input.setSession(MatchWindow.from(this.node.processor.getStreamId(),
+                        input.getWindow(), id, input.getData()).getSession()), whether)
                 .displayAs(SpecialDisplayNode.BRANCH.name()), this.node.getFlow());
         State<O, D, ?, F> branch = processor.process(branchStart);
         this.branches.add(branch);