Examples with SplitAssignment io.trino.execution.SplitAssignment used on opensource projects

Example 1 with SplitAssignment

use of io.trino.execution.SplitAssignment in project trino by trinodb.

the class HttpRemoteTask method sendUpdate.

private synchronized void sendUpdate() {
    TaskStatus taskStatus = getTaskStatus();
    // don't update if the task hasn't been started yet or if it is already finished
    if (!started.get() || !needsUpdate.get() || taskStatus.getState().isDone()) {
        return;
    }
    // currentRequest is always cleared when request is complete
    if (currentRequest != null) {
        return;
    }
    // if throttled due to error, asynchronously wait for timeout and try again
    ListenableFuture<Void> errorRateLimit = updateErrorTracker.acquireRequestPermit();
    if (!errorRateLimit.isDone()) {
        errorRateLimit.addListener(this::sendUpdate, executor);
        return;
    }
    List<SplitAssignment> splitAssignments = getSplitAssignments();
    VersionedDynamicFilterDomains dynamicFilterDomains = outboundDynamicFiltersCollector.acknowledgeAndGetNewDomains(sentDynamicFiltersVersion);
    // Workers don't need the embedded JSON representation when the fragment is sent
    Optional<PlanFragment> fragment = sendPlan.get() ? Optional.of(planFragment.withoutEmbeddedJsonRepresentation()) : Optional.empty();
    TaskUpdateRequest updateRequest = new TaskUpdateRequest(session.toSessionRepresentation(), session.getIdentity().getExtraCredentials(), fragment, splitAssignments, outputBuffers.get(), dynamicFilterDomains.getDynamicFilterDomains());
    byte[] taskUpdateRequestJson = taskUpdateRequestCodec.toJsonBytes(updateRequest);
    if (fragment.isPresent()) {
        stats.updateWithPlanBytes(taskUpdateRequestJson.length);
    }
    if (!dynamicFilterDomains.getDynamicFilterDomains().isEmpty()) {
        stats.updateWithDynamicFilterBytes(taskUpdateRequestJson.length);
    }
    HttpUriBuilder uriBuilder = getHttpUriBuilder(taskStatus);
    Request request = preparePost().setUri(uriBuilder.build()).setHeader(HttpHeaders.CONTENT_TYPE, MediaType.JSON_UTF_8.toString()).setBodyGenerator(createStaticBodyGenerator(taskUpdateRequestJson)).build();
    updateErrorTracker.startRequest();
    ListenableFuture<JsonResponse<TaskInfo>> future = httpClient.executeAsync(request, createFullJsonResponseHandler(taskInfoCodec));
    currentRequest = future;
    currentRequestStartNanos = System.nanoTime();
    // The needsUpdate flag needs to be set to false BEFORE adding the Future callback since callback might change the flag value
    // and does so without grabbing the instance lock.
    needsUpdate.set(false);
    Futures.addCallback(future, new SimpleHttpResponseHandler<>(new UpdateResponseHandler(splitAssignments, dynamicFilterDomains.getVersion()), request.getUri(), stats), executor);
}

Example 2 with SplitAssignment

use of io.trino.execution.SplitAssignment in project trino by trinodb.

the class HttpRemoteTask method getSplitAssignment.

private synchronized SplitAssignment getSplitAssignment(PlanNodeId planNodeId) {
    Set<ScheduledSplit> splits = pendingSplits.get(planNodeId);
    boolean pendingNoMoreSplits = Boolean.TRUE.equals(this.noMoreSplits.get(planNodeId));
    boolean noMoreSplits = this.noMoreSplits.containsKey(planNodeId);
    Set<Lifespan> noMoreSplitsForLifespan = pendingNoMoreSplitsForLifespan.get(planNodeId);
    SplitAssignment assignment = null;
    if (!splits.isEmpty() || !noMoreSplitsForLifespan.isEmpty() || pendingNoMoreSplits) {
        assignment = new SplitAssignment(planNodeId, splits, noMoreSplitsForLifespan, noMoreSplits);
    }
    return assignment;
}

Example 3 with SplitAssignment

use of io.trino.execution.SplitAssignment in project trino by trinodb.

the class Driver method processNewSources.

