@Benchmark @OperationsPerInvocation(SPLITS) public Object benchmark(BenchmarkData data) { List<RemoteTask> remoteTasks = ImmutableList.copyOf(data.getTaskMap().values()); Iterator<MockRemoteTaskFactory.MockRemoteTask> finishingTask = Iterators.cycle(data.getTaskMap().values()); Iterator<Split> splits = data.getSplits().iterator(); Set<Split> batch = new HashSet<>(); while (splits.hasNext() || !batch.isEmpty()) { Multimap<Node, Split> assignments = data.getNodeSelector().computeAssignments(batch, remoteTasks).getAssignments(); for (Node node : assignments.keySet()) { MockRemoteTaskFactory.MockRemoteTask remoteTask = data.getTaskMap().get(node); remoteTask.addSplits(ImmutableMultimap.<PlanNodeId, Split>builder() .putAll(new PlanNodeId("sourceId"), assignments.get(node)) .build()); remoteTask.startSplits(MAX_SPLITS_PER_NODE); } if (assignments.size() == batch.size()) { batch.clear(); } else { batch.removeAll(assignments.values()); } while (batch.size() < SPLIT_BATCH_SIZE && splits.hasNext()) { batch.add(splits.next()); } finishingTask.next().finishSplits((int) Math.ceil(MAX_SPLITS_PER_NODE / 50.0)); } return remoteTasks; }
public synchronized void clearSplits() { splits.clear(); partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); runningDrivers = 0; updateSplitQueueSpace(); }
private synchronized void updateSplitQueueSpace() { if (getQueuedPartitionedSplitCount() < 9) { if (!whenSplitQueueHasSpace.isDone()) { whenSplitQueueHasSpace.set(null); } } else { if (whenSplitQueueHasSpace.isDone()) { whenSplitQueueHasSpace = SettableFuture.create(); } } }
boolean summarizeTaskInfo) return new MockRemoteTask(taskId, fragment, node.getNodeIdentifier(), executor, scheduledExecutor, initialSplits, totalPartitions, partitionedSplitCountTracker);
boolean summarizeTaskInfo) return new MockRemoteTask(taskId, fragment, node.getNodeIdentifier(), executor, scheduledExecutor, initialSplits, totalPartitions, partitionedSplitCountTracker);
task++; MockRemoteTaskFactory.MockRemoteTask remoteTask = remoteTaskFactory.createTableScanTask(taskId, node, ImmutableList.copyOf(assignments.get(node)), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId)); remoteTask.startSplits(25); nodeTaskMap.addTask(node, remoteTask); taskMap.put(node, remoteTask);
task++; MockRemoteTaskFactory.MockRemoteTask remoteTask = remoteTaskFactory.createTableScanTask(taskId, node, ImmutableList.copyOf(assignments.get(node)), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId)); remoteTask.startSplits(25); nodeTaskMap.addTask(node, remoteTask); taskMap.put(node, remoteTask);
((MockRemoteTask) stage.getAllTasks().get(0)).clearSplits();
((MockRemoteTask) stage.getAllTasks().get(0)).clearSplits();
splits.putAll(initialSplits); this.partitionedSplitCountTracker = requireNonNull(partitionedSplitCountTracker, "partitionedSplitCountTracker is null"); partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); updateSplitQueueSpace();
splits.putAll(initialSplits); this.partitionedSplitCountTracker = requireNonNull(partitionedSplitCountTracker, "partitionedSplitCountTracker is null"); partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); updateSplitQueueSpace();
@Benchmark @OperationsPerInvocation(SPLITS) public Object benchmark(BenchmarkData data) { List<RemoteTask> remoteTasks = ImmutableList.copyOf(data.getTaskMap().values()); Iterator<MockRemoteTaskFactory.MockRemoteTask> finishingTask = Iterators.cycle(data.getTaskMap().values()); Iterator<Split> splits = data.getSplits().iterator(); Set<Split> batch = new HashSet<>(); while (splits.hasNext() || !batch.isEmpty()) { Multimap<Node, Split> assignments = data.getNodeSelector().computeAssignments(batch, remoteTasks).getAssignments(); for (Node node : assignments.keySet()) { MockRemoteTaskFactory.MockRemoteTask remoteTask = data.getTaskMap().get(node); remoteTask.addSplits(ImmutableMultimap.<PlanNodeId, Split>builder() .putAll(new PlanNodeId("sourceId"), assignments.get(node)) .build()); remoteTask.startSplits(MAX_SPLITS_PER_NODE); } if (assignments.size() == batch.size()) { batch.clear(); } else { batch.removeAll(assignments.values()); } while (batch.size() < SPLIT_BATCH_SIZE && splits.hasNext()) { batch.add(splits.next()); } finishingTask.next().finishSplits((int) Math.ceil(MAX_SPLITS_PER_NODE / 50.0)); } return remoteTasks; }
@Override public void addSplits(Multimap<PlanNodeId, Split> splits) { synchronized (this) { this.splits.putAll(splits); } partitionedSplitCountTracker.setPartitionedSplitCount(getPartitionedSplitCount()); updateSplitQueueSpace(); }
@Override public synchronized int getQueuedPartitionedSplitCount() { if (taskStateMachine.getState().isDone()) { return 0; } return getPartitionedSplitCount() - runningDrivers; } }
@Override public void addStateChangeListener(StateChangeListener<TaskStatus> stateChangeListener) { taskStateMachine.addStateChangeListener(newValue -> stateChangeListener.stateChanged(getTaskStatus())); }
@Override public void start() { taskStateMachine.addStateChangeListener(newValue -> { if (newValue.isDone()) { clearSplits(); } }); }
public synchronized void startSplits(int maxRunning) { runningDrivers = splits.size(); runningDrivers = Math.min(runningDrivers, maxRunning); updateSplitQueueSpace(); }
public synchronized void finishSplits(int splits) { List<Map.Entry<PlanNodeId, Split>> toRemove = new ArrayList<>(); Iterator<Map.Entry<PlanNodeId, Split>> iterator = this.splits.entries().iterator(); while (toRemove.size() < splits && iterator.hasNext()) { toRemove.add(iterator.next()); } for (Map.Entry<PlanNodeId, Split> entry : toRemove) { this.splits.remove(entry.getKey(), entry.getValue()); } updateSplitQueueSpace(); }
@Override public void addFinalTaskInfoListener(StateChangeListener<TaskInfo> stateChangeListener) { AtomicBoolean done = new AtomicBoolean(); StateChangeListener<TaskState> fireOnceStateChangeListener = state -> { if (state.isDone() && done.compareAndSet(false, true)) { stateChangeListener.stateChanged(getTaskInfo()); } }; taskStateMachine.addStateChangeListener(fireOnceStateChangeListener); fireOnceStateChangeListener.stateChanged(taskStateMachine.getState()); }
private synchronized void updateSplitQueueSpace() { if (getQueuedPartitionedSplitCount() < 9) { if (!whenSplitQueueHasSpace.isDone()) { whenSplitQueueHasSpace.set(null); } } else { if (whenSplitQueueHasSpace.isDone()) { whenSplitQueueHasSpace = SettableFuture.create(); } } }