StackTraceSampleCoordinator { @SuppressWarnings("unchecked") public Future<StackTraceSample> triggerStackTraceSample( ExecutionVertex[] tasksToSample, int numSamples, FiniteDuration 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]; for (int i = 0; i < triggerIds.length; i++) { Execution execution = tasksToSample[i].getCurrentExecutionAttempt(); if (execution != null && execution.getState() == ExecutionState.RUNNING) { triggerIds[i] = execution.getAttemptId(); } else { Promise failedPromise = new scala.concurrent.impl.Promise.DefaultPromise<>() .failure(new IllegalStateException("Task " + tasksToSample[i] .getTaskNameWithSubtaskIndex() + " is not running.")); return failedPromise.future(); } } synchronized (lock) { if (isShutDown) { Promise failedPromise = new scala.concurrent.impl.Promise.DefaultPromise<>() .failure(new IllegalStateException("Shut down")); return failedPromise.future(); } if (timer == null) { timer = new Timer("Stack trace sample coordinator timer"); } int sampleId = sampleIdCounter++; LOG.debug("Triggering stack trace sample {}", sampleId); final PendingStackTraceSample pending = new PendingStackTraceSample( sampleId, triggerIds); long expectedDuration = numSamples * delayBetweenSamples.toMillis(); long discardDelay = expectedDuration + sampleTimeout; TimerTask discardTask = new TimerTask() { @Override public void run() { try { synchronized (lock) { if (!pending.isDiscarded()) { LOG.info("Sample {} expired before completing", pending.getSampleId()); pending.discard(new RuntimeException("Time out")); if (pendingSamples.remove(pending.getSampleId()) != null) { rememberRecentSampleId(pending.getSampleId()); } } } } catch (Throwable t) { LOG.error("Exception while handling sample timeout", t); } } }; pendingSamples.put(sampleId, pending); timer.schedule(discardTask, discardDelay); boolean success = true; try { for (int i = 0; i < tasksToSample.length; i++) { TriggerStackTraceSample msg = new TriggerStackTraceSample( sampleId, triggerIds[i], numSamples, delayBetweenSamples, maxStackTraceDepth); if (!tasksToSample[i].sendMessageToCurrentExecution( msg, triggerIds[i], responseActor)) { success = false; break; } } return pending.getStackTraceSampleFuture(); } finally { if (!success) { pending.discard(new RuntimeException("Failed to trigger sample, " + "because task has been reset.")); pendingSamples.remove(sampleId); rememberRecentSampleId(sampleId); } } } } StackTraceSampleCoordinator(ActorSystem actorSystem, int sampleTimeout); @SuppressWarnings("unchecked") Future<StackTraceSample> triggerStackTraceSample(
			ExecutionVertex[] tasksToSample,
			int numSamples,
			FiniteDuration delayBetweenSamples,
			int maxStackTraceDepth); void cancelStackTraceSample(int sampleId, Exception 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; FiniteDuration delayBetweenSamples = new FiniteDuration(100, TimeUnit.MILLISECONDS); int maxStackTraceDepth = 0; Future<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); verify(vertex).sendMessageToCurrentExecution( eq(expectedMsg), eq(expectedExecutionId), any(AkkaActorGateway.class)); } assertFalse(sampleFuture.isCompleted()); 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) { assertTrue(sampleFuture.isCompleted()); } else { assertFalse(sampleFuture.isCompleted()); } } StackTraceSample sample = sampleFuture.value().get().get(); assertEquals(0, sample.getSampleId()); 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); assertNotNull("Task not found", sampleTraces); assertTrue(traces.equals(sampleTraces)); } assertEquals(0, coord.getNumberOfPendingSamples()); coord.collectStackTraces(0, vertices[0].getCurrentExecutionAttempt().getAttemptId(), traces); }
@Test public void testTriggerStackTraceSampleNotRunningTasks() throws Exception { ExecutionVertex[] vertices = new ExecutionVertex[] { mockExecutionVertex(new ExecutionAttemptID(), ExecutionState.RUNNING, true), mockExecutionVertex(new ExecutionAttemptID(), ExecutionState.DEPLOYING, true) }; Future<StackTraceSample> sampleFuture = coord.triggerStackTraceSample( vertices, 1, new FiniteDuration(100, TimeUnit.MILLISECONDS), 0); assertTrue(sampleFuture.isCompleted()); assertTrue(sampleFuture.failed().isCompleted()); assertTrue(sampleFuture.failed().value().get().get() instanceof IllegalStateException); }
@Test public void testTriggerStackTraceSampleResetRunningTasks() throws Exception { ExecutionVertex[] vertices = new ExecutionVertex[] { mockExecutionVertex(new ExecutionAttemptID(), ExecutionState.RUNNING, true), mockExecutionVertex(new ExecutionAttemptID(), ExecutionState.RUNNING, false) }; Future<StackTraceSample> sampleFuture = coord.triggerStackTraceSample( vertices, 1, new FiniteDuration(100, TimeUnit.MILLISECONDS), 0); assertTrue(sampleFuture.isCompleted()); assertTrue(sampleFuture.failed().isCompleted()); assertTrue(sampleFuture.failed().value().get().get().getCause() instanceof RuntimeException); }