diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/BlockEntry.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/BlockEntry.java index 3e040ef805e..6927ad44770 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/cache/BlockEntry.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/BlockEntry.java @@ -19,10 +19,6 @@ package org.apache.sysds.runtime.ooc.cache; -import org.apache.sysds.runtime.instructions.spark.data.IndexedMatrixValue; - -import java.util.List; - public final class BlockEntry { private final BlockKey _key; private final long _size; @@ -31,6 +27,19 @@ public final class BlockEntry { private Object _data; private int _retainHintCount; private int _referenceCount; // The number of references from different managing instances (e.g. CachingStream) + // Optional implementation-local cache metadata; null for cache implementations that do not need it. + private volatile Object _cacheMeta; + + public BlockEntry(BlockKey key) { + this._key = key; + this._size = -1; + this._pinCount = 0; + this._state = BlockState.COLD; + this._data = null; + this._retainHintCount = 0; + this._referenceCount = 0; + this._cacheMeta = null; + } public BlockEntry(BlockKey key, long size, Object data) { this._key = key; @@ -40,6 +49,18 @@ public BlockEntry(BlockKey key, long size, Object data) { this._data = data; this._retainHintCount = 0; this._referenceCount = 1; + this._cacheMeta = null; + } + + public BlockEntry(BlockKey key, long size, Object data, BlockState state) { + this._key = key; + this._size = size; + this._pinCount = 0; + this._state = state; + this._data = data; + this._retainHintCount = 0; + this._referenceCount = 1; + this._cacheMeta = null; } public BlockKey getKey() { @@ -56,18 +77,6 @@ public Object getData() { throw new IllegalStateException("Cannot get the data of an unpinned entry"); } - public int getGroupSize() { - if(_pinCount > 0) - return ((List)_data).size(); - throw new IllegalStateException("Cannot get the data of an unpinned entry"); - } - - public boolean isGrouped() { - if(_pinCount > 0) - return _data instanceof List; - throw new IllegalStateException("Cannot get the data of an unpinned entry"); - } - public Object getDataUnsafe() { return _data; } @@ -86,6 +95,10 @@ public boolean isPinned() { return _pinCount > 0; } + public synchronized int getPinCount() { + return _pinCount; + } + public synchronized int addReference() { return ++_referenceCount; } @@ -94,6 +107,18 @@ public synchronized int forget() { return --_referenceCount; } + public synchronized int getReferenceCount() { + return _referenceCount; + } + + public Object getCacheMeta() { + return _cacheMeta; + } + + public void setCacheMeta(Object meta) { + _cacheMeta = meta; + } + public synchronized void setState(BlockState state) { _state = state; } @@ -106,12 +131,6 @@ public synchronized void addRetainHint() { _retainHintCount++; } - public synchronized void removeRetainHint(int cnt) { - _retainHintCount -= cnt; - if(_retainHintCount < 0) - _retainHintCount = 0; - } - public synchronized void removeRetainHint() { if (_retainHintCount <= 0) return; @@ -129,8 +148,6 @@ public synchronized int getRetainHintCount() { public synchronized long clear() { if (_pinCount != 0 || _data == null) return 0; - if (_data instanceof IndexedMatrixValue) - ((IndexedMatrixValue)_data).setValue(null); // Explicitly clear _data = null; _retainHintCount = 0; return _size; diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCache.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCache.java new file mode 100644 index 00000000000..7f1fdb493d0 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCache.java @@ -0,0 +1,142 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.ooc.cache; + +import org.apache.sysds.runtime.ooc.memory.MemoryAllowance; + +import java.util.function.LongUnaryOperator; + +public interface OOCCache { + default OOCFuture pin(BlockKey key, MemoryAllowance allowance) { + return pin(key.getStreamId(), key.getSequenceNumber(), allowance); + } + + default OOCFuture pinAdmitted(BlockKey key, MemoryAllowance allowance) { + return pinAdmitted(key.getStreamId(), key.getSequenceNumber(), allowance); + } + + /** + * Adds a new pinned entry whose bytes are already owned by the given allowance. Ownership can later move only via + * pin/unpin. + */ + default BlockEntry putPinned(BlockKey key, Object data, long size, MemoryAllowance allowance) { + return putPinned(key.getStreamId(), key.getSequenceNumber(), data, size, allowance); + } + + /** + * Adds a new pinned entry whose bytes are already owned by the given allowance. Ownership can later move only via + * pin/unpin. + */ + BlockEntry putPinned(long sId, long tId, Object data, long size, MemoryAllowance allowance); + + /** + * Pins an item backed by an allowance. A successful pin transfers memory ownership from the cache to the owner of + * the allowance and guarantees data availability. While pinned, the bytes of the entry are not counted as + * cache-owned memory. + * + * @param sId + * @param tId + * @param allowance + * @return a non-null future of the pinned block entry; the future result is null if the required memory could not + * be reserved + */ + OOCFuture pin(long sId, long tId, MemoryAllowance allowance); + + default OOCFuture pinAdmitted(long sId, long tId, MemoryAllowance allowance) { + return pin(sId, tId, allowance); + } + + /** + * Pins an item backed by an allowance if it is already live in cache. A successful pin transfers memory ownership + * from the cache to the owner of the allowance and guarantees data availability. While pinned, the bytes of the + * entry are not counted as cache-owned memory. Implementations must reserve the required bytes from the allowance + * before making data available. + * + * @param sId + * @param tId + * @param allowance + * @return the pinned block entry if available. Null if the required memory could not be reserved or the block is + * not live + */ + BlockEntry pinIfLive(long sId, long tId, MemoryAllowance allowance); + + /** + * Unpins an item that is still backed by the given allowance. Unpinning tries to transfer memory ownership back to + * the cache. An ownership transfer may commit immediately only if this does not cause the cache to exceed its hard + * limit. Otherwise, the transfer is deferred and the allowance remains charged until the returned handle commits, + * is reclaimed, or is superseded by a later pin that transfers ownership to another allowance. Unpin can be viewed + * as an eventually resolving operation. + * + * @param entry + * @param allowance + * @return a handle describing the ownership transfer from allowance-owned memory back to cache-owned memory + */ + UnpinHandle unpin(BlockEntry entry, MemoryAllowance allowance); + + /** + * Referencing a pinned entry guarantees that its key remains in the cache until dereferenced. + * + * @param entry + * @return + */ + int reference(BlockEntry entry); + + /** + * Dereferencing allows an entry to be forgotten if no further reference is held. Dereferencing may not immediately + * cause entry removal if still pinned. + * + * @param entry + * @return + */ + int dereference(BlockEntry entry); + + /** + * Dereferencing allows an entry to be forgotten if no further reference is held. Dereferencing may not immediately + * cause entry removal if still pinned. + */ + int dereference(BlockKey key); + + void updateLimits(long hardLimit, long evictionLimit); + + /** + * Adds an eviction scoring policy for one logical cache stream. Larger scores are selected for eviction first. + * {@link Long#MAX_VALUE} remains reserved as "no policy score". + */ + void addEvictionPolicy(long streamId, LongUnaryOperator scoreFn); + + /** + * Returns the current cache-owned size in bytes. + */ + long getOwnedCacheSize(); + + void shutdown(); + + interface UnpinHandle { + BlockEntry entry(); + + MemoryAllowance allowance(); + + long bytes(); + + boolean isCommitted(); + + OOCFuture getCompletionFuture(); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheImpl.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheImpl.java new file mode 100644 index 00000000000..9b0008e84ab --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/OOCCacheImpl.java @@ -0,0 +1,745 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.ooc.cache; + +import org.apache.sysds.runtime.ooc.cache.collections.MaskedOnceArrayList; +import org.apache.sysds.runtime.ooc.cache.collections.SegmentedStreamTableList; +import org.apache.sysds.runtime.ooc.cache.eviction.EvictController; +import org.apache.sysds.runtime.ooc.cache.eviction.IndexedObjectPair; +import org.apache.sysds.runtime.ooc.cache.io.OOCIOHandler; +import org.apache.sysds.runtime.ooc.memory.MemoryAllowance; +import org.apache.sysds.utils.Statistics; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.PriorityQueue; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.Executor; +import java.util.concurrent.Executors; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.LongUnaryOperator; + +public class OOCCacheImpl implements OOCCache { + private static final int MIN_EVICTION_CANDIDATES = 1024; + private static final int MAX_EVICTION_CANDIDATES = 65536; + private static final long EVICTION_CANDIDATE_BYTE_FACTOR = 250_000; + + private final OOCIOHandler _ioHandler; + private final SegmentedStreamTableList _blocks; + private final SegmentedStreamTableList _evictControllers; + private final EvictController _defaultEvictController; + private final ConcurrentLinkedQueue _deferredUnpins; + private final Executor _collectorExecutor; + private final AtomicBoolean _evictionRunning; + + private long _hardLimit; + private long _evictionLimit; + private long _ownedBytes; + private long _evictingBytes; + private boolean _running; + + public OOCCacheImpl(OOCIOHandler ioHandler, long hardLimit, long evictionLimit) { + _ioHandler = ioHandler; + _hardLimit = hardLimit; + _evictionLimit = evictionLimit; + _ownedBytes = 0; + _evictingBytes = 0; + _running = true; + _blocks = new SegmentedStreamTableList<>(); + _evictControllers = new SegmentedStreamTableList<>(); + _defaultEvictController = new EvictController(); + _deferredUnpins = new ConcurrentLinkedQueue<>(); + _collectorExecutor = Executors.newSingleThreadExecutor(r -> { + Thread t = new Thread(r, "ooc-cache-collector"); + t.setDaemon(true); + return t; + }); + _evictionRunning = new AtomicBoolean(false); + } + + @Override + public BlockEntry putPinned(long sId, long tId, Object data, long size, MemoryAllowance allowance) { + BlockKey key = new BlockKey(sId, tId); + BlockEntry entry = new BlockEntry(key, size, data, BlockState.REMOVED); + entry.pin(); + EntryMeta meta = new EntryMeta(entry); + entry.setCacheMeta(meta); + synchronized(this) { + checkRunning(); + putEntry(entry); + } + Statistics.incrementOOCEvictionPut(); + return entry; + } + + @Override + public OOCFuture pin(long sId, long tId, MemoryAllowance allowance) { + return pinInternal(new BlockKey(sId, tId), allowance, false, false); + } + + @Override + public OOCFuture pinAdmitted(long sId, long tId, MemoryAllowance allowance) { + return pinInternal(new BlockKey(sId, tId), allowance, false, true); + } + + @Override + public BlockEntry pinIfLive(long sId, long tId, MemoryAllowance allowance) { + return pinInternal(new BlockKey(sId, tId), allowance, true, false).getNow(null); + } + + @Override + public UnpinHandle unpin(BlockEntry entry, MemoryAllowance allowance) { + if(entry.fastUnpin()) { + allowance.release(entry.getSize()); + return CacheUnpinHandle.committed(entry, allowance, entry.getSize()); + } + UnpinHandle result; + long releaseBytes; + synchronized(this) { + EntryMeta meta = getMeta(entry); + if(meta == null) + return CacheUnpinHandle.committed(entry, allowance, Math.max(0, entry.getSize())); + if(entry.getPinCount() > 1) { + entry.unpin(); + releaseBytes = entry.getSize(); + result = CacheUnpinHandle.committed(entry, allowance, releaseBytes); + } + else if(canAcceptOwnedBytes(entry.getSize())) { + releaseBytes = entry.getSize(); + result = commitLastUnpin(meta, allowance); + } + else { + CacheUnpinHandle handle = CacheUnpinHandle.deferred(entry, allowance); + meta.deferredUnpin = handle; + _deferredUnpins.offer(entry.getKey()); + return handle; + } + } + if(releaseBytes > 0) + allowance.release(releaseBytes); + return result; + } + + @Override + public synchronized int reference(BlockEntry entry) { + return entry.addReference(); + } + + @Override + public int dereference(BlockEntry entry) { + int refs; + synchronized(this) { + EntryMeta meta = getMeta(entry); + if(meta == null) + return 0; + refs = entry.forget(); + if(refs <= 0) + removeIfUnused(meta); + } + return refs; + } + + @Override + public int dereference(BlockKey key) { + BlockEntry entry = findEntry(key); + if(entry == null) + return 0; + return dereference(entry); + } + + @Override + public void updateLimits(long hardLimit, long evictionLimit) { + List completions; + synchronized(this) { + _hardLimit = hardLimit; + _evictionLimit = evictionLimit; + completions = processDeferredUnpins(); + scheduleEvictionIfNeeded(); + } + completions.forEach(this::completeDeferred); + } + + @Override + public synchronized void addEvictionPolicy(long streamId, LongUnaryOperator scoreFn) { + getOrCreateEvictController(streamId).addEvictionPolicy(scoreFn); + scheduleEvictionIfNeeded(); + } + + @Override + public synchronized long getOwnedCacheSize() { + return _ownedBytes; + } + + @Override + public synchronized void shutdown() { + _running = false; + _blocks.clear(); + _deferredUnpins.clear(); + _ownedBytes = 0; + _evictingBytes = 0; + _ioHandler.shutdown(); + } + + private OOCFuture pinInternal(BlockKey key, MemoryAllowance allowance, boolean liveOnly, + boolean waitForAdmission) { + BlockEntry deferredUnpinEntry = null; + CacheUnpinHandle deferredUnpinHandle = null; + long reserveBytes; + synchronized(this) { + checkRunning(); + BlockEntry entry = findEntry(key); + EntryMeta meta = getMeta(entry); + if(meta == null) + return OOCFuture.completed(null); + if(liveOnly && entry.getDataUnsafe() == null) + return OOCFuture.completed(null); + if(meta.deferredUnpin != null) { + if(meta.deferredUnpin.allowance == allowance) { + deferredUnpinHandle = meta.deferredUnpin; + meta.deferredUnpin = null; + deferredUnpinEntry = entry; + Statistics.incrementOOCEvictionGet(); + } + reserveBytes = entry.getSize(); + } + else if(isResidentForPin(entry)) + reserveBytes = entry.getSize(); + else if(liveOnly) + return OOCFuture.completed(null); + else + reserveBytes = entry.getSize(); + } + if(deferredUnpinEntry != null) { + deferredUnpinHandle.complete(false); + return OOCFuture.completed(deferredUnpinEntry); + } + if(!waitForAdmission) { + if(!allowance.tryReserve(reserveBytes)) + return OOCFuture.completed(null); + return pinReserved(key, allowance, reserveBytes, liveOnly); + } + + OOCFuture result = new OOCFuture<>(); + allowance.reserveAsync(reserveBytes).whenComplete((ignored, error) -> { + if(error != null) { + result.completeExceptionally(error); + return; + } + try { + pinReserved(key, allowance, reserveBytes, liveOnly).whenComplete((pinned, pinError) -> { + if(pinError != null) + result.completeExceptionally(pinError); + else + result.complete(pinned); + }); + } + catch(Throwable t) { + allowance.release(reserveBytes); + result.completeExceptionally(t); + } + }); + return result; + } + + private OOCFuture pinReserved(BlockKey key, MemoryAllowance allowance, long reservedBytes, + boolean liveOnly) { + EntryMeta meta = null; + BlockEntry deferredUnpinEntry = null; + CacheUnpinHandle deferredUnpinHandle = null; + MemoryAllowance releaseAllowance = null; + DeferredCompletion deferredCompletion = null; + BlockEntry resident = null; + long releaseBytes = 0; + boolean releaseReserved = false; + boolean returnNull = false; + synchronized(this) { + if(!_running) { + releaseReserved = true; + returnNull = true; + } + else { + BlockEntry entry = findEntry(key); + meta = getMeta(entry); + if(meta == null || (liveOnly && entry.getDataUnsafe() == null)) { + releaseReserved = true; + returnNull = true; + } + else if(meta.deferredUnpin != null) { + deferredUnpinHandle = meta.deferredUnpin; + meta.deferredUnpin = null; + deferredUnpinEntry = meta.entry; + if(deferredUnpinHandle.allowance == allowance) + releaseReserved = true; + else { + releaseAllowance = deferredUnpinHandle.allowance; + releaseBytes = meta.entry.getSize(); + } + Statistics.incrementOOCEvictionGet(); + } + else if(isResidentForPin(entry)) { + deferredCompletion = pinResident(meta); + Statistics.incrementOOCEvictionGet(); + resident = entry; + } + else if(liveOnly) { + releaseReserved = true; + returnNull = true; + } + } + } + if(releaseReserved) + allowance.release(reservedBytes); + if(releaseAllowance != null) { + releaseAllowance.release(releaseBytes); + deferredUnpinHandle.complete(false); + } + else if(deferredUnpinHandle != null) + deferredUnpinHandle.complete(false); + + completeDeferred(deferredCompletion); + if(resident != null) + return OOCFuture.completed(resident); + if(returnNull || deferredUnpinEntry != null) + return OOCFuture.completed(deferredUnpinEntry); + // Trigger read + return pinFromBackingReserved(meta, allowance, reservedBytes); + } + + private OOCFuture pinFromBackingReserved(EntryMeta meta, MemoryAllowance allowance, + long reservedBytes) { + OOCFuture readFuture; + boolean releaseReserved = false; + DeferredCompletion deferredCompletion = null; + BlockEntry resident = null; + synchronized(this) { + if(!_running || getMeta(meta.entry) != meta) { + releaseReserved = true; + readFuture = null; + } + else if(meta.entry.getDataUnsafe() != null) { + deferredCompletion = pinResident(meta); + Statistics.incrementOOCEvictionGet(); + resident = meta.entry; + readFuture = null; + } + else if(meta.readFuture == null) { + meta.entry.setState(BlockState.READING); + OOCFuture scheduled = _ioHandler.scheduleRead(meta.entry); + meta.readFuture = scheduled; + readFuture = scheduled; + scheduled.whenComplete((entry, ex) -> { + synchronized(OOCCacheImpl.this) { + if(meta.readFuture == scheduled) + meta.readFuture = null; + if(ex != null && meta.entry.getState() == BlockState.READING) + meta.entry.setState(BlockState.COLD); + } + }); + } + else + readFuture = meta.readFuture; + } + if(releaseReserved) { + allowance.release(reservedBytes); + return OOCFuture.completed(null); + } + completeDeferred(deferredCompletion); + if(resident != null) + return OOCFuture.completed(resident); + + OOCFuture result = new OOCFuture<>(); + readFuture.whenComplete((entry, ex) -> { + boolean release = false; + DeferredCompletion completion = null; + try { + if(ex != null) { + release = true; + allowance.release(reservedBytes); + result.completeExceptionally(ex); + return; + } + BlockEntry pinned; + synchronized(OOCCacheImpl.this) { + if(getMeta(meta.entry) != meta || meta.entry.getDataUnsafe() == null) { + release = true; + if(meta.entry.getState() == BlockState.READING) + meta.entry.setState(BlockState.COLD); + pinned = null; + } + else { + completion = pinResident(meta); + Statistics.incrementOOCEvictionGet(); + pinned = meta.entry; + } + } + if(release) + allowance.release(reservedBytes); + completeDeferred(completion); + result.complete(pinned); + } + catch(Throwable t) { + if(!release) + allowance.release(reservedBytes); + result.completeExceptionally(t); + } + }); + return result; + } + + private DeferredCompletion pinResident(EntryMeta meta) { + BlockEntry entry = meta.entry; + if(isCacheOwned(entry)) { + _ownedBytes -= entry.getSize(); + if(entry.getState() == BlockState.EVICTING) + _evictingBytes -= entry.getSize(); + clearLive(entry); + } + entry.setState(BlockState.REMOVED); + entry.pin(); + CacheUnpinHandle handle = meta.deferredUnpin; + if(handle == null) + return null; + long bytes = meta.entry.getSize(); + meta.deferredUnpin = null; + meta.entry.unpin(); + return new DeferredCompletion(handle, bytes, false); + } + + private UnpinHandle commitLastUnpin(EntryMeta meta, MemoryAllowance allowance) { + BlockEntry entry = meta.entry; + entry.unpin(); + if(entry.getReferenceCount() <= 0) { + removeEntry(entry.getKey()); + entry.clear(); + entry.setCacheMeta(null); + if(meta.backed) + _ioHandler.scheduleDeletion(entry); + return CacheUnpinHandle.committed(entry, allowance, entry.getSize()); + } + entry.setState(meta.backed ? BlockState.WARM : BlockState.HOT); + setLive(entry); + _ownedBytes += entry.getSize(); + scheduleEvictionIfNeeded(); + return CacheUnpinHandle.committed(entry, allowance, entry.getSize()); + } + + private List processDeferredUnpins() { + List completions = null; + while(true) { + BlockKey key = _deferredUnpins.peek(); + if(key == null) + return completions == null ? Collections.emptyList() : completions; + BlockEntry entry = findEntry(key); + EntryMeta meta = getMeta(entry); + if(meta == null || meta.deferredUnpin == null) { + _deferredUnpins.poll(); + continue; + } + if(!canAcceptOwnedBytes(meta.entry.getSize())) + return completions == null ? Collections.emptyList() : completions; + _deferredUnpins.poll(); + CacheUnpinHandle handle = meta.deferredUnpin; + meta.deferredUnpin = null; + long bytes = entry.getSize(); + entry.unpin(); + if(entry.getReferenceCount() <= 0) { + removeEntry(entry.getKey()); + entry.clear(); + entry.setCacheMeta(null); + if(meta.backed) + _ioHandler.scheduleDeletion(entry); + } + else { + entry.setState(meta.backed ? BlockState.WARM : BlockState.HOT); + setLive(entry); + _ownedBytes += entry.getSize(); + } + if(completions == null) + completions = new ArrayList<>(); + completions.add(new DeferredCompletion(handle, bytes, true)); + } + } + + private void completeDeferred(DeferredCompletion completion) { + if(completion == null) + return; + completion.handle.allowance.release(completion.bytes); + completion.handle.complete(completion.committed); + } + + private boolean canAcceptOwnedBytes(long bytes) { + return _ownedBytes + bytes <= _hardLimit; + } + + private void scheduleEvictionIfNeeded() { + if(evictionPressure() <= _evictionLimit || !_evictionRunning.compareAndSet(false, true)) + return; + _collectorExecutor.execute(this::runEviction); + } + + private void runEviction() { + try { + while(true) { + long bytes; + synchronized(this) { + bytes = evictionPressure() - _evictionLimit; + if(bytes <= 0) + return; + } + + List> candidates = collectEvictionCandidates(bytes); + if(candidates.isEmpty()) + return; + + List toWrite = new ArrayList<>(); + List completions; + boolean progress = false; + synchronized(this) { + for(IndexedObjectPair candidate : candidates) { + if(evictionPressure() <= _evictionLimit) + break; + EntryMeta meta = getMeta(candidate.obj()); + if(meta == null || candidate.obj().getPinCount() > 0 || meta.deferredUnpin != null) + continue; + BlockEntry entry = meta.entry; + if(entry.getState() == BlockState.WARM) { + entry.clear(); + entry.setState(BlockState.COLD); + clearLive(entry); + _ownedBytes -= entry.getSize(); + progress = true; + } + else if(entry.getState() == BlockState.HOT) { + entry.setState(BlockState.EVICTING); + _evictingBytes += entry.getSize(); + clearLive(entry); + toWrite.add(entry); + progress = true; + } + } + completions = processDeferredUnpins(); + } + completions.forEach(this::completeDeferred); + for(BlockEntry entry : toWrite) + _ioHandler.scheduleEviction(entry).whenComplete((ignored, ex) -> onEvicted(entry, ex)); + if(!progress) + return; + } + } + finally { + _evictionRunning.set(false); + synchronized(this) { + if(evictionPressure() > _evictionLimit) + scheduleEvictionIfNeeded(); + } + } + } + + private void onEvicted(BlockEntry entry, Throwable ex) { + List completions = null; + synchronized(this) { + EntryMeta meta = getMeta(entry); + if(meta == null) + return; + if(ex != null) { + if(entry.getState() == BlockState.EVICTING) { + entry.setState(BlockState.HOT); + _evictingBytes -= entry.getSize(); + setLive(entry); + scheduleEvictionIfNeeded(); + } + return; + } + meta.backed = true; + if(entry.getState() == BlockState.HOT) { + entry.setState(BlockState.WARM); + return; + } + if(entry.getState() != BlockState.EVICTING) + return; + entry.clear(); + entry.setState(BlockState.COLD); + _ownedBytes -= entry.getSize(); + _evictingBytes -= entry.getSize(); + removeIfUnused(meta); + completions = processDeferredUnpins(); + scheduleEvictionIfNeeded(); + } + completions.forEach(this::completeDeferred); + } + + private List> collectEvictionCandidates(long bytes) { + int k = evictionCandidateLimit(bytes); + PriorityQueue> queue = new PriorityQueue<>(); + _blocks.forEachStreamTable( + (streamId, stream) -> getEvictController(streamId).findEvictionCandidates(stream, queue, k, 0)); + + List> candidates = new ArrayList<>(queue.size()); + while(!queue.isEmpty()) + candidates.add(queue.poll()); + Collections.reverse(candidates); + return candidates; + } + + private int evictionCandidateLimit(long bytes) { + long limit = Math.max(MIN_EVICTION_CANDIDATES, + (bytes + EVICTION_CANDIDATE_BYTE_FACTOR - 1) / EVICTION_CANDIDATE_BYTE_FACTOR); + return (int) Math.min(MAX_EVICTION_CANDIDATES, limit); + } + + private EvictController getEvictController(long streamId) { + MaskedOnceArrayList controllers = _evictControllers.get(streamId); + if(controllers == null) + return _defaultEvictController; + EvictController controller = controllers.get(0); + return controller == null ? _defaultEvictController : controller; + } + + private EvictController getOrCreateEvictController(long streamId) { + MaskedOnceArrayList controllers = _evictControllers.getOrCreate(streamId); + EvictController controller = controllers.get(0); + if(controller != null) + return controller; + controller = new EvictController(); + controllers.put(0, controller); + return controller; + } + + private void removeIfUnused(EntryMeta meta) { + if(meta.entry.getReferenceCount() > 0 || meta.entry.getPinCount() > 0 || meta.deferredUnpin != null) + return; + BlockEntry entry = meta.entry; + if(isCacheOwned(entry)) + _ownedBytes -= entry.getSize(); + if(entry.getState() == BlockState.EVICTING) + _evictingBytes -= entry.getSize(); + removeEntry(entry.getKey()); + clearLive(entry); + entry.clear(); + entry.setCacheMeta(null); + if(meta.backed) + _ioHandler.scheduleDeletion(entry); + } + + private boolean isCacheOwned(BlockEntry entry) { + return entry.getState() == BlockState.HOT || entry.getState() == BlockState.WARM || + entry.getState() == BlockState.EVICTING; + } + + private boolean isResidentForPin(BlockEntry entry) { + return entry.getDataUnsafe() != null && entry.getState() != BlockState.COLD && + entry.getState() != BlockState.READING; + } + + private long evictionPressure() { + return _ownedBytes - _evictingBytes; + } + + private BlockEntry findEntry(BlockKey key) { + MaskedOnceArrayList stream = _blocks.get(key.getStreamId()); + return stream == null ? null : stream.get(blockIndex(key)); + } + + private void putEntry(BlockEntry entry) { + MaskedOnceArrayList stream = _blocks.getOrCreate(entry.getKey().getStreamId()); + int index = blockIndex(entry.getKey()); + if(stream.get(index) != null) + throw new IllegalStateException("Cache entry already exists: " + entry.getKey()); + stream.put(index, entry); + } + + private BlockEntry removeEntry(BlockKey key) { + MaskedOnceArrayList stream = _blocks.get(key.getStreamId()); + if(stream == null) + return null; + return stream.clear(blockIndex(key)) ? null : stream.get(blockIndex(key)); + } + + private void setLive(BlockEntry entry) { + MaskedOnceArrayList stream = _blocks.get(entry.getKey().getStreamId()); + if(stream != null) + stream.setLive(blockIndex(entry.getKey())); + } + + private void clearLive(BlockEntry entry) { + MaskedOnceArrayList stream = _blocks.get(entry.getKey().getStreamId()); + if(stream != null) + stream.clearLive(blockIndex(entry.getKey())); + } + + private int blockIndex(BlockKey key) { + long sequenceNumber = key.getSequenceNumber(); + if(sequenceNumber < 0 || sequenceNumber > Integer.MAX_VALUE) + throw new IndexOutOfBoundsException("Invalid block index: " + sequenceNumber); + return (int) sequenceNumber; + } + + private void checkRunning() { + if(!_running) + throw new IllegalStateException("Cache has been shut down."); + } + + private EntryMeta getMeta(BlockEntry entry) { + return entry == null ? null : (EntryMeta) entry.getCacheMeta(); + } + + private static class EntryMeta { + private final BlockEntry entry; + private boolean backed; + private OOCFuture readFuture; + private CacheUnpinHandle deferredUnpin; + + private EntryMeta(BlockEntry entry) { + this.entry = entry; + backed = entry.getState().isBackedByDisk(); + } + } + + private record DeferredCompletion(CacheUnpinHandle handle, long bytes, boolean committed) { + } + + private record CacheUnpinHandle(BlockEntry entry, MemoryAllowance allowance, long bytes, OOCFuture future) + implements UnpinHandle { + private static CacheUnpinHandle committed(BlockEntry entry, MemoryAllowance allowance, long bytes) { + return new CacheUnpinHandle(entry, allowance, bytes, OOCFuture.completed(true)); + } + + private static CacheUnpinHandle deferred(BlockEntry entry, MemoryAllowance allowance) { + return new CacheUnpinHandle(entry, allowance, entry.getSize(), new OOCFuture<>()); + } + + @Override + public boolean isCommitted() { + return future.getNow(false); + } + + @Override + public OOCFuture getCompletionFuture() { + return future; + } + + private void complete(boolean committed) { + if(future.isDone()) + return; + future.complete(committed); + } + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/collections/ConcurrentBitSet.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/collections/ConcurrentBitSet.java new file mode 100644 index 00000000000..5a1582d47d7 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/collections/ConcurrentBitSet.java @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.ooc.cache.collections; + +import java.lang.invoke.MethodHandles; +import java.lang.invoke.VarHandle; + +public class ConcurrentBitSet { + private static final VarHandle LONG_ARR = MethodHandles.arrayElementVarHandle(long[].class); + + private final long[] words; + + public ConcurrentBitSet(int bits) { + // (bits + 63) >>> 6 = ceil(bits / 64.0) + this.words = new long[(bits + 63) >>> 6]; + } + + public boolean get(int i) { + int w = i >>> 6; + long mask = 1L << (i & 63); + long word = (long) LONG_ARR.getAcquire(words, w); + return (word & mask) != 0; + } + + public boolean set(int i) { + int w = i >>> 6; + long mask = 1L << (i & 63); + + long prev = (long) LONG_ARR.getAndBitwiseOrRelease(words, w, mask); + return (prev & mask) == 0; // true if changed absent -> present + } + + public boolean clear(int i) { + int w = i >>> 6; + long mask = 1L << (i & 63); + + long prev = (long) LONG_ARR.getAndBitwiseAndRelease(words, w, ~mask); + return (prev & mask) != 0; // true if changed present -> absent + } + + public long getWord(int wordIndex) { + return (long) LONG_ARR.getAcquire(words, wordIndex); + } + + public int length() { + return words.length; + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/collections/IndexedObjectPredicate.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/collections/IndexedObjectPredicate.java new file mode 100644 index 00000000000..956a5d6498a --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/collections/IndexedObjectPredicate.java @@ -0,0 +1,24 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.ooc.cache.collections; + +public interface IndexedObjectPredicate { + boolean test(int idx, T value); +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/collections/MaskedOnceArray.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/collections/MaskedOnceArray.java new file mode 100644 index 00000000000..5d2d36f36ac --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/collections/MaskedOnceArray.java @@ -0,0 +1,165 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.ooc.cache.collections; + +import java.lang.invoke.MethodHandles; +import java.lang.invoke.VarHandle; +import java.util.concurrent.atomic.AtomicReferenceArray; +import java.util.function.Consumer; + +public class MaskedOnceArray { + private static final int RETIRED = Integer.MIN_VALUE; + private static final VarHandle NON_NULL_COUNT; + + static { + try { + NON_NULL_COUNT = MethodHandles.lookup().findVarHandle(MaskedOnceArray.class, "_nonNullCount", int.class); + } + catch(ReflectiveOperationException e) { + throw new ExceptionInInitializerError(e); + } + } + + private final AtomicReferenceArray _values; + protected final ConcurrentBitSet _liveState; + private volatile int _nonNullCount; + + public MaskedOnceArray(int length) { + _values = new AtomicReferenceArray<>(length); + _liveState = new ConcurrentBitSet(length); + _nonNullCount = 0; + } + + public boolean put(int i, T value) { + if(value == null) { + return clear(i); + } + if(!incrementNonNullCount()) + return false; + boolean changed = _values.getAndSet(i, value) == null; + if(!changed) + decrementNonNullCount(); + _liveState.set(i); + return changed; + } + + private boolean incrementNonNullCount() { + while(true) { + int count = (int) NON_NULL_COUNT.getAcquire(this); + if(count == RETIRED) + return false; + if(NON_NULL_COUNT.compareAndSet(this, count, count + 1)) + return true; + } + } + + private void decrementNonNullCount() { + while(true) { + int count = (int) NON_NULL_COUNT.getAcquire(this); + if(count <= 0) + return; + if(NON_NULL_COUNT.compareAndSet(this, count, count - 1)) + return; + } + } + + public boolean clear(int i) { + boolean changed = _values.getAndSet(i, null) != null; + if(changed) + decrementNonNullCount(); + _liveState.clear(i); + return changed; + } + + public T get(int i) { + return _values.get(i); + } + + public void forEachVisible(Consumer action) { + for(int i = 0; i < _values.length(); i++) { + T v = _values.get(i); + if(v != null) + action.accept(v); + } + } + + public boolean tryRetireIfEmpty() { + return NON_NULL_COUNT.compareAndSet(this, 0, RETIRED); + } + + public boolean isRetired() { + return (int) NON_NULL_COUNT.getAcquire(this) == RETIRED; + } + + public boolean isEmpty() { + return (int) NON_NULL_COUNT.getAcquire(this) == 0; + } + + public void setLive(int i) { + _liveState.set(i); + } + + public void clearLive(int i) { + _liveState.clear(i); + } + + public boolean forEachLive(IndexedObjectPredicate action, boolean reversed, int offset) { + if(reversed) + return forEachLiveBackward(action, offset); + else + return forEachLiveForward(action, offset); + } + + private boolean forEachLiveForward(IndexedObjectPredicate action, int offset) { + int len = _liveState.length(); + T data; + for(int word = 0; word < len; word++) { + if(_liveState.getWord(word) == 0) + continue; + int lower = word * 64; + int upper = (word + 1) * 64; + for(int i = lower; i < upper; i++) { + data = get(i); + if(data != null) + if(!action.test(offset + i, data)) + return false; + } + } + return true; + } + + private boolean forEachLiveBackward(IndexedObjectPredicate action, int offset) { + int len = _liveState.length(); + for(int word = len - 1; word >= 0; word--) { + if(_liveState.getWord(word) == 0) + continue; + int lower = word * 64; + int upper = (word + 1) * 64; + T data; + for(int i = upper - 1; i >= lower; i--) { + data = get(i); + if(data != null) + if(!action.test(offset + i, data)) + return false; + } + } + return true; + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/collections/MaskedOnceArrayList.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/collections/MaskedOnceArrayList.java new file mode 100644 index 00000000000..b4e0ed2368c --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/collections/MaskedOnceArrayList.java @@ -0,0 +1,225 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.ooc.cache.collections; + +import java.lang.invoke.MethodHandles; +import java.lang.invoke.VarHandle; +import java.util.function.Consumer; + +public class MaskedOnceArrayList { + private static final VarHandle PARTITIONS; + private static final VarHandle PARTITION = MethodHandles.arrayElementVarHandle(MaskedOnceArray[].class); + private static final int DEFAULT_PARTITION_SIZE = 1024; + + static { + try { + PARTITIONS = MethodHandles.lookup().findVarHandle(MaskedOnceArrayList.class, "_partitions", + MaskedOnceArray[].class); + } + catch(ReflectiveOperationException e) { + throw new ExceptionInInitializerError(e); + } + } + + private final int _partitionSize; + private final int _partitionBits; + private final int _partitionMask; + + @SuppressWarnings("rawtypes") + private volatile MaskedOnceArray[] _partitions; + + public MaskedOnceArrayList() { + this(DEFAULT_PARTITION_SIZE); + } + + public MaskedOnceArrayList(int partitionSize) { + validatePartitionSize(partitionSize); + _partitionSize = partitionSize; + _partitionBits = Integer.numberOfTrailingZeros(partitionSize); + _partitionMask = partitionSize - 1; + _partitions = new MaskedOnceArray[1]; + } + + @SuppressWarnings("rawtypes") + public boolean put(int i, T value) { + checkIndex(i); + if(value == null) + return clear(i); + int partitionIndex = partitionIndex(i); + int offset = offsetInPartition(i); + while(true) { + MaskedOnceArray[] partitions = ensurePartitionCapacity(partitionIndex); + MaskedOnceArray partition = partitionAt(partitions, partitionIndex); + boolean changed = partition.put(offset, value); + if(PARTITION.getAcquire(partitions, partitionIndex) == partition && !partition.isRetired()) + return changed; + } + } + + @SuppressWarnings("rawtypes") + public boolean clear(int i) { + checkIndex(i); + int partition = partitionIndex(i); + MaskedOnceArray[] partitions = (MaskedOnceArray[]) PARTITIONS.getAcquire(this); + if(partition < partitions.length) + return clear(partitions, partition, offsetInPartition(i)); + return false; + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + public T get(int i) { + checkIndex(i); + int partition = partitionIndex(i); + MaskedOnceArray[] partitions = (MaskedOnceArray[]) PARTITIONS.getAcquire(this); + if(partition >= partitions.length) + return null; + MaskedOnceArray p = (MaskedOnceArray) PARTITION.getAcquire(partitions, partition); + return p == null ? null : (T) p.get(offsetInPartition(i)); + } + + @SuppressWarnings("rawtypes") + public void setLive(int i) { + checkIndex(i); + int partitionIndex = partitionIndex(i); + MaskedOnceArray[] partitions = ensurePartitionCapacity(partitionIndex); + partitionAt(partitions, partitionIndex).setLive(offsetInPartition(i)); + } + + @SuppressWarnings("rawtypes") + public void clearLive(int i) { + checkIndex(i); + int partition = partitionIndex(i); + MaskedOnceArray[] partitions = (MaskedOnceArray[]) PARTITIONS.getAcquire(this); + if(partition < partitions.length) { + MaskedOnceArray p = (MaskedOnceArray) PARTITION.getAcquire(partitions, partition); + if(p != null) + p.clearLive(offsetInPartition(i)); + } + } + + @SuppressWarnings("rawtypes") + public int capacity() { + MaskedOnceArray[] partitions = (MaskedOnceArray[]) PARTITIONS.getAcquire(this); + return partitions.length * _partitionSize; + } + + @SuppressWarnings({"rawtypes", "unchecked"}) + public void forEachLive(IndexedObjectPredicate action, boolean reversed) { + MaskedOnceArray[] partitions = (MaskedOnceArray[]) PARTITIONS.getAcquire(this); + if(reversed) { + for(int i = partitions.length - 1; i >= 0; i--) { + MaskedOnceArray partition = (MaskedOnceArray) PARTITION.getAcquire(partitions, i); + if(partition != null) + partition.forEachLive(action, true, i * _partitionSize); + } + } + else { + for(int i = 0; i < partitions.length; i++) { + MaskedOnceArray partition = (MaskedOnceArray) PARTITION.getAcquire(partitions, i); + if(partition != null) + partition.forEachLive(action, false, i * _partitionSize); + } + } + } + + @SuppressWarnings({"rawtypes", "unchecked"}) + public void forEachVisible(Consumer action) { + MaskedOnceArray[] partitions = (MaskedOnceArray[]) PARTITIONS.getAcquire(this); + for(int i = 0; i < partitions.length; i++) { + MaskedOnceArray partition = (MaskedOnceArray) PARTITION.getAcquire(partitions, i); + if(partition != null) + partition.forEachVisible(action); + } + } + + @SuppressWarnings("rawtypes") + private MaskedOnceArray[] ensurePartitionCapacity(int partitionIndex) { + MaskedOnceArray[] partitions = (MaskedOnceArray[]) PARTITIONS.getAcquire(this); + while(partitionIndex >= partitions.length) { + MaskedOnceArray[] bigger = growPartitions(partitions, partitionIndex + 1); + if(PARTITIONS.compareAndSet(this, partitions, bigger)) + partitions = bigger; + else + partitions = (MaskedOnceArray[]) PARTITIONS.getAcquire(this); + } + return partitions; + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + private MaskedOnceArray partitionAt(MaskedOnceArray[] partitions, int partitionIndex) { + MaskedOnceArray partition; + while((partition = (MaskedOnceArray) PARTITION.getAcquire(partitions, partitionIndex)) == null || + partition.isRetired()) { + if(partition != null) { + PARTITION.compareAndSet(partitions, partitionIndex, partition, null); + continue; + } + MaskedOnceArray newPartition = new MaskedOnceArray<>(_partitionSize); + if(PARTITION.compareAndSet(partitions, partitionIndex, null, newPartition)) + return newPartition; + } + return partition; + } + + @SuppressWarnings("rawtypes") + private boolean clear(MaskedOnceArray[] partitions, int partitionIndex, int offset) { + MaskedOnceArray partition = (MaskedOnceArray) PARTITION.getAcquire(partitions, partitionIndex); + if(partition == null) + return false; + boolean changed = partition.clear(offset); + if(partition.tryRetireIfEmpty()) + PARTITION.compareAndSet(partitions, partitionIndex, partition, null); + return changed; + } + + @SuppressWarnings("rawtypes") + private MaskedOnceArray[] growPartitions(MaskedOnceArray[] partitions, int minLength) { + int newLength = partitions.length; + while(newLength < minLength) { + if(newLength > Integer.MAX_VALUE / 2) + throw new IllegalStateException("MaskedOnceArrayList capacity overflow"); + newLength <<= 1; + } + + MaskedOnceArray[] bigger = new MaskedOnceArray[newLength]; + System.arraycopy(partitions, 0, bigger, 0, partitions.length); + return bigger; + } + + private int partitionIndex(int index) { + return index >>> _partitionBits; + } + + private int offsetInPartition(int index) { + return index & _partitionMask; + } + + private static void validatePartitionSize(int partitionSize) { + if(partitionSize < 64 || (partitionSize & (partitionSize - 1)) != 0) { + throw new IllegalArgumentException( + "partitionSize must be a power of two and at least 64: " + partitionSize); + } + } + + private static void checkIndex(int i) { + if(i < 0) + throw new IndexOutOfBoundsException("Negative index: " + i); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/collections/SegmentedStreamTableList.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/collections/SegmentedStreamTableList.java new file mode 100644 index 00000000000..5dcf05bcd68 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/collections/SegmentedStreamTableList.java @@ -0,0 +1,206 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.ooc.cache.collections; + +import java.lang.invoke.MethodHandles; +import java.lang.invoke.VarHandle; +import java.util.function.BiConsumer; +import java.util.function.Consumer; + +public class SegmentedStreamTableList { + private static final VarHandle SEGMENTS; + private static final VarHandle ARRAY = MethodHandles.arrayElementVarHandle(Object[].class); + private static final int DEFAULT_SEGMENT_SIZE = 64; + + static { + try { + SEGMENTS = MethodHandles.lookup().findVarHandle(SegmentedStreamTableList.class, "_segments", + Object[].class); + } + catch(ReflectiveOperationException e) { + throw new ExceptionInInitializerError(e); + } + } + + private final int _segmentSize; + private final int _segmentBits; + private final int _segmentMask; + private final int _streamPartitionSize; + + private volatile Object[] _segments; + + public SegmentedStreamTableList() { + this(DEFAULT_SEGMENT_SIZE); + } + + public SegmentedStreamTableList(int segmentSize) { + this(segmentSize, 1024); + } + + public SegmentedStreamTableList(int segmentSize, int streamPartitionSize) { + validatePowerOfTwo(segmentSize, "segmentSize"); + _segmentSize = segmentSize; + _segmentBits = Integer.numberOfTrailingZeros(segmentSize); + _segmentMask = segmentSize - 1; + _streamPartitionSize = streamPartitionSize; + _segments = new Object[1]; + } + + public MaskedOnceArrayList get(int streamId) { + checkStreamId(streamId); + Object[] segments = (Object[]) SEGMENTS.getAcquire(this); + int segmentIndex = segmentIndex(streamId); + if(segmentIndex >= segments.length) + return null; + + Object[] segment = (Object[]) ARRAY.getAcquire(segments, segmentIndex); + if(segment == null) + return null; + + @SuppressWarnings("unchecked") + MaskedOnceArrayList streamTable = (MaskedOnceArrayList) ARRAY.getAcquire(segment, + offsetInSegment(streamId)); + return streamTable; + } + + public MaskedOnceArrayList get(long streamId) { + return get(asIntStreamId(streamId)); + } + + public MaskedOnceArrayList getOrCreate(int streamId) { + checkStreamId(streamId); + int segmentIndex = segmentIndex(streamId); + int offset = offsetInSegment(streamId); + + while(true) { + Object[] segments = ensureOuterCapacity(segmentIndex + 1); + Object[] segment = (Object[]) ARRAY.getAcquire(segments, segmentIndex); + if(segment == null) { + Object[] newSegment = new Object[_segmentSize]; + if(!ARRAY.compareAndSet(segments, segmentIndex, null, newSegment)) + continue; + segment = newSegment; + } + + @SuppressWarnings("unchecked") + MaskedOnceArrayList streamTable = (MaskedOnceArrayList) ARRAY.getAcquire(segment, offset); + if(streamTable != null) + return streamTable; + + MaskedOnceArrayList newTable = new MaskedOnceArrayList<>(_streamPartitionSize); + if(ARRAY.compareAndSet(segment, offset, null, newTable)) + return newTable; + } + } + + public MaskedOnceArrayList getOrCreate(long streamId) { + return getOrCreate(asIntStreamId(streamId)); + } + + public int capacity() { + Object[] segments = (Object[]) SEGMENTS.getAcquire(this); + return segments.length * _segmentSize; + } + + public void forEachLive(IndexedObjectPredicate action) { + forEachStreamTable(table -> table.forEachLive(action, false)); + } + + public void forEachVisible(Consumer action) { + forEachStreamTable(table -> table.forEachVisible(action)); + } + + public void forEachStreamTable(BiConsumer> action) { + Object[] segments = (Object[]) SEGMENTS.getAcquire(this); + for(int i = 0; i < segments.length; i++) { + Object[] segment = (Object[]) ARRAY.getAcquire(segments, i); + if(segment == null) + continue; + for(int j = 0; j < segment.length; j++) { + @SuppressWarnings("unchecked") + MaskedOnceArrayList table = (MaskedOnceArrayList) ARRAY.getAcquire(segment, j); + if(table != null) + action.accept((i << _segmentBits) | j, table); + } + } + } + + public void clear() { + SEGMENTS.setRelease(this, new Object[1]); + } + + private void forEachStreamTable(Consumer> action) { + Object[] segments = (Object[]) SEGMENTS.getAcquire(this); + for(int i = 0; i < segments.length; i++) { + Object[] segment = (Object[]) ARRAY.getAcquire(segments, i); + if(segment == null) + continue; + for(int j = 0; j < segment.length; j++) { + @SuppressWarnings("unchecked") + MaskedOnceArrayList table = (MaskedOnceArrayList) ARRAY.getAcquire(segment, j); + if(table != null) + action.accept(table); + } + } + } + + private Object[] ensureOuterCapacity(int minLength) { + Object[] segments = (Object[]) SEGMENTS.getAcquire(this); + while(minLength > segments.length) { + int newLength = segments.length; + while(newLength < minLength) { + if(newLength > Integer.MAX_VALUE / 2) + throw new IllegalStateException("SegmentedStreamTableList capacity overflow"); + newLength <<= 1; + } + + Object[] bigger = new Object[newLength]; + System.arraycopy(segments, 0, bigger, 0, segments.length); + if(SEGMENTS.compareAndSet(this, segments, bigger)) + return bigger; + segments = (Object[]) SEGMENTS.getAcquire(this); + } + return segments; + } + + private int segmentIndex(int streamId) { + return streamId >>> _segmentBits; + } + + private int offsetInSegment(int streamId) { + return streamId & _segmentMask; + } + + private static int asIntStreamId(long streamId) { + if(streamId < 0 || streamId > Integer.MAX_VALUE) + throw new IndexOutOfBoundsException("Invalid streamId: " + streamId); + return (int) streamId; + } + + private static void checkStreamId(int streamId) { + if(streamId < 0) + throw new IndexOutOfBoundsException("Invalid streamId: " + streamId); + } + + private static void validatePowerOfTwo(int value, String name) { + if(value <= 0 || (value & (value - 1)) != 0) + throw new IllegalArgumentException(name + " must be a power of two: " + value); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/eviction/EvictController.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/eviction/EvictController.java new file mode 100644 index 00000000000..64fefeb1160 --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/eviction/EvictController.java @@ -0,0 +1,86 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.ooc.cache.eviction; + +import org.apache.sysds.runtime.ooc.cache.BlockEntry; +import org.apache.sysds.runtime.ooc.cache.BlockState; +import org.apache.sysds.runtime.ooc.cache.collections.MaskedOnceArrayList; + +import java.util.PriorityQueue; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.function.LongUnaryOperator; + +public class EvictController { + private final CopyOnWriteArrayList _op = new CopyOnWriteArrayList<>(); + + public void addEvictionPolicy(LongUnaryOperator op) { + if(op == null) + throw new IllegalArgumentException("Eviction policy must not be null."); + _op.add(op); + } + + public void findEvictionCandidates(MaskedOnceArrayList list, + PriorityQueue> candidates, int k, long estimatedReuseTimestamp) { + if(_op.isEmpty()) { + list.forEachLive((idx, b) -> { + if(!isEvictionCandidate(b)) + return true; + var iop = new IndexedObjectPair<>(estimatedReuseTimestamp + idx, b); + if(candidates.size() < k) { + candidates.offer(iop); + } + else if(iop.compareTo(candidates.peek()) > 0) { + candidates.poll(); + candidates.offer(iop); + } + return true; + }, true); + return; + } + list.forEachLive((idx, b) -> { + if(!isEvictionCandidate(b)) + return true; + long score = computeScore(idx); + if(score == Long.MAX_VALUE) + score = idx + estimatedReuseTimestamp; + var iop = new IndexedObjectPair<>(score, b); + if(candidates.size() < k) { + candidates.offer(iop); + } + else if(iop.compareTo(candidates.peek()) > 0) { + candidates.poll(); + candidates.offer(iop); + } + return true; + }, true); + } + + private boolean isEvictionCandidate(BlockEntry entry) { + BlockState state = entry.getState(); + return state == BlockState.HOT || state == BlockState.WARM; + } + + private long computeScore(int idx) { + long out = Long.MAX_VALUE; + for(LongUnaryOperator uop : _op) + out = Math.min(out, uop.applyAsLong(idx)); + return out; + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/cache/eviction/IndexedObjectPair.java b/src/main/java/org/apache/sysds/runtime/ooc/cache/eviction/IndexedObjectPair.java new file mode 100644 index 00000000000..e04c2953ddf --- /dev/null +++ b/src/main/java/org/apache/sysds/runtime/ooc/cache/eviction/IndexedObjectPair.java @@ -0,0 +1,27 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.sysds.runtime.ooc.cache.eviction; + +public record IndexedObjectPair(long idx, T obj) implements Comparable> { + @Override + public int compareTo(IndexedObjectPair indexedObjectPair) { + return Long.compare(idx, indexedObjectPair.idx); + } +} diff --git a/src/main/java/org/apache/sysds/runtime/ooc/memory/MemoryAllowance.java b/src/main/java/org/apache/sysds/runtime/ooc/memory/MemoryAllowance.java index 64518ded4a3..b2db5ea7533 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/memory/MemoryAllowance.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/memory/MemoryAllowance.java @@ -19,15 +19,27 @@ package org.apache.sysds.runtime.ooc.memory; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; + public interface MemoryAllowance { boolean tryReserve(long bytes); + void reserveBlocking(long bytes); + + OOCFuture reserveAsync(long bytes); + void release(long bytes); + long getUsedMemory(); + long getGrantedMemory(); + long getTargetMemory(); + void setTargetMemory(long targetMemory); + void shutdown(); + boolean isShutdown(); default void destroy() { diff --git a/src/main/java/org/apache/sysds/runtime/ooc/memory/SyncMemoryAllowance.java b/src/main/java/org/apache/sysds/runtime/ooc/memory/SyncMemoryAllowance.java index 85d2cbfcd2b..2c4a1b7a0fa 100644 --- a/src/main/java/org/apache/sysds/runtime/ooc/memory/SyncMemoryAllowance.java +++ b/src/main/java/org/apache/sysds/runtime/ooc/memory/SyncMemoryAllowance.java @@ -20,22 +20,51 @@ package org.apache.sysds.runtime.ooc.memory; import org.apache.sysds.runtime.DMLRuntimeException; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; + +import java.util.ArrayDeque; +import java.util.concurrent.ExecutionException; public class SyncMemoryAllowance implements MemoryAllowance { + private static final long RELEASE_TRIM_BUFFER_BYTES = 20_000_000L; + protected final MemoryBroker _broker; + protected final long _consumptionLimit; + protected final long _minimumOperatingBytes; protected volatile long _usedBytes; protected volatile long _grantedBytes; protected volatile long _targetBytes; protected volatile boolean _shutdown; protected volatile boolean _destroyed; + private final ArrayDeque _reservationWaiters; + private boolean _drainingReservationWaiters; + private boolean _reservationDrainRequested; public SyncMemoryAllowance(MemoryBroker broker) { + this(broker, Long.MAX_VALUE); + } + + public SyncMemoryAllowance(MemoryBroker broker, long consumptionLimit) { + this(broker, consumptionLimit, 0); + } + + public SyncMemoryAllowance(MemoryBroker broker, long consumptionLimit, long minimumOperatingBytes) { + if(consumptionLimit < 0) + throw new IllegalArgumentException("Consumption limit must not be negative: " + consumptionLimit); + if(minimumOperatingBytes < 0) + throw new IllegalArgumentException( + "Minimum operating memory must not be negative: " + minimumOperatingBytes); _broker = broker; + _consumptionLimit = consumptionLimit; + _minimumOperatingBytes = Math.min(minimumOperatingBytes, consumptionLimit); _usedBytes = 0; _grantedBytes = 0; _targetBytes = 0; _shutdown = false; _destroyed = false; + _reservationWaiters = new ArrayDeque<>(); + _drainingReservationWaiters = false; + _reservationDrainRequested = false; broker.attachAllowance(this); } @@ -46,19 +75,23 @@ public boolean tryReserve(long bytes) { synchronized(this) { if(_shutdown || _destroyed) return false; - if(_usedBytes + bytes > _targetBytes) - return false; if(_usedBytes + bytes <= _grantedBytes) { _usedBytes += bytes; return true; } + if(_usedBytes + bytes > _targetBytes) + return false; minRequest = _usedBytes + bytes - _grantedBytes; maxRequest = Math.max(minRequest, Math.max(_grantedBytes, bytes) * 2); } + if(bytes > _consumptionLimit) + throw new IllegalArgumentException("Cannot reserve more memory than the consumption limit"); + long granted = _broker.requestMemory(this, minRequest, maxRequest); long refund = 0; boolean success = false; + boolean drainWaiters = false; synchronized(this) { if(_shutdown || _destroyed) refund = granted; @@ -68,36 +101,52 @@ public boolean tryReserve(long bytes) { _usedBytes += bytes; success = true; } + drainWaiters = success && !_reservationWaiters.isEmpty(); notifyAll(); } } if(refund > 0) _broker.freeMemory(this, refund); + if(drainWaiters) + requestReservationDrain(); return success; } @Override public void reserveBlocking(long bytes) { - if(_shutdown || _destroyed) - throw new IllegalStateException("Cannot reserve memory on closed allowance."); - while(true) { - if(tryReserve(bytes)) { - synchronized(this) { - notifyAll(); - } - return; - } - synchronized(this) { - if(_shutdown || _destroyed) - throw new IllegalStateException("Cannot reserve memory on closed allowance."); - try { - wait(); - } - catch(InterruptedException e) { - throw new DMLRuntimeException(e); - } + try { + reserveAsync(bytes).get(); + } + catch(InterruptedException e) { + Thread.currentThread().interrupt(); + throw new DMLRuntimeException(e); + } + catch(ExecutionException e) { + throw DMLRuntimeException.of(e.getCause()); + } + } + + @Override + public OOCFuture reserveAsync(long bytes) { + if(bytes < 0) + throw new IllegalArgumentException("Cannot reserve negative bytes: " + bytes); + if(bytes == 0) + return OOCFuture.completed(null); + if(bytes > _consumptionLimit) + return OOCFuture + .failed(new IllegalArgumentException("Cannot reserve more memory than the consumption limit")); + if(tryReserve(bytes)) + return OOCFuture.completed(null); + OOCFuture future = new OOCFuture<>(); + synchronized(this) { + if(_shutdown || _destroyed) { + future.completeExceptionally(new IllegalStateException("Cannot reserve memory on closed allowance.")); + return future; } + _reservationWaiters.addLast(new ReservationWaiter(bytes, future)); } + requestReservationDrain(); + return future; } @Override @@ -105,11 +154,23 @@ public void release(long bytes) { long freedMemory = 0; long destroyFreedMemory = 0; boolean destroy = false; + boolean drainWaiters; synchronized(this) { + if(bytes < 0) + throw new IllegalArgumentException("Cannot release negative bytes: " + bytes); + if(_usedBytes < bytes) { + throw new IllegalArgumentException("Memory allowance underflow in " + getClass().getSimpleName() + + ": release=" + bytes + ", used=" + _usedBytes + ", granted=" + _grantedBytes + ", target=" + + _targetBytes + ", shutdown=" + _shutdown + ", destroyed=" + _destroyed); + } _usedBytes -= bytes; if(_shutdown) { long oldGrantedBytes = _grantedBytes; _grantedBytes = _usedBytes; + if(_grantedBytes < 0) { + throw new IllegalArgumentException("Granted memory underflow in " + getClass().getSimpleName() + + ": granted=" + _grantedBytes + ", used=" + _usedBytes + ", released=" + bytes); + } if(_usedBytes == 0) { _destroyed = true; destroy = true; @@ -124,12 +185,20 @@ else if(_grantedBytes > _targetBytes) { _grantedBytes = Math.max(_usedBytes, _targetBytes); freedMemory = oldGrantedBytes - _grantedBytes; } + else if(_usedBytes * 3 < _grantedBytes * 2) { + long oldGrantedBytes = _grantedBytes; + _grantedBytes = Math.max(_usedBytes, Math.min(_grantedBytes, _usedBytes + RELEASE_TRIM_BUFFER_BYTES)); + freedMemory = oldGrantedBytes - _grantedBytes; + } + drainWaiters = !_reservationWaiters.isEmpty() && !_shutdown && !_destroyed; notifyAll(); } if(destroy) _broker.destroyAllowance(this, destroyFreedMemory); else if(freedMemory > 0) _broker.freeMemory(this, freedMemory); + if(drainWaiters) + requestReservationDrain(); } @Override @@ -148,11 +217,25 @@ public long getTargetMemory() { } @Override - public synchronized void setTargetMemory(long targetMemory) { - if(_shutdown || _destroyed) - return; - _targetBytes = targetMemory; - notifyAll(); + public void setTargetMemory(long targetMemory) { + long freedMemory = 0; + boolean drainWaiters = false; + synchronized(this) { + if(_shutdown || _destroyed) + return; + _targetBytes = Math.min(Math.max(targetMemory, _minimumOperatingBytes), _consumptionLimit); + if(_grantedBytes > _targetBytes) { + long oldGrantedBytes = _grantedBytes; + _grantedBytes = Math.max(_usedBytes, _targetBytes); + freedMemory = oldGrantedBytes - _grantedBytes; + } + drainWaiters = !_reservationWaiters.isEmpty(); + notifyAll(); + } + if(freedMemory > 0) + _broker.freeMemory(this, freedMemory); + if(drainWaiters) + requestReservationDrain(); } @Override @@ -160,6 +243,7 @@ public void shutdown() { long freedMemory = 0; long destroyFreedMemory = 0; boolean destroy = false; + ArrayDeque waiters; synchronized(this) { if(_shutdown || _destroyed) return; @@ -175,6 +259,8 @@ public void shutdown() { else { freedMemory = oldGrantedBytes - _grantedBytes; } + waiters = new ArrayDeque<>(_reservationWaiters); + _reservationWaiters.clear(); notifyAll(); } _broker.shutdownAllowance(this); @@ -182,10 +268,81 @@ public void shutdown() { _broker.destroyAllowance(this, destroyFreedMemory); else if(freedMemory > 0) _broker.freeMemory(this, freedMemory); + IllegalStateException ex = new IllegalStateException("Cannot reserve memory on closed allowance."); + while(!waiters.isEmpty()) + waiters.removeFirst().future.completeExceptionally(ex); } @Override public boolean isShutdown() { return _shutdown || _destroyed; } + + private void requestReservationDrain() { + synchronized(this) { + _reservationDrainRequested = true; + if(_drainingReservationWaiters) + return; + _drainingReservationWaiters = true; + } + try { + while(true) { + synchronized(this) { + _reservationDrainRequested = false; + } + drainReservationWaitersOnce(); + synchronized(this) { + if(!_reservationDrainRequested) { + _drainingReservationWaiters = false; + return; + } + } + } + } + catch(RuntimeException | Error t) { + synchronized(this) { + _drainingReservationWaiters = false; + } + throw t; + } + } + + private void drainReservationWaitersOnce() { + while(true) { + ReservationWaiter waiter; + synchronized(this) { + if(_shutdown || _destroyed) + return; + waiter = _reservationWaiters.peekFirst(); + if(waiter == null) + return; + } + boolean admitted; + try { + admitted = tryReserve(waiter.bytes); + } + catch(Throwable t) { + removeReservationWaiter(waiter); + waiter.future.completeExceptionally(t); + continue; + } + if(!admitted) + return; + if(removeReservationWaiter(waiter)) + waiter.future.complete(null); + else + release(waiter.bytes); + } + } + + private synchronized boolean removeReservationWaiter(ReservationWaiter waiter) { + if(_reservationWaiters.peekFirst() == waiter) { + _reservationWaiters.removeFirst(); + return true; + } + return _reservationWaiters.remove(waiter); + } + + private record ReservationWaiter(long bytes, OOCFuture future) { + } } diff --git a/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCCacheImplTest.java b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCCacheImplTest.java new file mode 100644 index 00000000000..283cebe4667 --- /dev/null +++ b/src/test/java/org/apache/sysds/test/component/ooc/cache/OOCCacheImplTest.java @@ -0,0 +1,334 @@ +/* + * 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.cache; + +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BooleanSupplier; + +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.OOCCacheImpl; +import org.apache.sysds.runtime.ooc.cache.OOCFuture; +import org.apache.sysds.runtime.ooc.cache.io.OOCIOHandler; +import org.apache.sysds.runtime.ooc.memory.GlobalMemoryBroker; +import org.apache.sysds.runtime.ooc.memory.SyncMemoryAllowance; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +public class OOCCacheImplTest { + private static final long STREAM_ID = 7; + private static final long BLOCK_ID = 3; + private static final long BYTES = 1_000; + private static final long WAIT_TIMEOUT_SEC = 10; + + private RecordingIOHandler _io; + private GlobalMemoryBroker _broker; + private SyncMemoryAllowance _producer; + private SyncMemoryAllowance _reader; + private OOCCacheImpl _cache; + + @Before + public void setUp() { + _io = new RecordingIOHandler(); + _broker = new GlobalMemoryBroker(8 * BYTES); + _producer = new SyncMemoryAllowance(_broker, 4 * BYTES); + _reader = new SyncMemoryAllowance(_broker, 4 * BYTES); + _cache = new OOCCacheImpl(_io, 4 * BYTES, 4 * BYTES); + } + + @After + public void tearDown() { + if(_cache != null) + _cache.shutdown(); + if(_producer != null) + _producer.destroy(); + if(_reader != null) + _reader.destroy(); + } + + @Test + public void testPinMissingEntryReturnsNullWithoutReservation() throws Exception { + BlockEntry pinned = _cache.pin(new BlockKey(STREAM_ID, BLOCK_ID), _reader).get(WAIT_TIMEOUT_SEC, + TimeUnit.SECONDS); + + Assert.assertNull(pinned); + Assert.assertNull(_cache.pinIfLive(STREAM_ID, BLOCK_ID, _reader)); + Assert.assertEquals(0, _reader.getUsedMemory()); + Assert.assertEquals(0, _io.readCount()); + } + + @Test + public void testResidentPinTransfersOwnershipBetweenCacheAndAllowance() throws Exception { + BlockKey key = new BlockKey(STREAM_ID, BLOCK_ID); + String payload = "resident"; + + _producer.reserveBlocking(BYTES); + BlockEntry entry = _cache.putPinned(key, payload, BYTES, _producer); + Assert.assertEquals(0, _cache.getOwnedCacheSize()); + Assert.assertEquals(BYTES, _producer.getUsedMemory()); + + await(_cache.unpin(entry, _producer)); + Assert.assertEquals(BYTES, _cache.getOwnedCacheSize()); + Assert.assertEquals(0, _producer.getUsedMemory()); + + BlockEntry pinned = _cache.pin(key, _reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + Assert.assertSame(entry, pinned); + Assert.assertEquals(payload, pinned.getData()); + Assert.assertEquals(0, _cache.getOwnedCacheSize()); + Assert.assertEquals(BYTES, _reader.getUsedMemory()); + Assert.assertEquals(0, _io.readCount()); + + await(_cache.unpin(pinned, _reader)); + Assert.assertEquals(BYTES, _cache.getOwnedCacheSize()); + Assert.assertEquals(0, _reader.getUsedMemory()); + } + + @Test + public void testPinReloadsColdBackedEntry() throws Exception { + useEvictingCache(); + BlockKey key = new BlockKey(STREAM_ID, BLOCK_ID); + String payload = "payload"; + + _producer.reserveBlocking(BYTES); + BlockEntry entry = _cache.putPinned(key, payload, BYTES, _producer); + await(_cache.unpin(entry, _producer)); + waitFor(() -> _io.evictionCount() == 1 && BlockEntryTestAccess.getDataUnsafe(entry) == null); + Assert.assertEquals(0, _producer.getUsedMemory()); + + BlockEntry pinned = _cache.pin(key, _reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + + Assert.assertNotNull(pinned); + Assert.assertSame(entry, pinned); + Assert.assertEquals(payload, pinned.getData()); + Assert.assertEquals(1, _io.readCount()); + Assert.assertEquals(BYTES, _reader.getUsedMemory()); + + await(_cache.unpin(pinned, _reader)); + Assert.assertEquals(0, _reader.getUsedMemory()); + } + + @Test + public void testPinIfLiveDoesNotReadColdBackedEntry() throws Exception { + useEvictingCache(); + BlockKey key = new BlockKey(STREAM_ID, BLOCK_ID); + String payload = "cold"; + + _producer.reserveBlocking(BYTES); + BlockEntry entry = _cache.putPinned(key, payload, BYTES, _producer); + await(_cache.unpin(entry, _producer)); + waitFor(() -> _io.evictionCount() == 1 && BlockEntryTestAccess.getDataUnsafe(entry) == null); + + BlockEntry pinned = _cache.pinIfLive(STREAM_ID, BLOCK_ID, _reader); + + Assert.assertNull(pinned); + Assert.assertEquals(0, _io.readCount()); + Assert.assertEquals(0, _reader.getUsedMemory()); + } + + @Test + public void testDeferredUnpinCommitsWhenLimitsGrow() throws Exception { + useZeroHardLimitCache(); + BlockKey key = new BlockKey(STREAM_ID, BLOCK_ID); + + _producer.reserveBlocking(BYTES); + BlockEntry entry = _cache.putPinned(key, "deferred", BYTES, _producer); + OOCCache.UnpinHandle deferred = _cache.unpin(entry, _producer); + + Assert.assertFalse(deferred.isCommitted()); + Assert.assertFalse(deferred.getCompletionFuture().isDone()); + Assert.assertEquals(BYTES, _producer.getUsedMemory()); + Assert.assertEquals(0, _cache.getOwnedCacheSize()); + + _cache.updateLimits(BYTES, BYTES); + deferred.getCompletionFuture().get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + + Assert.assertTrue(deferred.isCommitted()); + Assert.assertEquals(0, _producer.getUsedMemory()); + Assert.assertEquals(BYTES, _cache.getOwnedCacheSize()); + } + + @Test + public void testDeferredUnpinCanBeAdoptedBySameAllowance() throws Exception { + useZeroHardLimitCache(); + BlockKey key = new BlockKey(STREAM_ID, BLOCK_ID); + + _producer.reserveBlocking(BYTES); + BlockEntry entry = _cache.putPinned(key, "adopt", BYTES, _producer); + OOCCache.UnpinHandle deferred = _cache.unpin(entry, _producer); + + BlockEntry repinned = _cache.pin(key, _producer).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + + Assert.assertSame(entry, repinned); + Assert.assertTrue(deferred.getCompletionFuture().isDone()); + Assert.assertFalse(deferred.isCommitted()); + Assert.assertEquals(BYTES, _producer.getUsedMemory()); + Assert.assertEquals(0, _cache.getOwnedCacheSize()); + + OOCCache.UnpinHandle cleanup = _cache.unpin(repinned, _producer); + _cache.updateLimits(BYTES, BYTES); + cleanup.getCompletionFuture().get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + Assert.assertEquals(0, _producer.getUsedMemory()); + } + + @Test + public void testDereferenceRemovesEntryAfterLastUnpin() throws Exception { + BlockKey key = new BlockKey(STREAM_ID, BLOCK_ID); + + _producer.reserveBlocking(BYTES); + BlockEntry entry = _cache.putPinned(key, "drop", BYTES, _producer); + + Assert.assertEquals(0, _cache.dereference(entry)); + await(_cache.unpin(entry, _producer)); + + Assert.assertEquals(0, _producer.getUsedMemory()); + Assert.assertNull(_cache.pin(key, _reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS)); + Assert.assertEquals(0, _reader.getUsedMemory()); + } + + @Test + public void testBackingReadFailureReleasesReservedBytes() throws Exception { + useEvictingCache(); + BlockKey key = new BlockKey(STREAM_ID, BLOCK_ID); + + _producer.reserveBlocking(BYTES); + BlockEntry entry = _cache.putPinned(key, "fail-read", BYTES, _producer); + await(_cache.unpin(entry, _producer)); + waitFor(() -> _io.evictionCount() == 1 && BlockEntryTestAccess.getDataUnsafe(entry) == null); + + _io.failReads(true); + try { + _cache.pin(key, _reader).get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + Assert.fail("A failed backing read must fail the pin future."); + } + catch(ExecutionException expected) { + // expected + } + + Assert.assertEquals(1, _io.readCount()); + Assert.assertEquals(0, _reader.getUsedMemory()); + } + + private void useEvictingCache() { + _cache.shutdown(); + _io.reset(); + _cache = new OOCCacheImpl(_io, 4 * BYTES, 0); + } + + private void useZeroHardLimitCache() { + _cache.shutdown(); + _io.reset(); + _cache = new OOCCacheImpl(_io, 0, 0); + } + + private static void await(OOCCache.UnpinHandle handle) throws Exception { + if(!handle.isCommitted()) + handle.getCompletionFuture().get(WAIT_TIMEOUT_SEC, TimeUnit.SECONDS); + } + + private static void waitFor(BooleanSupplier condition) throws Exception { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(WAIT_TIMEOUT_SEC); + while(!condition.getAsBoolean() && System.nanoTime() < deadline) + Thread.sleep(1); + Assert.assertTrue(condition.getAsBoolean()); + } + + private static final class RecordingIOHandler implements OOCIOHandler { + private final Map _spilled = new ConcurrentHashMap<>(); + private final AtomicInteger _evictions = new AtomicInteger(); + private final AtomicInteger _reads = new AtomicInteger(); + private volatile boolean _failReads; + + @Override + public void shutdown() { + _spilled.clear(); + } + + @Override + public CompletableFuture scheduleEviction(BlockEntry block) { + _spilled.put(block.getKey(), BlockEntryTestAccess.getDataUnsafe(block)); + _evictions.incrementAndGet(); + return CompletableFuture.completedFuture(null); + } + + @Override + public OOCFuture scheduleRead(BlockEntry block) { + _reads.incrementAndGet(); + if(_failReads) + return OOCFuture.failed(new IllegalStateException("read failed")); + Object data = _spilled.get(block.getKey()); + if(data == null) + return OOCFuture.completed(null); + BlockEntryTestAccess.setDataUnsafe(block, data); + return OOCFuture.completed(block); + } + + @Override + public void prioritizeRead(BlockKey key, double priority) { + } + + @Override + public CompletableFuture scheduleDeletion(BlockEntry block) { + _spilled.remove(block.getKey()); + return CompletableFuture.completedFuture(true); + } + + @Override + public void registerSourceLocation(BlockKey key, SourceBlockDescriptor descriptor) { + } + + @Override + public CompletableFuture scheduleSourceRead(SourceReadRequest request) { + return CompletableFuture.failedFuture(new UnsupportedOperationException()); + } + + @Override + public CompletableFuture continueSourceRead(SourceReadContinuation continuation, + long maxBytesInFlight) { + return CompletableFuture.failedFuture(new UnsupportedOperationException()); + } + + private int evictionCount() { + return _evictions.get(); + } + + private int readCount() { + return _reads.get(); + } + + private void failReads(boolean failReads) { + _failReads = failReads; + } + + private void reset() { + _spilled.clear(); + _evictions.set(0); + _reads.set(0); + _failReads = false; + } + } +}