Examples with CountDownLatchInjection org.apache.drill.exec.testing.CountDownLatchInjection used on opensource projects

Example 1 with CountDownLatchInjection

use of org.apache.drill.exec.testing.CountDownLatchInjection in project drill by axbaretto.

the class PartitionerDecorator method executeMethodLogic.

/**
 * Helper to execute the different methods wrapped into same logic
 * @param iface
 * @throws IOException
 */
protected void executeMethodLogic(final GeneralExecuteIface iface) throws IOException {
    if (partitioners.size() == 1) {
        // no need for threads
        final OperatorStats localStatsSingle = partitioners.get(0).getStats();
        localStatsSingle.clear();
        localStatsSingle.startProcessing();
        try {
            iface.execute(partitioners.get(0));
        } finally {
            localStatsSingle.stopProcessing();
            stats.mergeMetrics(localStatsSingle);
            // since main stats did not have any wait time - adjust based of partitioner stats wait time
            // main stats processing time started recording in BaseRootExec
            stats.adjustWaitNanos(localStatsSingle.getWaitNanos());
        }
        return;
    }
    long maxProcessTime = 0l;
    // start waiting on main stats to adjust by sum(max(processing)) at the end
    stats.startWait();
    final CountDownLatch latch = new CountDownLatch(partitioners.size());
    final List<CustomRunnable> runnables = Lists.newArrayList();
    final List<Future<?>> taskFutures = Lists.newArrayList();
    CountDownLatchInjection testCountDownLatch = null;
    try {
        // To simulate interruption of main fragment thread and interrupting the partition threads, create a
        // CountDownInject patch. Partitioner threads await on the latch and main fragment thread counts down or
        // interrupts waiting threads. This makes sures that we are actually interrupting the blocked partitioner threads.
        testCountDownLatch = injector.getLatch(context.getExecutionControls(), "partitioner-sender-latch");
        testCountDownLatch.initialize(1);
        for (final Partitioner part : partitioners) {
            final CustomRunnable runnable = new CustomRunnable(childThreadPrefix, latch, iface, part, testCountDownLatch);
            runnables.add(runnable);
            taskFutures.add(executor.submit(runnable));
        }
        while (true) {
            try {
                // Wait for main fragment interruption.
                injector.injectInterruptiblePause(context.getExecutionControls(), "wait-for-fragment-interrupt", logger);
                // If there is no pause inserted at site "wait-for-fragment-interrupt", release the latch.
                injector.getLatch(context.getExecutionControls(), "partitioner-sender-latch").countDown();
                latch.await();
                break;
            } catch (final InterruptedException e) {
                // If the fragment state says we shouldn't continue, cancel or interrupt partitioner threads
                if (!context.getExecutorState().shouldContinue()) {
                    logger.debug("Interrupting partioner threads. Fragment thread {}", tName);
                    for (Future<?> f : taskFutures) {
                        f.cancel(true);
                    }
                    break;
                }
            }
        }
        IOException excep = null;
        for (final CustomRunnable runnable : runnables) {
            IOException myException = runnable.getException();
            if (myException != null) {
                if (excep == null) {
                    excep = myException;
                } else {
                    excep.addSuppressed(myException);
                }
            }
            final OperatorStats localStats = runnable.getPart().getStats();
            long currentProcessingNanos = localStats.getProcessingNanos();
            // find out max Partitioner processing time
            maxProcessTime = (currentProcessingNanos > maxProcessTime) ? currentProcessingNanos : maxProcessTime;
            stats.mergeMetrics(localStats);
        }
        if (excep != null) {
            throw excep;
        }
    } finally {
        stats.stopWait();
        // scale down main stats wait time based on calculated processing time
        // since we did not wait for whole duration of above execution
        stats.adjustWaitNanos(-maxProcessTime);
        // Done with the latch, close it.
        if (testCountDownLatch != null) {
            testCountDownLatch.close();
        }
    }
}

Example 2 with CountDownLatchInjection

use of org.apache.drill.exec.testing.CountDownLatchInjection in project drill by apache.

the class PartitionerDecorator method executeMethodLogic.

/**
 * Helper to execute the different methods wrapped into same logic
 * @param iface
 * @throws ExecutionException
 */
@VisibleForTesting
void executeMethodLogic(final GeneralExecuteIface iface) throws ExecutionException {
    // interrupts waiting threads. This makes sure that we are actually interrupting the blocked partitioner threads.
    try (CountDownLatchInjection testCountDownLatch = injector.getLatch(context.getExecutionControls(), "partitioner-sender-latch")) {
        testCountDownLatch.initialize(1);
        final AtomicInteger count = new AtomicInteger();
        List<PartitionerTask> partitionerTasks = new ArrayList<>(partitioners.size());
        ExecutionException executionException = null;
        // start waiting on main stats to adjust by sum(max(processing)) at the end
        startWait();
        try {
            partitioners.forEach(partitioner -> createAndExecute(iface, testCountDownLatch, count, partitionerTasks, partitioner));
            // Wait for main fragment interruption.
            injector.injectInterruptiblePause(context.getExecutionControls(), "wait-for-fragment-interrupt", logger);
            testCountDownLatch.countDown();
        } catch (InterruptedException e) {
            logger.warn("fragment thread interrupted", e);
            throw new QueryCancelledException();
        } catch (RejectedExecutionException e) {
            logger.warn("Failed to execute partitioner tasks. Execution service down?", e);
            executionException = new ExecutionException(e);
        } finally {
            await(count, partitionerTasks);
            stopWait();
            processPartitionerTasks(partitionerTasks, executionException);
        }
    }
}