예제 #1
0
  private synchronized void enqueueDrivers(boolean forceRunSplit, List<DriverSplitRunner> runners) {
    // schedule driver to be executed
    List<ListenableFuture<?>> finishedFutures =
        taskExecutor.enqueueSplits(taskHandle, forceRunSplit, runners);
    checkState(
        finishedFutures.size() == runners.size(),
        "Expected %s futures but got %s",
        runners.size(),
        finishedFutures.size());

    // record new driver
    remainingDrivers.addAndGet(finishedFutures.size());

    // when driver completes, update state and fire events
    for (int i = 0; i < finishedFutures.size(); i++) {
      ListenableFuture<?> finishedFuture = finishedFutures.get(i);
      final DriverSplitRunner splitRunner = runners.get(i);
      Futures.addCallback(
          finishedFuture,
          new FutureCallback<Object>() {
            @Override
            public void onSuccess(Object result) {
              try (SetThreadName setThreadName = new SetThreadName("Task-%s", taskId)) {
                // record driver is finished
                remainingDrivers.decrementAndGet();

                checkTaskCompletion();

                queryMonitor.splitCompletionEvent(
                    taskId, splitRunner.getDriverContext().getDriverStats());
              }
            }

            @Override
            public void onFailure(Throwable cause) {
              try (SetThreadName setThreadName = new SetThreadName("Task-%s", taskId)) {
                taskStateMachine.failed(cause);

                // record driver is finished
                remainingDrivers.decrementAndGet();

                DriverContext driverContext = splitRunner.getDriverContext();
                DriverStats driverStats;
                if (driverContext != null) {
                  driverStats = driverContext.getDriverStats();
                } else {
                  // split runner did not start successfully
                  driverStats = new DriverStats();
                }

                // fire failed event with cause
                queryMonitor.splitFailedEvent(taskId, driverStats, cause);
              }
            }
          },
          notificationExecutor);
    }
  }
예제 #2
0
  public TestSqlTask() {
    taskExecutor = new TaskExecutor(8);
    taskExecutor.start();

    taskNotificationExecutor = newScheduledThreadPool(5, threadsNamed("task-notification-%d"));

    LocalExecutionPlanner planner = createTestingPlanner();

    sqlTaskExecutionFactory =
        new SqlTaskExecutionFactory(
            taskNotificationExecutor,
            taskExecutor,
            planner,
            new QueryMonitor(
                new ObjectMapperProvider().get(), new NullEventClient(), new NodeInfo("test")),
            new TaskManagerConfig());
  }
예제 #3
0
  private SqlTaskExecution(
      Session session,
      TaskId taskId,
      URI location,
      PlanFragment fragment,
      OutputBuffers outputBuffers,
      LocalExecutionPlanner planner,
      DataSize maxBufferSize,
      TaskExecutor taskExecutor,
      DataSize maxTaskMemoryUsage,
      DataSize operatorPreAllocatedMemory,
      QueryMonitor queryMonitor,
      Executor notificationExecutor,
      boolean cpuTimerEnabled) {
    try (SetThreadName setThreadName = new SetThreadName("Task-%s", taskId)) {
      this.taskId = checkNotNull(taskId, "taskId is null");
      this.location = checkNotNull(location, "location is null");
      this.taskExecutor = checkNotNull(taskExecutor, "driverExecutor is null");
      this.notificationExecutor =
          checkNotNull(notificationExecutor, "notificationExecutor is null");

      this.taskStateMachine = new TaskStateMachine(taskId, notificationExecutor);
      taskStateMachine.addStateChangeListener(
          new StateChangeListener<TaskState>() {
            @Override
            public void stateChanged(TaskState taskState) {
              if (taskState.isDone()) {
                SqlTaskExecution.this.taskExecutor.removeTask(taskHandle);
                // make sure buffers are cleaned up
                if (taskState != TaskState.FAILED) {
                  // don't close buffers for a failed query
                  // closed buffers signal to upstream tasks that everything finished cleanly
                  sharedBuffer.destroy();
                }
              }
            }
          });

      this.taskContext =
          new TaskContext(
              taskStateMachine,
              notificationExecutor,
              session,
              checkNotNull(maxTaskMemoryUsage, "maxTaskMemoryUsage is null"),
              checkNotNull(operatorPreAllocatedMemory, "operatorPreAllocatedMemory is null"),
              cpuTimerEnabled);

      this.sharedBuffer =
          new SharedBuffer(
              taskId,
              notificationExecutor,
              checkNotNull(maxBufferSize, "maxBufferSize is null"),
              outputBuffers);
      sharedBuffer.addStateChangeListener(
          new StateChangeListener<QueueState>() {
            @Override
            public void stateChanged(QueueState taskState) {
              if (taskState == QueueState.FINISHED) {
                checkTaskCompletion();
              }
            }
          });

      this.queryMonitor = checkNotNull(queryMonitor, "queryMonitor is null");

      taskHandle = taskExecutor.addTask(taskId);

      LocalExecutionPlan localExecutionPlan =
          planner.plan(
              session,
              fragment.getRoot(),
              fragment.getSymbols(),
              new TaskOutputFactory(sharedBuffer));
      List<DriverFactory> driverFactories = localExecutionPlan.getDriverFactories();

      // index driver factories
      DriverSplitRunnerFactory partitionedDriverFactory = null;
      ImmutableList.Builder<DriverSplitRunnerFactory> unpartitionedDriverFactories =
          ImmutableList.builder();
      for (DriverFactory driverFactory : driverFactories) {
        if (driverFactory.getSourceIds().contains(fragment.getPartitionedSource())) {
          checkState(
              partitionedDriverFactory == null, "multiple partitioned sources are not supported");
          partitionedDriverFactory = new DriverSplitRunnerFactory(driverFactory);
        } else {
          unpartitionedDriverFactories.add(new DriverSplitRunnerFactory(driverFactory));
        }
      }
      this.unpartitionedDriverFactories = unpartitionedDriverFactories.build();

      if (fragment.getDistribution() == PlanDistribution.SOURCE) {
        checkArgument(
            partitionedDriverFactory != null,
            "Fragment is partitioned, but no partitioned driver found");
      }
      this.partitionedSourceId = fragment.getPartitionedSource();
      this.partitionedDriverFactory = partitionedDriverFactory;
    }
  }
