From bba24e6e5b98613556d69706b076443932c509b6 Mon Sep 17 00:00:00 2001 From: Jannik Lindemann Date: Wed, 15 Jul 2026 16:39:54 +0200 Subject: [PATCH] [OOC] Add StateTable --- .../sysds/runtime/ooc/cache/OOCFuture.java | 35 +- .../sysds/runtime/ooc/store/StateTable.java | 414 ++++++++++++++++++ .../sysds/runtime/ooc/store/StoreLease.java | 4 +- .../runtime/ooc/util/StateTableUtils.java | 102 +++++ .../component/ooc/StateTableUtilsTest.java | 143 ++++++ 5 files changed, 693 insertions(+), 5 deletions(-) create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/store/StateTable.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/util/StateTableUtils.java create mode 100644 src/test/java/org/apache/sysds/test/component/ooc/StateTableUtilsTest.java 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 491fefaad87..d6796e35a87 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 T get(long timeout, TimeUnit unit) throws InterruptedException, Execution } } + public OOCFuture thenCompose(Function> mapper) { + OOCFuture result = new OOCFuture<>(); + whenComplete((value, error) -> { + if(error != null) { + result.completeExceptionally(error); + return; + } + + final OOCFuture 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 void subscribe(Function mapper, Consumer action, BiConsumer completion) { T value; @@ -167,7 +197,6 @@ else if(resultError == null) action.accept(result); } catch(Throwable ignored) { - // Subscribers are independent; one failed callback must not prevent the remaining notifications. } } @@ -181,8 +210,8 @@ private static final class Subscriber { private Subscriber(Function mapper, Consumer action, BiConsumer completion, Subscriber next) { this.mapper = mapper; - this.action = (Consumer)action; - this.completion = (BiConsumer)(BiConsumer)completion; + this.action = (Consumer) action; + this.completion = (BiConsumer) 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 00000000000..80eb57dcfe4 --- /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 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 _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 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 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> putOrTake(int index, ManagedPayload payload, MemoryAllowance leaseAllowance) { + return putSlotOrTake(index, leaseAllowance, slot -> finalizeOwnedPut(index, slot, payload)); + } + + public OOCFuture> putReferenceOrTake(int index, BlockEntry pinned, MemoryAllowance leaseAllowance) { + checkPinned(pinned); + return putSlotOrTake(index, leaseAllowance, slot -> finalizeReferencePut(index, slot, pinned)); + } + + private OOCFuture> putSlotOrTake(int index, MemoryAllowance leaseAllowance, + Consumer finalizer) { + Slot putting = null; + Slot taken = null; + OOCFuture 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> take(int index, MemoryAllowance leaseAllowance) { + Slot taken = null; + OOCFuture 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> 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 pinned = OOCUtils.pinAdmitted(_cache, key.getStreamId(), key.getSequenceNumber(), + leaseAllowance, () -> _closed); + OOCFuture> 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 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 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 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 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 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 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> pinTaken(Slot slot, MemoryAllowance leaseAllowance) { + OOCFuture pinned = OOCUtils.pinAdmitted(_cache, slot._key.getStreamId(), + slot._key.getSequenceNumber(), leaseAllowance, () -> _closed); + OOCFuture> 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 _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 814c95cb56a..c9d4ca6f8cd 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 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 00000000000..a6aecd40c71 --- /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 putOrTake(StateTable table, int slot, + OOCStream.QueueCallback tile, MemoryAllowance allowance) { + if(tile instanceof MaterializedCallback pinned && pinned.pinnedEntry() != null) + return putReferenceOrTake(table, slot, pinned, allowance); + ManagedPayload 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 result = new OOCFuture<>(); + OOCFuture> 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 putReferenceOrTake(StateTable table, int slot, + MaterializedCallback pinned, MemoryAllowance allowance) { + OOCFuture result = new OOCFuture<>(); + OOCFuture> 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 left, + OOCStream.QueueCallback 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 00000000000..7debcd04e3a --- /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 _source; + private StateTable _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 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 left = referenced.left(); + OOCStream.QueueCallback 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 own = copied.left(); + OOCStream.QueueCallback 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 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 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)); + } +}