Examples with FlinkCompletableFuture org.apache.flink.runtime.concurrent.impl.FlinkCompletableFuture used on opensource projects

Example 11 with FlinkCompletableFuture

use of org.apache.flink.runtime.concurrent.impl.FlinkCompletableFuture in project flink by apache.

the class FlinkFutureTest method testAcceptAsync.

@Test(timeout = 10000L)
public void testAcceptAsync() throws ExecutionException, InterruptedException {
    CompletableFuture<Integer> initialFuture = new FlinkCompletableFuture<>();
    final AtomicInteger atomicInteger = new AtomicInteger(0);
    int expectedValue = 42;
    Future<Void> result = initialFuture.thenAcceptAsync(new AcceptFunction<Integer>() {

        @Override
        public void accept(Integer value) {
            atomicInteger.set(value);
        }
    }, executor);
    initialFuture.complete(expectedValue);
    result.get();
    assertEquals(expectedValue, atomicInteger.get());
}

Example 12 with FlinkCompletableFuture

use of org.apache.flink.runtime.concurrent.impl.FlinkCompletableFuture in project flink by apache.

the class SlotPool method internalAllocateSlot.

Future<SimpleSlot> internalAllocateSlot(ScheduledUnit task, ResourceProfile resources, Iterable<TaskManagerLocation> locationPreferences) {
    // (1) do we have a slot available already?
    SlotAndLocality slotFromPool = availableSlots.poll(resources, locationPreferences);
    if (slotFromPool != null) {
        SimpleSlot slot = createSimpleSlot(slotFromPool.slot(), slotFromPool.locality());
        allocatedSlots.add(slot);
        return FlinkCompletableFuture.completed(slot);
    }
    // the request will be completed by a future
    final AllocationID allocationID = new AllocationID();
    final FlinkCompletableFuture<SimpleSlot> future = new FlinkCompletableFuture<>();
    if (resourceManagerGateway == null) {
        // no slot available, and no resource manager connection
        stashRequestWaitingForResourceManager(allocationID, resources, future);
    } else {
        // we have a resource manager connection, so let's ask it for more resources
        requestSlotFromResourceManager(allocationID, future, resources);
    }
    return future;
}

Example 13 with FlinkCompletableFuture

use of org.apache.flink.runtime.concurrent.impl.FlinkCompletableFuture in project flink by apache.

the class BackPressureStatsTrackerTest method testTriggerStackTraceSample.

