From 9c439a816602e1c5a435129b3c7d584ec9be6655 Mon Sep 17 00:00:00 2001 From: Jannik Lindemann Date: Wed, 15 Jul 2026 14:05:57 +0200 Subject: [PATCH] [OOC] Add MaterializedStore Structures --- .../ooc/memory/InMemoryQueueCallback.java | 22 ++ .../runtime/ooc/memory/ManagedPayload.java | 68 ++++ .../store/IndexedMaterializedStoreReader.java | 117 ++++++ .../ooc/store/MaterializedCallback.java | 84 ++++ .../runtime/ooc/store/MaterializedStore.java | 233 ++++++++++++ .../ooc/store/OOCStreamMaterializer.java | 151 ++++++++ .../store/OrderedMaterializedStoreReader.java | 253 ++++++++++++ .../sysds/runtime/ooc/store/StoreLease.java | 78 ++++ .../sysds/runtime/ooc/util/OOCUtils.java | 36 ++ .../component/ooc/MaterializedStoreTest.java | 359 ++++++++++++++++++ 10 files changed, 1401 insertions(+) create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/memory/ManagedPayload.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/store/IndexedMaterializedStoreReader.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedCallback.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStore.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/store/OOCStreamMaterializer.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/store/OrderedMaterializedStoreReader.java create mode 100644 src/main/java/org/apache/sysds/runtime/ooc/store/StoreLease.java create mode 100644 src/test/java/org/apache/sysds/test/component/ooc/MaterializedStoreTest.java diff --git a/src/main/java/org/apache/sysds/runtime/ooc/memory/InMemoryQueueCallback.java b/src/main/java/org/apache/sysds/runtime/ooc/memory/InMemoryQueueCallback.java index 7496279d3e3..7fafc042e41 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/memory/InMemoryQueueCallback.java +++ b/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 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); + } + } } diff --git a/src/main/java/org/apache/sysds/runtime/ooc/memory/ManagedPayload.java b/src/main/java/org/apache/sysds/runtime/ooc/memory/ManagedPayload.java new file mode 100644 index 00000000000..494ced159a0 --- /dev/null +++ b/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 { + 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(); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/IndexedMaterializedStoreReader.java b/src/main/java/org/apache/sysds/runtime/ooc/store/IndexedMaterializedStoreReader.java new file mode 100644 index 00000000000..3106705a34a --- /dev/null +++ b/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 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> request(int index, MemoryAllowance requestAllowance) { + checkReady(index); + reserve(index); + OOCFuture pinned = OOCUtils.pinAdmitted(_cache, _streamId, index, requestAllowance, () -> _closed); + OOCFuture> 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 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); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedCallback.java b/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedCallback.java new file mode 100644 index 00000000000..829e60b32ea --- /dev/null +++ b/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 { + private final StoreLease _lease; + private final AtomicReference _failure; + private boolean _closed; + + public MaterializedCallback(StoreLease lease) { + this(lease, new AtomicReference<>()); + } + + private MaterializedCallback(StoreLease lease, AtomicReference 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 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; + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStore.java b/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStore.java new file mode 100644 index 00000000000..d1ed01ffa3e --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/MaterializedStore.java @@ -0,0 +1,233 @@ +/* + * 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.BitSet; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; + +public final class MaterializedStore { + private final OOCCache _cache; + private final long _streamId; + private final ArrayList _registeredReaders; + private final BitSet _forgotten; + private final AtomicInteger _published; + private final AtomicInteger _publishedCount; + + private volatile List _readers; + private volatile int _completedSize; + private volatile boolean _complete; + private volatile boolean _readersSealed; + private volatile boolean _closed; + + public MaterializedStore(OOCCache cache, long streamId) { + _cache = cache; + _streamId = streamId; + _registeredReaders = new ArrayList<>(); + _forgotten = new BitSet(); + _published = new AtomicInteger(); + _publishedCount = new AtomicInteger(); + _readers = Collections.emptyList(); + } + + StoreLease publishPinnedLive(int index, T value, long bytes, MemoryAllowance allowance) { + BlockEntry entry; + try { + if(_complete || _closed) + throw new IllegalStateException("Store no longer accepts published items"); + if(index < 0 || index == Integer.MAX_VALUE) + throw new IndexOutOfBoundsException("Invalid index: " + index); + entry = _cache.putPinned(_streamId, index, value, bytes, allowance); + } + catch(RuntimeException ex) { + if(bytes > 0) + allowance.release(bytes); + throw ex; + } + _publishedCount.incrementAndGet(); + updatePublished(index + 1); + return new StoreLease<>(entry, () -> { + _cache.unpin(entry, allowance); + tryForget(index); + }); + } + + StoreLease publishPinnedLive(int index, ManagedPayload payload) { + payload.transfer(); + return publishPinnedLive(index, payload.value(), payload.bytes(), payload.owner()); + } + + public synchronized void complete() { + if(_complete) + return; + _completedSize = _published.get(); + if(_publishedCount.get() != _completedSize) + throw new IllegalStateException("Incomplete publication: " + _publishedCount.get() + + " published items for logical range [0, " + _completedSize + ")"); + _complete = true; + } + + public synchronized OrderedMaterializedStoreReader openReader(AccessPattern pattern, MemoryAllowance allowance, + int maxPrefetch) { + return openReader(pattern, allowance, maxPrefetch, true); + } + + public synchronized OrderedMaterializedStoreReader openReader(AccessPattern pattern, MemoryAllowance allowance, + int maxPrefetch, boolean softOrdering) { + if(!_complete || _closed) + throw new IllegalStateException("Readers require a completed store"); + if(_readersSealed) + throw new IllegalStateException("Store no longer accepts new readers"); + OrderedMaterializedStoreReader reader = new OrderedMaterializedStoreReader<>(_cache, _streamId, pattern, + allowance, Math.max(1, maxPrefetch), softOrdering, this::forgetAfterReaderClose, this::tryForget); + _registeredReaders.add(reader); + return reader; + } + + public synchronized IndexedMaterializedStoreReader openIndexedReader(Liveness liveness) { + if(!_complete || _closed) + throw new IllegalStateException("Readers require a completed store"); + if(_readersSealed) + throw new IllegalStateException("Store no longer accepts new readers"); + IndexedMaterializedStoreReader reader = new IndexedMaterializedStoreReader<>(_cache, _streamId, + () -> _completedSize, liveness, this::forgetAfterReaderClose, this::tryForget); + _registeredReaders.add(reader); + return reader; + } + + public OOCFuture> requestPublished(int index, MemoryAllowance allowance) { + if(_closed) + throw new IllegalStateException("Store is closed"); + if(index < 0 || index >= _published.get()) + throw new IndexOutOfBoundsException("Invalid requested index: " + index); + OOCFuture pinned = OOCUtils.pinAdmitted(_cache, _streamId, index, allowance, () -> _closed); + OOCFuture> result = new OOCFuture<>(); + pinned.whenComplete((entry, error) -> { + if(error != null) + result.completeExceptionally(error); + else if(entry == null) + result.complete(null); + else + result.complete(new StoreLease<>(entry, () -> _cache.unpin(entry, allowance))); + }); + return result; + } + + public synchronized void sealReaders() { + if(_closed) + throw new IllegalStateException("Cannot seal readers for a closed store"); + if(_readersSealed) + return; + _readers = new ArrayList<>(_registeredReaders); + _readersSealed = true; + int publishedSize = _complete ? _completedSize : _published.get(); + for(int i = 0; i < publishedSize; i++) + tryForget(i); + } + + public int size() { + return _complete ? _completedSize : _published.get(); + } + + public void close() { + List localReaders; + synchronized(this) { + if(_closed) + return; + _closed = true; + localReaders = _readersSealed ? _readers : new ArrayList<>(_registeredReaders); + } + for(StoreReader localReader : localReaders) + localReader.close(); + for(int i = 0; i < size(); i++) + if(markForgotten(i)) + _cache.dereference(new BlockKey(_streamId, i)); + } + + private void tryForget(int index) { + if(!_readersSealed) + return; + List localReaders = _readers; + for(StoreReader reader : localReaders) + if(!reader.isClosed() && reader.liveness().needs(index)) + return; + if(markForgotten(index)) + _cache.dereference(new BlockKey(_streamId, index)); + } + + private synchronized boolean markForgotten(int index) { + if(_forgotten.get(index)) + return false; + _forgotten.set(index); + return true; + } + + private void forgetAfterReaderClose() { + if(_closed || !_readersSealed) + return; + for(int i = 0; i < _completedSize; i++) + tryForget(i); + } + + private void updatePublished(int size) { + int current = _published.get(); + while(current < size && !_published.compareAndSet(current, size)) + current = _published.get(); + } + + public interface Liveness { + boolean needs(int index); + + void consumed(int index); + + default boolean reserve(int index) { + return needs(index); + } + + default void unreserve(int index) { + } + } + + public interface AccessPattern extends Liveness { + boolean hasNext(); + + int next(); + } + + public interface StoreReader extends AutoCloseable { + Liveness liveness(); + + boolean isClosed(); + + @Override + void close(); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/OOCStreamMaterializer.java b/src/main/java/org/apache/sysds/runtime/ooc/store/OOCStreamMaterializer.java new file mode 100644 index 00000000000..2cede71f804 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/OOCStreamMaterializer.java @@ -0,0 +1,151 @@ +/* + * 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.matrix.data.MatrixBlock; +import org.apache.sysds.runtime.matrix.data.MatrixIndexes; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; +import org.apache.sysds.runtime.ooc.memory.InMemoryQueueCallback; +import org.apache.sysds.runtime.ooc.memory.MemoryAllowance; + +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Consumer; +import java.util.function.ToIntFunction; + +public final class OOCStreamMaterializer implements Consumer> { + private final MaterializedStore _store; + private final ToIntFunction _linearize; + private final MemoryAllowance _allowance; + private final List>> _liveConsumers; + private final OOCFuture _completion; + private final AtomicBoolean _done; + + public OOCStreamMaterializer(MaterializedStore store, ToIntFunction linearize, + MemoryAllowance allowance) { + this(store, linearize, allowance, List.of()); + } + + public OOCStreamMaterializer(MaterializedStore store, ToIntFunction linearize, + MemoryAllowance allowance, List>> liveConsumers) { + _store = store; + _linearize = linearize; + _allowance = allowance; + _liveConsumers = List.copyOf(liveConsumers); + _completion = new OOCFuture<>(); + _done = new AtomicBoolean(false); + } + + public void attach(OOCStream source) { + source.setSubscriber(this); + } + + public OOCFuture completion() { + return _completion; + } + + @Override + public void accept(OOCStream.QueueCallback callback) { + if(_done.get()) { + callback.close(); + return; + } + try(callback) { + if(callback.isFailure()) { + try { + callback.get(); + } + catch(DMLRuntimeException ex) { + fail(ex); + } + return; + } + if(callback.isEos()) { + finish(); + return; + } + publish(callback); + } + catch(RuntimeException ex) { + fail(DMLRuntimeException.of(ex)); + } + } + + private void publish(OOCStream.QueueCallback callback) { + IndexedMatrixValue value = callback.get(); + int index = _linearize.applyAsInt(value.getIndexes()); + StoreLease lease; + if(callback instanceof InMemoryQueueCallback managed && managed.getManagedBytes() > 0) { + lease = _store.publishPinnedLive(index, managed.extractManagedPayload()); + } + else { + // To handle not-yet managed callbacks + long bytes = serializedSize(value); + _allowance.reserveBlocking(bytes); + lease = _store.publishPinnedLive(index, value, bytes, _allowance); + } + try(lease) { + for(Consumer> liveConsumer : _liveConsumers) { + try(OOCStream.QueueCallback alias = new MaterializedCallback(lease.retain())) { + liveConsumer.accept(alias); + } + } + } + } + + private void finish() { + if(!_done.compareAndSet(false, true)) + return; + try { + _store.complete(); + } + catch(RuntimeException ex) { + deliverEos(DMLRuntimeException.of(ex)); + _completion.completeExceptionally(ex); + return; + } + deliverEos(null); + _completion.complete(null); + } + + private void fail(DMLRuntimeException failure) { + if(!_done.compareAndSet(false, true)) + return; + deliverEos(failure); + _completion.completeExceptionally(failure); + } + + private void deliverEos(DMLRuntimeException failure) { + for(Consumer> liveConsumer : _liveConsumers) { + try { + liveConsumer.accept(OOCStream.eos(failure)); + } + catch(RuntimeException ignored) { + } + } + } + + private static long serializedSize(IndexedMatrixValue value) { + return ((MatrixBlock) value.getValue()).getExactSerializedSize(); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/store/OrderedMaterializedStoreReader.java b/src/main/java/org/apache/sysds/runtime/ooc/store/OrderedMaterializedStoreReader.java new file mode 100644 index 00000000000..2dd048750c1 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/OrderedMaterializedStoreReader.java @@ -0,0 +1,253 @@ +/* + * 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.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.ArrayDeque; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.IntConsumer; + +public final class OrderedMaterializedStoreReader implements MaterializedStore.StoreReader { + private static final Request CLOSED = new Request(-1, OOCFuture.completed(null)); + + private final OOCCache _cache; + private final long _streamId; + private final MaterializedStore.AccessPattern _pattern; + private final MemoryAllowance _allowance; + private final int _maxPrefetch; + private final boolean _softOrdering; + private final Runnable _afterClose; + private final IntConsumer _afterRelease; + private final BlockingQueue _requests; + private final AtomicInteger _inFlightRequests; + private volatile boolean _closed; + + OrderedMaterializedStoreReader(OOCCache cache, long streamId, MaterializedStore.AccessPattern pattern, + MemoryAllowance allowance, int maxPrefetch, boolean softOrdering, Runnable afterClose, + IntConsumer afterRelease) { + _cache = cache; + _streamId = streamId; + _pattern = pattern; + _allowance = allowance; + _maxPrefetch = maxPrefetch; + _softOrdering = softOrdering; + _afterClose = afterClose; + _afterRelease = afterRelease; + _requests = new LinkedBlockingQueue<>(); + _inFlightRequests = new AtomicInteger(); + } + + @Override + public MaterializedStore.Liveness liveness() { + return _pattern; + } + + @Override + public boolean isClosed() { + return _closed; + } + + @Override + public void close() { + ArrayDeque pending = new ArrayDeque<>(); + synchronized(this) { + if(_closed) + return; + _closed = true; + _requests.drainTo(pending); + } + for(Request request : pending) { + releaseWhenReady(request); + if(_softOrdering) + _inFlightRequests.decrementAndGet(); + } + if(_softOrdering) + _requests.offer(CLOSED); + _afterClose.run(); + } + + public boolean hasNext() { + if(_softOrdering) { + checkReady(); + fillSoft(); + return _inFlightRequests.get() > 0; + } + synchronized(this) { + checkReady(); + fillStrict(); + return !_requests.isEmpty(); + } + } + + public StoreLease next() throws InterruptedException { + if(_softOrdering) { + checkReady(); + return nextSoft(); + } + Request request; + synchronized(this) { + checkReady(); + fillStrict(); + if(_requests.isEmpty()) + throw new IllegalStateException("No remaining item"); + request = _requests.remove(); + } + BlockEntry entry; + try { + entry = awaitEntry(request); + } + catch(InterruptedException | RuntimeException ex) { + releaseWhenReady(request); + throw ex; + } + synchronized(this) { + if(_closed) { + _cache.unpin(entry, _allowance); + throw new IllegalStateException("Reader is closed"); + } + try { + fillStrict(); + } + catch(RuntimeException ex) { + _cache.unpin(entry, _allowance); + throw ex; + } + } + return new StoreLease<>(entry, () -> release(request._index, entry)); + } + + public void release(int index, BlockEntry entry) { + _cache.unpin(entry, _allowance); + _pattern.consumed(index); + _afterRelease.accept(index); + } + + private void checkReady() { + if(_closed) + throw new IllegalStateException("Reader is closed"); + } + + private StoreLease nextSoft() throws InterruptedException { + fillSoft(); + if(_inFlightRequests.get() <= 0) + throw new IllegalStateException("No remaining item"); + Request request = _requests.take(); + if(request == CLOSED) + throw new IllegalStateException("Reader is closed"); + _inFlightRequests.decrementAndGet(); + BlockEntry entry; + try { + entry = request._future.get(); + } + catch(ExecutionException e) { + throw DMLRuntimeException.of(e.getCause()); + } + if(entry == null) + throw new IllegalStateException("Reader is closed"); + fillSoft(); + return new StoreLease<>(entry, () -> release(request._index, entry)); + } + + private void fillStrict() { + while(_requests.size() < _maxPrefetch && _pattern.hasNext()) { + int index = _pattern.next(); + OOCFuture future = _cache.pin(_streamId, index, _allowance); + _requests.offer(new Request(index, future)); + } + } + + private void fillSoft() { + while(_inFlightRequests.get() < _maxPrefetch && _pattern.hasNext()) { + int index = _pattern.next(); + OOCFuture future = _cache.pin(_streamId, index, _allowance); + _inFlightRequests.incrementAndGet(); + registerSoftRequest(new Request(index, future)); + } + } + + private void registerSoftRequest(Request request) { + request._future.whenComplete((entry, error) -> { + if(error != null || entry != null) { + completeSoft(request); + return; + } + request._future = OOCUtils.pinAdmitted(_cache, _streamId, request._index, _allowance, () -> _closed); + request._future.whenComplete((admittedEntry, admittedError) -> completeSoft(request)); + }); + } + + private void completeSoft(Request request) { + if(_closed) { + releaseWhenReady(request); + _inFlightRequests.decrementAndGet(); + return; + } + _requests.offer(request); + if(_closed && _requests.remove(request)) { + releaseWhenReady(request); + _inFlightRequests.decrementAndGet(); + } + } + + private BlockEntry awaitEntry(Request request) throws InterruptedException { + try { + BlockEntry entry = request._future.get(); + if(entry == null) { + OOCFuture retried = OOCUtils.pinAdmitted(_cache, _streamId, request._index, _allowance, + () -> _closed); + request._future = retried; + entry = retried.get(); + if(entry == null) + throw new IllegalStateException("Reader is closed"); + } + return entry; + } + catch(ExecutionException e) { + throw DMLRuntimeException.of(e.getCause()); + } + } + + private void releaseWhenReady(Request request) { + request._future.whenComplete((entry, error) -> { + if(entry != null) + _cache.unpin(entry, _allowance); + }); + } + + private static final class Request { + private final int _index; + private OOCFuture _future; + + private Request(int index, OOCFuture future) { + _index = index; + _future = future; + } + } +} 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 new file mode 100644 index 00000000000..814c95cb56a --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/store/StoreLease.java @@ -0,0 +1,78 @@ +/* + * 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.io.SpillableObject; + +import java.util.concurrent.atomic.AtomicInteger; + +public final class StoreLease implements AutoCloseable { + private final Runnable _releaser; + private final T _value; + private final BlockEntry _entry; + private final AtomicInteger _shared; + private boolean _open; + + StoreLease(BlockEntry entry, Runnable releaser) { + this(null, entry, releaser, new AtomicInteger(1)); + } + + StoreLease(T value, Runnable releaser) { + this(value, null, releaser, new AtomicInteger(1)); + } + + private StoreLease(T value, BlockEntry entry, Runnable releaser, AtomicInteger shared) { + _releaser = releaser; + _value = value; + _entry = entry; + _shared = shared; + _open = true; + } + + @SuppressWarnings("unchecked") + public synchronized T value() { + if(!_open) + throw new IllegalStateException("Lease is closed"); + return _entry == null ? _value : (T) _entry.getData(); + } + + synchronized BlockEntry entry() { + if(!_open) + throw new IllegalStateException("Lease is closed"); + return _entry; + } + + public synchronized StoreLease retain() { + if(!_open) + throw new IllegalStateException("Lease is closed"); + _shared.incrementAndGet(); + return new StoreLease<>(_value, _entry, _releaser, _shared); + } + + @Override + public synchronized void close() { + if(!_open) + return; + _open = false; + if(_shared.decrementAndGet() == 0) + _releaser.run(); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/util/OOCUtils.java b/src/main/java/org/apache/sysds/runtime/ooc/util/OOCUtils.java index f33d17ea132..74e6ed30a8f 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/util/OOCUtils.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/util/OOCUtils.java @@ -21,14 +21,50 @@ import org.apache.sysds.runtime.matrix.data.MatrixIndexes; import org.apache.sysds.runtime.meta.DataCharacteristics; +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.memory.MemoryAllowance; import org.apache.sysds.runtime.util.IndexRange; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; import java.util.List; +import java.util.function.BooleanSupplier; public class OOCUtils { + public static OOCFuture pinAdmitted(OOCCache cache, long streamId, long sequenceNumber, + MemoryAllowance allowance, BooleanSupplier cancelled) { + if(cancelled.getAsBoolean()) + return OOCFuture.completed(null); + if(allowance.isShutdown()) + return OOCFuture + .failed(new IllegalStateException("Allowance was shut down while a pin admission was pending.")); + + OOCFuture admitted; + try { + admitted = cache.pinAdmitted(streamId, sequenceNumber, allowance); + } + catch(RuntimeException ex) { + return OOCFuture.failed(ex); + } + OOCFuture result = new OOCFuture<>(); + admitted.whenComplete((entry, error) -> { + if(error != null) { + result.completeExceptionally(error); + return; + } + if(entry != null && cancelled.getAsBoolean()) { + cache.unpin(entry, allowance); + result.complete(null); + return; + } + result.complete(entry); + }); + return result; + } + public static IndexRange getRangeOfTile(MatrixIndexes tileIdx, long blen) { long rs = 1 + tileIdx.getRowIndex() * blen; long re = (tileIdx.getRowIndex() + 1) * blen; diff --git a/src/test/java/org/apache/sysds/test/component/ooc/MaterializedStoreTest.java b/src/test/java/org/apache/sysds/test/component/ooc/MaterializedStoreTest.java new file mode 100644 index 00000000000..d8ea0e3986f --- /dev/null +++ b/src/test/java/org/apache/sysds/test/component/ooc/MaterializedStoreTest.java @@ -0,0 +1,359 @@ +/* + * 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.List; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import org.apache.sysds.runtime.DMLRuntimeException; +import org.apache.sysds.runtime.instructions.ooc.CachingStream; +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.SyncMemoryAllowance; +import org.apache.sysds.runtime.ooc.store.IndexedMaterializedStoreReader; +import org.apache.sysds.runtime.ooc.store.MaterializedStore; +import org.apache.sysds.runtime.ooc.store.OOCStreamMaterializer; +import org.apache.sysds.runtime.ooc.store.OrderedMaterializedStoreReader; +import org.apache.sysds.runtime.ooc.store.StoreLease; +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 MaterializedStoreTest { + 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 GlobalMemoryBroker _broker; + private SyncMemoryAllowance _producer; + private SyncMemoryAllowance _materializerAllowance; + private SyncMemoryAllowance _readerAllowance; + private OOCCacheImpl _cache; + private MaterializedStore _store; + + @Before + public void setUp() { + _broker = new GlobalMemoryBroker(1_000_000_000); + _producer = new SyncMemoryAllowance(_broker); + _materializerAllowance = new SyncMemoryAllowance(_broker); + _readerAllowance = new SyncMemoryAllowance(_broker); + _producer.setTargetMemory(MEMORY_LIMIT); + _materializerAllowance.setTargetMemory(MEMORY_LIMIT); + _readerAllowance.setTargetMemory(MEMORY_LIMIT); + _cache = new OOCCacheImpl(new OOCCacheTestUtils.RecordingOOCIOHandler(), MEMORY_LIMIT, MEMORY_LIMIT); + _store = new MaterializedStore<>(_cache, CachingStream._streamSeq.getNextID()); + } + + @After + public void tearDown() { + _store.close(); + _cache.shutdown(); + _producer.destroy(); + _materializerAllowance.destroy(); + _readerAllowance.destroy(); + } + + @Test + public void testMaterializationReadersAndForgetting() throws Exception { + OOCStreamMaterializer materializer = new OOCStreamMaterializer(_store, + indexes -> (int) indexes.getRowIndex() - 1, _materializerAllowance); + _producer.reserveBlocking(TILE_BYTES); + materializer.accept(new InMemoryQueueCallback(tile(0, 1.0), null, _producer, TILE_BYTES)); + materializer.accept(new OOCStream.SimpleQueueCallback<>(tile(1, 2.0), null)); + _producer.reserveBlocking(TILE_BYTES); + materializer.accept(new InMemoryQueueCallback(tile(2, 3.0), null, _producer, TILE_BYTES)); + materializer.accept(OOCStream.eos(null)); + materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); + + Assert.assertEquals(0, _producer.getUsedMemory()); + Assert.assertEquals(0, _materializerAllowance.getUsedMemory()); + Assert.assertEquals(3, _store.size()); + + MaterializedStore.AccessPattern pattern = sequentialPattern(3); + boolean[] indexedConsumed = new boolean[3]; + MaterializedStore.Liveness liveness = new MaterializedStore.Liveness() { + @Override + public boolean needs(int index) { + return !indexedConsumed[index]; + } + + @Override + public void consumed(int index) { + indexedConsumed[index] = true; + } + }; + + OrderedMaterializedStoreReader ordered = _store.openReader(pattern, _readerAllowance, 2); + IndexedMaterializedStoreReader indexed = _store.openIndexedReader(liveness); + _store.sealReaders(); + + int index = 0; + while(ordered.hasNext()) { + try(StoreLease lease = ordered.next()) { + Assert.assertEquals(index + 1L, lease.value().getIndexes().getRowIndex()); + index++; + } + } + ordered.close(); + Assert.assertEquals(3, index); + Assert.assertTrue(_cache.getOwnedCacheSize() > 0); + + for(index = 2; index > 0; index--) { + try(StoreLease lease = indexed.request(index, _readerAllowance).get(WAIT_SECONDS, + TimeUnit.SECONDS)) { + Assert.assertEquals(index + 1L, lease.value().getIndexes().getRowIndex()); + } + } + Assert.assertTrue(_cache.getOwnedCacheSize() > 0); + indexed.close(); + OOCCacheTestUtils.await(() -> _cache.getOwnedCacheSize() == 0, WAIT_SECONDS); + Assert.assertEquals(0, _readerAllowance.getUsedMemory()); + } + + @Test + public void testOrderedReaderRetries() throws Exception { + OOCStreamMaterializer materializer = new OOCStreamMaterializer(_store, + indexes -> (int) indexes.getRowIndex() - 1, _materializerAllowance); + for(int i = 0; i < 2; i++) { + _producer.reserveBlocking(TILE_BYTES); + materializer.accept(new InMemoryQueueCallback(tile(i, i + 1.0), null, _producer, TILE_BYTES)); + } + materializer.accept(OOCStream.eos(null)); + materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); + + _readerAllowance.destroy(); + _readerAllowance = new SyncMemoryAllowance(_broker, TILE_BYTES); + _readerAllowance.setTargetMemory(TILE_BYTES); + OrderedMaterializedStoreReader reader = _store.openReader(sequentialPattern(2), + _readerAllowance, 2, false); + _store.sealReaders(); + + Assert.assertTrue(reader.hasNext()); + try(StoreLease first = reader.next()) { + Assert.assertEquals(1L, first.value().getIndexes().getRowIndex()); + } + try(StoreLease second = reader.next()) { + Assert.assertEquals(2L, second.value().getIndexes().getRowIndex()); + } + reader.close(); + Assert.assertEquals(0, _readerAllowance.getUsedMemory()); + } + + @Test + public void testSoftOrderingReturnsReadyRequestFirst() throws Exception { + MatrixBlock largeBlock = new MatrixBlock(16, 16, 1.0); + long largeBytes = largeBlock.getExactSerializedSize(); + OOCStreamMaterializer materializer = new OOCStreamMaterializer(_store, + indexes -> (int) indexes.getRowIndex() - 1, _materializerAllowance); + _producer.reserveBlocking(largeBytes); + materializer.accept(new InMemoryQueueCallback(new IndexedMatrixValue(new MatrixIndexes(1, 1), largeBlock), null, + _producer, largeBytes)); + _producer.reserveBlocking(TILE_BYTES); + materializer.accept(new InMemoryQueueCallback(tile(1, 2.0), null, _producer, TILE_BYTES)); + materializer.accept(OOCStream.eos(null)); + materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); + + _readerAllowance.destroy(); + _readerAllowance = new SyncMemoryAllowance(_broker, largeBytes); + _readerAllowance.setTargetMemory(largeBytes); + long heldBytes = largeBytes - TILE_BYTES; + _readerAllowance.reserveBlocking(heldBytes); + OrderedMaterializedStoreReader reader = _store.openReader(sequentialPattern(2), + _readerAllowance, 2); + _store.sealReaders(); + + Assert.assertTrue(reader.hasNext()); + try(StoreLease first = reader.next()) { + Assert.assertEquals(2L, first.value().getIndexes().getRowIndex()); + } + _readerAllowance.release(heldBytes); + try(StoreLease second = reader.next()) { + Assert.assertEquals(1L, second.value().getIndexes().getRowIndex()); + } + Assert.assertFalse(reader.hasNext()); + reader.close(); + Assert.assertEquals(0, _readerAllowance.getUsedMemory()); + } + + @Test + public void testDirectRequests() throws Exception { + OOCStreamMaterializer materializer = new OOCStreamMaterializer(_store, + indexes -> (int) indexes.getRowIndex() - 1, _materializerAllowance); + _producer.reserveBlocking(TILE_BYTES); + materializer.accept(new InMemoryQueueCallback(tile(0, 1.0), null, _producer, TILE_BYTES)); + materializer.accept(OOCStream.eos(null)); + materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); + + boolean[] needed = {true}; + IndexedMaterializedStoreReader reader = _store + .openIndexedReader(new MaterializedStore.Liveness() { + @Override + public boolean needs(int index) { + return needed[index]; + } + + @Override + public void consumed(int index) { + needed[index] = false; + } + }); + _store.sealReaders(); + + try(StoreLease published = _store.requestPublished(0, _readerAllowance).get(WAIT_SECONDS, + TimeUnit.SECONDS)) { + Assert.assertEquals(1L, published.value().getIndexes().getRowIndex()); + } + StoreLease live = reader.requestIfLive(0, _readerAllowance); + Assert.assertNotNull(live); + StoreLease retained = live.retain(); + live.close(); + Assert.assertEquals(TILE_BYTES, _readerAllowance.getUsedMemory()); + Assert.assertEquals(1L, retained.value().getIndexes().getRowIndex()); + retained.close(); + Assert.assertEquals(0, _readerAllowance.getUsedMemory()); + reader.close(); + Assert.assertTrue(reader.isClosed()); + } + + @Test + public void testCompletionMissingPublications() throws Exception { + AtomicInteger failures = new AtomicInteger(); + OOCStreamMaterializer materializer = new OOCStreamMaterializer(_store, + indexes -> (int) indexes.getRowIndex() - 1, _materializerAllowance, List.of(callback -> { + if(callback.isFailure()) + failures.incrementAndGet(); + })); + for(int index : new int[] {0, 2}) { + _producer.reserveBlocking(TILE_BYTES); + materializer.accept(new InMemoryQueueCallback(tile(index, 1.0), null, _producer, TILE_BYTES)); + } + materializer.accept(OOCStream.eos(null)); + + try { + materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); + Assert.fail("Completion must reject a missing publication index"); + } + catch(ExecutionException ex) { + Assert.assertTrue(ex.getCause() instanceof IllegalStateException); + } + Assert.assertEquals(1, failures.get()); + Assert.assertEquals(0, _producer.getUsedMemory()); + } + + @Test + public void testLiveCallbackKeepsPublicationPinned() throws Exception { + AtomicReference> retained = new AtomicReference<>(); + AtomicInteger eos = new AtomicInteger(); + OOCStreamMaterializer materializer = new OOCStreamMaterializer(_store, + indexes -> (int) indexes.getRowIndex() - 1, _materializerAllowance, List.of(callback -> { + if(callback.isEos()) + eos.incrementAndGet(); + else + retained.set(callback.keepOpen()); + })); + + _producer.reserveBlocking(TILE_BYTES); + materializer.accept(new InMemoryQueueCallback(tile(0, 1.0), null, _producer, TILE_BYTES)); + Assert.assertEquals(TILE_BYTES, _producer.getUsedMemory()); + Assert.assertNotNull(retained.get()); + retained.get().close(); + retained.get().close(); + Assert.assertEquals(0, _producer.getUsedMemory()); + + materializer.accept(OOCStream.eos(null)); + materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); + Assert.assertEquals(1, eos.get()); + } + + @Test + public void testFailurePropagation() throws Exception { + DMLRuntimeException sourceFailure = new DMLRuntimeException("injected failure"); + AtomicInteger failures = new AtomicInteger(); + OOCStreamMaterializer materializer = new OOCStreamMaterializer(_store, + indexes -> (int) indexes.getRowIndex() - 1, _materializerAllowance, List.of(callback -> { + if(callback.isFailure()) + failures.incrementAndGet(); + })); + materializer.accept(OOCStream.eos(sourceFailure)); + try { + materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); + Assert.fail("The source failure must reach materializer completion"); + } + catch(ExecutionException ex) { + Assert.assertSame(sourceFailure, ex.getCause()); + } + Assert.assertEquals(1, failures.get()); + + _store.close(); + _store = new MaterializedStore<>(_cache, CachingStream._streamSeq.getNextID()); + _store.close(); + materializer = new OOCStreamMaterializer(_store, indexes -> (int) indexes.getRowIndex() - 1, + _materializerAllowance); + _producer.reserveBlocking(TILE_BYTES); + materializer.accept(new InMemoryQueueCallback(tile(0, 1.0), null, _producer, TILE_BYTES)); + try { + materializer.completion().get(WAIT_SECONDS, TimeUnit.SECONDS); + Assert.fail("Publishing into a closed store must fail"); + } + catch(ExecutionException ex) { + Assert.assertTrue(ex.getCause() instanceof DMLRuntimeException); + } + Assert.assertEquals(0, _producer.getUsedMemory()); + } + + private static IndexedMatrixValue tile(int index, double value) { + return new IndexedMatrixValue(new MatrixIndexes(index + 1L, 1), new MatrixBlock(4, 4, value)); + } + + private static MaterializedStore.AccessPattern sequentialPattern(int to) { + return new MaterializedStore.AccessPattern() { + private int _next; + + @Override + public boolean hasNext() { + return _next < to; + } + + @Override + public int next() { + return _next++; + } + + @Override + public boolean needs(int index) { + return true; + } + + @Override + public void consumed(int index) { + } + }; + } +}