@Test public void testSingleNode() { FixedCountScheduler nodeScheduler = new FixedCountScheduler( (node, partition, totalPartitions) -> Optional.of(taskFactory.createTableScanTask( new TaskId("test", 1, 1), node, ImmutableList.of(), new PartitionedSplitCountTracker(delta -> {}))), generateRandomNodes(1)); ScheduleResult result = nodeScheduler.schedule(); assertTrue(result.isFinished()); assertTrue(result.getBlocked().isDone()); assertEquals(result.getNewTasks().size(), 1); assertTrue(result.getNewTasks().iterator().next().getNodeId().equals("other 0")); }
partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); updateSplitQueueSpace(); partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); updateSplitQueueSpace();
continuousTaskInfoFetcher = new ContinuousTaskInfoFetcher(refreshMaxWait); partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount());
partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount());
partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount());
@Override public synchronized void addSplits(Multimap<PlanNodeId, Split> splitsBySource) { try (SetThreadName ignored = new SetThreadName("HttpRemoteTask-%s", taskId)) { requireNonNull(splitsBySource, "splitsBySource is null"); // only add pending split if not done if (getTaskInfo().getState().isDone()) { return; } for (Entry<PlanNodeId, Collection<Split>> entry : splitsBySource.asMap().entrySet()) { PlanNodeId sourceId = entry.getKey(); Collection<Split> splits = entry.getValue(); checkState(!noMoreSplits.contains(sourceId), "noMoreSplits has already been set for %s", sourceId); checkState(!noMoreSplits.contains(sourceId), "noMoreSplits has already been set for %s", sourceId); int added = 0; for (Split split : splits) { if (pendingSplits.put(sourceId, new ScheduledSplit(nextSplitId.getAndIncrement(), split))) { added++; } } if (sourceId.equals(planFragment.getPartitionedSource())) { pendingSourceSplitCount += added; partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); } needsUpdate.set(true); } scheduleUpdate(); } }
public MockRemoteTask(TaskId taskId, PlanFragment fragment, String nodeId, int partition, Executor executor, Multimap<PlanNodeId, Split> initialSplits, PartitionedSplitCountTracker partitionedSplitCountTracker) { this.taskStateMachine = new TaskStateMachine(requireNonNull(taskId, "taskId is null"), requireNonNull(executor, "executor is null")); MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)); MemoryPool memorySystemPool = new MemoryPool(new MemoryPoolId("testSystem"), new DataSize(1, GIGABYTE)); this.taskContext = new QueryContext(taskId.getQueryId(), new DataSize(1, MEGABYTE), memoryPool, memorySystemPool, executor).addTaskContext(taskStateMachine, TEST_SESSION, new DataSize(1, MEGABYTE), true, true); this.location = URI.create("fake://task/" + taskId); this.sharedBuffer = new SharedBuffer(taskId, TASK_INSTANCE_ID, executor, requireNonNull(new DataSize(1, BYTE), "maxBufferSize is null")); this.fragment = requireNonNull(fragment, "fragment is null"); this.nodeId = requireNonNull(nodeId, "nodeId is null"); checkArgument(partition >= 0, "partition is negative"); this.partition = partition; splits.putAll(initialSplits); this.partitionedSplitCountTracker = requireNonNull(partitionedSplitCountTracker, "partitionedSplitCountTracker is null"); partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); }
@Test public void testSingleNode() throws Exception { FixedCountScheduler nodeScheduler = new FixedCountScheduler( (node, partition) -> taskFactory.createTableScanTask( new TaskId(new StageId("test", "1"), "1"), (Node) node, ImmutableList.of(), new PartitionedSplitCountTracker(delta -> { })), generateRandomNodes(1)); ScheduleResult result = nodeScheduler.schedule(); assertTrue(result.isFinished()); assertTrue(result.getBlocked().isDone()); assertEquals(result.getNewTasks().size(), 1); result.getNewTasks().iterator().next().getNodeId().equals("other 0"); }
@Test public void testMultipleNodes() throws Exception { FixedCountScheduler nodeScheduler = new FixedCountScheduler( (node, partition) -> taskFactory.createTableScanTask( new TaskId(new StageId("test", "1"), "1"), (Node) node, ImmutableList.of(), new PartitionedSplitCountTracker(delta -> { })), generateRandomNodes(5)); ScheduleResult result = nodeScheduler.schedule(); assertTrue(result.isFinished()); assertTrue(result.getBlocked().isDone()); assertEquals(result.getNewTasks().size(), 5); assertEquals(result.getNewTasks().stream().map(RemoteTask::getNodeId).collect(toImmutableSet()).size(), 5); }
public synchronized void clearSplits() { splits.clear(); partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); runningDrivers = 0; }
public PartitionedSplitCountTracker createPartitionedSplitCountTracker(TaskId taskId) { requireNonNull(taskId, "taskId is null"); TaskPartitionedSplitCountTracker tracker = new TaskPartitionedSplitCountTracker(taskId); PartitionedSplitCountTracker partitionedSplitCountTracker = new PartitionedSplitCountTracker(tracker::setPartitionedSplitCount); // when partitionedSplitCountTracker is garbage collected, run the cleanup method on the tracker // Note: tracker can not have a reference to partitionedSplitCountTracker finalizerService.addFinalizer(partitionedSplitCountTracker, tracker::cleanup); return partitionedSplitCountTracker; }
@Override public void addSplits(Multimap<PlanNodeId, Split> splits) { synchronized (this) { this.splits.putAll(splits); } partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); }
splits.putAll(initialSplits); this.partitionedSplitCountTracker = requireNonNull(partitionedSplitCountTracker, "partitionedSplitCountTracker is null"); partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); updateSplitQueueSpace();
partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount());
private synchronized void cleanUpTask() { checkState(getTaskStatus().getState().isDone(), "attempt to clean up a task that is not done yet"); // clear pending splits to free memory pendingSplits.clear(); pendingSourceSplitCount = 0; partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); splitQueueHasSpace = true; whenSplitQueueHasSpace.complete(null, executor); // cancel pending request if (currentRequest != null) { currentRequest.cancel(true); currentRequest = null; currentRequestStartNanos = 0; } taskStatusFetcher.stop(); // The remote task is likely to get a delete from the PageBufferClient first. // We send an additional delete anyway to get the final TaskInfo HttpUriBuilder uriBuilder = getHttpUriBuilder(getTaskStatus()); Request request = prepareDelete() .setUri(uriBuilder.build()) .build(); scheduleAsyncCleanupRequest(createCleanupBackoff(), request, "cleanup"); }
private synchronized void processTaskUpdate(TaskInfo newValue, List<TaskSource> sources) { updateTaskInfo(newValue); // remove acknowledged splits, which frees memory for (TaskSource source : sources) { PlanNodeId planNodeId = source.getPlanNodeId(); int removed = 0; for (ScheduledSplit split : source.getSplits()) { if (pendingSplits.remove(planNodeId, split)) { removed++; } } if (source.isNoMoreSplits()) { noMoreSplits.put(planNodeId, false); } for (Lifespan lifespan : source.getNoMoreSplitsForLifespan()) { pendingNoMoreSplitsForLifespan.remove(planNodeId, lifespan); } if (planFragment.isPartitionedSources(planNodeId)) { pendingSourceSplitCount -= removed; } } updateSplitQueueSpace(); partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); }
@Override public void addSplits(Multimap<PlanNodeId, Split> splits) { synchronized (this) { this.splits.putAll(splits); } partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); updateSplitQueueSpace(); }
@Test public void testMultipleNodes() { FixedCountScheduler nodeScheduler = new FixedCountScheduler( (node, partition, totalPartitions) -> Optional.of(taskFactory.createTableScanTask( new TaskId("test", 1, 1), node, ImmutableList.of(), new PartitionedSplitCountTracker(delta -> {}))), generateRandomNodes(5)); ScheduleResult result = nodeScheduler.schedule(); assertTrue(result.isFinished()); assertTrue(result.getBlocked().isDone()); assertEquals(result.getNewTasks().size(), 5); assertEquals(result.getNewTasks().stream().map(RemoteTask::getNodeId).collect(toImmutableSet()).size(), 5); }
private RemoteTask createRemoteTask(HttpRemoteTaskFactory httpRemoteTaskFactory) { return httpRemoteTaskFactory.createRemoteTask( TEST_SESSION, new TaskId("test", 1, 2), new PrestoNode("node-id", URI.create("http://fake.invalid/"), new NodeVersion("version"), false), TaskTestUtils.PLAN_FRAGMENT, ImmutableMultimap.of(), OptionalInt.empty(), createInitialEmptyOutputBuffers(OutputBuffers.BufferType.BROADCAST), new NodeTaskMap.PartitionedSplitCountTracker(i -> {}), true); }
public synchronized void clearSplits() { splits.clear(); partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); runningDrivers = 0; updateSplitQueueSpace(); }