@GuardedBy("exclusiveLock")
private void processNewSources() {
    checkLockHeld("Lock must be held to call processNewSources");
    // only update if the driver is still alive
    if (state.get() != State.ALIVE) {
        return;
    }
    SplitAssignment splitAssignment = pendingSplitAssignmentUpdates.getAndSet(null);
    if (splitAssignment == null) {
        return;
    }
    // merge the current assignment and the specified assignment
    SplitAssignment newAssignment = currentSplitAssignment.update(splitAssignment);
    // if the update contains no new data, just return
    if (newAssignment == currentSplitAssignment) {
        return;
    }
    // determine new splits to add
    Set<ScheduledSplit> newSplits = Sets.difference(newAssignment.getSplits(), currentSplitAssignment.getSplits());
    // add new splits
    SourceOperator sourceOperator = this.sourceOperator.orElseThrow(VerifyException::new);
    for (ScheduledSplit newSplit : newSplits) {
        Split split = newSplit.getSplit();
        Supplier<Optional<UpdatablePageSource>> pageSource = sourceOperator.addSplit(split);
        deleteOperator.ifPresent(deleteOperator -> deleteOperator.setPageSource(pageSource));
        updateOperator.ifPresent(updateOperator -> updateOperator.setPageSource(pageSource));
    }
    // set no more splits
    if (newAssignment.isNoMoreSplits()) {
        sourceOperator.noMoreSplits();
    }
    currentSplitAssignment = newAssignment;
}

Example 4 with SplitAssignment

use of io.trino.execution.SplitAssignment in project trino by trinodb.

the class TestMemoryBlocking method testTableScanMemoryBlocking.

@Test
public void testTableScanMemoryBlocking() {
    PlanNodeId sourceId = new PlanNodeId("source");
    List<Type> types = ImmutableList.of(VARCHAR);
    TableScanOperator source = new TableScanOperator(driverContext.addOperatorContext(1, new PlanNodeId("test"), "values"), sourceId, (session, split, table, columns, dynamicFilter) -> new FixedPageSource(rowPagesBuilder(types).addSequencePage(10, 1).addSequencePage(10, 1).addSequencePage(10, 1).addSequencePage(10, 1).addSequencePage(10, 1).build()), TEST_TABLE_HANDLE, ImmutableList.of(), DynamicFilter.EMPTY);
    PageConsumerOperator sink = createSinkOperator(types);
    Driver driver = Driver.createDriver(driverContext, source, sink);
    assertSame(driver.getDriverContext(), driverContext);
    assertFalse(driver.isFinished());
    Split testSplit = new Split(new CatalogName("test"), new TestSplit(), Lifespan.taskWide());
    driver.updateSplitAssignment(new SplitAssignment(sourceId, ImmutableSet.of(new ScheduledSplit(0, sourceId, testSplit)), true));
    ListenableFuture<Void> blocked = driver.processFor(new Duration(1, NANOSECONDS));
    // the driver shouldn't block in the first call as it will be able to move a page between source and the sink operator
    // but the operator should be blocked
    assertTrue(blocked.isDone());
    assertFalse(source.getOperatorContext().isWaitingForMemory().isDone());
    // and they should stay blocked until more memory becomes available
    for (int i = 0; i < 10; i++) {
        blocked = driver.processFor(new Duration(1, NANOSECONDS));
        assertFalse(blocked.isDone());
        assertFalse(source.getOperatorContext().isWaitingForMemory().isDone());
    }
    // free up some memory
    memoryPool.free(QUERY_ID, "test", memoryPool.getReservedBytes());
    // the operator should be unblocked
    assertTrue(source.getOperatorContext().isWaitingForMemory().isDone());
    // the driver shouldn't be blocked
    blocked = driver.processFor(new Duration(1, NANOSECONDS));
    assertTrue(blocked.isDone());
}

Example 5 with SplitAssignment

use of io.trino.execution.SplitAssignment in project trino by trinodb.

the class HttpRemoteTask method processTaskUpdate.

private synchronized void processTaskUpdate(TaskInfo newValue, List<SplitAssignment> splitAssignments) {
    updateTaskInfo(newValue);
    // remove acknowledged splits, which frees memory
    for (SplitAssignment assignment : splitAssignments) {
        PlanNodeId planNodeId = assignment.getPlanNodeId();
        boolean isPartitionedSource = planFragment.isPartitionedSources(planNodeId);
        int removed = 0;
        long removedWeight = 0;
        for (ScheduledSplit split : assignment.getSplits()) {
            if (pendingSplits.remove(planNodeId, split)) {
                if (isPartitionedSource) {
                    removed++;
                    removedWeight = addExact(removedWeight, split.getSplit().getSplitWeight().getRawValue());
                }
            }
        }
        if (assignment.isNoMoreSplits()) {
            noMoreSplits.put(planNodeId, false);
        }
        for (Lifespan lifespan : assignment.getNoMoreSplitsForLifespan()) {
            pendingNoMoreSplitsForLifespan.remove(planNodeId, lifespan);
        }
        if (isPartitionedSource) {
            pendingSourceSplitCount -= removed;
            pendingSourceSplitsWeight -= removedWeight;
        }
    }
    // Update node level split tracker before split queue space to ensure it's up to date before waking up the scheduler
    partitionedSplitCountTracker.setPartitionedSplits(getPartitionedSplitsInfo());
    updateSplitQueueSpace();
}