You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@systemds.apache.org by ja...@apache.org on 2026/07/16 13:34:01 UTC
(systemds) branch main updated: [OOC] Add StateTable
This is an automated email from the ASF dual-hosted git repository.
janniklinde pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/systemds.git
The following commit(s) were added to refs/heads/main by this push:
new 14a586fa67 [OOC] Add StateTable
14a586fa67 is described below
commit 14a586fa672ab5156e51d1f5288f31b7a1524c96
Author: Jannik Lindemann <52...@users.noreply.github.com>
AuthorDate: Thu Jul 16 15:33:55 2026 +0200
[OOC] Add StateTable
---
.../apache/sysds/runtime/ooc/cache/OOCFuture.java | 35 +-
.../apache/sysds/runtime/ooc/store/StateTable.java | 414 +++++++++++++++++++++
.../apache/sysds/runtime/ooc/store/StoreLease.java | 4 +-
.../sysds/runtime/ooc/util/StateTableUtils.java | 102 +++++
.../test/component/ooc/StateTableUtilsTest.java | 143 +++++++
5 files changed, 693 insertions(+), 5 deletions(-)
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCFuture.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCFuture.java
index 491fefaad8..d6796e35a8 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCFuture.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCFuture.java
@@ -111,6 +111,36 @@ public class OOCFuture<T> {
}
}
+ public <R> OOCFuture<R> thenCompose(Function<? super T, ? extends OOCFuture<? extends R>> mapper) {
+ OOCFuture<R> result = new OOCFuture<>();
+ whenComplete((value, error) -> {
+ if(error != null) {
+ result.completeExceptionally(error);
+ return;
+ }
+
+ final OOCFuture<? extends R> next;
+ try {
+ next = mapper.apply(value);
+ if(next == null)
+ throw new NullPointerException("thenCompose mapper returned null");
+ }
+ catch(Throwable t) {
+ result.completeExceptionally(t);
+ return;
+ }
+
+ next.whenComplete((nextValue, nextError) -> {
+ if(nextError != null)
+ result.completeExceptionally(nextError);
+ else
+ result.complete(nextValue);
+ });
+ });
+
+ return result;
+ }
+
private <R> void subscribe(Function<? super T, ? extends R> mapper, Consumer<? super R> action,
BiConsumer<? super R, ? super Throwable> completion) {
T value;
@@ -167,7 +197,6 @@ public class OOCFuture<T> {
action.accept(result);
}
catch(Throwable ignored) {
- // Subscribers are independent; one failed callback must not prevent the remaining notifications.
}
}
@@ -181,8 +210,8 @@ public class OOCFuture<T> {
private <R> Subscriber(Function<? super T, ? extends R> mapper, Consumer<? super R> action,
BiConsumer<? super R, ? super Throwable> completion, Subscriber<T> next) {
this.mapper = mapper;
- this.action = (Consumer<Object>)action;
- this.completion = (BiConsumer<Object, Throwable>)(BiConsumer<?, ?>)completion;
+ this.action = (Consumer<Object>) action;
+ this.completion = (BiConsumer<Object, Throwable>) completion;
this.next = next;
}
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/StateTable.java b/src/main/java/org/apache/sysds/runtime/ooc/store/StateTable.java
new file mode 100644
index 0000000000..80eb57dcfe
--- /dev/null
+++ b/src/main/java/org/apache/sysds/runtime/ooc/store/StateTable.java
@@ -0,0 +1,414 @@
+/*
+ * 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.BlockKey;
+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.ManagedPayload;
+import org.apache.sysds.runtime.ooc.memory.MemoryAllowance;
+import org.apache.sysds.runtime.ooc.util.OOCUtils;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicIntegerArray;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.IntToLongFunction;
+
+public final class StateTable<T extends SpillableObject> implements AutoCloseable {
+ private static final int INITIAL_SLOTS = 64;
+
+ private final OOCCache _cache;
+ private final long _streamId;
+ private final AtomicLong _nextGeneration = new AtomicLong();
+ private final CopyOnWriteArrayList<IntToLongFunction> _evictionPolicies = new CopyOnWriteArrayList<>();
+ private final AtomicBoolean _evictionPolicyInstalled = new AtomicBoolean(false);
+ private Slot[] _slots;
+ private volatile AtomicIntegerArray _generationSlots;
+ private volatile boolean _closed;
+
+ public StateTable(OOCCache cache, long streamId) {
+ this(cache, streamId, INITIAL_SLOTS);
+ }
+
+ public StateTable(OOCCache cache, long streamId, int numSlots) {
+ int capacity = Math.max(1, numSlots);
+ _cache = cache;
+ _streamId = streamId;
+ _generationSlots = new AtomicIntegerArray(capacity);
+ _slots = new Slot[capacity];
+ }
+
+ public void addEvictionPolicy(IntToLongFunction slotPolicy) {
+ _evictionPolicies.add(slotPolicy);
+ if(_evictionPolicyInstalled.compareAndSet(false, true))
+ _cache.addEvictionPolicy(_streamId, this::scoreTableEntry);
+ }
+
+ public void put(int index, ManagedPayload<T> payload) {
+ putSlot(index, slot -> finalizeOwnedPut(index, slot, payload));
+ }
+
+ public void putReference(int index, BlockEntry pinned) {
+ checkPinned(pinned);
+ putSlot(index, slot -> finalizeReferencePut(index, slot, pinned));
+ }
+
+ private void putSlot(int index, Consumer<Slot> finalizer) {
+ Slot slot;
+ synchronized(this) {
+ checkOpen();
+ ensureCapacity(index);
+ if(_slots[index] != null)
+ throw new IllegalStateException("State table slot " + index + " is already occupied.");
+ slot = new Slot();
+ _slots[index] = slot;
+ }
+ finalizer.accept(slot);
+ }
+
+ public OOCFuture<StoreLease<T>> putOrTake(int index, ManagedPayload<T> payload, MemoryAllowance leaseAllowance) {
+ return putSlotOrTake(index, leaseAllowance, slot -> finalizeOwnedPut(index, slot, payload));
+ }
+
+ public OOCFuture<StoreLease<T>> putReferenceOrTake(int index, BlockEntry pinned, MemoryAllowance leaseAllowance) {
+ checkPinned(pinned);
+ return putSlotOrTake(index, leaseAllowance, slot -> finalizeReferencePut(index, slot, pinned));
+ }
+
+ private OOCFuture<StoreLease<T>> putSlotOrTake(int index, MemoryAllowance leaseAllowance,
+ Consumer<Slot> finalizer) {
+ Slot putting = null;
+ Slot taken = null;
+ OOCFuture<Void> waitFor = null;
+ synchronized(this) {
+ checkOpen();
+ ensureCapacity(index);
+ Slot existing = _slots[index];
+ if(existing == null) {
+ putting = new Slot();
+ _slots[index] = putting;
+ }
+ else if(existing._putFuture == null) {
+ _slots[index] = null;
+ taken = existing;
+ }
+ else {
+ waitFor = existing._putFuture;
+ }
+ }
+ if(putting != null) {
+ finalizer.accept(putting);
+ return OOCFuture.completed(null);
+ }
+ if(taken != null)
+ return pinTaken(taken, leaseAllowance);
+ return waitFor.thenCompose(ignored -> putSlotOrTake(index, leaseAllowance, finalizer));
+ }
+
+ public OOCFuture<StoreLease<T>> take(int index, MemoryAllowance leaseAllowance) {
+ Slot taken = null;
+ OOCFuture<Void> waitFor = null;
+ synchronized(this) {
+ checkOpen();
+ if(index < 0 || index >= _slots.length)
+ return OOCFuture.completed(null);
+ Slot existing = _slots[index];
+ if(existing == null)
+ return OOCFuture.completed(null);
+ if(existing._putFuture == null) {
+ _slots[index] = null;
+ taken = existing;
+ }
+ else {
+ waitFor = existing._putFuture;
+ }
+ }
+ if(taken != null)
+ return pinTaken(taken, leaseAllowance);
+ return waitFor.thenCompose(ignored -> take(index, leaseAllowance));
+ }
+
+ public OOCFuture<StoreLease<T>> acquire(int index, MemoryAllowance leaseAllowance) {
+ BlockKey key;
+ synchronized(this) {
+ checkOpen();
+ if(index < 0 || index >= _slots.length)
+ return OOCFuture.completed(null);
+ Slot slot = _slots[index];
+ if(slot == null || slot._putFuture != null)
+ return OOCFuture.completed(null);
+ key = slot._key;
+ }
+ OOCFuture<BlockEntry> pinned = OOCUtils.pinAdmitted(_cache, key.getStreamId(), key.getSequenceNumber(),
+ leaseAllowance, () -> _closed);
+ OOCFuture<StoreLease<T>> result = new OOCFuture<>();
+ pinned.whenComplete((entry, error) -> {
+ if(error != null)
+ result.completeExceptionally(error);
+ else
+ result.complete(
+ entry == null ? null : new StoreLease<>(entry, () -> _cache.unpin(entry, leaseAllowance)));
+ });
+ return result;
+ }
+
+ public StoreLease<T> peek(int index, MemoryAllowance leaseAllowance) {
+ BlockKey key;
+ synchronized(this) {
+ checkOpen();
+ if(index < 0 || index >= _slots.length)
+ return null;
+ Slot slot = _slots[index];
+ if(slot == null || slot._putFuture != null)
+ return null;
+ key = slot._key;
+ }
+ BlockEntry entry = _cache.pinIfLive(key.getStreamId(), key.getSequenceNumber(), leaseAllowance);
+ return entry == null ? null : new StoreLease<>(entry, () -> _cache.unpin(entry, leaseAllowance));
+ }
+
+ public void clear(int index) {
+ Slot removed = null;
+ synchronized(this) {
+ if(index < 0 || index >= _slots.length)
+ return;
+ Slot slot = _slots[index];
+ if(slot == null)
+ return;
+ _slots[index] = null;
+ if(slot._putFuture == null)
+ removed = slot;
+ else
+ slot._cleared = true;
+ }
+ if(removed != null)
+ releaseSlot(removed);
+ }
+
+ @Override
+ public void close() {
+ List<Slot> toRelease = new ArrayList<>();
+ synchronized(this) {
+ if(_closed)
+ return;
+ _closed = true;
+ for(int i = 0; i < _slots.length; i++) {
+ Slot slot = _slots[i];
+ if(slot == null)
+ continue;
+ _slots[i] = null;
+ if(slot._putFuture == null)
+ toRelease.add(slot);
+ else
+ slot._cleared = true;
+ }
+ }
+ for(Slot slot : toRelease)
+ releaseSlot(slot);
+ }
+
+ private void finalizeOwnedPut(int index, Slot slot, ManagedPayload<T> payload) {
+ BlockKey key = new BlockKey(_streamId, _nextGeneration.getAndIncrement());
+ BlockEntry entry;
+ try {
+ payload.transfer();
+ }
+ catch(RuntimeException ex) {
+ failPut(index, slot, ex);
+ throw ex;
+ }
+ try {
+ entry = _cache.putPinned(key.getStreamId(), key.getSequenceNumber(), payload.value(), payload.bytes(),
+ payload.owner());
+ }
+ catch(RuntimeException ex) {
+ if(payload.bytes() > 0)
+ payload.owner().release(payload.bytes());
+ failPut(index, slot, ex);
+ throw ex;
+ }
+ boolean cleared;
+ OOCFuture<Void> putFuture;
+ synchronized(this) {
+ slot._key = key;
+ slot._tableOwnedKey = true;
+ int generation = blockIndex(key.getSequenceNumber());
+ ensureGenerationCapacity(generation);
+ _generationSlots.set(generation, index + 1);
+ cleared = slot._cleared;
+ putFuture = slot._putFuture;
+ slot._putFuture = null;
+ }
+ _cache.unpin(entry, payload.owner());
+ if(cleared)
+ releaseSlot(slot);
+ putFuture.complete(null);
+ }
+
+ private void finalizeReferencePut(int index, Slot slot, BlockEntry pinned) {
+ try {
+ _cache.reference(pinned);
+ }
+ catch(RuntimeException ex) {
+ failPut(index, slot, ex);
+ throw ex;
+ }
+
+ boolean cleared;
+ OOCFuture<Void> putFuture;
+ synchronized(this) {
+ slot._key = pinned.getKey();
+ slot._tableOwnedKey = false;
+ cleared = slot._cleared;
+ putFuture = slot._putFuture;
+ slot._putFuture = null;
+ }
+ if(cleared)
+ _cache.dereference(pinned.getKey());
+ putFuture.complete(null);
+ }
+
+ private void failPut(int index, Slot slot, RuntimeException ex) {
+ OOCFuture<Void> putFuture;
+ synchronized(this) {
+ if(index < _slots.length && _slots[index] == slot)
+ _slots[index] = null;
+ putFuture = slot._putFuture;
+ slot._putFuture = null;
+ }
+ if(putFuture != null)
+ putFuture.completeExceptionally(ex);
+ }
+
+ private OOCFuture<StoreLease<T>> pinTaken(Slot slot, MemoryAllowance leaseAllowance) {
+ OOCFuture<BlockEntry> pinned = OOCUtils.pinAdmitted(_cache, slot._key.getStreamId(),
+ slot._key.getSequenceNumber(), leaseAllowance, () -> _closed);
+ OOCFuture<StoreLease<T>> result = new OOCFuture<>();
+ pinned.whenComplete((entry, error) -> {
+ Throwable completionError = error;
+ try {
+ releaseSlot(slot);
+ }
+ catch(RuntimeException releaseError) {
+ if(completionError == null)
+ completionError = releaseError;
+ }
+ if(completionError == null && entry == null)
+ completionError = new IllegalStateException("State table closed while a take was pending.");
+ if(completionError != null) {
+ if(entry != null) {
+ try {
+ _cache.unpin(entry, leaseAllowance);
+ }
+ catch(RuntimeException ignored) {
+ }
+ }
+ result.completeExceptionally(completionError);
+ return;
+ }
+ result.complete(new StoreLease<>(entry, () -> _cache.unpin(entry, leaseAllowance)));
+ });
+ return result;
+ }
+
+ private void releaseSlot(Slot slot) {
+ if(slot._tableOwnedKey) {
+ int generation = blockIndex(slot._key.getSequenceNumber());
+ AtomicIntegerArray slots = _generationSlots;
+ if(generation < slots.length())
+ slots.set(generation, 0);
+ }
+ _cache.dereference(slot._key);
+ }
+
+ private long scoreTableEntry(long generation) {
+ int index = blockIndex(generation);
+ AtomicIntegerArray slots = _generationSlots;
+ if(index >= slots.length())
+ return Long.MAX_VALUE;
+ int encodedSlot = slots.get(index);
+ if(encodedSlot == 0)
+ return Long.MAX_VALUE;
+ int slot = encodedSlot - 1;
+ long score = Long.MAX_VALUE;
+ for(IntToLongFunction policy : _evictionPolicies)
+ score = Math.min(score, policy.applyAsLong(slot));
+ return score;
+ }
+
+ private void ensureGenerationCapacity(int index) {
+ AtomicIntegerArray slots = _generationSlots;
+ if(index < slots.length())
+ return;
+ int newLength = slots.length();
+ while(index >= newLength) {
+ if(newLength > Integer.MAX_VALUE / 2)
+ throw new IllegalStateException("State table generation map capacity overflow");
+ newLength *= 2;
+ }
+ AtomicIntegerArray grown = new AtomicIntegerArray(newLength);
+ for(int i = 0; i < slots.length(); i++)
+ grown.set(i, slots.get(i));
+ _generationSlots = grown;
+ }
+
+ private static int blockIndex(long sequenceNumber) {
+ if(sequenceNumber < 0 || sequenceNumber > Integer.MAX_VALUE)
+ throw new IndexOutOfBoundsException("Invalid block index: " + sequenceNumber);
+ return (int) sequenceNumber;
+ }
+
+ private void checkOpen() {
+ if(_closed)
+ throw new IllegalStateException("State table is closed.");
+ }
+
+ private static void checkPinned(BlockEntry pinned) {
+ if(!pinned.isPinned())
+ throw new IllegalArgumentException(
+ "Reference install requires the supplied entry to be pinned: " + pinned.getKey());
+ }
+
+ private void ensureCapacity(int index) {
+ if(index < 0)
+ throw new IndexOutOfBoundsException("Invalid slot index: " + index);
+ if(index < _slots.length)
+ return;
+ int newLength = _slots.length;
+ while(index >= newLength)
+ newLength *= 2;
+ Slot[] grown = new Slot[newLength];
+ System.arraycopy(_slots, 0, grown, 0, _slots.length);
+ _slots = grown;
+ }
+
+ private static final class Slot {
+ private boolean _cleared;
+ private boolean _tableOwnedKey;
+ private BlockKey _key;
+ private OOCFuture<Void> _putFuture = new OOCFuture<>();
+ }
+}
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/StoreLease.java b/src/main/java/org/apache/sysds/runtime/ooc/store/StoreLease.java
index 814c95cb56..c9d4ca6f8c 100644
--- a/src/main/java/org/apache/sysds/runtime/ooc/store/StoreLease.java
+++ b/src/main/java/org/apache/sysds/runtime/ooc/store/StoreLease.java
@@ -31,11 +31,11 @@ public final class StoreLease<T extends SpillableObject> implements AutoCloseabl
private final AtomicInteger _shared;
private boolean _open;
- StoreLease(BlockEntry entry, Runnable releaser) {
+ public StoreLease(BlockEntry entry, Runnable releaser) {
this(null, entry, releaser, new AtomicInteger(1));
}
- StoreLease(T value, Runnable releaser) {
+ public StoreLease(T value, Runnable releaser) {
this(value, null, releaser, new AtomicInteger(1));
}
diff --git a/src/main/java/org/apache/sysds/runtime/ooc/util/StateTableUtils.java b/src/main/java/org/apache/sysds/runtime/ooc/util/StateTableUtils.java
new file mode 100644
index 0000000000..a6aecd40c7
--- /dev/null
+++ b/src/main/java/org/apache/sysds/runtime/ooc/util/StateTableUtils.java
@@ -0,0 +1,102 @@
+/*
+ * 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.util;
+
+import org.apache.sysds.runtime.instructions.ooc.OOCStream;
+import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue;
+import org.apache.sysds.runtime.matrix.data.MatrixBlock;
+import org.apache.sysds.runtime.ooc.cache.OOCFuture;
+import org.apache.sysds.runtime.ooc.memory.InMemoryQueueCallback;
+import org.apache.sysds.runtime.ooc.memory.ManagedPayload;
+import org.apache.sysds.runtime.ooc.memory.MemoryAllowance;
+import org.apache.sysds.runtime.ooc.store.MaterializedCallback;
+import org.apache.sysds.runtime.ooc.store.StateTable;
+import org.apache.sysds.runtime.ooc.store.StoreLease;
+
+public final class StateTableUtils {
+ public static OOCFuture<Match> putOrTake(StateTable<IndexedMatrixValue> table, int slot,
+ OOCStream.QueueCallback<IndexedMatrixValue> tile, MemoryAllowance allowance) {
+ if(tile instanceof MaterializedCallback pinned && pinned.pinnedEntry() != null)
+ return putReferenceOrTake(table, slot, pinned, allowance);
+ ManagedPayload<IndexedMatrixValue> payload;
+ if(tile instanceof InMemoryQueueCallback managed && managed.getManagedBytes() > 0) {
+ payload = managed.extractManagedPayload();
+ managed.close();
+ }
+ else {
+ IndexedMatrixValue value = tile.get();
+ long bytes = ((MatrixBlock) value.getValue()).getExactSerializedSize();
+ allowance.reserveBlocking(bytes);
+ payload = new ManagedPayload<>(value, bytes, allowance);
+ tile.close();
+ }
+ OOCFuture<Match> result = new OOCFuture<>();
+ OOCFuture<StoreLease<IndexedMatrixValue>> matched;
+ try {
+ matched = table.putOrTake(slot, payload, allowance);
+ }
+ catch(RuntimeException ex) {
+ payload.release();
+ return OOCFuture.failed(ex);
+ }
+ matched.whenComplete((lease, error) -> {
+ if(error != null) {
+ payload.release();
+ result.completeExceptionally(error);
+ }
+ else if(lease == null)
+ result.complete(null);
+ else
+ result.complete(new Match(new MaterializedCallback(new StoreLease<>(payload.value(), payload::release)),
+ new MaterializedCallback(lease)));
+ });
+ return result;
+ }
+
+ private static OOCFuture<Match> putReferenceOrTake(StateTable<IndexedMatrixValue> table, int slot,
+ MaterializedCallback pinned, MemoryAllowance allowance) {
+ OOCFuture<Match> result = new OOCFuture<>();
+ OOCFuture<StoreLease<IndexedMatrixValue>> matched;
+ try {
+ matched = table.putReferenceOrTake(slot, pinned.pinnedEntry(), allowance);
+ }
+ catch(RuntimeException ex) {
+ pinned.close();
+ return OOCFuture.failed(ex);
+ }
+ matched.whenComplete((lease, error) -> {
+ if(error != null) {
+ pinned.close();
+ result.completeExceptionally(error);
+ }
+ else if(lease == null) {
+ pinned.close();
+ result.complete(null);
+ }
+ else
+ result.complete(new Match(pinned, new MaterializedCallback(lease)));
+ });
+ return result;
+ }
+
+ public record Match(OOCStream.QueueCallback<IndexedMatrixValue> left,
+ OOCStream.QueueCallback<IndexedMatrixValue> right) {
+ }
+}
diff --git a/src/test/java/org/apache/sysds/test/component/ooc/StateTableUtilsTest.java b/src/test/java/org/apache/sysds/test/component/ooc/StateTableUtilsTest.java
new file mode 100644
index 0000000000..7debcd04e3
--- /dev/null
+++ b/src/test/java/org/apache/sysds/test/component/ooc/StateTableUtilsTest.java
@@ -0,0 +1,143 @@
+/*
+ * 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.test.component.ooc;
+
+import java.util.concurrent.TimeUnit;
+
+import org.apache.sysds.runtime.instructions.ooc.OOCStream;
+import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue;
+import org.apache.sysds.runtime.matrix.data.MatrixBlock;
+import org.apache.sysds.runtime.matrix.data.MatrixIndexes;
+import org.apache.sysds.runtime.ooc.cache.OOCCacheImpl;
+import org.apache.sysds.runtime.ooc.memory.GlobalMemoryBroker;
+import org.apache.sysds.runtime.ooc.memory.InMemoryQueueCallback;
+import org.apache.sysds.runtime.ooc.memory.ManagedPayload;
+import org.apache.sysds.runtime.ooc.memory.SyncMemoryAllowance;
+import org.apache.sysds.runtime.ooc.store.MaterializedCallback;
+import org.apache.sysds.runtime.ooc.store.StateTable;
+import org.apache.sysds.runtime.ooc.store.StoreLease;
+import org.apache.sysds.runtime.ooc.util.StateTableUtils;
+import org.apache.sysds.test.component.ooc.cache.OOCCacheTestUtils;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+
+public class StateTableUtilsTest {
+ private static final long MEMORY_LIMIT = 100_000_000;
+ private static final long WAIT_SECONDS = 10;
+ private static final long TILE_BYTES = new MatrixBlock(4, 4, 1.0).getExactSerializedSize();
+
+ private SyncMemoryAllowance _producer;
+ private SyncMemoryAllowance _reader;
+ private OOCCacheImpl _cache;
+ private StateTable<IndexedMatrixValue> _source;
+ private StateTable<IndexedMatrixValue> _table;
+
+ @Before
+ public void setUp() {
+ GlobalMemoryBroker broker = new GlobalMemoryBroker(1_000_000_000);
+ _producer = new SyncMemoryAllowance(broker);
+ _reader = new SyncMemoryAllowance(broker);
+ _producer.setTargetMemory(MEMORY_LIMIT);
+ _reader.setTargetMemory(MEMORY_LIMIT);
+ _cache = new OOCCacheImpl(new OOCCacheTestUtils.RecordingOOCIOHandler(), MEMORY_LIMIT, MEMORY_LIMIT);
+ _source = new StateTable<>(_cache, 1);
+ _table = new StateTable<>(_cache, 2);
+ }
+
+ @After
+ public void tearDown() {
+ _source.close();
+ _table.close();
+ _cache.shutdown();
+ _producer.destroy();
+ _reader.destroy();
+ }
+
+ @Test
+ public void testCallbackPutOrTake() throws Exception {
+ _producer.reserveBlocking(TILE_BYTES);
+ _source.put(0, new ManagedPayload<>(tile(1.0), TILE_BYTES, _producer));
+ StoreLease<IndexedMatrixValue> pinned = _source.peek(0, _reader);
+ Assert.assertNotNull(pinned);
+ Assert.assertNull(StateTableUtils.putOrTake(_table, 0, new MaterializedCallback(pinned), _reader)
+ .get(WAIT_SECONDS, TimeUnit.SECONDS));
+ Assert.assertEquals(0, _reader.getUsedMemory());
+
+ _producer.reserveBlocking(TILE_BYTES);
+ StateTableUtils.Match referenced = StateTableUtils
+ .putOrTake(_table, 0, new InMemoryQueueCallback(tile(2.0), null, _producer, TILE_BYTES), _reader)
+ .get(WAIT_SECONDS, TimeUnit.SECONDS);
+ Assert.assertNotNull(referenced);
+ try(OOCStream.QueueCallback<IndexedMatrixValue> left = referenced.left();
+ OOCStream.QueueCallback<IndexedMatrixValue> right = referenced.right()) {
+ Assert.assertEquals(2.0, left.get().getValue().get(0, 0), 0.0);
+ Assert.assertEquals(1.0, right.get().getValue().get(0, 0), 0.0);
+ }
+
+ _producer.reserveBlocking(TILE_BYTES);
+ Assert.assertNull(StateTableUtils
+ .putOrTake(_table, 1, new InMemoryQueueCallback(tile(3.0), null, _producer, TILE_BYTES), _reader)
+ .get(WAIT_SECONDS, TimeUnit.SECONDS));
+ StateTableUtils.Match copied = StateTableUtils
+ .putOrTake(_table, 1, new OOCStream.SimpleQueueCallback<>(tile(4.0), null), _reader)
+ .get(WAIT_SECONDS, TimeUnit.SECONDS);
+ Assert.assertNotNull(copied);
+ try(OOCStream.QueueCallback<IndexedMatrixValue> own = copied.left();
+ OOCStream.QueueCallback<IndexedMatrixValue> partner = copied.right()) {
+ Assert.assertEquals(4.0, own.get().getValue().get(0, 0), 0.0);
+ Assert.assertEquals(3.0, partner.get().getValue().get(0, 0), 0.0);
+ }
+
+ Assert.assertEquals(0, _producer.getUsedMemory());
+ Assert.assertEquals(0, _reader.getUsedMemory());
+ _source.close();
+ _table.close();
+ OOCCacheTestUtils.await(() -> _cache.getOwnedCacheSize() == 0, WAIT_SECONDS);
+ }
+
+ @Test
+ public void testStateTableLifecycle() throws Exception {
+ _producer.reserveBlocking(TILE_BYTES);
+ _table.put(0, new ManagedPayload<>(tile(5.0), TILE_BYTES, _producer));
+ try(StoreLease<IndexedMatrixValue> lease = _table.acquire(0, _reader).get(WAIT_SECONDS, TimeUnit.SECONDS)) {
+ Assert.assertNotNull(lease);
+ Assert.assertEquals(5.0, lease.value().getValue().get(0, 0), 0.0);
+ }
+ try(StoreLease<IndexedMatrixValue> lease = _table.take(0, _reader).get(WAIT_SECONDS, TimeUnit.SECONDS)) {
+ Assert.assertNotNull(lease);
+ Assert.assertEquals(5.0, lease.value().getValue().get(0, 0), 0.0);
+ }
+ Assert.assertNull(_table.take(0, _reader).get(WAIT_SECONDS, TimeUnit.SECONDS));
+
+ _producer.reserveBlocking(TILE_BYTES);
+ _table.put(1, new ManagedPayload<>(tile(6.0), TILE_BYTES, _producer));
+ _table.clear(1);
+ Assert.assertNull(_table.take(1, _reader).get(WAIT_SECONDS, TimeUnit.SECONDS));
+ Assert.assertEquals(0, _producer.getUsedMemory());
+ Assert.assertEquals(0, _reader.getUsedMemory());
+ OOCCacheTestUtils.await(() -> _cache.getOwnedCacheSize() == 0, WAIT_SECONDS);
+ }
+
+ private static IndexedMatrixValue tile(double value) {
+ return new IndexedMatrixValue(new MatrixIndexes(1, 1), new MatrixBlock(4, 4, value));
+ }
+}