org.modeshape.jcr.query.engine.process.RestartableSequence Maven / Gradle / Ivy
/*
* ModeShape (http://www.modeshape.org)
*
* Licensed 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.modeshape.jcr.query.engine.process;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Queue;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import org.mapdb.Serializer;
import org.modeshape.common.collection.CloseableSupplier;
import org.modeshape.common.collection.EmptyIterator;
import org.modeshape.jcr.cache.CachedNode;
import org.modeshape.jcr.cache.CachedNodeSupplier;
import org.modeshape.jcr.query.BufferManager;
import org.modeshape.jcr.query.BufferManager.QueueBuffer;
import org.modeshape.jcr.query.NodeSequence;
import org.modeshape.jcr.query.NodeSequence.Restartable;
import org.modeshape.jcr.query.engine.process.BufferedRows.BufferedRow;
import org.modeshape.jcr.query.engine.process.BufferedRows.BufferedRowFactory;
/**
* A {@link NodeSequence} that captures the buffers as they are used (or all at once) so that the sequence can be restarted.
*
* @author Randall Hauch ([email protected])
*/
public class RestartableSequence extends NodeSequence implements Restartable {
protected final NodeSequence original;
private final BufferedRowFactory rowFactory;
protected final Serializer rowSerializer;
protected final Queue inMemoryBatches;
protected final QueueBufferSupplier offHeapBatchesSupplier;
protected final String workspaceName;
protected final AtomicLong remainingRowCount = new AtomicLong();
private final int targetNumRowsInMemory;
protected final int width;
private BatchSequence batches;
protected final AtomicLong batchSize = new AtomicLong();
protected int actualNumRowsInMemory = 0;
protected long totalSize = 0L;
protected boolean loadedAll = false;
protected boolean usedOffHeap = false;
@SuppressWarnings( "unchecked" )
public RestartableSequence( String workspaceName,
NodeSequence original,
final BufferManager bufferMgr,
CachedNodeSupplier nodeCache,
final int numRowsInMemory ) {
this.original = original;
this.workspaceName = workspaceName;
this.width = original.width();
assert !original.isEmpty();
assert original.width() != 0;
// Create the row factory ...
this.rowFactory = BufferedRows.serializer(nodeCache, width);
// Create the buffer into which we'll place the rows with null keys ...
rowSerializer = (Serializer)BufferedRows.serializer(nodeCache, width);
// Create the in-memory storage ...
this.targetNumRowsInMemory = numRowsInMemory;
inMemoryBatches = new LinkedList<>();
// Create the supplier for the off-heap storage ...
offHeapBatchesSupplier = new QueueBufferSupplier(bufferMgr);
// Create the batch sequence that will copy and load each batch before returning it ...
batches = new BatchSequence() {
private final AtomicReference copiedBatch = new AtomicReference<>();
@Override
public Batch nextBatch() {
Batch batch = RestartableSequence.this.original.nextBatch();
if (batch == null) {
// There are no more batches, so we have seen, copied, and loaded all of the batches ...
loadedAll = true;
return null;
}
// Otherwise, make a copy of it and load it into the storage ...
boolean loadIntoMemory = inMemoryBatches != null && actualNumRowsInMemory < numRowsInMemory;
totalSize += loadBatch(batch, loadIntoMemory, copiedBatch);
if (batchSize.get() == 0L) batchSize.set(copiedBatch.get().rowCount());
return copiedBatch.get();
}
};
}
@Override
public int width() {
return width;
}
@Override
public long getRowCount() {
if (batches == null) return 0L;// closed
loadRemaining();
return totalSize;
}
@Override
public boolean isEmpty() {
return false;
}
@Override
public Batch nextBatch() {
if (batches == null) return null; // closed
return batches.nextBatch();
}
@Override
public void restart() {
loadRemaining();
restartBatches();
}
protected void restartBatches() {
if (batches == null) return; // closed
remainingRowCount.set(totalSize);
batches = new BatchSequence() {
private Iterator inMemory = inMemoryBatches.iterator();
private Iterator persisted;
@Override
public Batch nextBatch() {
if (inMemory.hasNext()) {
Batch result = inMemory.next();
if (result instanceof Restartable) {
((Restartable)result).restart();
}
return result;
}
if (persisted == null) {
persisted = offHeapBatchesSupplier.iterator();
}
if (!persisted.hasNext()) return null;
return batchFrom(persisted, batchSize.get());
}
};
}
@Override
public void close() {
RuntimeException error = null;
try {
original.close();
} catch (RuntimeException e) {
error = e;
} finally {
try {
inMemoryBatches.clear();
totalSize = 0L;
remainingRowCount.set(0L);
offHeapBatchesSupplier.close();
loadedAll = true;
} catch (RuntimeException e) {
if (error == null) error = e;
} finally {
batches = null;
if (error != null) throw error;
}
}
}
/**
* Load all of the remaining rows from the supplied sequence into the buffer.
*/
protected void loadRemaining() {
if (!loadedAll) {
// Put all of the batches from the sequence into the buffer
assert targetNumRowsInMemory >= 0L;
assert batchSize != null;
Batch batch = original.nextBatch();
boolean loadIntoMemory = inMemoryBatches != null && actualNumRowsInMemory < targetNumRowsInMemory;
while (batch != null) {
long rows = loadBatch(batch, loadIntoMemory, null);
if (batchSize.get() == 0L) batchSize.set(rows);
if (loadIntoMemory) {
assert inMemoryBatches != null;
if (actualNumRowsInMemory >= targetNumRowsInMemory) loadIntoMemory = false;
}
batch = original.nextBatch();
}
long numInMemory = inMemoryBatches != null ? actualNumRowsInMemory : 0L;
totalSize = offHeapBatchesSupplier.size() + numInMemory;
loadedAll = true;
restartBatches();
}
}
protected long loadBatch( Batch batch,
boolean loadIntoMemory,
AtomicReference copyOutput ) {
assert batch != null;
if (batch.isEmpty()) {
if (copyOutput != null) copyOutput.set(batch);
return 0L;
}
// Put all of the batches from the sequence into the buffer
if (loadIntoMemory) {
Batch copy = NodeSequence.copy(batch);
inMemoryBatches.add(copy);
if (copyOutput != null) copyOutput.set(copy);
long numRows = copy.rowCount();
actualNumRowsInMemory += numRows;
return numRows;
}
if (copyOutput != null) {
batch = NodeSequence.copy(batch);
copyOutput.set(batch);
}
long batchSize = 0L;
QueueBuffer persistedBatches = offHeapBatchesSupplier.get();
while (batch.hasNext()) {
batch.nextRow();
persistedBatches.append(createRow(batch));
++batchSize;
}
if (batch instanceof Restartable) {
((Restartable)batch).restart();
}
return batchSize;
}
protected BufferedRow createRow( Batch currentRow ) {
return rowFactory.createRow(currentRow);
}
protected Batch batchFrom( final Iterator rows,
final long maxBatchSize ) {
if (remainingRowCount.get() <= 0 || !rows.hasNext()) return null;
if (maxBatchSize == 0) return NodeSequence.emptyBatch(workspaceName, this.width);
final long rowsInBatch = Math.min(maxBatchSize, remainingRowCount.get());
return new Batch() {
private BufferedRow current;
@Override
public int width() {
return width;
}
@Override
public long rowCount() {
return rowsInBatch;
}
@Override
public String getWorkspaceName() {
return workspaceName;
}
@Override
public boolean isEmpty() {
return false;
}
@Override
public boolean hasNext() {
return remainingRowCount.get() > 0 && rows.hasNext();
}
@Override
public void nextRow() {
current = rows.next();
remainingRowCount.decrementAndGet();
}
@Override
public CachedNode getNode() {
return current.getNode();
}
@Override
public CachedNode getNode( int index ) {
return current.getNode(index);
}
@Override
public float getScore() {
return current.getScore();
}
@Override
public float getScore( int index ) {
return current.getScore(index);
}
@Override
public String toString() {
return "(restartable-batch width=" + width + " rows=" + rowCount() + ")";
}
};
}
@Override
public String toString() {
return "(restartable total-size=" + totalSize + " " + original + " )";
}
protected static interface BatchSequence {
Batch nextBatch();
}
protected class QueueBufferSupplier implements CloseableSupplier>, Iterable {
private QueueBuffer buffer;
private final BufferManager bufferMgr;
protected QueueBufferSupplier( BufferManager bufferMgr ) {
this.bufferMgr = bufferMgr;
}
@Override
public Iterator iterator() {
return buffer != null ? buffer.iterator() : new EmptyIterator();
}
protected long size() {
return buffer != null ? buffer.size() : 0L;
}
@Override
public QueueBuffer get() {
if (buffer == null) {
buffer = bufferMgr.createQueueBuffer(rowSerializer).useHeap(false).make();
}
return buffer;
}
@Override
public void close() {
if (buffer != null) {
try {
buffer.close();
} finally {
buffer = null;
}
}
}
}
}
© 2015 - 2024 Weber Informatics LLC | Privacy Policy