How to use

getDriverContext

method

in

com.facebook.presto.operator.Driver

public synchronized DriverContext getDriverContext()
{
  if (driver == null) {
    return null;
  }
  return driver.getDriverContext();
}

private static boolean isOperatorBlocked(List<Driver> drivers, Predicate<OperatorContext> reason)
{
  for (Driver driver : drivers) {
    for (OperatorContext operatorContext : driver.getDriverContext().getOperatorContexts()) {
      if (reason.apply(operatorContext)) {
        return true;
      }
    }
  }
  return false;
}

private static void processRow(final Driver joinDriver, final TaskStateMachine taskStateMachine)
{
  joinDriver.getDriverContext().getYieldSignal().setWithDelay(TimeUnit.SECONDS.toNanos(1), joinDriver.getDriverContext().getYieldExecutor());
  joinDriver.process();
  joinDriver.getDriverContext().getYieldSignal().reset();
  checkErrors(taskStateMachine);
}

for (Driver driver : drivers) {
  if (alwaysRevokeMemory) {
    driver.getDriverContext().getOperatorContexts().stream()
        .filter(operatorContext -> operatorContext.getOperatorStats().getRevocableMemoryReservation().getValue() > 0)
        .forEach(OperatorContext::requestMemoryRevoking);

/**
 * Runs Driver in another thread until it is finished
 */
private static void runDriverInThread(ExecutorService executor, Driver driver)
{
  executor.execute(() -> {
    if (!driver.isFinished()) {
      try {
        driver.process();
      }
      catch (PrestoException e) {
        driver.getDriverContext().failed(e);
        throw e;
      }
      runDriverInThread(executor, driver);
    }
  });
}

/**
 * Runs Driver in another thread until it is finished
 */
private static void runDriverInThread(ExecutorService executor, Driver driver)
{
  executor.execute(() -> {
    if (!driver.isFinished()) {
      try {
        driver.process();
      }
      catch (PrestoException e) {
        driver.getDriverContext().failed(e);
        throw e;
      }
      runDriverInThread(executor, driver);
    }
  });
}

@Test
public void testBrokenOperatorCloseWhileProcessing()
    throws Exception
{
  BrokenOperator brokenOperator = new BrokenOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "source"), false);
  final Driver driver = Driver.createDriver(driverContext, brokenOperator, createSinkOperator(ImmutableList.of()));
  assertSame(driver.getDriverContext(), driverContext);
  // block thread in operator processing
  Future<Boolean> driverProcessFor = executor.submit(new Callable<Boolean>()
  {
    @Override
    public Boolean call()
    {
      return driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone();
    }
  });
  brokenOperator.waitForLocked();
  driver.close();
  assertTrue(driver.isFinished());
  try {
    driverProcessFor.get(1, TimeUnit.SECONDS);
    fail("Expected InterruptedException");
  }
  catch (ExecutionException e) {
    assertDriverInterrupted(e.getCause());
  }
}

@Test
public void testBrokenOperatorProcessWhileClosing()
    throws Exception
{
  BrokenOperator brokenOperator = new BrokenOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "source"), true);
  final Driver driver = Driver.createDriver(driverContext, brokenOperator, createSinkOperator(ImmutableList.of()));
  assertSame(driver.getDriverContext(), driverContext);
  // block thread in operator close
  Future<Boolean> driverClose = executor.submit(new Callable<Boolean>()
  {
    @Override
    public Boolean call()
    {
      driver.close();
      return true;
    }
  });
  brokenOperator.waitForLocked();
  assertTrue(driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone());
  assertTrue(driver.isFinished());
  brokenOperator.unlock();
  assertTrue(driverClose.get());
}

@Test
public void testNormalFinish()
{
  List<Type> types = ImmutableList.of(VARCHAR, BIGINT, BIGINT);
  ValuesOperator source = new ValuesOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "values"), rowPagesBuilder(types)
      .addSequencePage(10, 20, 30, 40)
      .build());
  Operator sink = createSinkOperator(types);
  Driver driver = Driver.createDriver(driverContext, source, sink);
  assertSame(driver.getDriverContext(), driverContext);
  assertFalse(driver.isFinished());
  ListenableFuture<?> blocked = driver.processFor(new Duration(1, TimeUnit.SECONDS));
  assertTrue(blocked.isDone());
  assertTrue(driver.isFinished());
  assertTrue(sink.isFinished());
  assertTrue(source.isFinished());
}

