[MINOR] OOC Bugfix Cache Reference Management + Return Right BlockKey on Externally Managed Grouped Callbacks

Closes #2454.

Author: janniklinde

src/main/java/org/apache/sysds/runtime/instructions/ooc/CachingStream.java

@@ -130,6 +130,7 @@ public CachingStream(OOCStream<IndexedMatrixValue> source, long streamId) {
 							boolean ownsEntry = true;
 							if(tmp instanceof OOCCacheManager.CachedGroupCallback<?> cachedGroup) {
 								baseKey = cachedGroup.getBlockKey();
+								ensureReferencedOrRematerialize(baseKey, cachedGroup);
 								ownsEntry = false;
 								if(mSubscribers != null && mSubscribers.length > 0)
 									mCallback = tmp.keepOpen();
@@ -183,12 +184,14 @@ public CachingStream(OOCStream<IndexedMatrixValue> source, long streamId) {
 
 							if(tmp instanceof OOCCacheManager.CachedQueueCallback<?> cachedQueue) {
 								blockKey = cachedQueue.getBlockKey();
+								ensureReferencedOrRematerialize(blockKey, task);
 								ownsEntry = false;
 								if(mSubscribers != null && mSubscribers.length > 0)
 									mCallback = tmp.keepOpen();
 							}
 							else if(tmp instanceof OOCCacheManager.CachedSubCallback<?> cachedSub) {
 								BlockKey parent = cachedSub.getParent().getBlockKey();
+								ensureReferencedOrRematerialize(parent, cachedSub.getParent());
 								blockKey = new GroupedBlockKey(parent.getStreamId(), (int) parent.getSequenceNumber(),
 									cachedSub.getGroupIndex());
 								ownsEntry = false;
@@ -297,6 +300,49 @@ else if(tmp instanceof OOCCacheManager.CachedSubCallback<?> cachedSub) {
 		});
 	}
 
+
+	private void ensureReferencedOrRematerialize(BlockKey key, IndexedMatrixValue value) {
+		try {
+			OOCCacheManager.getCache().addReference(key);
+		}
+		catch(IllegalArgumentException ex) {
+			try {
+				OOCCacheManager.putRaw(key, value, ((MatrixBlock) value.getValue()).getExactSerializedSize());
+			}
+			catch(IllegalStateException putEx) {
+				// Another downstream stream may have re-materialized the same entry first.
+				OOCCacheManager.getCache().addReference(key);
+			}
+		}
+	}
+
+	private void ensureReferencedOrRematerialize(BlockKey key, OOCCacheManager.CachedGroupCallback<?> group) {
+		try {
+			OOCCacheManager.getCache().addReference(key);
+		}
+		catch(IllegalArgumentException ex) {
+			try {
+				List<IndexedMatrixValue> values = new ArrayList<>(group.size());
+				long totalSize = 0;
+				for(int gi = 0; gi < group.size(); gi++) {
+					@SuppressWarnings("unchecked")
+					OOCStream.QueueCallback<IndexedMatrixValue> sub =
+						(OOCStream.QueueCallback<IndexedMatrixValue>) group.getCallback(gi);
+					try(sub) {
+						IndexedMatrixValue imv = sub.get();
+						values.add(imv);
+						totalSize += ((MatrixBlock) imv.getValue()).getExactSerializedSize();
+					}
+				}
+				OOCCacheManager.putRaw(key, values, totalSize);
+			}
+			catch(IllegalStateException putEx) {
+				// Another downstream stream may have re-materialized the same entry first.
+				OOCCacheManager.getCache().addReference(key);
+			}
+		}
+	}
+
 	private String getCtxMsg() {
 		StackTraceElement[] st = new Exception().getStackTrace();
 		// Skip the first few frames (constructor, createWritableStream, etc.)
@@ -687,7 +733,7 @@ public void setSubscriber(Consumer<OOCStream.QueueCallback<IndexedMatrixValue>>
 			if(groupIdx > 0)
 				continue; // only replay grouped blocks once at the base index
 
-			BlockKey replayKey = (groupSize > 1 && groupIdx == 0) ? new BlockKey(_streamId, idx) : getBlockKey(i);
+			BlockKey replayKey = (groupSize > 1 && groupIdx == 0) ? getEntryBlockKey(idx) : getBlockKey(i);
 			OOCCacheManager.requestBlock(replayKey).whenComplete((cb, r) -> {
 				if(r != null) {
 					subscriber.accept(OOCStream.eos(DMLRuntimeException.of(r)));
@@ -697,7 +743,6 @@ public void setSubscriber(Consumer<OOCStream.QueueCallback<IndexedMatrixValue>>
 					synchronized(CachingStream.this) {
 						if(_index != null) {
 							if(cb instanceof OOCStream.GroupQueueCallback<?> && groupSize > 1) {
-								@SuppressWarnings("unchecked")
 								OOCStream.GroupQueueCallback<IndexedMatrixValue> group =
 									(OOCStream.GroupQueueCallback<IndexedMatrixValue>) cb;
 								for(int gi = 0; gi < groupSize; gi++) {

src/main/java/org/apache/sysds/runtime/ooc/cache/BlockEntry.java

@@ -30,6 +30,7 @@ public final class BlockEntry {
 	private volatile BlockState _state;
 	private Object _data;
 	private int _retainHintCount;
+	private int _referenceCount; // The number of references from different managing instances (e.g. CachingStream)
 
 	BlockEntry(BlockKey key, long size, Object data) {
 		this._key = key;
@@ -38,6 +39,7 @@ public final class BlockEntry {
 		this._state = BlockState.HOT;
 		this._data = data;
 		this._retainHintCount = 0;
+		this._referenceCount = 1;
 	}
 
 	public BlockKey getKey() {
@@ -84,6 +86,14 @@ public boolean isPinned() {
 		return _pinCount > 0;
 	}
 
+	synchronized int addReference() {
+		return ++_referenceCount;
+	}
+
+	synchronized int forget() {
+		return --_referenceCount;
+	}
+
 	synchronized void setState(BlockState state) {
 		_state = state;
 	}

src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheScheduler.java

@@ -103,6 +103,14 @@ public interface OOCCacheScheduler {
 	BlockEntry putAndPinSourceBacked(BlockKey key, Object data, long size,
 		OOCIOHandler.SourceBlockDescriptor descriptor);
 
+	/**
+	 * Notifies the cache that there is another reference to the same block key.
+	 * This will prevent forget(key) from removing the block from cache.
+	 * A block will only be forgotten after all referencing instances called forget(key).
+	 * @param key
+	 */
+	void addReference(BlockKey key);
+
 	/**
 	 * Forgets a block from the cache.
 	 * @param key the associated key of the block

src/main/java/org/apache/sysds/runtime/ooc/cache/OOCLRUCacheScheduler.java

@@ -19,6 +19,7 @@
 
 package org.apache.sysds.runtime.ooc.cache;
 
+import org.apache.commons.lang3.mutable.MutableObject;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.sysds.api.DMLScript;
@@ -291,6 +292,18 @@ public BlockEntry putAndPinSourceBacked(BlockKey key, Object data, long size, OO
 		return put(key, data, size, true, descriptor);
 	}
 
+	@Override
+	public void addReference(BlockKey key) {
+		synchronized(this) {
+			BlockEntry entry = _cache.get(key);
+			if(entry == null)
+				entry = _evictionCache.get(key);
+			if(entry == null)
+				throw new IllegalArgumentException("Could not find requested block with key " + key);
+			entry.addReference();
+		}
+	}
+
 	private BlockEntry put(BlockKey key, Object data, long size, boolean pin, OOCIOHandler.SourceBlockDescriptor descriptor) {
 		if (!this._running)
 			throw new IllegalStateException();
@@ -324,14 +337,34 @@ private BlockEntry put(BlockKey key, Object data, long size, boolean pin, OOCIOH
 	public void forget(BlockKey key) {
 		if (!this._running)
 			return;
+		final MutableObject<BlockEntry> mEntry = new MutableObject<>();
 		BlockEntry entry;
 		boolean shouldScheduleDeletion = false;
 		long cacheSizeDelta = 0;
 		synchronized(this) {
-			entry = _cache.remove(key);
+			_cache.compute(key, (k, e) -> {
+				if(e == null)
+					return null;
+				if(e.forget() == 0) {
+					mEntry.setValue(e);
+					return null;
+				}
+				return e;
+			});
 
-			if (entry == null)
-				entry = _evictionCache.remove(key);
+			if (mEntry.getValue() == null) {
+				_evictionCache.compute(key, (k, e) -> {
+					if(e == null)
+						return null;
+					if(e.forget() == 0) {
+						mEntry.setValue(e);
+						return null;
+					}
+					return e;
+				});
+			}
+
+			entry = mEntry.getValue();
 
 			if (entry != null) {
 				synchronized(entry) {