Merge pull request #628 from stariy95/FIX-4.2-CAY-2876

CAY-2876 Memory leak in the ObjectStore

Author: stariy95

RELEASE-NOTES.txt

@@ -22,6 +22,7 @@ Bug Fixes:
 CAY-2866 DefaultDataDomainFlushAction breaks on circular relationship update
 CAY-2868 Regression: DefaultDbRowOpSorter shouldn't sort update operations
 CAY-2871 QualifierTranslator breaks on a relationship with a compound FK
+CAY-2876 Memory leak in the ObjectStore
 CAY-2879 Negative number for non parameterized ObjectSelect query not processed correctly
 
 ----------------------------------

cayenne-server/src/main/java/org/apache/cayenne/access/ObjectStore.java

@@ -44,6 +44,8 @@
 import org.apache.cayenne.reflect.PropertyVisitor;
 import org.apache.cayenne.reflect.ToManyProperty;
 import org.apache.cayenne.reflect.ToOneProperty;
+import org.apache.cayenne.util.SoftValueMap;
+import org.apache.cayenne.util.WeakValueMap;
 
 import java.io.Serializable;
 import java.util.ArrayList;
@@ -53,7 +55,6 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 
 /**
@@ -109,13 +110,23 @@ public class ObjectStore implements Serializable, SnapshotEventListener, GraphMa
      */
     public ObjectStore(DataRowStore dataRowCache, Map<Object, Persistent> objectMap) {
         setDataRowCache(dataRowCache);
-        if (objectMap != null) {
-            this.objectMap = objectMap;
-        }
-        else {
+        setObjectMap(objectMap);
+        this.changes = new HashMap<>();
+    }
+
+    /**
+     * @since 4.2.2
+     */
+    void setObjectMap(Map<Object, Persistent> objectMap) {
+        if(objectMap == null) {
             throw new CayenneRuntimeException("Object map is null.");
         }
-        this.changes = new HashMap<>();
+        this.objectMap = objectMap;
+        if(objectMap instanceof SoftValueMap) {
+            ((SoftValueMap<Object, Persistent>) objectMap).setKeyCleanupCallback(this::onObjectKeyCleanup);
+        } else if(objectMap instanceof WeakValueMap) {
+            ((WeakValueMap<Object, Persistent>) objectMap).setKeyCleanupCallback(this::onObjectKeyCleanup);
+        }
     }
 
     /**
@@ -138,7 +149,7 @@ void childContextSyncStopped() {
     Collection<GraphDiff> getLifecycleEventInducedChanges() {
         return lifecycleEventInducedChanges != null
                 ? lifecycleEventInducedChanges
-                : Collections.<GraphDiff>emptyList();
+                : Collections.emptyList();
     }
 
     void registerLifecycleEventInducedChange(GraphDiff diff) {
@@ -416,6 +427,9 @@ public void postprocessAfterCommit(GraphDiff parentChanges) {
             switch (object.getPersistenceState()) {
                 case PersistenceState.DELETED:
                     objectMap.remove(id);
+                    if(trackedFlattenedPaths != null) {
+                        trackedFlattenedPaths.remove(id);
+                    }
                     object.setObjectContext(null);
                     object.setPersistenceState(PersistenceState.TRANSIENT);
                     break;
@@ -640,9 +654,11 @@ void processDeletedID(ObjectId nodeId) {
                     if (dataObject == null || delegate.shouldProcessDelete(dataObject)) {
                         objectMap.remove(nodeId);
                         changes.remove(nodeId);
+                        if(trackedFlattenedPaths != null) {
+                            trackedFlattenedPaths.remove(nodeId);
+                        }
 
-                        // setting DataContext to null will also set
-                        // state to transient
+                        // setting DataContext to null will also set state to transient
                         object.setObjectContext(null);
 
                         if (dataObject != null) {
@@ -1049,6 +1065,16 @@ public void markFlattenedPath(ObjectId objectId, String path, ObjectId id) {
                 .put(path, id);
     }
 
+    /**
+     * @param key object id that was removed from the {@link #objectMap}
+     * @since 4.2.2
+     */
+    void onObjectKeyCleanup(Object key) {
+        if(trackedFlattenedPaths != null) {
+            trackedFlattenedPaths.remove(key);
+        }
+    }
+
     // an ObjectIdQuery optimized for retrieval of multiple snapshots - it can be reset
     // with the new id
     final class CachedSnapshotQuery extends ObjectIdQuery {

cayenne-server/src/main/java/org/apache/cayenne/util/ReferenceMap.java

@@ -35,6 +35,7 @@
 import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.Set;
+import java.util.function.Consumer;
 
 /**
  * Map that transparently stores values as references and resolves them as needed.
@@ -82,6 +83,11 @@ abstract class ReferenceMap<K, V, R extends Reference<V>> extends AbstractMap<K,
      */
     protected transient Set<Entry<K, V>> entrySet;
 
+    /**
+     * @since 4.2.2
+     */
+    protected transient Consumer<K> keyCleanupCallback;
+
     public ReferenceMap() {
         map = new HashMap<>();
         referenceQueue = new ReferenceQueue<>();
@@ -221,6 +227,17 @@ public Set<Entry<K, V>> entrySet() {
         return es;
     }
 
+    /**
+     * Set callback that will be notified with a key on each value removal
+     * due to the corresponding value reference cleaned up by the GC.
+     *
+     * @param keyCleanupCallback callback to set
+     * @since 4.2.2
+     */
+    public void setKeyCleanupCallback(Consumer<K> keyCleanupCallback) {
+        this.keyCleanupCallback = keyCleanupCallback;
+    }
+
     /**
      * Cleanup all references collected by GC so far
      */
@@ -247,6 +264,9 @@ protected void checkReferenceQueue() {
 
         for(K keyToRemove : keysToRemove) {
             map.remove(keyToRemove);
+            if(keyCleanupCallback != null) {
+                keyCleanupCallback.accept(keyToRemove);
+            }
         }
     }