[DROOLS-7480] Persist info to identify which activation is fired or not

drools-reliability/drools-reliability-core/src/main/java/org/drools/reliability/core/BaseStoredEvent.java

@@ -41,9 +41,10 @@ public long getDuration() {
     }
 
     @Override
-    public void repropagate(InternalWorkingMemoryEntryPoint ep) {
+    public long repropagate(InternalWorkingMemoryEntryPoint ep) {
         FactHandleFactory fhFactory = ep.getHandleFactory();
         DefaultEventHandle eFh = fhFactory.createEventFactHandle(fhFactory.getNextId(), getObject(), fhFactory.getNextRecency(), ep, timestamp, duration);
         ep.insert(eFh);
+        return eFh.getId();
     }
 }

drools-reliability/drools-reliability-core/src/main/java/org/drools/reliability/core/BaseStoredObject.java

@@ -18,6 +18,7 @@
 import java.io.Serializable;
 
 import org.drools.core.common.InternalWorkingMemoryEntryPoint;
+import org.kie.api.runtime.rule.FactHandle;
 
 public abstract class BaseStoredObject implements StoredObject,
                                                   Serializable {
@@ -34,7 +35,8 @@ public boolean isPropagated() {
     }
 
     @Override
-    public void repropagate(InternalWorkingMemoryEntryPoint ep) {
-        ep.insert(getObject());
+    public long repropagate(InternalWorkingMemoryEntryPoint ep) {
+        FactHandle factHandle = ep.insert(getObject());
+        return factHandle.getId();
     }
 }

drools-reliability/drools-reliability-core/src/main/java/org/drools/reliability/core/ReliableKieSession.java

@@ -15,8 +15,13 @@
 
 package org.drools.reliability.core;
 
+import org.drools.core.common.Storage;
 import org.kie.api.runtime.KieSession;
 
 public interface ReliableKieSession extends KieSession {
     void safepoint();
+
+    Storage<String, Object> getActivationsStorage();
+
+    void setActivationsStorage(Storage<String, Object> activationsStorage);
 }

drools-reliability/drools-reliability-core/src/main/java/org/drools/reliability/core/ReliableSessionInitializer.java

@@ -15,24 +15,29 @@
 
 package org.drools.reliability.core;
 
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
 import org.drools.core.SessionConfiguration;
 import org.drools.core.WorkingMemoryEntryPoint;
 import org.drools.core.common.InternalFactHandle;
 import org.drools.core.common.InternalWorkingMemory;
 import org.drools.core.common.InternalWorkingMemoryEntryPoint;
 import org.drools.core.common.Storage;
 import org.drools.core.phreak.PropagationEntry;
+import org.drools.reliability.core.util.ReliabilityUtils;
+import org.kie.api.event.rule.AfterMatchFiredEvent;
+import org.kie.api.event.rule.DefaultAgendaEventListener;
+import org.kie.api.event.rule.MatchCancelledEvent;
+import org.kie.api.event.rule.MatchCreatedEvent;
 import org.kie.api.event.rule.ObjectDeletedEvent;
 import org.kie.api.event.rule.ObjectInsertedEvent;
 import org.kie.api.event.rule.ObjectUpdatedEvent;
 import org.kie.api.event.rule.RuleRuntimeEventListener;
 import org.kie.api.runtime.conf.PersistedSessionOption;
 import org.kie.api.runtime.rule.EntryPoint;
 
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-
 import static org.drools.reliability.core.ReliablePropagationList.PROPAGATION_LIST;
 
 public class ReliableSessionInitializer {
@@ -54,13 +59,24 @@ static class StoresOnlySessionInitializer implements SessionInitializer {
 
         @Override
         public InternalWorkingMemory init(InternalWorkingMemory session, PersistedSessionOption persistedSessionOption) {
+
+            if (persistedSessionOption.getActivationStrategy() == PersistedSessionOption.ActivationStrategy.ACTIVATION_KEY) {
+                Storage<String, Object> activationsStorage = StorageManagerFactory.get().getStorageManager().getOrCreateStorageForSession(session, "activations");
+                ((ReliableKieSession)session).setActivationsStorage(activationsStorage);
+            }
+
             if (!persistedSessionOption.isNewSession()) {
                 // re-propagate objects from the storage to the new session
                 populateSessionFromStorage(session);
             }
 
+            // These listeners should be added after populateSessionFromStorage()
             session.setWorkingMemoryActionListener(entry -> onWorkingMemoryAction(session, entry));
             session.getRuleRuntimeEventSupport().addEventListener(new SimpleStoreRuntimeEventListener(session));
+            if (persistedSessionOption.getActivationStrategy() == PersistedSessionOption.ActivationStrategy.ACTIVATION_KEY) {
+                ((ReliableKieSession)session).getActivationsStorage().clear();
+                session.getAgendaEventSupport().addEventListener(new SimpleStoreAgendaEventListener((ReliableKieSession)session));
+            }
 
             return session;
         }
@@ -108,6 +124,30 @@ public void objectUpdated(ObjectUpdatedEvent ev) {
                 store.putIntoPersistedStorage(fh, false);
             }
         }
+
+        static class SimpleStoreAgendaEventListener extends DefaultAgendaEventListener {
+
+            private final ReliableKieSession session;
+
+            public SimpleStoreAgendaEventListener(ReliableKieSession session) {
+                this.session = session;
+            }
+
+            @Override
+            public void matchCreated(MatchCreatedEvent event) {
+                session.getActivationsStorage().put(ReliabilityUtils.getActivationKey(event.getMatch()), true);
+            }
+
+            @Override
+            public void matchCancelled(MatchCancelledEvent event) {
+                session.getActivationsStorage().remove(ReliabilityUtils.getActivationKey(event.getMatch()));
+            }
+
+            @Override
+            public void afterMatchFired(AfterMatchFiredEvent event) {
+                session.getActivationsStorage().remove(ReliabilityUtils.getActivationKey(event.getMatch()));
+            }
+        }
     }
 
     static class FullReliableSessionInitializer implements SessionInitializer {

drools-reliability/drools-reliability-core/src/main/java/org/drools/reliability/core/ReliableStatefulKnowledgeSessionImpl.java

@@ -16,7 +16,7 @@
 package org.drools.reliability.core;
 
 import org.drools.core.SessionConfiguration;
-import org.drools.core.common.InternalAgenda;
+import org.drools.core.common.Storage;
 import org.drools.core.rule.accessor.FactHandleFactory;
 import org.drools.kiesession.rulebase.InternalKnowledgeBase;
 import org.drools.kiesession.session.StatefulKnowledgeSessionImpl;
@@ -27,6 +27,8 @@
 
 public class ReliableStatefulKnowledgeSessionImpl extends StatefulKnowledgeSessionImpl implements ReliableKieSession {
 
+    private transient Storage<String, Object> activationsStorage;
+
     public ReliableStatefulKnowledgeSessionImpl() {
     }
 
@@ -73,8 +75,22 @@ public void endOperation(InternalOperationType operationType) {
         }
     }
 
+    @Override
+    public Storage<String, Object> getActivationsStorage() {
+        return activationsStorage;
+    }
+
+    @Override
+    public void setActivationsStorage(Storage<String, Object> activationsStorage) {
+        this.activationsStorage = activationsStorage;
+    }
+
     @Override
     public void safepoint() {
         getEntryPoints().stream().map(ReliableNamedEntryPoint.class::cast).forEach(ReliableNamedEntryPoint::safepoint);
+        if (getSessionConfiguration().getPersistedSessionOption().getActivationStrategy() == PersistedSessionOption.ActivationStrategy.ACTIVATION_KEY
+                && activationsStorage.requiresFlush()) {
+            activationsStorage.flush();
+        }
     }
 }

drools-reliability/drools-reliability-core/src/main/java/org/drools/reliability/core/SimpleReliableObjectStore.java

@@ -15,13 +15,13 @@
 
 package org.drools.reliability.core;
 
+import java.util.List;
+
 import org.drools.core.common.InternalFactHandle;
 import org.drools.core.common.InternalWorkingMemory;
 import org.drools.core.common.InternalWorkingMemoryEntryPoint;
 import org.drools.core.common.ObjectStore;
 
-import java.util.List;
-
 public interface SimpleReliableObjectStore extends ObjectStore {
 
     List<StoredObject> reInit(InternalWorkingMemory session, InternalWorkingMemoryEntryPoint ep);

drools-reliability/drools-reliability-core/src/main/java/org/drools/reliability/core/SimpleSerializationReliableObjectStore.java

@@ -15,17 +15,22 @@
 
 package org.drools.reliability.core;
 
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+
 import org.drools.core.ClockType;
 import org.drools.core.common.DefaultEventHandle;
 import org.drools.core.common.IdentityObjectStore;
 import org.drools.core.common.InternalFactHandle;
 import org.drools.core.common.InternalWorkingMemory;
 import org.drools.core.common.InternalWorkingMemoryEntryPoint;
 import org.drools.core.common.Storage;
-
-import java.util.ArrayList;
-import java.util.List;
-import java.util.concurrent.TimeUnit;
+import org.drools.reliability.core.util.ReliabilityUtils;
+import org.kie.api.runtime.conf.PersistedSessionOption;
 
 public class SimpleSerializationReliableObjectStore extends IdentityObjectStore implements SimpleReliableObjectStore {
 
@@ -57,23 +62,23 @@ public void removeHandle(InternalFactHandle handle) {
     @Override
     public List<StoredObject> reInit(InternalWorkingMemory session, InternalWorkingMemoryEntryPoint ep) {
         reInitPropagated = true;
-        List<StoredObject> propagated = new ArrayList<>();
+        Map<Long, StoredObject> propagated = new HashMap<>();
         List<StoredObject> notPropagated = new ArrayList<>();
-        for (StoredObject entry : storage.values()) {
-            if (entry.isPropagated()) {
-                propagated.add(entry);
+        for (Long factHandleId : storage.keySet()) {
+            StoredObject storedObject = storage.get(factHandleId);
+            if (storedObject.isPropagated()) {
+                propagated.put(factHandleId, storedObject);
             } else {
-                notPropagated.add(entry);
+                notPropagated.add(storedObject);
             }
         }
+
         storage.clear();
 
         if (session.getSessionConfiguration().getClockType() == ClockType.PSEUDO_CLOCK) {
             repropagateWithPseudoClock(session, ep, propagated);
         } else {
-            // fact handles with a match have been already propagated in the original session, so they shouldn't fire
-            propagated.forEach(obj -> obj.repropagate(ep));
-            session.fireAllRules(match -> false);
+            repropagate(session, ep, propagated);
         }
 
         reInitPropagated = false;
@@ -82,9 +87,41 @@ public List<StoredObject> reInit(InternalWorkingMemory session, InternalWorkingM
         return notPropagated;
     }
 
-    private void repropagateWithPseudoClock(InternalWorkingMemory session, InternalWorkingMemoryEntryPoint ep, List<StoredObject> propagated) {
+    private void repropagate(InternalWorkingMemory session, InternalWorkingMemoryEntryPoint ep, Map<Long, StoredObject> propagated) {
+        Map<Long, Long> factHandleIdMap = new HashMap<>();
+        propagated.forEach((oldFactHandleId, storedObject) -> {
+            long newFactHandleId = storedObject.repropagate(ep);
+            factHandleIdMap.put(newFactHandleId, oldFactHandleId);
+        });
+
+        fireOnlyWhenActivationRemaining(session, factHandleIdMap);
+    }
+
+    private void fireOnlyWhenActivationRemaining(InternalWorkingMemory session, Map<Long, Long> factHandleIdMap) {
+        if (session.getSessionConfiguration().getPersistedSessionOption().getActivationStrategy() == PersistedSessionOption.ActivationStrategy.ACTIVATION_KEY) {
+            // fact handles with a match have been already propagated in the original session, so they shouldn't fire unless remained in activationsStorage
+            Storage<String, Object> activationsStorage = ((ReliableKieSession)session).getActivationsStorage();
+            Set<String> activationKeySet = activationsStorage.keySet();
+            session.fireAllRules(match -> {
+                String activationKey = ReliabilityUtils.getActivationKeyReplacingNewIdWithOldId(match, factHandleIdMap);
+                if (activationKeySet.contains(activationKey)) {
+                    // If there is a remaining activation, it can fire
+                    activationsStorage.remove(activationKey);
+                    return true;
+                } else {
+                    return false;
+                }
+            });
+        } else {
+            session.fireAllRules(match -> false);
+        }
+    }
+
+    private void repropagateWithPseudoClock(InternalWorkingMemory session, InternalWorkingMemoryEntryPoint ep, Map<Long, StoredObject> propagated) {
         ReliablePseudoClockScheduler clock = (ReliablePseudoClockScheduler) session.getSessionClock();
-        for (StoredObject storedObject : propagated) {
+        Map<Long, Long> factHandleIdMap = new HashMap<>();
+        for (Map.Entry<Long, StoredObject> entry : propagated.entrySet()) {
+            StoredObject storedObject = entry.getValue();
             if (storedObject.isEvent()) {
                 StoredEvent storedEvent = (StoredEvent) storedObject;
                 long currentTime = clock.getCurrentTime();
@@ -93,10 +130,11 @@ private void repropagateWithPseudoClock(InternalWorkingMemory session, InternalW
                     clock.advanceTime(timestamp - currentTime, TimeUnit.MILLISECONDS);
                 }
             }
-            storedObject.repropagate(ep); // This may schedule an expiration
+            long newFactHandleId = storedObject.repropagate(ep); // This may schedule an expiration
+            factHandleIdMap.put(newFactHandleId, entry.getKey());
         }
-        // fact handles with a match have been already propagated in the original session, so they shouldn't fire
-        session.fireAllRules(match -> false);
+
+        fireOnlyWhenActivationRemaining(session, factHandleIdMap);
 
         // Finally, meet with the persistedTime
         long currentTime = clock.getCurrentTime();

drools-reliability/drools-reliability-core/src/main/java/org/drools/reliability/core/StoredObject.java

@@ -27,5 +27,5 @@ default boolean isEvent() {
 
     Object getObject();
 
-    void repropagate(InternalWorkingMemoryEntryPoint ep);
+    long repropagate(InternalWorkingMemoryEntryPoint ep);
 }

drools-reliability/drools-reliability-core/src/main/java/org/drools/reliability/core/util/ReliabilityUtils.java

@@ -0,0 +1,68 @@
+/*
+ * Copyright 2023 Red Hat, Inc. and/or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.drools.reliability.core.util;
+
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import org.drools.core.reteoo.RuleTerminalNodeLeftTuple;
+import org.drools.reliability.core.ReliabilityRuntimeException;
+import org.kie.api.runtime.rule.FactHandle;
+import org.kie.api.runtime.rule.Match;
+
+public class ReliabilityUtils {
+
+    private ReliabilityUtils() {
+        // no constructor
+    }
+
+    /**
+     * Returns a String representation of the activation.
+     */
+    public static String getActivationKey(Match match) {
+        return getActivationKey(match, null);
+    }
+
+    /**
+     * Returns a String representation of the activation, replacing the new fact handle id with the old fact handle id.
+     * Used to find an activation key in the persisted storage.
+     */
+    public static String getActivationKeyReplacingNewIdWithOldId(Match match, Map<Long, Long> factHandleIdMap) {
+        return getActivationKey(match, factHandleIdMap);
+    }
+
+    private static String getActivationKey(Match match, Map<Long, Long> factHandleIdMap) {
+        if (!(match instanceof RuleTerminalNodeLeftTuple)) {
+            throw new ReliabilityRuntimeException("getActivationKey doesn't support " + match.getClass());
+        }
+        RuleTerminalNodeLeftTuple ruleTerminalNodeLeftTuple = (RuleTerminalNodeLeftTuple) match;
+        String packageName = ruleTerminalNodeLeftTuple.getRule().getPackageName();
+        String ruleName = ruleTerminalNodeLeftTuple.getRule().getName();
+        List<FactHandle> factHandles = ruleTerminalNodeLeftTuple.getFactHandles();
+        List<Long> factHandleIdList = factHandles.stream()
+                                                 .map(FactHandle::getId)
+                                                 .map(handleId -> {
+                                                     if (factHandleIdMap != null) {
+                                                         return factHandleIdMap.get(handleId); // replace new id with old id
+                                                     } else {
+                                                         return handleId; // don't replace
+                                                     }
+                                                 })
+                                                 .collect(Collectors.toList());
+        return "ActivationKey [packageName=" + packageName + ", ruleName=" + ruleName + ", factHandleIdList=" + factHandleIdList + "]";
+    }
+}