StackTraceSampleCoordinator { @SuppressWarnings("unchecked") public CompletableFuture<StackTraceSample> triggerStackTraceSample( ExecutionVertex[] tasksToSample, int numSamples, Time delayBetweenSamples, int maxStackTraceDepth) { checkNotNull(tasksToSample, "Tasks to sample"); checkArgument(tasksToSample.length >= 1, "No tasks to sample"); checkArgument(numSamples >= 1, "No number of samples"); checkArgument(maxStackTraceDepth >= 0, "Negative maximum stack trace depth"); ExecutionAttemptID[] triggerIds = new ExecutionAttemptID[tasksToSample.length]; Execution[] executions = new Execution[tasksToSample.length]; for (int i = 0; i < triggerIds.length; i++) { Execution execution = tasksToSample[i].getCurrentExecutionAttempt(); if (execution != null && execution.getState() == ExecutionState.RUNNING) { executions[i] = execution; triggerIds[i] = execution.getAttemptId(); } else { return FutureUtils.completedExceptionally(new IllegalStateException("Task " + tasksToSample[i] .getTaskNameWithSubtaskIndex() + " is not running.")); } } synchronized (lock) { if (isShutDown) { return FutureUtils.completedExceptionally(new IllegalStateException("Shut down")); } final int sampleId = sampleIdCounter++; LOG.debug("Triggering stack trace sample {}", sampleId); final PendingStackTraceSample pending = new PendingStackTraceSample( sampleId, triggerIds); long expectedDuration = numSamples * delayBetweenSamples.toMilliseconds(); Time timeout = Time.milliseconds(expectedDuration + sampleTimeout); pendingSamples.put(sampleId, pending); for (Execution execution: executions) { final CompletableFuture<StackTraceSampleResponse> stackTraceSampleFuture = execution.requestStackTraceSample( sampleId, numSamples, delayBetweenSamples, maxStackTraceDepth, timeout); stackTraceSampleFuture.handleAsync( (StackTraceSampleResponse stackTraceSampleResponse, Throwable throwable) -> { if (stackTraceSampleResponse != null) { collectStackTraces( stackTraceSampleResponse.getSampleId(), stackTraceSampleResponse.getExecutionAttemptID(), stackTraceSampleResponse.getSamples()); } else { cancelStackTraceSample(sampleId, throwable); } return null; }, executor); } return pending.getStackTraceSampleFuture(); } } StackTraceSampleCoordinator(Executor executor, long sampleTimeout); @SuppressWarnings("unchecked") CompletableFuture<StackTraceSample> triggerStackTraceSample(
			ExecutionVertex[] tasksToSample,
			int numSamples,
			Time delayBetweenSamples,
			int maxStackTraceDepth); void cancelStackTraceSample(int sampleId, Throwable cause); void shutDown(); void collectStackTraces(
			int sampleId,
			ExecutionAttemptID executionId,
			List<StackTraceElement[]> stackTraces);  }
@Test public void testTriggerStackTraceSample() throws Exception { ExecutionVertex[] vertices = new ExecutionVertex[] { mockExecutionVertex(new ExecutionAttemptID(), ExecutionState.RUNNING, true), mockExecutionVertex(new ExecutionAttemptID(), ExecutionState.RUNNING, true), mockExecutionVertex(new ExecutionAttemptID(), ExecutionState.RUNNING, true), mockExecutionVertex(new ExecutionAttemptID(), ExecutionState.RUNNING, true) }; int numSamples = 1; Time delayBetweenSamples = Time.milliseconds(100L); int maxStackTraceDepth = 0; CompletableFuture<StackTraceSample> sampleFuture = coord.triggerStackTraceSample( vertices, numSamples, delayBetweenSamples, maxStackTraceDepth); for (ExecutionVertex vertex : vertices) { ExecutionAttemptID expectedExecutionId = vertex .getCurrentExecutionAttempt().getAttemptId(); TriggerStackTraceSample expectedMsg = new TriggerStackTraceSample( 0, expectedExecutionId, numSamples, delayBetweenSamples, maxStackTraceDepth); Mockito.verify(vertex.getCurrentExecutionAttempt()) .requestStackTraceSample(Matchers.eq(0), Matchers.eq(numSamples), Matchers.eq(delayBetweenSamples), Matchers.eq(maxStackTraceDepth), Matchers.any(Time.class)); } Assert.assertFalse(sampleFuture.isDone()); StackTraceElement[] stackTraceSample = Thread.currentThread().getStackTrace(); List<StackTraceElement[]> traces = new ArrayList<>(); traces.add(stackTraceSample); traces.add(stackTraceSample); traces.add(stackTraceSample); for (int i = 0; i < vertices.length; i++) { ExecutionAttemptID executionId = vertices[i].getCurrentExecutionAttempt().getAttemptId(); coord.collectStackTraces(0, executionId, traces); if (i == vertices.length - 1) { Assert.assertTrue(sampleFuture.isDone()); } else { Assert.assertFalse(sampleFuture.isDone()); } } StackTraceSample sample = sampleFuture.get(); Assert.assertEquals(0, sample.getSampleId()); Assert.assertTrue(sample.getEndTime() >= sample.getStartTime()); Map<ExecutionAttemptID, List<StackTraceElement[]>> tracesByTask = sample.getStackTraces(); for (ExecutionVertex vertex : vertices) { ExecutionAttemptID executionId = vertex.getCurrentExecutionAttempt().getAttemptId(); List<StackTraceElement[]> sampleTraces = tracesByTask.get(executionId); Assert.assertNotNull("Task not found", sampleTraces); Assert.assertTrue(traces.equals(sampleTraces)); } Assert.assertEquals(0, coord.getNumberOfPendingSamples()); coord.collectStackTraces(0, vertices[0].getCurrentExecutionAttempt().getAttemptId(), traces); }