/**
* 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.drill.exec.work.batch;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.DrillBuf;
import java.io.EOFException;
import java.io.IOException;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;
import org.apache.drill.exec.ExecConstants;
import org.apache.drill.exec.memory.BufferAllocator;
import org.apache.drill.exec.ops.FragmentContext;
import org.apache.drill.exec.proto.BitData;
import org.apache.drill.exec.proto.ExecProtos;
import org.apache.drill.exec.proto.helper.QueryIdHelper;
import org.apache.drill.exec.record.RawFragmentBatch;
import org.apache.drill.exec.store.LocalSyncableFileSystem;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import com.google.common.base.Joiner;
import com.google.common.base.Preconditions;
import com.google.common.base.Stopwatch;
import com.google.common.collect.Queues;
/**
* This implementation of RawBatchBuffer starts writing incoming batches to disk once the buffer size reaches a threshold.
* The order of the incoming buffers is maintained.
*/
public class SpoolingRawBatchBuffer extends BaseRawBatchBuffer<SpoolingRawBatchBuffer.RawFragmentBatchWrapper> {
static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(SpoolingRawBatchBuffer.class);
private static String DRILL_LOCAL_IMPL_STRING = "fs.drill-local.impl";
private static final float STOP_SPOOLING_FRACTION = (float) 0.5;
public static final long ALLOCATOR_INITIAL_RESERVATION = 1*1024*1024;
public static final long ALLOCATOR_MAX_RESERVATION = 20L*1000*1000*1000;
private enum SpoolingState {
NOT_SPOOLING,
SPOOLING,
PAUSE_SPOOLING,
STOP_SPOOLING
}
private final BufferAllocator allocator;
private final long threshold;
private final int oppositeId;
private final int bufferIndex;
private volatile SpoolingState spoolingState;
private volatile long currentSizeInMemory = 0;
private volatile Spooler spooler;
private FileSystem fs;
private Path path;
private FSDataOutputStream outputStream;
public SpoolingRawBatchBuffer(FragmentContext context, int fragmentCount, int oppositeId, int bufferIndex) {
super(context, fragmentCount);
this.allocator = context.getNewChildAllocator(
"SpoolingRawBatchBufer", 100, ALLOCATOR_INITIAL_RESERVATION, ALLOCATOR_MAX_RESERVATION);
this.threshold = context.getConfig().getLong(ExecConstants.SPOOLING_BUFFER_MEMORY);
this.oppositeId = oppositeId;
this.bufferIndex = bufferIndex;
this.bufferQueue = new SpoolingBufferQueue();
}
private class SpoolingBufferQueue implements BufferQueue<RawFragmentBatchWrapper> {
private final LinkedBlockingDeque<RawFragmentBatchWrapper> buffer = Queues.newLinkedBlockingDeque();
@Override
public void addOomBatch(RawFragmentBatch batch) {
RawFragmentBatchWrapper batchWrapper = new RawFragmentBatchWrapper(batch, true);
batchWrapper.setOutOfMemory(true);
buffer.addFirst(batchWrapper);
}
@Override
public RawFragmentBatch poll() throws IOException {
RawFragmentBatchWrapper batchWrapper = buffer.poll();
if (batchWrapper != null) {
try {
return batchWrapper.get();
} catch (InterruptedException e) {
return null;
}
}
return null;
}
@Override
public RawFragmentBatch take() throws IOException, InterruptedException {
return buffer.take().get();
}
@Override
public boolean checkForOutOfMemory() {
return buffer.peek().isOutOfMemory();
}
@Override
public int size() {
return buffer.size();
}
@Override
public boolean isEmpty() {
return buffer.size() == 0;
}
public void add(RawFragmentBatchWrapper batchWrapper) {
buffer.add(batchWrapper);
}
}
private synchronized void setSpoolingState(SpoolingState newState) {
SpoolingState currentState = spoolingState;
if (newState == SpoolingState.NOT_SPOOLING ||
currentState == SpoolingState.STOP_SPOOLING) {
return;
}
spoolingState = newState;
}
private boolean isCurrentlySpooling() {
return spoolingState == SpoolingState.SPOOLING;
}
private void startSpooling() {
setSpoolingState(SpoolingState.SPOOLING);
}
private void pauseSpooling() {
setSpoolingState(SpoolingState.PAUSE_SPOOLING);
}
private boolean isSpoolingStopped() {
return spoolingState == SpoolingState.STOP_SPOOLING;
}
private void stopSpooling() {
setSpoolingState(SpoolingState.STOP_SPOOLING);
}
public String getDir() {
List<String> dirs = context.getConfig().getStringList(ExecConstants.TEMP_DIRECTORIES);
return dirs.get(ThreadLocalRandom.current().nextInt(dirs.size()));
}
private synchronized void initSpooler() throws IOException {
if (spooler != null) {
return;
}
Configuration conf = new Configuration();
conf.set(FileSystem.FS_DEFAULT_NAME_KEY, context.getConfig().getString(ExecConstants.TEMP_FILESYSTEM));
conf.set(DRILL_LOCAL_IMPL_STRING, LocalSyncableFileSystem.class.getName());
fs = FileSystem.get(conf);
path = getPath();
outputStream = fs.create(path);
final String spoolingThreadName = QueryIdHelper.getExecutorThreadName(context.getHandle()).concat(
":Spooler-" + oppositeId + "-" + bufferIndex);
spooler = new Spooler(spoolingThreadName);
spooler.start();
}
@Override
protected void enqueueInner(RawFragmentBatch batch) throws IOException {
assert batch.getHeader().getSendingMajorFragmentId() == oppositeId;
logger.debug("Enqueue batch. Current buffer size: {}. Last batch: {}. Sending fragment: {}", bufferQueue.size(), batch.getHeader().getIsLastBatch(), batch.getHeader().getSendingMajorFragmentId());
RawFragmentBatchWrapper wrapper;
boolean spoolCurrentBatch = isCurrentlySpooling();
wrapper = new RawFragmentBatchWrapper(batch, !spoolCurrentBatch);
currentSizeInMemory += wrapper.getBodySize();
if (spoolCurrentBatch) {
if (spooler == null) {
initSpooler();
}
spooler.addBatchForSpooling(wrapper);
}
bufferQueue.add(wrapper);
if (!spoolCurrentBatch && currentSizeInMemory > threshold) {
logger.debug("Buffer size {} greater than threshold {}. Start spooling to disk", currentSizeInMemory, threshold);
startSpooling();
}
}
@Override
public void kill(FragmentContext context) {
allocator.close();
if (spooler != null) {
spooler.terminate();
}
}
@Override
protected void upkeep(RawFragmentBatch batch) {
if (context.isOverMemoryLimit()) {
outOfMemory.set(true);
}
DrillBuf body = batch.getBody();
if (body != null) {
currentSizeInMemory -= body.capacity();
}
if (isCurrentlySpooling() && currentSizeInMemory < threshold * STOP_SPOOLING_FRACTION) {
logger.debug("buffer size {} less than {}x threshold. Stop spooling.", currentSizeInMemory, STOP_SPOOLING_FRACTION);
pauseSpooling();
}
logger.debug("Got batch. Current buffer size: {}", bufferQueue.size());
}
@Override
public void close() {
if (spooler != null) {
spooler.terminate();
while (spooler.isAlive()) {
try {
spooler.join();
} catch (InterruptedException e) {
logger.warn("Interrupted while waiting for spooling thread to exit");
continue;
}
}
}
allocator.close();
try {
if (outputStream != null) {
outputStream.close();
}
} catch (IOException e) {
logger.warn("Failed to cleanup I/O streams", e);
}
if (context.getConfig().getBoolean(ExecConstants.SPOOLING_BUFFER_DELETE)) {
try {
if (fs != null) {
fs.delete(path, false);
logger.debug("Deleted file {}", path.toString());
}
} catch (IOException e) {
logger.warn("Failed to delete temporary files", e);
}
}
super.close();
}
private class Spooler extends Thread {
private final LinkedBlockingDeque<RawFragmentBatchWrapper> spoolingQueue;
private volatile boolean shouldContinue = true;
private Thread spoolingThread;
public Spooler(String name) {
setDaemon(true);
setName(name);
spoolingQueue = Queues.newLinkedBlockingDeque();
}
public void run() {
try {
while (shouldContinue) {
RawFragmentBatchWrapper batch;
try {
batch = spoolingQueue.take();
} catch (InterruptedException e) {
if (shouldContinue) {
continue;
} else {
break;
}
}
try {
batch.writeToStream(outputStream);
} catch (IOException e) {
context.fail(e);
}
}
} catch (Throwable e) {
context.fail(e);
} finally {
logger.info("Spooler thread exiting");
}
}
public void addBatchForSpooling(RawFragmentBatchWrapper batchWrapper) {
if (isSpoolingStopped()) {
spoolingQueue.add(batchWrapper);
} else {
// will not spill this batch
batchWrapper.available = true;
batchWrapper.batch.sendOk();
batchWrapper.latch.countDown();
}
}
public void terminate() {
stopSpooling();
shouldContinue = false;
if (spoolingThread.isAlive()) {
spoolingThread.interrupt();
}
}
}
class RawFragmentBatchWrapper {
private RawFragmentBatch batch;
private volatile boolean available;
private CountDownLatch latch;
private volatile int bodyLength;
private volatile boolean outOfMemory = false;
private long start = -1;
private long check;
public RawFragmentBatchWrapper(RawFragmentBatch batch, boolean available) {
Preconditions.checkNotNull(batch);
this.batch = batch;
this.available = available;
this.latch = new CountDownLatch(available ? 0 : 1);
if (available) {
batch.sendOk();
}
}
public boolean isNull() {
return batch == null;
}
public RawFragmentBatch get() throws InterruptedException, IOException {
if (available) {
assert batch.getHeader() != null : "batch header null";
return batch;
} else {
latch.await();
readFromStream();
available = true;
return batch;
}
}
public long getBodySize() {
if (batch.getBody() == null) {
return 0;
}
assert batch.getBody().readableBytes() >= 0;
return batch.getBody().readableBytes();
}
public void writeToStream(FSDataOutputStream stream) throws IOException {
Stopwatch watch = Stopwatch.createStarted();
available = false;
check = ThreadLocalRandom.current().nextLong();
start = stream.getPos();
logger.debug("Writing check value {} at position {}", check, start);
stream.writeLong(check);
batch.getHeader().writeDelimitedTo(stream);
ByteBuf buf = batch.getBody();
if (buf != null) {
bodyLength = buf.capacity();
} else {
bodyLength = 0;
}
if (bodyLength > 0) {
buf.getBytes(0, stream, bodyLength);
}
stream.hsync();
FileStatus status = fs.getFileStatus(path);
long len = status.getLen();
logger.debug("After spooling batch, stream at position {}. File length {}", stream.getPos(), len);
batch.sendOk();
latch.countDown();
long t = watch.elapsed(TimeUnit.MICROSECONDS);
logger.debug("Took {} us to spool {} to disk. Rate {} mb/s", t, bodyLength, bodyLength / t);
if (buf != null) {
buf.release();
}
}
public void readFromStream() throws IOException, InterruptedException {
long pos = start;
boolean tryAgain = true;
int duration = 0;
while (tryAgain) {
// Sometimes, the file isn't quite done writing when we attempt to read it. As such, we need to wait and retry.
Thread.sleep(duration);
try(final FSDataInputStream stream = fs.open(path);
final DrillBuf buf = allocator.buffer(bodyLength)) {
stream.seek(start);
final long currentPos = stream.getPos();
final long check = stream.readLong();
pos = stream.getPos();
assert check == this.check : String.format("Check values don't match: %d %d, Position %d", this.check, check, currentPos);
Stopwatch watch = Stopwatch.createStarted();
BitData.FragmentRecordBatch header = BitData.FragmentRecordBatch.parseDelimitedFrom(stream);
pos = stream.getPos();
assert header != null : "header null after parsing from stream";
buf.writeBytes(stream, bodyLength);
pos = stream.getPos();
batch = new RawFragmentBatch(header, buf, null);
available = true;
latch.countDown();
long t = watch.elapsed(TimeUnit.MICROSECONDS);
logger.debug("Took {} us to read {} from disk. Rate {} mb/s", t, bodyLength, bodyLength / t);
tryAgain = false;
} catch (EOFException e) {
FileStatus status = fs.getFileStatus(path);
logger.warn("EOF reading from file {} at pos {}. Current file size: {}", path, pos, status.getLen());
duration = Math.max(1, duration * 2);
if (duration < 60000) {
continue;
} else {
throw e;
}
} finally {
if (tryAgain) {
// we had a premature exit, release batch memory so we don't leak it.
if (batch != null) {
batch.getBody().release();
}
}
}
}
}
private boolean isOutOfMemory() {
return outOfMemory;
}
private void setOutOfMemory(boolean outOfMemory) {
this.outOfMemory = outOfMemory;
}
}
private Path getPath() {
ExecProtos.FragmentHandle handle = context.getHandle();
String qid = QueryIdHelper.getQueryId(handle.getQueryId());
int majorFragmentId = handle.getMajorFragmentId();
int minorFragmentId = handle.getMinorFragmentId();
String fileName = Joiner.on(Path.SEPARATOR).join(getDir(), qid, majorFragmentId, minorFragmentId, oppositeId, bufferIndex);
return new Path(fileName);
}
}