exec/java-exec/src/main/java/org/apache/drill/exec/work/filter/RuntimeFilterSink.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.filter;

 import io.netty.buffer.DrillBuf;
 import org.apache.drill.common.exceptions.DrillRuntimeException;

 import org.apache.drill.exec.ops.AccountingDataTunnel;
 import org.apache.drill.exec.ops.Consumer;
 import org.apache.drill.exec.ops.DataTunnelStatusHandler;
 import org.apache.drill.exec.ops.SendingAccountor;
 import org.apache.drill.exec.proto.BitData;
 import org.apache.drill.exec.proto.CoordinationProtos;
 import org.apache.drill.exec.proto.UserBitShared;
 import org.apache.drill.exec.rpc.RpcException;
 import org.apache.drill.exec.rpc.RpcOutcomeListener;
 import org.apache.drill.exec.rpc.data.DataTunnel;
 import org.apache.drill.exec.server.DrillbitContext;
 import com.google.common.base.Stopwatch;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import java.io.Closeable;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;

 /**
  * This sink receives the RuntimeFilters from the netty thread,
  * aggregates them in an async thread, broadcast the final aggregated
  * one to the RuntimeFilterRecordBatch.
  */
 public class RuntimeFilterSink implements Closeable
 {

   private BlockingQueue<RuntimeFilterWritable> rfQueue = new LinkedBlockingQueue<>();

   private Map<Integer, Integer> joinMjId2rfNumber;

   //HashJoin node's major fragment id to its corresponding probe side nodes's endpoints
   private Map<Integer, List<CoordinationProtos.DrillbitEndpoint>> joinMjId2probeScanEps = new HashMap<>();

   //HashJoin node's major fragment id to its corresponding probe side scan node's belonging major fragment id
   private Map<Integer, Integer> joinMjId2ScanMjId = new HashMap<>();

   //HashJoin node's major fragment id to its aggregated RuntimeFilterWritable
   private Map<Integer, RuntimeFilterWritable> joinMjId2AggregatedRF = new HashMap<>();
   //for debug usage
   private Map<Integer, Stopwatch> joinMjId2Stopwatch = new HashMap<>();

   private DrillbitContext drillbitContext;

   private SendingAccountor sendingAccountor;

   private  AsyncAggregateWorker asyncAggregateWorker;

   private AtomicBoolean running = new AtomicBoolean(true);

   private static final Logger logger = LoggerFactory.getLogger(RuntimeFilterSink.class);


   public RuntimeFilterSink(DrillbitContext drillbitContext, SendingAccountor sendingAccountor)
   {
     this.drillbitContext = drillbitContext;
     this.sendingAccountor = sendingAccountor;
     asyncAggregateWorker = new AsyncAggregateWorker();
     drillbitContext.getExecutor().submit(asyncAggregateWorker);
   }

   public void add(RuntimeFilterWritable runtimeFilterWritable)
   {
     if (!running.get()) {
       runtimeFilterWritable.close();
       return;
     }
     runtimeFilterWritable.retainBuffers(1);
     int joinMjId = runtimeFilterWritable.getRuntimeFilterBDef().getMajorFragmentId();
     if (joinMjId2Stopwatch.get(joinMjId) == null) {
       Stopwatch stopwatch = Stopwatch.createStarted();
       joinMjId2Stopwatch.put(joinMjId, stopwatch);
     }
     synchronized (rfQueue) {
       if (!running.get()) {
         runtimeFilterWritable.close();
         return;
       }
       rfQueue.add(runtimeFilterWritable);
       rfQueue.notify();
     }
   }

   public void close() {
     running.set(false);
     if (asyncAggregateWorker != null) {
       synchronized (rfQueue) {
         rfQueue.notify();
       }
     }
     while (!asyncAggregateWorker.over.get()) {
       try {
         Thread.sleep(100);
       } catch (InterruptedException e) {
         logger.error("interrupted while sleeping to wait for the aggregating worker thread to exit", e);
       }
     }
     for (RuntimeFilterWritable runtimeFilterWritable : joinMjId2AggregatedRF.values()) {
       runtimeFilterWritable.close();
     }
   }

   private void aggregate(RuntimeFilterWritable srcRuntimeFilterWritable)
   {
     BitData.RuntimeFilterBDef runtimeFilterB = srcRuntimeFilterWritable.getRuntimeFilterBDef();
     int joinMajorId = runtimeFilterB.getMajorFragmentId();
     int buildSideRfNumber;
     RuntimeFilterWritable toAggregated = null;
     buildSideRfNumber = joinMjId2rfNumber.get(joinMajorId);
     buildSideRfNumber--;
     joinMjId2rfNumber.put(joinMajorId, buildSideRfNumber);
     toAggregated = joinMjId2AggregatedRF.get(joinMajorId);
     if (toAggregated == null) {
       toAggregated = srcRuntimeFilterWritable;
       toAggregated.retainBuffers(1);
     } else {
       toAggregated.aggregate(srcRuntimeFilterWritable);
     }
     joinMjId2AggregatedRF.put(joinMajorId, toAggregated);
     if (buildSideRfNumber == 0) {
       joinMjId2AggregatedRF.remove(joinMajorId);
       route(toAggregated);
       joinMjId2rfNumber.remove(joinMajorId);
       Stopwatch stopwatch = joinMjId2Stopwatch.get(joinMajorId);
       logger.info(
           "received all the RFWs belonging to the majorId {}'s HashJoin nodes and flushed aggregated RFW out elapsed {} ms",
           joinMajorId,
           stopwatch.elapsed(TimeUnit.MILLISECONDS)
       );
     }
   }

   private void route(RuntimeFilterWritable srcRuntimeFilterWritable)
   {
     BitData.RuntimeFilterBDef runtimeFilterB = srcRuntimeFilterWritable.getRuntimeFilterBDef();
     int joinMajorId = runtimeFilterB.getMajorFragmentId();
     UserBitShared.QueryId queryId = runtimeFilterB.getQueryId();
     List<String> probeFields = runtimeFilterB.getProbeFieldsList();
     List<Integer> sizeInBytes = runtimeFilterB.getBloomFilterSizeInBytesList();
     long rfIdentifier = runtimeFilterB.getRfIdentifier();
     DrillBuf[] data = srcRuntimeFilterWritable.getData();
     List<CoordinationProtos.DrillbitEndpoint> scanNodeEps = joinMjId2probeScanEps.get(joinMajorId);
     int scanNodeSize = scanNodeEps.size();
     srcRuntimeFilterWritable.retainBuffers(scanNodeSize - 1);
     int scanNodeMjId = joinMjId2ScanMjId.get(joinMajorId);
     for (int minorId = 0; minorId < scanNodeEps.size(); minorId++) {
       BitData.RuntimeFilterBDef.Builder builder = BitData.RuntimeFilterBDef.newBuilder();
       for (String probeField : probeFields) {
         builder.addProbeFields(probeField);
       }
       BitData.RuntimeFilterBDef runtimeFilterBDef = builder.setQueryId(queryId)
                                                            .setMajorFragmentId(scanNodeMjId)
                                                            .setMinorFragmentId(minorId)
                                                            .setToForeman(false)
                                                            .setRfIdentifier(rfIdentifier)
                                                            .addAllBloomFilterSizeInBytes(sizeInBytes)
                                                            .build();
       RuntimeFilterWritable runtimeFilterWritable = new RuntimeFilterWritable(runtimeFilterBDef, data);
       CoordinationProtos.DrillbitEndpoint drillbitEndpoint = scanNodeEps.get(minorId);

       DataTunnel dataTunnel = drillbitContext.getDataConnectionsPool().getTunnel(drillbitEndpoint);
       Consumer<RpcException> exceptionConsumer = new Consumer<RpcException>()
       {
         @Override
         public void accept(final RpcException e)
         {
           logger.warn("fail to broadcast a runtime filter to the probe side scan node", e);
         }

         @Override
         public void interrupt(final InterruptedException e)
         {
           logger.warn("fail to broadcast a runtime filter to the probe side scan node", e);
         }
       };
       RpcOutcomeListener<BitData.AckWithCredit> statusHandler = new DataTunnelStatusHandler(exceptionConsumer, sendingAccountor);
       AccountingDataTunnel accountingDataTunnel = new AccountingDataTunnel(dataTunnel, sendingAccountor, statusHandler);
       accountingDataTunnel.sendRuntimeFilter(runtimeFilterWritable);
     }
   }

   public void setJoinMjId2rfNumber(Map<Integer, Integer> joinMjId2rfNumber)
   {
     this.joinMjId2rfNumber = joinMjId2rfNumber;
   }

   public void setJoinMjId2probeScanEps(Map<Integer, List<CoordinationProtos.DrillbitEndpoint>> joinMjId2probeScanEps)
   {
     this.joinMjId2probeScanEps = joinMjId2probeScanEps;
   }

   public void setJoinMjId2ScanMjId(Map<Integer, Integer> joinMjId2ScanMjId)
   {
     this.joinMjId2ScanMjId = joinMjId2ScanMjId;
   }

   private class AsyncAggregateWorker implements Runnable
   {
     private AtomicBoolean over = new AtomicBoolean(false);

     @Override
     public void run()
     {
       while ((joinMjId2rfNumber == null || !joinMjId2rfNumber.isEmpty() ) && running.get()) {
         RuntimeFilterWritable toAggregate = null;
         synchronized (rfQueue) {
           try {
             toAggregate = rfQueue.poll();
             while (toAggregate == null && running.get()) {
               rfQueue.wait();
               toAggregate = rfQueue.poll();
             }
           } catch (InterruptedException ex) {
             logger.error("RFW_Aggregator thread being interrupted", ex);
             continue;
           }
         }
         if (toAggregate == null) {
           continue;
         }
         // perform aggregate outside the sync block.
         try {
           aggregate(toAggregate);
         } catch (Exception ex) {
           logger.error("Failed to aggregate or route the RFW", ex);

           // Set running to false and cleanup pending RFW in queue. This will make sure producer
           // thread is also indicated to stop and queue is cleaned up properly in failure cases
           synchronized (rfQueue) {
             running.set(false);
           }
           cleanupQueue();
           throw new DrillRuntimeException(ex);
         } finally {
             toAggregate.close();
         }
       }
       cleanupQueue();
     }

     private void cleanupQueue() {
       if (!running.get()) {
         RuntimeFilterWritable toClose;
         while ((toClose = rfQueue.poll()) != null) {
           toClose.close();
         }
       }
       over.set(true);
     }
   }
 }
	/*
	* 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.filter;

	import io.netty.buffer.DrillBuf;
	import org.apache.drill.common.exceptions.DrillRuntimeException;

	import org.apache.drill.exec.ops.AccountingDataTunnel;
	import org.apache.drill.exec.ops.Consumer;
	import org.apache.drill.exec.ops.DataTunnelStatusHandler;
	import org.apache.drill.exec.ops.SendingAccountor;
	import org.apache.drill.exec.proto.BitData;
	import org.apache.drill.exec.proto.CoordinationProtos;
	import org.apache.drill.exec.proto.UserBitShared;
	import org.apache.drill.exec.rpc.RpcException;
	import org.apache.drill.exec.rpc.RpcOutcomeListener;
	import org.apache.drill.exec.rpc.data.DataTunnel;
	import org.apache.drill.exec.server.DrillbitContext;
	import com.google.common.base.Stopwatch;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import java.io.Closeable;
	import java.util.HashMap;
	import java.util.List;
	import java.util.Map;
	import java.util.concurrent.BlockingQueue;
	import java.util.concurrent.LinkedBlockingQueue;
	import java.util.concurrent.TimeUnit;
	import java.util.concurrent.atomic.AtomicBoolean;

	/**
	* This sink receives the RuntimeFilters from the netty thread,
	* aggregates them in an async thread, broadcast the final aggregated
	* one to the RuntimeFilterRecordBatch.
	*/
	public class RuntimeFilterSink implements Closeable
	{

	private BlockingQueue<RuntimeFilterWritable> rfQueue = new LinkedBlockingQueue<>();

	private Map<Integer, Integer> joinMjId2rfNumber;

	//HashJoin node's major fragment id to its corresponding probe side nodes's endpoints
	private Map<Integer, List<CoordinationProtos.DrillbitEndpoint>> joinMjId2probeScanEps = new HashMap<>();

	//HashJoin node's major fragment id to its corresponding probe side scan node's belonging major fragment id
	private Map<Integer, Integer> joinMjId2ScanMjId = new HashMap<>();

	//HashJoin node's major fragment id to its aggregated RuntimeFilterWritable
	private Map<Integer, RuntimeFilterWritable> joinMjId2AggregatedRF = new HashMap<>();
	//for debug usage
	private Map<Integer, Stopwatch> joinMjId2Stopwatch = new HashMap<>();

	private DrillbitContext drillbitContext;

	private SendingAccountor sendingAccountor;

	private AsyncAggregateWorker asyncAggregateWorker;

	private AtomicBoolean running = new AtomicBoolean(true);

	private static final Logger logger = LoggerFactory.getLogger(RuntimeFilterSink.class);


	public RuntimeFilterSink(DrillbitContext drillbitContext, SendingAccountor sendingAccountor)
	{
	this.drillbitContext = drillbitContext;
	this.sendingAccountor = sendingAccountor;
	asyncAggregateWorker = new AsyncAggregateWorker();
	drillbitContext.getExecutor().submit(asyncAggregateWorker);
	}

	public void add(RuntimeFilterWritable runtimeFilterWritable)
	{
	if (!running.get()) {
	runtimeFilterWritable.close();
	return;
	}
	runtimeFilterWritable.retainBuffers(1);
	int joinMjId = runtimeFilterWritable.getRuntimeFilterBDef().getMajorFragmentId();
	if (joinMjId2Stopwatch.get(joinMjId) == null) {
	Stopwatch stopwatch = Stopwatch.createStarted();
	joinMjId2Stopwatch.put(joinMjId, stopwatch);
	}
	synchronized (rfQueue) {
	if (!running.get()) {
	runtimeFilterWritable.close();
	return;
	}
	rfQueue.add(runtimeFilterWritable);
	rfQueue.notify();
	}
	}

	public void close() {
	running.set(false);
	if (asyncAggregateWorker != null) {
	synchronized (rfQueue) {
	rfQueue.notify();
	}
	}
	while (!asyncAggregateWorker.over.get()) {
	try {
	Thread.sleep(100);
	} catch (InterruptedException e) {
	logger.error("interrupted while sleeping to wait for the aggregating worker thread to exit", e);
	}
	}
	for (RuntimeFilterWritable runtimeFilterWritable : joinMjId2AggregatedRF.values()) {
	runtimeFilterWritable.close();
	}
	}

	private void aggregate(RuntimeFilterWritable srcRuntimeFilterWritable)
	{
	BitData.RuntimeFilterBDef runtimeFilterB = srcRuntimeFilterWritable.getRuntimeFilterBDef();
	int joinMajorId = runtimeFilterB.getMajorFragmentId();
	int buildSideRfNumber;
	RuntimeFilterWritable toAggregated = null;
	buildSideRfNumber = joinMjId2rfNumber.get(joinMajorId);
	buildSideRfNumber--;
	joinMjId2rfNumber.put(joinMajorId, buildSideRfNumber);
	toAggregated = joinMjId2AggregatedRF.get(joinMajorId);
	if (toAggregated == null) {
	toAggregated = srcRuntimeFilterWritable;
	toAggregated.retainBuffers(1);
	} else {
	toAggregated.aggregate(srcRuntimeFilterWritable);
	}
	joinMjId2AggregatedRF.put(joinMajorId, toAggregated);
	if (buildSideRfNumber == 0) {
	joinMjId2AggregatedRF.remove(joinMajorId);
	route(toAggregated);
	joinMjId2rfNumber.remove(joinMajorId);
	Stopwatch stopwatch = joinMjId2Stopwatch.get(joinMajorId);
	logger.info(
	"received all the RFWs belonging to the majorId {}'s HashJoin nodes and flushed aggregated RFW out elapsed {} ms",
	joinMajorId,
	stopwatch.elapsed(TimeUnit.MILLISECONDS)
	);
	}
	}

	private void route(RuntimeFilterWritable srcRuntimeFilterWritable)
	{
	BitData.RuntimeFilterBDef runtimeFilterB = srcRuntimeFilterWritable.getRuntimeFilterBDef();
	int joinMajorId = runtimeFilterB.getMajorFragmentId();
	UserBitShared.QueryId queryId = runtimeFilterB.getQueryId();
	List<String> probeFields = runtimeFilterB.getProbeFieldsList();
	List<Integer> sizeInBytes = runtimeFilterB.getBloomFilterSizeInBytesList();
	long rfIdentifier = runtimeFilterB.getRfIdentifier();
	DrillBuf[] data = srcRuntimeFilterWritable.getData();
	List<CoordinationProtos.DrillbitEndpoint> scanNodeEps = joinMjId2probeScanEps.get(joinMajorId);
	int scanNodeSize = scanNodeEps.size();
	srcRuntimeFilterWritable.retainBuffers(scanNodeSize - 1);
	int scanNodeMjId = joinMjId2ScanMjId.get(joinMajorId);
	for (int minorId = 0; minorId < scanNodeEps.size(); minorId++) {
	BitData.RuntimeFilterBDef.Builder builder = BitData.RuntimeFilterBDef.newBuilder();
	for (String probeField : probeFields) {
	builder.addProbeFields(probeField);
	}
	BitData.RuntimeFilterBDef runtimeFilterBDef = builder.setQueryId(queryId)
	.setMajorFragmentId(scanNodeMjId)
	.setMinorFragmentId(minorId)
	.setToForeman(false)
	.setRfIdentifier(rfIdentifier)
	.addAllBloomFilterSizeInBytes(sizeInBytes)
	.build();
	RuntimeFilterWritable runtimeFilterWritable = new RuntimeFilterWritable(runtimeFilterBDef, data);
	CoordinationProtos.DrillbitEndpoint drillbitEndpoint = scanNodeEps.get(minorId);

	DataTunnel dataTunnel = drillbitContext.getDataConnectionsPool().getTunnel(drillbitEndpoint);
	Consumer<RpcException> exceptionConsumer = new Consumer<RpcException>()
	{
	@Override
	public void accept(final RpcException e)
	{
	logger.warn("fail to broadcast a runtime filter to the probe side scan node", e);
	}

	@Override
	public void interrupt(final InterruptedException e)
	{
	logger.warn("fail to broadcast a runtime filter to the probe side scan node", e);
	}
	};
	RpcOutcomeListener<BitData.AckWithCredit> statusHandler = new DataTunnelStatusHandler(exceptionConsumer, sendingAccountor);
	AccountingDataTunnel accountingDataTunnel = new AccountingDataTunnel(dataTunnel, sendingAccountor, statusHandler);
	accountingDataTunnel.sendRuntimeFilter(runtimeFilterWritable);
	}
	}

	public void setJoinMjId2rfNumber(Map<Integer, Integer> joinMjId2rfNumber)
	{
	this.joinMjId2rfNumber = joinMjId2rfNumber;
	}

	public void setJoinMjId2probeScanEps(Map<Integer, List<CoordinationProtos.DrillbitEndpoint>> joinMjId2probeScanEps)
	{
	this.joinMjId2probeScanEps = joinMjId2probeScanEps;
	}

	public void setJoinMjId2ScanMjId(Map<Integer, Integer> joinMjId2ScanMjId)
	{
	this.joinMjId2ScanMjId = joinMjId2ScanMjId;
	}

	private class AsyncAggregateWorker implements Runnable
	{
	private AtomicBoolean over = new AtomicBoolean(false);

	@Override
	public void run()
	{
	while ((joinMjId2rfNumber == null \|\| !joinMjId2rfNumber.isEmpty() ) && running.get()) {
	RuntimeFilterWritable toAggregate = null;
	synchronized (rfQueue) {
	try {
	toAggregate = rfQueue.poll();
	while (toAggregate == null && running.get()) {
	rfQueue.wait();
	toAggregate = rfQueue.poll();
	}
	} catch (InterruptedException ex) {
	logger.error("RFW_Aggregator thread being interrupted", ex);
	continue;
	}
	}
	if (toAggregate == null) {
	continue;
	}
	// perform aggregate outside the sync block.
	try {
	aggregate(toAggregate);
	} catch (Exception ex) {
	logger.error("Failed to aggregate or route the RFW", ex);

	// Set running to false and cleanup pending RFW in queue. This will make sure producer
	// thread is also indicated to stop and queue is cleaned up properly in failure cases
	synchronized (rfQueue) {
	running.set(false);
	}
	cleanupQueue();
	throw new DrillRuntimeException(ex);
	} finally {
	toAggregate.close();
	}
	}
	cleanupQueue();
	}

	private void cleanupQueue() {
	if (!running.get()) {
	RuntimeFilterWritable toClose;
	while ((toClose = rfQueue.poll()) != null) {
	toClose.close();
	}
	}
	over.set(true);
	}
	}
	}