예제 #4
0
 @Override
 public void stateChanged(TaskState taskState) {
   if (taskState.isDone()) {
     taskExecutor.removeTask(taskHandle);
   }
 }
예제 #5
0
  private SqlTaskExecution(
      TaskStateMachine taskStateMachine,
      TaskContext taskContext,
      SharedBuffer sharedBuffer,
      PlanFragment fragment,
      LocalExecutionPlanner planner,
      TaskExecutor taskExecutor,
      QueryMonitor queryMonitor,
      Executor notificationExecutor) {
    this.taskStateMachine = checkNotNull(taskStateMachine, "taskStateMachine is null");
    this.taskId = taskStateMachine.getTaskId();
    this.taskContext = checkNotNull(taskContext, "taskContext is null");
    this.sharedBuffer = checkNotNull(sharedBuffer, "sharedBuffer is null");

    this.taskExecutor = checkNotNull(taskExecutor, "driverExecutor is null");
    this.notificationExecutor = checkNotNull(notificationExecutor, "notificationExecutor is null");

    this.queryMonitor = checkNotNull(queryMonitor, "queryMonitor is null");

    try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) {
      List<DriverFactory> driverFactories;
      try {
        OutputFactory outputOperatorFactory;
        if (fragment.getOutputPartitioning() == OutputPartitioning.NONE) {
          outputOperatorFactory = new TaskOutputFactory(sharedBuffer);
        } else if (fragment.getOutputPartitioning() == OutputPartitioning.HASH) {
          outputOperatorFactory = new PartitionedOutputFactory(sharedBuffer);
        } else {
          throw new PrestoException(
              NOT_SUPPORTED,
              format("OutputPartitioning %s is not supported", fragment.getOutputPartitioning()));
        }

        LocalExecutionPlan localExecutionPlan =
            planner.plan(
                taskContext.getSession(),
                fragment.getRoot(),
                fragment.getOutputLayout(),
                fragment.getSymbols(),
                fragment.getDistribution(),
                outputOperatorFactory);
        driverFactories = localExecutionPlan.getDriverFactories();
      } catch (Throwable e) {
        // planning failed
        taskStateMachine.failed(e);
        throw Throwables.propagate(e);
      }

      // index driver factories
      DriverSplitRunnerFactory partitionedDriverFactory = null;
      ImmutableList.Builder<DriverSplitRunnerFactory> unpartitionedDriverFactories =
          ImmutableList.builder();
      for (DriverFactory driverFactory : driverFactories) {
        if (driverFactory.getSourceIds().contains(fragment.getPartitionedSource())) {
          checkState(
              partitionedDriverFactory == null, "multiple partitioned sources are not supported");
          partitionedDriverFactory = new DriverSplitRunnerFactory(driverFactory);
        } else {
          unpartitionedDriverFactories.add(new DriverSplitRunnerFactory(driverFactory));
        }
      }
      this.unpartitionedDriverFactories = unpartitionedDriverFactories.build();

      if (fragment.getDistribution() == PlanDistribution.SOURCE) {
        checkArgument(
            partitionedDriverFactory != null,
            "Fragment is partitioned, but no partitioned driver found");
      }
      this.partitionedSourceId = fragment.getPartitionedSource();
      this.partitionedDriverFactory = partitionedDriverFactory;

      // don't register the task if it is already completed (most likely failed during planning
      // above)
      if (!taskStateMachine.getState().isDone()) {
        taskHandle = taskExecutor.addTask(taskId);
        taskStateMachine.addStateChangeListener(
            new RemoveTaskHandleWhenDone(taskExecutor, taskHandle));
      } else {
        taskHandle = null;
      }

      sharedBuffer.addStateChangeListener(
          new CheckTaskCompletionOnBufferFinish(SqlTaskExecution.this));
    }
  }
