public void register(OpChain operatorChain) {
    Future<?> scheduledFuture = _executorService.submit(new TraceRunnable() {
      @Override
      public void runJob() {
        boolean isFinished = false;
        TransferableBlock returnedErrorBlock = null;
        Throwable thrown = null;
        try {
          LOGGER.trace("({}): Executing", operatorChain);
          TransferableBlock result = operatorChain.getRoot().nextBlock();
          while (!result.isEndOfStreamBlock()) {
            result = operatorChain.getRoot().nextBlock();
          }
          isFinished = true;
          if (result.isErrorBlock()) {
            returnedErrorBlock = result;
            LOGGER.error("({}): Completed erroneously {} {}", operatorChain, result.getQueryStats(),
                result.getExceptions());
          } else {
            LOGGER.debug("({}): Completed {}", operatorChain, result.getQueryStats());
          }
        } catch (Exception e) {
          LOGGER.error("({}): Failed to execute operator chain!", operatorChain, e);
          thrown = e;
        } finally {
          _submittedOpChainMap.remove(operatorChain.getId());
          if (returnedErrorBlock != null || thrown != null) {
            if (thrown == null) {
              thrown = new RuntimeException("Error block " + returnedErrorBlock.getExceptions());
            }
            operatorChain.cancel(thrown);
          } else if (isFinished) {
            operatorChain.close();
          }
        }
      }
    });
    _submittedOpChainMap.put(operatorChain.getId(), scheduledFuture);
  }
@Test
  public void shouldCallCloseOnOperatorsThatFinishSuccessfully()
      throws InterruptedException {
    OpChain opChain = getChain(_operatorA);
    OpChainSchedulerService schedulerService = new OpChainSchedulerService(_executor);

    CountDownLatch latch = new CountDownLatch(1);
    Mockito.when(_operatorA.nextBlock()).thenReturn(TransferableBlockTestUtils.getEndOfStreamTransferableBlock(0));
    Mockito.doAnswer(inv -> {
      latch.countDown();
      return null;
    }).when(_operatorA).close();

    schedulerService.register(opChain);

    Assert.assertTrue(latch.await(10, TimeUnit.SECONDS), "expected await to be called in less than 10 seconds");
  }