exec/java-exec/src/main/java/org/apache/drill/exec/work/batch/BaseRawBatchBuffer.java - drill - Git at Google

 /*
  * 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 java.io.IOException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;

 import org.apache.drill.common.exceptions.DrillRuntimeException;
 import org.apache.drill.exec.ExecConstants;
 import org.apache.drill.exec.ops.FragmentContext;
 import org.apache.drill.exec.record.RawFragmentBatch;

 public abstract class BaseRawBatchBuffer<T> implements RawBatchBuffer {
   private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(BaseRawBatchBuffer.class);

   private enum BufferState {
     INIT,
     STREAMS_FINISHED,
     KILLED
   }

   protected interface BufferQueue<T> {
     void addOomBatch(RawFragmentBatch batch);
     RawFragmentBatch poll() throws IOException, InterruptedException;
     RawFragmentBatch take() throws IOException, InterruptedException;
     RawFragmentBatch poll(long timeout, TimeUnit timeUnit) throws InterruptedException, IOException;
     boolean checkForOutOfMemory();
     int size();
     boolean isEmpty();
     void add(T obj);
   }

   protected BufferQueue<T> bufferQueue;
   private volatile BufferState state = BufferState.INIT;
   protected final int bufferSizePerSocket;
   protected final AtomicBoolean outOfMemory = new AtomicBoolean(false);
   private int streamCounter;
   private final int fragmentCount;
   protected final FragmentContext context;

   public BaseRawBatchBuffer(final FragmentContext context, final int fragmentCount) {
     bufferSizePerSocket = context.getConfig().getInt(ExecConstants.INCOMING_BUFFER_SIZE);

     this.fragmentCount = fragmentCount;
     this.streamCounter = fragmentCount;
     this.context = context;
   }

   /**
    * Return the fragment count from construction time.
    *
    * @return the fragment count
    */
   protected int getFragmentCount() {
     return fragmentCount;
   }

   @Override
   public synchronized void enqueue(final RawFragmentBatch batch) throws IOException {

     // if this fragment is already canceled or failed, we shouldn't need any or more stuff. We do the null check to
     // ensure that tests run.
     if (context != null && !context.getExecutorState().shouldContinue()) {
       this.kill(context);
     }

     if (isTerminated()) {
       if (state == BufferState.KILLED) {
         // do not even enqueue just release and send ack back
         batch.release();
         batch.sendOk();
         return;
       } else {
         throw new IOException("Attempted to enqueue batch after finished");
       }
     }
     enqueueInner(batch);
   }

   /**
    * implementation specific method to enqueue batch
    *
    * @param batch
    * @throws IOException
    */
   protected abstract void enqueueInner(final RawFragmentBatch batch) throws IOException;

 //  ## Add assertion that all acks have been sent. TODO
   @Override
   public void close() {
     if (!isTerminated() && context.getExecutorState().shouldContinue()) {
       final String msg = String.format("Cleanup before finished. %d out of %d streams have finished", completedStreams(), fragmentCount);
       throw  new IllegalStateException(msg);
     }

     if (!bufferQueue.isEmpty()) {
       if (context.getExecutorState().shouldContinue()) {
         context.getExecutorState().fail(new IllegalStateException("Batches still in queue during cleanup"));
         logger.error("{} Batches in queue.", bufferQueue.size());
       }
       clearBufferWithBody();
     }
   }

   @Override
   public synchronized void kill(final FragmentContext context) {
     state = BufferState.KILLED;
     clearBufferWithBody();
   }

   /**
    * Helper method to clear buffer with request bodies release also flushes ack queue - in case there are still
    * responses pending
    */
   private void clearBufferWithBody() {
     RawFragmentBatch batch;
     while (!bufferQueue.isEmpty()) {
       batch = null;
       try {
         batch = bufferQueue.poll();
         assertAckSent(batch);
       } catch (IOException e) {
         context.getExecutorState().fail(e);
         continue;
       } catch (InterruptedException e) {
         context.getExecutorState().fail(e);
         // keep the state that the thread is interrupted
         Thread.currentThread().interrupt();
         continue;
       } finally {
         if (batch != null && batch.getBody() != null) {
           batch.getBody().release();
         }
       }
     }
   }

   private void allStreamsFinished() {
     if (state != BufferState.KILLED) {
       state = BufferState.STREAMS_FINISHED;
     }

     if (!bufferQueue.isEmpty()) {
       throw new IllegalStateException("buffer not empty when finished");
     }
   }

   @Override
   public RawFragmentBatch getNext() throws IOException {

     if (outOfMemory.get()) {
       if (bufferQueue.size() < 10) {
         outOfMemory.set(false);
       }
     }

     RawFragmentBatch b;
     try {
       b = bufferQueue.poll();

       // if we didn't get a batch, block on waiting for queue.
       if (b == null && (!isTerminated() || !bufferQueue.isEmpty())) {
         // We shouldn't block infinitely here. There can be a condition such that due to a failure FragmentExecutor
         // state is changed to FAILED and queue is empty. Because of this the minor fragment main thread will block
         // here waiting for next batch to arrive. Meanwhile when next batch arrived and was enqueued it sees
         // FragmentExecutor failure state and doesn't enqueue the batch and cleans up the buffer queue. Hence this
         // thread will stuck forever. So we pool for 5 seconds until we get a batch or FragmentExecutor state is in
         // error condition.
         while (b == null) {
           b = bufferQueue.poll(5, TimeUnit.SECONDS);
           if (!context.getExecutorState().shouldContinue()) {
             kill(context);
             if (b != null) {
               assertAckSent(b);
               if (b.getBody() != null) {
                 b.getBody().release();
               }
               b = null;
             }
           } // else b will be assigned a valid batch
         }
       }
     } catch (final InterruptedException e) {

       // We expect that the interrupt means the fragment is canceled or failed, so we should kill this buffer
       if (!context.getExecutorState().shouldContinue()) {
         kill(context);
       } else {
         throw new DrillRuntimeException("Interrupted but context.shouldContinue() is true", e);
       }
       // Preserve evidence that the interruption occurred so that code higher up on the call stack can learn of the
       // interruption and respond to it if it wants to.
       Thread.currentThread().interrupt();

       return null;
     }

     if (context.getAllocator().isOverLimit()) {
       outOfMemory.set(true);
     }

     if (b != null) {

       upkeep(b);

       if (b.getHeader().getIsLastBatch()) {
         logger.debug("Got last batch from {}:{}", b.getHeader().getSendingMajorFragmentId(), b.getHeader()
             .getSendingMinorFragmentId());
         final int remainingStreams = decrementStreamCounter();
         if (remainingStreams == 0) {
           logger.debug("Streams finished");
           allStreamsFinished();
         }
       }
     } else {
       if (!bufferQueue.isEmpty()) {
         throw new IllegalStateException("Returning null when there are batches left in queue");
       }
       if (!isTerminated()) {
         throw new IllegalStateException("Returning null when not finished");
       }
     }

     assertAckSent(b);
     return b;
   }

   private void assertAckSent(RawFragmentBatch batch) {
     assert batch == null || batch.isAckSent() : "Ack not sent for batch";
   }

   private int decrementStreamCounter() {
     streamCounter--;
     return streamCounter;
   }

   private int completedStreams() {
     return fragmentCount - streamCounter;
   }

   /**
    * Handle miscellaneous tasks after batch retrieval
    */
   protected abstract void upkeep(RawFragmentBatch batch);

   protected boolean isTerminated() {
     return (state == BufferState.KILLED || state == BufferState.STREAMS_FINISHED);
   }
 }
	/*
	* 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 java.io.IOException;
	import java.util.concurrent.TimeUnit;
	import java.util.concurrent.atomic.AtomicBoolean;

	import org.apache.drill.common.exceptions.DrillRuntimeException;
	import org.apache.drill.exec.ExecConstants;
	import org.apache.drill.exec.ops.FragmentContext;
	import org.apache.drill.exec.record.RawFragmentBatch;

	public abstract class BaseRawBatchBuffer<T> implements RawBatchBuffer {
	private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(BaseRawBatchBuffer.class);

	private enum BufferState {
	INIT,
	STREAMS_FINISHED,
	KILLED
	}

	protected interface BufferQueue<T> {
	void addOomBatch(RawFragmentBatch batch);
	RawFragmentBatch poll() throws IOException, InterruptedException;
	RawFragmentBatch take() throws IOException, InterruptedException;
	RawFragmentBatch poll(long timeout, TimeUnit timeUnit) throws InterruptedException, IOException;
	boolean checkForOutOfMemory();
	int size();
	boolean isEmpty();
	void add(T obj);
	}

	protected BufferQueue<T> bufferQueue;
	private volatile BufferState state = BufferState.INIT;
	protected final int bufferSizePerSocket;
	protected final AtomicBoolean outOfMemory = new AtomicBoolean(false);
	private int streamCounter;
	private final int fragmentCount;
	protected final FragmentContext context;

	public BaseRawBatchBuffer(final FragmentContext context, final int fragmentCount) {
	bufferSizePerSocket = context.getConfig().getInt(ExecConstants.INCOMING_BUFFER_SIZE);

	this.fragmentCount = fragmentCount;
	this.streamCounter = fragmentCount;
	this.context = context;
	}

	/**
	* Return the fragment count from construction time.
	*
	* @return the fragment count
	*/
	protected int getFragmentCount() {
	return fragmentCount;
	}

	@Override
	public synchronized void enqueue(final RawFragmentBatch batch) throws IOException {

	// if this fragment is already canceled or failed, we shouldn't need any or more stuff. We do the null check to
	// ensure that tests run.
	if (context != null && !context.getExecutorState().shouldContinue()) {
	this.kill(context);
	}

	if (isTerminated()) {
	if (state == BufferState.KILLED) {
	// do not even enqueue just release and send ack back
	batch.release();
	batch.sendOk();
	return;
	} else {
	throw new IOException("Attempted to enqueue batch after finished");
	}
	}
	enqueueInner(batch);
	}

	/**
	* implementation specific method to enqueue batch
	*
	* @param batch
	* @throws IOException
	*/
	protected abstract void enqueueInner(final RawFragmentBatch batch) throws IOException;

	// ## Add assertion that all acks have been sent. TODO
	@Override
	public void close() {
	if (!isTerminated() && context.getExecutorState().shouldContinue()) {
	final String msg = String.format("Cleanup before finished. %d out of %d streams have finished", completedStreams(), fragmentCount);
	throw new IllegalStateException(msg);
	}

	if (!bufferQueue.isEmpty()) {
	if (context.getExecutorState().shouldContinue()) {
	context.getExecutorState().fail(new IllegalStateException("Batches still in queue during cleanup"));
	logger.error("{} Batches in queue.", bufferQueue.size());
	}
	clearBufferWithBody();
	}
	}

	@Override
	public synchronized void kill(final FragmentContext context) {
	state = BufferState.KILLED;
	clearBufferWithBody();
	}

	/**
	* Helper method to clear buffer with request bodies release also flushes ack queue - in case there are still
	* responses pending
	*/
	private void clearBufferWithBody() {
	RawFragmentBatch batch;
	while (!bufferQueue.isEmpty()) {
	batch = null;
	try {
	batch = bufferQueue.poll();
	assertAckSent(batch);
	} catch (IOException e) {
	context.getExecutorState().fail(e);
	continue;
	} catch (InterruptedException e) {
	context.getExecutorState().fail(e);
	// keep the state that the thread is interrupted
	Thread.currentThread().interrupt();
	continue;
	} finally {
	if (batch != null && batch.getBody() != null) {
	batch.getBody().release();
	}
	}
	}
	}

	private void allStreamsFinished() {
	if (state != BufferState.KILLED) {
	state = BufferState.STREAMS_FINISHED;
	}

	if (!bufferQueue.isEmpty()) {
	throw new IllegalStateException("buffer not empty when finished");
	}
	}

	@Override
	public RawFragmentBatch getNext() throws IOException {

	if (outOfMemory.get()) {
	if (bufferQueue.size() < 10) {
	outOfMemory.set(false);
	}
	}

	RawFragmentBatch b;
	try {
	b = bufferQueue.poll();

	// if we didn't get a batch, block on waiting for queue.
	if (b == null && (!isTerminated() \|\| !bufferQueue.isEmpty())) {
	// We shouldn't block infinitely here. There can be a condition such that due to a failure FragmentExecutor
	// state is changed to FAILED and queue is empty. Because of this the minor fragment main thread will block
	// here waiting for next batch to arrive. Meanwhile when next batch arrived and was enqueued it sees
	// FragmentExecutor failure state and doesn't enqueue the batch and cleans up the buffer queue. Hence this
	// thread will stuck forever. So we pool for 5 seconds until we get a batch or FragmentExecutor state is in
	// error condition.
	while (b == null) {
	b = bufferQueue.poll(5, TimeUnit.SECONDS);
	if (!context.getExecutorState().shouldContinue()) {
	kill(context);
	if (b != null) {
	assertAckSent(b);
	if (b.getBody() != null) {
	b.getBody().release();
	}
	b = null;
	}
	} // else b will be assigned a valid batch
	}
	}
	} catch (final InterruptedException e) {

	// We expect that the interrupt means the fragment is canceled or failed, so we should kill this buffer
	if (!context.getExecutorState().shouldContinue()) {
	kill(context);
	} else {
	throw new DrillRuntimeException("Interrupted but context.shouldContinue() is true", e);
	}
	// Preserve evidence that the interruption occurred so that code higher up on the call stack can learn of the
	// interruption and respond to it if it wants to.
	Thread.currentThread().interrupt();

	return null;
	}

	if (context.getAllocator().isOverLimit()) {
	outOfMemory.set(true);
	}

	if (b != null) {

	upkeep(b);

	if (b.getHeader().getIsLastBatch()) {
	logger.debug("Got last batch from {}:{}", b.getHeader().getSendingMajorFragmentId(), b.getHeader()
	.getSendingMinorFragmentId());
	final int remainingStreams = decrementStreamCounter();
	if (remainingStreams == 0) {
	logger.debug("Streams finished");
	allStreamsFinished();
	}
	}
	} else {
	if (!bufferQueue.isEmpty()) {
	throw new IllegalStateException("Returning null when there are batches left in queue");
	}
	if (!isTerminated()) {
	throw new IllegalStateException("Returning null when not finished");
	}
	}

	assertAckSent(b);
	return b;
	}

	private void assertAckSent(RawFragmentBatch batch) {
	assert batch == null \|\| batch.isAckSent() : "Ack not sent for batch";
	}

	private int decrementStreamCounter() {
	streamCounter--;
	return streamCounter;
	}

	private int completedStreams() {
	return fragmentCount - streamCounter;
	}

	/**
	* Handle miscellaneous tasks after batch retrieval
	*/
	protected abstract void upkeep(RawFragmentBatch batch);

	protected boolean isTerminated() {
	return (state == BufferState.KILLED \|\| state == BufferState.STREAMS_FINISHED);
	}
	}