[SYSTEMDS-3891][OOC] Add MaterializedStore Structures

src/main/java/org/apache/sysds/runtime/ooc/memory/InMemoryQueueCallback.java

@@ -191,4 +191,26 @@ private void closeFinal() {
 	public IndexedMatrixValue takeManagedResultForHandover() {
 		return _handle.takeManagedResultForHandover();
 	}
+
+	public synchronized ManagedPayload<IndexedMatrixValue> extractManagedPayload() {
+		if(_closed)
+			throw new IllegalStateException("Cannot extract a managed payload from a closed callback.");
+		CallbackHandle handle = _handle;
+		synchronized(handle) {
+			if(handle._failure != null)
+				throw handle._failure;
+			if(!handle.isExclusiveToRoot())
+				throw new IllegalStateException("Cannot extract a managed payload while callback aliases exist.");
+			if(handle._cacheIdx >= 0)
+				throw new IllegalStateException("Cannot extract a managed payload from a cached-slot callback.");
+			IndexedMatrixValue result = handle._result;
+			if(result == null)
+				throw new IllegalStateException("Cannot extract a managed payload from an empty callback.");
+			long bytes = handle._reservedBytes;
+			MemoryAllowance owner = handle._allow;
+			handle._result = null;
+			handle._reservedBytes = 0;
+			return new ManagedPayload<>(result, bytes, owner);
+		}
+	}
 }

src/main/java/org/apache/sysds/runtime/ooc/memory/ManagedPayload.java