예제 #6
0
 @AfterClass
 public void destroy() throws Exception {
   taskExecutor.stop();
   taskNotificationExecutor.shutdownNow();
 }
예제 #7
0
  @BeforeMethod
  public void setUp() throws Exception {
    DualMetadata dualMetadata = new DualMetadata();
    TableHandle tableHandle =
        dualMetadata.getTableHandle(new SchemaTableName("default", DualMetadata.NAME));
    assertNotNull(tableHandle, "tableHandle is null");

    ColumnHandle columnHandle = dualMetadata.getColumnHandle(tableHandle, DualMetadata.COLUMN_NAME);
    assertNotNull(columnHandle, "columnHandle is null");
    Symbol symbol = new Symbol(DualMetadata.COLUMN_NAME);

    MetadataManager metadata = new MetadataManager(new FeaturesConfig());
    metadata.addInternalSchemaMetadata(MetadataManager.INTERNAL_CONNECTOR_ID, dualMetadata);

    DualSplitManager dualSplitManager = new DualSplitManager(new InMemoryNodeManager());
    PartitionResult partitionResult =
        dualSplitManager.getPartitions(tableHandle, TupleDomain.all());

    SplitSource splitSource =
        dualSplitManager.getPartitionSplits(tableHandle, partitionResult.getPartitions());
    split = Iterables.getOnlyElement(splitSource.getNextBatch(1));
    assertTrue(splitSource.isFinished());

    LocalExecutionPlanner planner =
        new LocalExecutionPlanner(
            new NodeInfo("test"),
            metadata,
            new DataStreamManager(new DualDataStreamProvider()),
            new MockLocalStorageManager(new File("target/temp")),
            new RecordSinkManager(),
            new MockExchangeClientSupplier(),
            new ExpressionCompiler(metadata));

    taskExecutor = new TaskExecutor(8);
    taskExecutor.start();

    tableScanNodeId = new PlanNodeId("tableScan");
    PlanFragment testFragment =
        new PlanFragment(
            new PlanFragmentId("fragment"),
            new TableScanNode(
                tableScanNodeId,
                tableHandle,
                ImmutableList.of(symbol),
                ImmutableMap.of(symbol, columnHandle),
                null,
                Optional.<GeneratedPartitions>absent()),
            ImmutableMap.of(symbol, Type.VARCHAR),
            PlanDistribution.SOURCE,
            tableScanNodeId,
            OutputPartitioning.NONE,
            ImmutableList.<Symbol>of());

    TaskId taskId = new TaskId("query", "stage", "task");
    Session session = new Session("user", "test", "default", "default", "test", "test");

    taskNotificationExecutor = Executors.newCachedThreadPool(threadsNamed("task-notification-%d"));

    outputBuffers = OutputBuffers.INITIAL_EMPTY_OUTPUT_BUFFERS;

    taskExecution =
        SqlTaskExecution.createSqlTaskExecution(
            session,
            taskId,
            URI.create("fake://task/" + taskId),
            testFragment,
            ImmutableList.<TaskSource>of(),
            outputBuffers,
            planner,
            new DataSize(32, Unit.MEGABYTE),
            taskExecutor,
            taskNotificationExecutor,
            new DataSize(256, Unit.MEGABYTE),
            new DataSize(8, Unit.MEGABYTE),
            new QueryMonitor(
                new ObjectMapperProvider().get(), new NullEventClient(), new NodeInfo("test")),
            false);
  }
예제 #8
0
 @AfterMethod
 public void tearDown() throws Exception {
   taskExecutor.stop();
   taskNotificationExecutor.shutdownNow();
 }