/** Tests simple statistics with fake stack traces. */
@Test
@SuppressWarnings("unchecked")
public void testTriggerStackTraceSample() throws Exception {
    CompletableFuture<StackTraceSample> sampleFuture = new FlinkCompletableFuture<>();
    StackTraceSampleCoordinator sampleCoordinator = mock(StackTraceSampleCoordinator.class);
    when(sampleCoordinator.triggerStackTraceSample(any(ExecutionVertex[].class), anyInt(), any(Time.class), anyInt())).thenReturn(sampleFuture);
    ExecutionGraph graph = mock(ExecutionGraph.class);
    when(graph.getState()).thenReturn(JobStatus.RUNNING);
    // Same Thread execution context
    when(graph.getFutureExecutor()).thenReturn(new Executor() {

        @Override
        public void execute(Runnable runnable) {
            runnable.run();
        }
    });
    ExecutionVertex[] taskVertices = new ExecutionVertex[4];
    ExecutionJobVertex jobVertex = mock(ExecutionJobVertex.class);
    when(jobVertex.getJobId()).thenReturn(new JobID());
    when(jobVertex.getJobVertexId()).thenReturn(new JobVertexID());
    when(jobVertex.getGraph()).thenReturn(graph);
    when(jobVertex.getTaskVertices()).thenReturn(taskVertices);
    taskVertices[0] = mockExecutionVertex(jobVertex, 0);
    taskVertices[1] = mockExecutionVertex(jobVertex, 1);
    taskVertices[2] = mockExecutionVertex(jobVertex, 2);
    taskVertices[3] = mockExecutionVertex(jobVertex, 3);
    int numSamples = 100;
    Time delayBetweenSamples = Time.milliseconds(100L);
    BackPressureStatsTracker tracker = new BackPressureStatsTracker(sampleCoordinator, 9999, numSamples, delayBetweenSamples);
    // Trigger
    assertTrue("Failed to trigger", tracker.triggerStackTraceSample(jobVertex));
    verify(sampleCoordinator).triggerStackTraceSample(eq(taskVertices), eq(numSamples), eq(delayBetweenSamples), eq(BackPressureStatsTracker.MAX_STACK_TRACE_DEPTH));
    // Trigger again for pending request, should not fire
    assertFalse("Unexpected trigger", tracker.triggerStackTraceSample(jobVertex));
    assertTrue(tracker.getOperatorBackPressureStats(jobVertex).isEmpty());
    verify(sampleCoordinator).triggerStackTraceSample(eq(taskVertices), eq(numSamples), eq(delayBetweenSamples), eq(BackPressureStatsTracker.MAX_STACK_TRACE_DEPTH));
    assertTrue(tracker.getOperatorBackPressureStats(jobVertex).isEmpty());
    // Complete the future
    Map<ExecutionAttemptID, List<StackTraceElement[]>> traces = new HashMap<>();
    for (ExecutionVertex vertex : taskVertices) {
        List<StackTraceElement[]> taskTraces = new ArrayList<>();
        for (int i = 0; i < taskVertices.length; i++) {
            // Traces until sub task index are back pressured
            taskTraces.add(createStackTrace(i <= vertex.getParallelSubtaskIndex()));
        }
        traces.put(vertex.getCurrentExecutionAttempt().getAttemptId(), taskTraces);
    }
    int sampleId = 1231;
    int endTime = 841;
    StackTraceSample sample = new StackTraceSample(sampleId, 0, endTime, traces);
    // Succeed the promise
    sampleFuture.complete(sample);
    assertTrue(tracker.getOperatorBackPressureStats(jobVertex).isDefined());
    OperatorBackPressureStats stats = tracker.getOperatorBackPressureStats(jobVertex).get();
    // Verify the stats
    assertEquals(sampleId, stats.getSampleId());
    assertEquals(endTime, stats.getEndTimestamp());
    assertEquals(taskVertices.length, stats.getNumberOfSubTasks());
    for (int i = 0; i < taskVertices.length; i++) {
        double ratio = stats.getBackPressureRatio(i);
        // Traces until sub task index are back pressured
        assertEquals((i + 1) / ((double) 4), ratio, 0.0);
    }
}

Example 14 with FlinkCompletableFuture

use of org.apache.flink.runtime.concurrent.impl.FlinkCompletableFuture in project flink by apache.

the class TaskManagerLogHandlerTest method testLogFetchingFailure.