@Test
public void testAbruptFinish()
{
  List<Type> types = ImmutableList.of(VARCHAR, BIGINT, BIGINT);
  ValuesOperator source = new ValuesOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "values"), rowPagesBuilder(types)
      .addSequencePage(10, 20, 30, 40)
      .build());
  PageConsumerOperator sink = createSinkOperator(types);
  Driver driver = Driver.createDriver(driverContext, source, sink);
  assertSame(driver.getDriverContext(), driverContext);
  assertFalse(driver.isFinished());
  driver.close();
  assertTrue(driver.isFinished());
  // finish is only called in normal operations
  assertFalse(source.isFinished());
  assertFalse(sink.isFinished());
  // close is always called (values operator doesn't have a closed state)
  assertTrue(sink.isClosed());
}

Driver driver = Driver.createDriver(driverContext, source, sink);
assertSame(driver.getDriverContext(), driverContext);

brokenOperator.waitForLocked();
assertSame(driver.getDriverContext(), driverContext);

PageConsumerOperator sink = createSinkOperator(types);
Driver driver = Driver.createDriver(driverContext, source, sink);
assertSame(driver.getDriverContext(), driverContext);
assertFalse(driver.isFinished());
Split testSplit = new Split(new ConnectorId("test"), TestingTransactionHandle.create(), new TestSplit());

public synchronized DriverContext getDriverContext()
{
  if (driver == null) {
    return null;
  }
  return driver.getDriverContext();
}

/**
 * Runs Driver in another thread until it is finished
 */
private static void runDriverInThread(ExecutorService executor, Driver driver)
{
  executor.execute(() -> {
    if (!driver.isFinished()) {
      try {
        driver.process();
      }
      catch (PrestoException e) {
        driver.getDriverContext().failed(e);
        throw e;
      }
      runDriverInThread(executor, driver);
    }
  });
}

@Test
public void testBrokenOperatorCloseWhileProcessing()
    throws Exception
{
  BrokenOperator brokenOperator = new BrokenOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "source"), false);
  final Driver driver = new Driver(driverContext, brokenOperator, createSinkOperator(brokenOperator));
  assertSame(driver.getDriverContext(), driverContext);
  // block thread in operator processing
  Future<Boolean> driverProcessFor = executor.submit(new Callable<Boolean>()
  {
    @Override
    public Boolean call()
        throws Exception
    {
      return driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone();
    }
  });
  brokenOperator.waitForLocked();
  driver.close();
  assertTrue(driver.isFinished());
  try {
    driverProcessFor.get(1, TimeUnit.SECONDS);
    fail("Expected InterruptedException");
  }
  catch (ExecutionException e) {
    checkArgument(getRootCause(e) instanceof InterruptedException, "Expected root cause exception to be an instance of InterruptedException");
  }
}

@Test
public void testBrokenOperatorProcessWhileClosing()
    throws Exception
{
  BrokenOperator brokenOperator = new BrokenOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "source"), true);
  final Driver driver = new Driver(driverContext, brokenOperator, createSinkOperator(brokenOperator));
  assertSame(driver.getDriverContext(), driverContext);
  // block thread in operator close
  Future<Boolean> driverClose = executor.submit(new Callable<Boolean>()
  {
    @Override
    public Boolean call()
        throws Exception
    {
      driver.close();
      return true;
    }
  });
  brokenOperator.waitForLocked();
  assertTrue(driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone());
  assertTrue(driver.isFinished());
  brokenOperator.unlock();
  assertTrue(driverClose.get());
}

@Test
public void testAbruptFinish()
{
  List<Type> types = ImmutableList.<Type>of(VARCHAR, BIGINT, BIGINT);
  ValuesOperator source = new ValuesOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "values"), types, rowPagesBuilder(types)
      .addSequencePage(10, 20, 30, 40)
      .build());
  MaterializingOperator sink = createSinkOperator(source);
  Driver driver = new Driver(driverContext, source, sink);
  assertSame(driver.getDriverContext(), driverContext);
  assertFalse(driver.isFinished());
  driver.close();
  assertTrue(driver.isFinished());
  // finish is only called in normal operations
  assertFalse(source.isFinished());
  assertFalse(sink.isFinished());
  // close is always called (values operator doesn't have a closed state)
  assertTrue(sink.isClosed());
}

@Test
public void testNormalFinish()
{
  List<Type> types = ImmutableList.<Type>of(VARCHAR, BIGINT, BIGINT);
  ValuesOperator source = new ValuesOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "values"), types, rowPagesBuilder(types)
      .addSequencePage(10, 20, 30, 40)
      .build());
  MaterializingOperator sink = createSinkOperator(source);
  Driver driver = new Driver(driverContext, source, sink);
  assertSame(driver.getDriverContext(), driverContext);
  assertFalse(driver.isFinished());
  ListenableFuture<?> blocked = driver.processFor(new Duration(1, TimeUnit.SECONDS));
  assertTrue(blocked.isDone());
  assertTrue(driver.isFinished());
  assertTrue(sink.isFinished());
  assertTrue(source.isFinished());
}

Driver driver = new Driver(driverContext, source, sink);
assertSame(driver.getDriverContext(), driverContext);