@@ -0,0 +1,68 @@
+/*
+ * 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.memory;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+public final class ManagedPayload<T> {
+	private final T _value;
+	private final long _bytes;
+	private final MemoryAllowance _owner;
+	private final AtomicBoolean _transferred;
+
+	public ManagedPayload(T value, long bytes, MemoryAllowance owner) {
+		if(value == null)
+			throw new IllegalArgumentException("Managed payload requires a value.");
+		if(bytes < 0)
+			throw new IllegalArgumentException("Managed payload bytes must not be negative: " + bytes);
+		if(bytes > 0 && owner == null)
+			throw new IllegalArgumentException("Managed payload with charged bytes requires an owning allowance.");
+		_value = value;
+		_bytes = bytes;
+		_owner = owner;
+		_transferred = new AtomicBoolean(false);
+	}
+
+	public T value() {
+		return _value;
+	}
+
+	public long bytes() {
+		return _bytes;
+	}
+
+	public MemoryAllowance owner() {
+		return _owner;
+	}
+
+	public void transfer() {
+		if(!_transferred.compareAndSet(false, true))
+			throw new IllegalStateException("Managed payload was already settled.");
+	}
+
+	public void release() {
+		if(_transferred.compareAndSet(false, true) && _bytes > 0)
+			_owner.release(_bytes);
+	}
+
+	public boolean isTransferred() {
+		return _transferred.get();
+	}
+}

src/main/java/org/apache/sysds/runtime/ooc/store/IndexedMaterializedStoreReader.java

@@ -0,0 +1,117 @@
+/*
+ * 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.store;
+
+import org.apache.sysds.runtime.ooc.cache.BlockEntry;
+import org.apache.sysds.runtime.ooc.cache.OOCCache;
+import org.apache.sysds.runtime.ooc.cache.OOCFuture;
+import org.apache.sysds.runtime.ooc.cache.io.SpillableObject;
+import org.apache.sysds.runtime.ooc.memory.MemoryAllowance;
+import org.apache.sysds.runtime.ooc.util.OOCUtils;
+
+import java.util.function.IntConsumer;
+import java.util.function.IntSupplier;
+
+public final class IndexedMaterializedStoreReader<T extends SpillableObject> implements MaterializedStore.StoreReader {
+	private final OOCCache _cache;
+	private final long _streamId;
+	private final IntSupplier _completedSize;
+	private final MaterializedStore.Liveness _liveness;
+	private final Runnable _afterClose;
+	private final IntConsumer _afterRelease;
+	private volatile boolean _closed;
+
+	IndexedMaterializedStoreReader(OOCCache cache, long streamId, IntSupplier completedSize,
+		MaterializedStore.Liveness liveness, Runnable afterClose, IntConsumer afterRelease) {
+		_cache = cache;
+		_streamId = streamId;
+		_completedSize = completedSize;
+		_liveness = liveness;
+		_afterClose = afterClose;
+		_afterRelease = afterRelease;
+	}
+
+	@Override
+	public MaterializedStore.Liveness liveness() {
+		return _liveness;
+	}
+
+	@Override
+	public boolean isClosed() {
+		return _closed;
+	}
+
+	@Override
+	public void close() {
+		if(_closed)
+			return;
+		_closed = true;
+		_afterClose.run();
+	}
+
+	public OOCFuture<StoreLease<T>> request(int index, MemoryAllowance requestAllowance) {
+		checkReady(index);
+		reserve(index);
+		OOCFuture<BlockEntry> pinned = OOCUtils.pinAdmitted(_cache, _streamId, index, requestAllowance, () -> _closed);
+		OOCFuture<StoreLease<T>> result = new OOCFuture<>();
+		pinned.whenComplete((entry, error) -> {
+			if(error != null) {
+				_liveness.unreserve(index);
+				result.completeExceptionally(error);
+			}
+			else if(entry == null) {
+				_liveness.unreserve(index);
+				result.complete(null);
+			}
+			else
+				result.complete(new StoreLease<>(entry, () -> release(index, entry, requestAllowance)));
+		});
+		return result;
+	}
+
+	public StoreLease<T> requestIfLive(int index, MemoryAllowance requestAllowance) {
+		checkReady(index);
+		reserve(index);
+		BlockEntry entry = _cache.pinIfLive(_streamId, index, requestAllowance);
+		if(entry == null) {
+			_liveness.unreserve(index);
+			return null;
+		}
+		return new StoreLease<>(entry, () -> release(index, entry, requestAllowance));
+	}
+
+	private void release(int index, BlockEntry entry, MemoryAllowance requestAllowance) {
+		_cache.unpin(entry, requestAllowance);
+		_liveness.consumed(index);
+		_afterRelease.accept(index);
+	}
+
+	private void reserve(int index) {
+		if(!_liveness.reserve(index))
+			throw new IllegalStateException("Index is no longer live for this reader: " + index);
+	}
+
+	private void checkReady(int index) {
+		if(_closed)
+			throw new IllegalStateException("Reader is closed");
+		if(index < 0 || index >= _completedSize.getAsInt())
+			throw new IndexOutOfBoundsException("Invalid requested index: " + index);
+	}
+}

src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedCallback.java

@@ -0,0 +1,84 @@
+/*
+ * 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.store;
+
+import org.apache.sysds.runtime.DMLRuntimeException;
+import org.apache.sysds.runtime.instructions.ooc.OOCStream;
+import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue;
+import org.apache.sysds.runtime.ooc.cache.BlockEntry;
+
+import java.util.concurrent.atomic.AtomicReference;
+
+public final class MaterializedCallback implements OOCStream.QueueCallback<IndexedMatrixValue> {
+	private final StoreLease<IndexedMatrixValue> _lease;
+	private final AtomicReference<DMLRuntimeException> _failure;
+	private boolean _closed;
+
+	public MaterializedCallback(StoreLease<IndexedMatrixValue> lease) {
+		this(lease, new AtomicReference<>());
+	}
+
+	private MaterializedCallback(StoreLease<IndexedMatrixValue> lease, AtomicReference<DMLRuntimeException> failure) {
+		_lease = lease;
+		_failure = failure;
+	}
+
+	public BlockEntry pinnedEntry() {
+		return _lease != null ? _lease.entry() : null;
+	}
+
+	@Override
+	public IndexedMatrixValue get() {
+		DMLRuntimeException failure = _failure.get();
+		if(failure != null)
+			throw failure;
+		return _lease.value();
+	}
+
+	@Override
+	public synchronized OOCStream.QueueCallback<IndexedMatrixValue> keepOpen() {
+		if(_closed)
+			throw new IllegalStateException("Cannot keep open a closed callback");
+		return new MaterializedCallback(_lease.retain(), _failure);
+	}
+
+	@Override
+	public synchronized void close() {
+		if(_closed)
+			return;
+		_closed = true;
+		_lease.close();
+	}
+
+	@Override
+	public void fail(DMLRuntimeException failure) {
+		_failure.set(failure);
+	}
+
+	@Override
+	public boolean isEos() {
+		return false;
+	}
+
+	@Override
+	public boolean isFailure() {
+		return _failure.get() != null;
+	}
+}