Examples with DataTunnel org.apache.drill.exec.rpc.data.DataTunnel used on opensource projects

Example 1 with DataTunnel

use of org.apache.drill.exec.rpc.data.DataTunnel in project drill by axbaretto.

the class TestBitRpc method testConnectionBackpressure.

@Test
public void testConnectionBackpressure(@Injectable WorkerBee bee, @Injectable final WorkEventBus workBus) throws Exception {
    DrillConfig config1 = DrillConfig.create();
    final BootStrapContext c = new BootStrapContext(config1, ClassPathScanner.fromPrescan(config1));
    DrillConfig config2 = DrillConfig.create();
    BootStrapContext c2 = new BootStrapContext(config2, ClassPathScanner.fromPrescan(config2));
    final FragmentContext fcon = new MockUp<FragmentContext>() {

        BufferAllocator getAllocator() {
            return c.getAllocator();
        }
    }.getMockInstance();
    final FragmentManager fman = new MockUp<FragmentManager>() {

        int v = 0;

        @Mock
        boolean handle(IncomingDataBatch batch) throws FragmentSetupException, IOException {
            try {
                v++;
                if (v % 10 == 0) {
                    System.out.println("sleeping.");
                    Thread.sleep(3000);
                }
            } catch (InterruptedException e) {
            }
            RawFragmentBatch rfb = batch.newRawFragmentBatch(c.getAllocator());
            rfb.sendOk();
            rfb.release();
            return true;
        }

        public FragmentContext getFragmentContext() {
            return fcon;
        }
    }.getMockInstance();
    new NonStrictExpectations() {

        {
            workBus.getFragmentManagerIfExists((FragmentHandle) any);
            result = fman;
            workBus.getFragmentManager((FragmentHandle) any);
            result = fman;
        }
    };
    int port = 1234;
    DataServer server = new DataServer(c, c.getAllocator(), workBus, null);
    port = server.bind(port, true);
    DrillbitEndpoint ep = DrillbitEndpoint.newBuilder().setAddress("localhost").setDataPort(port).build();
    DataConnectionManager manager = new DataConnectionManager(ep, c2);
    DataTunnel tunnel = new DataTunnel(manager);
    AtomicLong max = new AtomicLong(0);
    for (int i = 0; i < 40; i++) {
        long t1 = System.currentTimeMillis();
        tunnel.sendRecordBatch(new TimingOutcome(max), new FragmentWritableBatch(false, QueryId.getDefaultInstance(), 1, 1, 1, 1, getRandomBatch(c.getAllocator(), 5000)));
        System.out.println(System.currentTimeMillis() - t1);
    // System.out.println("sent.");
    }
    System.out.println(String.format("Max time: %d", max.get()));
    assertTrue(max.get() > 2700);
    Thread.sleep(5000);
}

Example 2 with DataTunnel

use of org.apache.drill.exec.rpc.data.DataTunnel in project drill by apache.

the class RuntimeFilterSink method route.

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);
    }
}