package org.infinispan.persistence.util;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Executor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import org.infinispan.container.InternalEntryFactory;
import org.infinispan.container.entries.CacheEntry;
import org.infinispan.container.entries.InternalCacheEntry;
import org.infinispan.filter.KeyFilter;
import org.infinispan.marshall.core.MarshalledEntry;
import org.infinispan.persistence.PersistenceUtil;
import org.infinispan.persistence.manager.PersistenceManager;
import org.infinispan.persistence.spi.AdvancedCacheLoader;
import org.infinispan.util.concurrent.TimeoutException;
import org.infinispan.util.concurrent.WithinThreadExecutor;
import org.infinispan.util.function.CloseableSupplier;
/**
* A closeable supplier that provides a way to supply cache entries from a given persistence manager. On the first
* call to get this class will submit a task to collect all of the entries from the loader (or optionally a subset
* provided a given {@link org.infinispan.filter.KeyFilter}). A timeout value is required so that if a get blocks
* for the given timeout it will throw a {@link TimeoutException}.
* @author William Burns
* @since 8.0
*/
public class PersistenceManagerCloseableSupplier<K, V> implements CloseableSupplier<CacheEntry<K, V>> {
private final Executor executor;
private final PersistenceManager manager;
private final KeyFilter<K> filter;
private final InternalEntryFactory factory;
private final BlockingQueue<CacheEntry<K, V>> queue;
private final long timeout;
private final TimeUnit unit;
private final Lock closeLock = new ReentrantLock();
private final Condition closeCondition = closeLock.newCondition();
private boolean closed = false;
private final AtomicReference<AdvancedCacheLoader.CacheLoaderTask<K, V>> taskRef = new AtomicReference<>();
public PersistenceManagerCloseableSupplier(Executor executor, PersistenceManager manager,
InternalEntryFactory factory, KeyFilter<K> filter, long timeout,
TimeUnit unit, int maxQueue) {
this.executor = executor;
this.manager = manager;
this.factory = factory;
this.filter = filter;
this.timeout = timeout;
this.unit = unit;
this.queue = new ArrayBlockingQueue<>(maxQueue);
}
class SupplierCacheLoaderTask implements AdvancedCacheLoader.CacheLoaderTask<K, V> {
@Override
public void processEntry(MarshalledEntry<K, V> marshalledEntry, AdvancedCacheLoader.TaskContext taskContext)
throws InterruptedException {
if (!taskContext.isStopped()) {
closeLock.lock();
try {
if (closed) {
taskContext.stop();
return;
}
} finally {
closeLock.unlock();
}
InternalCacheEntry<K, V> ice = PersistenceUtil.convert(marshalledEntry, factory);
// We do a read without acquiring lock
boolean stop = closed;
while (!stop) {
// If we were able to offer a value this means someone took from the queue so let us come back around main
// loop to offer all values we can
// TODO: do some sort of batching here - to reduce wakeup contention ?
if (queue.offer(ice, 100, TimeUnit.MILLISECONDS)) {
closeLock.lock();
try {
// Wake up anyone waiting for a value
closeCondition.signalAll();
} finally {
closeLock.unlock();
}
break;
}
// If we couldn't offer an entry check if we were completed concurrently
// We have to do in lock to ensure we see updated value properly
closeLock.lock();
try {
stop = closed;
} finally {
closeLock.unlock();
}
}
}
}
}
@Override
public CacheEntry<K, V> get() throws TimeoutException {
// We do a regular get first just so subsequent get calls don't allocate the task. If not we update the
// ref to the new task atomically - the one who sets it starts the task
if (taskRef.get() == null && taskRef.getAndUpdate((t) -> t == null ? new SupplierCacheLoaderTask() : t) == null) {
AdvancedCacheLoader.CacheLoaderTask<K, V> task = taskRef.get();
// TODO: unfortunately processOnAllStores requires 2 threads minimum unless using within thread executor - We
// can't really use the persistence executor since we will block while waiting for additional work
executor.execute(() -> {
try {
manager.processOnAllStores(new WithinThreadExecutor(), filter, task, true, true);
} finally {
close();
}
});
}
CacheEntry<K, V> entry;
boolean interrupted = false;
// TODO: replace this with ForkJoinPool.ManagedBlocker
while ((entry = queue.poll()) == null) {
closeLock.lock();
try {
if (closed) {
// If is possible that this supplier was closed right after we polled the queue. In that case they may
// have also inserted a value. We need to double check that the queue is REALLY empty after being closed.
if ((entry = queue.poll()) != null) {
break;
}
break;
}
long targetTime = System.nanoTime() + unit.toNanos(timeout);
try {
if (!closeCondition.await(targetTime - System.nanoTime(), TimeUnit.NANOSECONDS)) {
throw new TimeoutException("Couldn't retrieve entry an entry from store in allotted timeout: " + timeout
+ " unit: " + unit);
}
} catch (InterruptedException e) {
interrupted = true;
}
} finally {
closeLock.unlock();
}
}
if (interrupted) {
Thread.currentThread().interrupt();
}
return entry;
}
@Override
public void close() {
closeLock.lock();
try {
closed = true;
closeCondition.signalAll();
} finally {
closeLock.unlock();
}
}
}