@Test
public void testLogFetchingFailure() throws Exception {
    // ========= setup TaskManager =================================================================================
    InstanceID tmID = new InstanceID();
    ResourceID tmRID = new ResourceID(tmID.toString());
    TaskManagerGateway taskManagerGateway = mock(TaskManagerGateway.class);
    when(taskManagerGateway.getAddress()).thenReturn("/tm/address");
    Instance taskManager = mock(Instance.class);
    when(taskManager.getId()).thenReturn(tmID);
    when(taskManager.getTaskManagerID()).thenReturn(tmRID);
    when(taskManager.getTaskManagerGateway()).thenReturn(taskManagerGateway);
    CompletableFuture<BlobKey> future = new FlinkCompletableFuture<>();
    future.completeExceptionally(new IOException("failure"));
    when(taskManagerGateway.requestTaskManagerLog(any(Time.class))).thenReturn(future);
    // ========= setup JobManager ==================================================================================
    ActorGateway jobManagerGateway = mock(ActorGateway.class);
    Object registeredTaskManagersAnswer = new JobManagerMessages.RegisteredTaskManagers(JavaConverters.collectionAsScalaIterableConverter(Collections.singletonList(taskManager)).asScala());
    when(jobManagerGateway.ask(isA(JobManagerMessages.RequestRegisteredTaskManagers$.class), any(FiniteDuration.class))).thenReturn(Future$.MODULE$.successful(registeredTaskManagersAnswer));
    when(jobManagerGateway.ask(isA(JobManagerMessages.getRequestBlobManagerPort().getClass()), any(FiniteDuration.class))).thenReturn(Future$.MODULE$.successful((Object) 5));
    when(jobManagerGateway.ask(isA(JobManagerMessages.RequestTaskManagerInstance.class), any(FiniteDuration.class))).thenReturn(Future$.MODULE$.successful((Object) new JobManagerMessages.TaskManagerInstance(Option.apply(taskManager))));
    when(jobManagerGateway.path()).thenReturn("/jm/address");
    JobManagerRetriever retriever = mock(JobManagerRetriever.class);
    when(retriever.getJobManagerGatewayAndWebPort()).thenReturn(Option.apply(new scala.Tuple2<ActorGateway, Integer>(jobManagerGateway, 0)));
    TaskManagerLogHandler handler = new TaskManagerLogHandler(retriever, ExecutionContext$.MODULE$.fromExecutor(Executors.directExecutor()), Future$.MODULE$.successful("/jm/address"), AkkaUtils.getDefaultClientTimeout(), TaskManagerLogHandler.FileMode.LOG, new Configuration(), false);
    final AtomicReference<String> exception = new AtomicReference<>();
    ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
    when(ctx.write(isA(ByteBuf.class))).thenAnswer(new Answer<Object>() {

        @Override
        public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
            ByteBuf data = invocationOnMock.getArgumentAt(0, ByteBuf.class);
            exception.set(new String(data.array(), ConfigConstants.DEFAULT_CHARSET));
            return null;
        }
    });
    Map<String, String> pathParams = new HashMap<>();
    pathParams.put(TaskManagersHandler.TASK_MANAGER_ID_KEY, tmID.toString());
    Routed routed = mock(Routed.class);
    when(routed.pathParams()).thenReturn(pathParams);
    when(routed.request()).thenReturn(new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/taskmanagers/" + tmID + "/log"));
    handler.respondAsLeader(ctx, routed, jobManagerGateway);
    Assert.assertEquals("Fetching TaskManager log failed.", exception.get());
}

Example 15 with FlinkCompletableFuture

use of org.apache.flink.runtime.concurrent.impl.FlinkCompletableFuture in project flink by apache.

the class StackTraceSampleCoordinatorTest method mockExecutionVertexWithTimeout.

private ExecutionVertex mockExecutionVertexWithTimeout(ExecutionAttemptID executionId, ExecutionState state, ScheduledExecutorService scheduledExecutorService, int timeout) {
    final CompletableFuture<StackTraceSampleResponse> future = new FlinkCompletableFuture<>();
    Execution exec = mock(Execution.class);
    when(exec.getAttemptId()).thenReturn(executionId);
    when(exec.getState()).thenReturn(state);
    when(exec.requestStackTraceSample(anyInt(), anyInt(), any(Time.class), anyInt(), any(Time.class))).thenReturn(future);
    scheduledExecutorService.schedule(new Runnable() {

        @Override
        public void run() {
            future.completeExceptionally(new TimeoutException("Timeout"));
        }
    }, timeout, TimeUnit.MILLISECONDS);
    ExecutionVertex vertex = mock(ExecutionVertex.class);
    when(vertex.getJobvertexId()).thenReturn(new JobVertexID());
    when(vertex.getCurrentExecutionAttempt()).thenReturn(exec);
    return vertex;
}