[OOC] Add New OOCCache Implementation

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

@@ -19,10 +19,6 @@
 
 package org.apache.sysds.runtime.ooc.cache;
 
-import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue;
-
-import java.util.List;
-
 public final class BlockEntry {
 	private final BlockKey _key;
 	private final long _size;
@@ -31,6 +27,19 @@ public final class BlockEntry {
 	private Object _data;
 	private int _retainHintCount;
 	private int _referenceCount; // The number of references from different managing instances (e.g. CachingStream)
+	// Optional implementation-local cache metadata; null for cache implementations that do not need it.
+	private volatile Object _cacheMeta;
+
+	public BlockEntry(BlockKey key) {
+		this._key = key;
+		this._size = -1;
+		this._pinCount = 0;
+		this._state = BlockState.COLD;
+		this._data = null;
+		this._retainHintCount = 0;
+		this._referenceCount = 0;
+		this._cacheMeta = null;
+	}
 
 	public BlockEntry(BlockKey key, long size, Object data) {
 		this._key = key;
@@ -40,6 +49,18 @@ public BlockEntry(BlockKey key, long size, Object data) {
 		this._data = data;
 		this._retainHintCount = 0;
 		this._referenceCount = 1;
+		this._cacheMeta = null;
+	}
+
+	public BlockEntry(BlockKey key, long size, Object data, BlockState state) {
+		this._key = key;
+		this._size = size;
+		this._pinCount = 0;
+		this._state = state;
+		this._data = data;
+		this._retainHintCount = 0;
+		this._referenceCount = 1;
+		this._cacheMeta = null;
 	}
 
 	public BlockKey getKey() {
@@ -56,18 +77,6 @@ public Object getData() {
 		throw new IllegalStateException("Cannot get the data of an unpinned entry");
 	}
 
-	public int getGroupSize() {
-		if(_pinCount > 0)
-			return ((List<?>)_data).size();
-		throw new IllegalStateException("Cannot get the data of an unpinned entry");
-	}
-
-	public boolean isGrouped() {
-		if(_pinCount > 0)
-			return _data instanceof List;
-		throw new IllegalStateException("Cannot get the data of an unpinned entry");
-	}
-
 	public Object getDataUnsafe() {
 		return _data;
 	}
@@ -86,6 +95,10 @@ public boolean isPinned() {
 		return _pinCount > 0;
 	}
 
+	public synchronized int getPinCount() {
+		return _pinCount;
+	}
+
 	public synchronized int addReference() {
 		return ++_referenceCount;
 	}
@@ -94,6 +107,18 @@ public synchronized int forget() {
 		return --_referenceCount;
 	}
 
+	public synchronized int getReferenceCount() {
+		return _referenceCount;
+	}
+
+	public Object getCacheMeta() {
+		return _cacheMeta;
+	}
+
+	public void setCacheMeta(Object meta) {
+		_cacheMeta = meta;
+	}
+
 	public synchronized void setState(BlockState state) {
 		_state = state;
 	}
@@ -106,12 +131,6 @@ public synchronized void addRetainHint() {
 		_retainHintCount++;
 	}
 
-	public synchronized void removeRetainHint(int cnt) {
-		_retainHintCount -= cnt;
-		if(_retainHintCount < 0)
-			_retainHintCount = 0;
-	}
-
 	public synchronized void removeRetainHint() {
 		if (_retainHintCount <= 0)
 			return;
@@ -129,8 +148,6 @@ public synchronized int getRetainHintCount() {
 	public synchronized long clear() {
 		if (_pinCount != 0 || _data == null)
 			return 0;
-		if (_data instanceof IndexedMatrixValue)
-			((IndexedMatrixValue)_data).setValue(null); // Explicitly clear
 		_data = null;
 		_retainHintCount = 0;
 		return _size;

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

@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   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.apache.sysds.runtime.ooc.cache;
+
+import org.apache.sysds.runtime.ooc.memory.MemoryAllowance;
+
+import java.util.function.LongUnaryOperator;
+
+public interface OOCCache {
+	default OOCFuture<BlockEntry> pin(BlockKey key, MemoryAllowance allowance) {
+		return pin(key.getStreamId(), key.getSequenceNumber(), allowance);
+	}
+
+	default OOCFuture<BlockEntry> pinAdmitted(BlockKey key, MemoryAllowance allowance) {
+		return pinAdmitted(key.getStreamId(), key.getSequenceNumber(), allowance);
+	}
+
+	/**
+	 * Adds a new pinned entry whose bytes are already owned by the given allowance. Ownership can later move only via
+	 * pin/unpin.
+	 */
+	default BlockEntry putPinned(BlockKey key, Object data, long size, MemoryAllowance allowance) {
+		return putPinned(key.getStreamId(), key.getSequenceNumber(), data, size, allowance);
+	}
+
+	/**
+	 * Adds a new pinned entry whose bytes are already owned by the given allowance. Ownership can later move only via
+	 * pin/unpin.
+	 */
+	BlockEntry putPinned(long sId, long tId, Object data, long size, MemoryAllowance allowance);
+
+	/**
+	 * Pins an item backed by an allowance. A successful pin transfers memory ownership from the cache to the owner of
+	 * the allowance and guarantees data availability. While pinned, the bytes of the entry are not counted as
+	 * cache-owned memory.
+	 *
+	 * @param sId
+	 * @param tId
+	 * @param allowance
+	 * @return a non-null future of the pinned block entry; the future result is null if the required memory could not
+	 * be reserved
+	 */
+	OOCFuture<BlockEntry> pin(long sId, long tId, MemoryAllowance allowance);
+
+	default OOCFuture<BlockEntry> pinAdmitted(long sId, long tId, MemoryAllowance allowance) {
+		return pin(sId, tId, allowance);
+	}
+
+	/**
+	 * Pins an item backed by an allowance if it is already live in cache. A successful pin transfers memory ownership
+	 * from the cache to the owner of the allowance and guarantees data availability. While pinned, the bytes of the
+	 * entry are not counted as cache-owned memory. Implementations must reserve the required bytes from the allowance
+	 * before making data available.
+	 *
+	 * @param sId
+	 * @param tId
+	 * @param allowance
+	 * @return the pinned block entry if available. Null if the required memory could not be reserved or the block is
+	 * not live
+	 */
+	BlockEntry pinIfLive(long sId, long tId, MemoryAllowance allowance);
+
+	/**
+	 * Unpins an item that is still backed by the given allowance. Unpinning tries to transfer memory ownership back to
+	 * the cache. An ownership transfer may commit immediately only if this does not cause the cache to exceed its hard
+	 * limit. Otherwise, the transfer is deferred and the allowance remains charged until the returned handle commits,
+	 * is reclaimed, or is superseded by a later pin that transfers ownership to another allowance. Unpin can be viewed
+	 * as an eventually resolving operation.
+	 *
+	 * @param entry
+	 * @param allowance
+	 * @return a handle describing the ownership transfer from allowance-owned memory back to cache-owned memory
+	 */
+	UnpinHandle unpin(BlockEntry entry, MemoryAllowance allowance);
+
+	/**
+	 * Referencing a pinned entry guarantees that its key remains in the cache until dereferenced.
+	 *
+	 * @param entry
+	 * @return
+	 */
+	int reference(BlockEntry entry);
+
+	/**
+	 * Dereferencing allows an entry to be forgotten if no further reference is held. Dereferencing may not immediately
+	 * cause entry removal if still pinned.
+	 *
+	 * @param entry
+	 * @return
+	 */
+	int dereference(BlockEntry entry);
+
+	/**
+	 * Dereferencing allows an entry to be forgotten if no further reference is held. Dereferencing may not immediately
+	 * cause entry removal if still pinned.
+	 */
+	int dereference(BlockKey key);
+
+	void updateLimits(long hardLimit, long evictionLimit);
+
+	/**
+	 * Adds an eviction scoring policy for one logical cache stream. Larger scores are selected for eviction first.
+	 * {@link Long#MAX_VALUE} remains reserved as "no policy score".
+	 */
+	void addEvictionPolicy(long streamId, LongUnaryOperator scoreFn);
+
+	/**
+	 * Returns the current cache-owned size in bytes.
+	 */
+	long getOwnedCacheSize();
+
+	void shutdown();
+
+	interface UnpinHandle {
+		BlockEntry entry();
+
+		MemoryAllowance allowance();
+
+		long bytes();
+
+		boolean isCommitted();
+
+		OOCFuture<Boolean> getCompletionFuture();
+	}
+}