akka.pattern.AskTimeoutException Java Examples

/**
 * Retrieve for response of a query command and limit the response according to a policy enforcer.
 *
 * @param commandWithReadSubjects the command to ask.
 * @param enforcer enforcer to build JsonView with.
 * @return always {@code true}.
 */
private CompletionStage<WithDittoHeaders> askThingsShardRegionAndBuildJsonView(
        final ThingQueryCommand commandWithReadSubjects, final Enforcer enforcer) {

    return Patterns.ask(thingsShardRegion, commandWithReadSubjects, getAskTimeout())
            .handle((response, error) -> {
                if (response instanceof ThingQueryCommandResponse) {
                    return reportJsonViewForThingQuery((ThingQueryCommandResponse) response, enforcer);
                } else if (response instanceof DittoRuntimeException) {
                    return (DittoRuntimeException) response;
                } else if (isAskTimeoutException(response, error)) {
                    final AskTimeoutException askTimeoutException = error instanceof AskTimeoutException
                            ? (AskTimeoutException) error
                            : (AskTimeoutException) response;
                    return reportTimeoutForThingQuery(commandWithReadSubjects, askTimeoutException);
                } else if (error != null) {
                    return reportUnexpectedError("before building JsonView", error);
                } else {
                    return reportUnknownResponse("before building JsonView", response);
                }
            });
}
 

private CompletionStage<WithDittoHeaders> askPoliciesShardRegionAndBuildJsonView(
        final PolicyQueryCommand commandWithReadSubjects,
        final Enforcer enforcer) {

    return Patterns.ask(policiesShardRegion, commandWithReadSubjects, getAskTimeout())
            .handle((response, error) -> {
                if (response instanceof PolicyQueryCommandResponse) {
                    return reportJsonViewForPolicyQuery((PolicyQueryCommandResponse<?>) response, enforcer);
                } else if (response instanceof DittoRuntimeException) {
                    throw (DittoRuntimeException) response;
                } else if (isAskTimeoutException(response, error)) {
                    throw reportTimeoutForPolicyQuery(commandWithReadSubjects, (AskTimeoutException) response);
                } else if (error != null) {
                    throw reportUnexpectedError("before building JsonView", error);
                } else {
                    throw reportUnknownResponse("before building JsonView", response);
                }
            });
}
 

@Test
public void timeout() {
    DittoTestSystem.run(this, kit -> {
        // GIVEN: SignalEnrichmentFacade.retrievePartialThing() with a short timeout
        final SignalEnrichmentFacade underTest =
                createSignalEnrichmentFacadeUnderTest(kit, Duration.ofMillis(1L));
        final ThingId thingId = ThingId.dummy();
        final DittoHeaders headers = DittoHeaders.newBuilder().correlationId(UUID.randomUUID().toString()).build();
        final CompletionStage<JsonObject> askResult =
                underTest.retrievePartialThing(thingId, SELECTOR, headers, THING_EVENT);

        // WHEN: Command handler does not respond
        final RetrieveThing retrieveThing = kit.expectMsgClass(RetrieveThing.class);
        assertThat(retrieveThing.getSelectedFields()).contains(actualSelectedFields(SELECTOR));

        // THEN: The result future fails with an AskTimeoutException.
        askResult.toCompletableFuture().exceptionally(e -> null).join();
        assertThat(askResult).hasFailedWithThrowableThat()
                .isInstanceOf(AskTimeoutException.class);
    });
}
 

/**
 * Handles AskTimeoutException. Created to encapsulate errors with more detail than akka.pattern.AskTimeoutException
 *
 * @param ex the AskTimeoutException
 * @return the ApiError object
 */
@ExceptionHandler(AskTimeoutException.class)
protected ResponseEntity<Object> handleAkkaAskTimeout(
		AskTimeoutException ex) {
	ApiError apiError = new ApiError(NOT_FOUND);
	apiError.setMessage(ex.getMessage());
	return buildResponseEntity(apiError);
}
 

/**
 * Report timeout of {@code ThingQueryCommand}.
 *
 * @param command the original command.
 * @param askTimeoutException the timeout exception.
 */
private ThingUnavailableException reportTimeoutForThingQuery(final ThingQueryCommand command,
        final AskTimeoutException askTimeoutException) {

    log(command).error(askTimeoutException, "Timeout before building JsonView");
    return ThingUnavailableException.newBuilder(command.getThingEntityId())
            .dittoHeaders(command.getDittoHeaders())
            .build();
}
 

private PolicyUnavailableException reportTimeoutForPolicyQuery(
        final PolicyQueryCommand command,
        final AskTimeoutException askTimeoutException) {
    log(command).error(askTimeoutException, "Timeout before building JsonView");
    return PolicyUnavailableException.newBuilder(command.getEntityId())
            .dittoHeaders(command.getDittoHeaders())
            .build();
}
 

@Test
public void retrieveThingErrorOnTimeout() {
    new TestKit(system) {{
        final ActorRef underTest = newEnforcerActor(getRef());

        // GIVEN: Thing and its Policy exist in cache
        final ThingId thingId = ThingId.of("thing", UUID.randomUUID().toString());
        final PolicyId policyId = PolicyId.of("policy", UUID.randomUUID().toString());
        final Thing thing = emptyThing(thingId, policyId.toString());
        final SudoRetrieveThingResponse sudoRetrieveThingResponse =
                SudoRetrieveThingResponse.of(thing.toJson(FieldType.all()), DittoHeaders.empty());
        final Policy policy = defaultPolicy(policyId);
        final SudoRetrievePolicyResponse sudoRetrievePolicyResponse =
                SudoRetrievePolicyResponse.of(policyId, policy, DittoHeaders.empty());
        mockThingsActor.underlyingActor()
                .setReply(THING_SUDO, sudoRetrieveThingResponse)
                .setReply(THING, new AskTimeoutException("thing timeout"));
        mockPoliciesActor.underlyingActor()
                .setReply(POLICY_SUDO, sudoRetrievePolicyResponse)
                .setReply(POLICY, new AskTimeoutException("policy timeout"));

        // WHEN: received RetrieveThing but both shard regions time out
        final JsonFieldSelector selectedFields = JsonFieldSelector.newInstance("_policy", "thingId");
        final RetrieveThing retrieveThing = RetrieveThing.getBuilder(thingId, DEFAULT_HEADERS)
                .withSelectedFields(selectedFields)
                .build();

        underTest.tell(retrieveThing, getRef());

        // THEN: initial requester receives error
        fishForMsgClass(this, ThingUnavailableException.class);

        // WHEN: Thing is responsive but Policy times out
        mockThingsActor.underlyingActor()
                .setReply(THING, RetrieveThingResponse.of(thingId, thing, DEFAULT_HEADERS));
        underTest.tell(retrieveThing, getRef());

        // THEN: initial requester receives Thing without Policy
        final RetrieveThingResponse response = expectMsgClass(RetrieveThingResponse.class);
        assertThat((CharSequence) response.getThingEntityId()).isEqualTo(thingId);
        assertThat(response.getEntity()).isObject();
        assertThat(response.getEntity().asObject()).doesNotContain(JsonKey.of("_policy"));
    }};
}
 

public Future<Object> queryStateFuture(final QueryState<K> queryState) {
	LOG.debug("Try to get ActorRef future for key {}.", queryState.getKey());

	Future<ActorRef> actorRefFuture = getActorRefFuture(queryState.getKey());

	@SuppressWarnings("unchecked")
	Future<Object> result =  actorRefFuture.flatMap(new Mapper<ActorRef, Future<Object>>() {
		public Future<Object> apply(ActorRef actorRef) {
			LOG.debug("Ask response actor for state for key {}.", queryState.getKey());
			return Patterns.ask(actorRef, queryState, new Timeout(askTimeout));
		}
	}, executor).recoverWith(new Recover<Future<Object>>() {
		@Override
		public Future<Object> recover(final Throwable failure) throws Throwable {
			if (failure instanceof WrongKeyPartitionException || failure instanceof ActorNotFound) {
				// wait askTimeout because we communicated with the wrong actor. This usually
				// indicates that not all actors have registered at the registry.
				return Patterns.after(
					askTimeout,
					getContext().system().scheduler(),
					executor,
					new Callable<Future<Object>>() {
						@Override
						public Future<Object> call() throws Exception {
							refreshCache();
							return Futures.failed(failure);
						}
					});
			} else if (failure instanceof AskTimeoutException) {
				LOG.debug("Ask timed out.", failure);
				handleAskTimeout();
				return Futures.failed(failure);
			} else {
				LOG.debug("State query failed with.", failure);
				refreshCache();
				return Futures.failed(failure);
			}
		}
	}, executor);

	return result;
}
 

@Test
public void testConsumerRecreationFailureWhenConnected() throws JMSException {
    new TestKit(actorSystem) {{
        final Props props =
                AmqpClientActor.propsForTests(singleConsumerConnection(), getRef(), getRef(),
                        (ac, el) -> mockConnection);
        final TestActorRef<AmqpClientActor> amqpClientActorRef = TestActorRef.apply(props, actorSystem);
        final AmqpClientActor amqpClientActor = amqpClientActorRef.underlyingActor();

        amqpClientActorRef.tell(OpenConnection.of(CONNECTION_ID, DittoHeaders.empty()), getRef());
        expectMsg(CONNECTED_SUCCESS);

        // GIVEN: JMS session fails, but the JMS connection can create a new functional session
        final Session mockSession2 = Mockito.mock(Session.class);
        final MessageConsumer mockConsumer2 = Mockito.mock(JmsMessageConsumer.class);
        when(mockSession.createConsumer(any()))
                .thenThrow(new IllegalStateException("expected exception"));
        doReturn(mockSession2).when(mockConnection).createSession(anyInt());
        doReturn(mockConsumer2).when(mockSession2).createConsumer(any());

        // WHEN: consumer is closed and cannot be recreated
        final ActorRef amqpConsumerActor = amqpClientActor.context().children().toStream()
                .find(child -> child.path().name().startsWith(AmqpConsumerActor.ACTOR_NAME_PREFIX))
                .get();
        final Throwable error = new IllegalStateException("Forcibly detached");
        final Status.Failure failure = new Status.Failure(new AskTimeoutException("Consumer creation timeout"));
        amqpClientActor.connectionListener.onConsumerClosed(mockConsumer, error);
        verify(mockSession, atLeastOnce()).createConsumer(any());
        amqpConsumerActor.tell(failure, amqpConsumerActor);

        // THEN: connection gets restarted
        verify(mockConnection).createSession(anyInt());
        final ArgumentCaptor<MessageListener> captor = ArgumentCaptor.forClass(MessageListener.class);
        verify(mockConsumer2, timeout(1000).atLeastOnce()).setMessageListener(captor.capture());
        final MessageListener messageListener = captor.getValue();

        // THEN: recreated connection is working
        messageListener.onMessage(mockMessage());
        expectMsgClass(Command.class);
    }};
}
 

public static OneForOneStrategy createStrategy(final LoggingAdapter log) {
    return new OneForOneStrategy(true, DeciderBuilder
            .match(NullPointerException.class, e -> {
                log.error(e, "NullPointer in child actor: {}", e.getMessage());
                log.info(RESTARTING_CHILD_MSG);
                return SupervisorStrategy.restart();
            }).match(IllegalArgumentException.class, e -> {
                log.warning("Illegal Argument in child actor: {}", e.getMessage());
                return SupervisorStrategy.resume();
            }).match(IndexOutOfBoundsException.class, e -> {
                log.warning("IndexOutOfBounds in child actor: {}", e.getMessage());
                return SupervisorStrategy.resume();
            }).match(IllegalStateException.class, e -> {
                log.warning("Illegal State in child actor: {}", e.getMessage());
                return SupervisorStrategy.resume();
            }).match(NoSuchElementException.class, e -> {
                log.warning("NoSuchElement in child actor: {}", e.getMessage());
                return SupervisorStrategy.resume();
            }).match(AskTimeoutException.class, e -> {
                log.warning("AskTimeoutException in child actor: {}", e.getMessage());
                return SupervisorStrategy.resume();
            }).match(ConnectException.class, e -> {
                log.warning("ConnectException in child actor: {}", e.getMessage());
                log.info(RESTARTING_CHILD_MSG);
                return SupervisorStrategy.restart();
            }).match(InvalidActorNameException.class, e -> {
                log.warning("InvalidActorNameException in child actor: {}", e.getMessage());
                return SupervisorStrategy.resume();
            }).match(ActorKilledException.class, e -> {
                log.error(e, "ActorKilledException in child actor: {}", e.message());
                log.info(RESTARTING_CHILD_MSG);
                return SupervisorStrategy.restart();
            }).match(DittoRuntimeException.class, e -> {
                log.error(e,
                        "DittoRuntimeException '{}' should not be escalated to RootActor. Simply resuming Actor.",
                        e.getErrorCode());
                return SupervisorStrategy.resume();
            }).match(Throwable.class, e -> {
                log.error(e, "Escalating above root actor!");
                return SupervisorStrategy.escalate();
            }).matchAny(e -> {
                log.error("Unknown message:'{}'! Escalating above root actor!", e);
                return SupervisorStrategy.escalate();
            }).build());
}
 

private void becomeStatisticsDetailsAwaiting(final Collection<ClusterRoleStatus> relevantRoles,
        final Consumer<StatisticsDetails> statisticsDetailsConsumer) {

    final Map<String, ShardStatisticsWrapper> shardStatisticsMap = new HashMap<>();
    final AskTimeoutException askTimeoutException = new AskTimeoutException("Timed out");

    getTimers().startSingleTimer(askTimeoutException, askTimeoutException, statisticsConfig.getAskTimeout());

    final int reachableMembers = relevantRoles.stream()
            .mapToInt(clusterRoleStatus -> clusterRoleStatus.getReachable().size())
            .sum();

    final ShardStatisticsWrapper messageCounter = new ShardStatisticsWrapper();
    messageCounter.count = 0L;

    getContext().become(ReceiveBuilder.create()
                    .match(RetrieveStatistics.class, this::respondWithCachedStatistics)
                    .match(RetrieveStatisticsDetailsResponse.class, retrieveStatisticsDetailsResponse -> {
                        final String shardRegion = retrieveStatisticsDetailsResponse.getStatisticsDetails()
                                .stream()
                                .findFirst()
                                .map(JsonField::getKeyName)
                                .orElse(null);

                        if (shardRegion != null) {
                            final ShardStatisticsWrapper shardStatistics =
                                    shardStatisticsMap.computeIfAbsent(shardRegion, s -> new ShardStatisticsWrapper());

                            retrieveStatisticsDetailsResponse.getStatisticsDetails().getValue(shardRegion)
                                    .map(JsonValue::asObject)
                                    .map(JsonObject::stream)
                                    .ifPresent(namespaceEntries -> namespaceEntries.forEach(field ->
                                            shardStatistics.hotnessMap
                                                    .merge(field.getKeyName(), field.getValue().asLong(), Long::sum)
                                    ));

                            // all reachable members sent reply; stop waiting.
                            if (++messageCounter.count >= reachableMembers) {
                                getTimers().cancel(askTimeoutException);
                                getSelf().tell(askTimeoutException, getSelf());
                            }
                        }
                    })
                    .match(AskTimeoutException.class, askTimeout -> {
                        currentStatisticsDetails = StatisticsDetails.fromShardStatisticsWrappers(shardStatisticsMap);
                        statisticsDetailsConsumer.accept(currentStatisticsDetails);
                        scheduleResetStatisticsDetails();
                        unbecome();
                    })
                    .matchEquals(InternalResetStatisticsDetails.INSTANCE, this::resetStatisticsDetails)
                    .match(DistributedPubSubMediator.SubscribeAck.class, this::logSubscribeAck)
                    .matchAny(m -> {
                        log.info("Stashing message during 'statisticsDetailsAwaiting': {}", m);
                        stash();
                    })
                    .build(),
            false);
}
 

private void becomeStatisticsAwaiting() {

        final Map<String, ShardStatisticsWrapper> shardStatisticsMap = new HashMap<>();

        final AskTimeoutException askTimeoutException = new AskTimeoutException("Timed out");
        getTimers().startSingleTimer(askTimeoutException, askTimeoutException, statisticsConfig.getAskTimeout());

        getContext().become(ReceiveBuilder.create()
                        .match(RetrieveStatistics.class, this::respondWithCachedStatistics)
                        .match(ShardRegion.ClusterShardingStats.class, clusterShardingStats -> {
                            final Optional<ShardStatisticsWrapper> shardStatistics =
                                    getShardStatistics(shardStatisticsMap, getSender());

                            if (shardStatistics.isPresent()) {
                                final Map<Address, ShardRegion.ShardRegionStats> regions = clusterShardingStats.getRegions();
                                shardStatistics.get().count = regions.isEmpty() ? 0 : regions.values().stream()
                                        .mapToInt(shardRegionStats -> shardRegionStats.getStats().isEmpty() ? 0 :
                                                shardRegionStats.getStats().values().stream()
                                                        .mapToInt(o -> (Integer) o)
                                                        .sum())
                                        .sum();
                            } else {
                                log.warning("Got stats from unknown shard <{}>: <{}>", getSender(), clusterShardingStats);
                            }

                            // all shard statistics are present; no need to wait more.
                            if (shardStatisticsMap.size() >= statisticsConfig.getShards().size()) {
                                getTimers().cancel(askTimeoutException);
                                getSelf().tell(askTimeoutException, getSelf());
                            }
                        })
                        .matchEquals(askTimeoutException, unit -> {
                            updateGauges(shardStatisticsMap);
                            currentStatistics = Statistics.fromGauges(gauges);
                            unbecome();
                        })
                        .matchEquals(InternalResetStatisticsDetails.INSTANCE, this::resetStatisticsDetails)
                        .match(DistributedPubSubMediator.SubscribeAck.class, this::logSubscribeAck)
                        .matchAny(m -> {
                            log.info("Stashing message during 'statisticsAwaiting': {}", m);
                            stash();
                        })
                        .build(),
                false);
    }
 

private OneForOneStrategy buildDefaultSupervisorStrategy() {
    return new OneForOneStrategy(true, DeciderBuilder
            .match(NullPointerException.class, e -> {
                log.error(e, "NullPointer in singleton actor: {}", e.getMessage());
                return restartChild();
            }).match(IllegalArgumentException.class, e -> {
                log.warning("Illegal Argument in singleton actor: {}", e.getMessage());

                StringWriter sw = new StringWriter();
                PrintWriter pw = new PrintWriter(sw);
                e.printStackTrace(pw);

                log.warning("Illegal Argument in singleton actor: {}", sw.toString());
                return SupervisorStrategy.resume();
            }).match(IllegalStateException.class, e -> {
                log.warning("Illegal State in singleton actor: {}", e.getMessage());
                return SupervisorStrategy.resume();
            }).match(IndexOutOfBoundsException.class, e -> {
                log.warning("IndexOutOfBounds in singleton actor: {}", e.getMessage());
                return SupervisorStrategy.resume();
            }).match(NoSuchElementException.class, e -> {
                log.warning("NoSuchElement in singleton actor: {}", e.getMessage());
                return SupervisorStrategy.resume();
            }).match(AskTimeoutException.class, e -> {
                log.warning("AskTimeoutException in singleton actor: {}", e.getMessage());
                return SupervisorStrategy.resume();
            }).match(ConnectException.class, e -> {
                log.warning("ConnectException in singleton actor: {}", e.getMessage());
                return restartChild();
            }).match(InvalidActorNameException.class, e -> {
                log.warning("InvalidActorNameException in singleton actor: {}", e.getMessage());
                return SupervisorStrategy.resume();
            }).match(ActorKilledException.class, e -> {
                log.error(e, "ActorKilledException in singleton actor: {}", e.message());
                return restartChild();
            }).match(DittoRuntimeException.class, e -> {
                log.error(e,
                        "DittoRuntimeException '{}' should not be escalated to SupervisorActor. Simply resuming Actor.",
                        e.getErrorCode());
                return SupervisorStrategy.resume();
            }).match(UnsupportedOperationException.class, e -> {
                log.error(e, "UnsupportedOperationException in singleton actor: {}",
                        e.getMessage());
                terminateActorSystem();
                return SupervisorStrategy.stop(); // only stopping as ActorSystem is terminated anyways
            }).match(Throwable.class, e -> {
                log.error(e, "Escalating above root actor!");
                terminateActorSystem();
                return SupervisorStrategy.stop(); // only stopping as ActorSystem is terminated anyways
            }).matchAny(e -> {
                log.error("Unknown message:'{}'! Escalating above root actor!", e);
                terminateActorSystem();
                return SupervisorStrategy.stop(); // only stopping as ActorSystem is terminated anyways
            }).build());
}
 

@Test
public void parallelSessions() {
    // Suppressing logs due to stacktrace for test4. Comment out to see logs.
    actorSystem.eventStream().setLogLevel(Attributes.logLevelOff());

    final ActorRef underTest = createSubscriptionManager();

    final TestProbe probe1 = TestProbe.apply("test1", actorSystem);
    final TestProbe probe2 = TestProbe.apply("test2", actorSystem);
    final TestProbe probe3 = TestProbe.apply("test3", actorSystem);
    final TestProbe probe4 = TestProbe.apply("test4", actorSystem);

    // expect success with results
    underTest.tell(createSubscription(1), probe1.ref());
    // expect success with prefix
    final String prefix = "ABCDEFGHIJKLMNOPQRSTUVWXYZ";
    underTest.tell(createSubscription(2).setPrefix(prefix), probe2.ref());
    // expect cancellation
    underTest.tell(createSubscription(3), probe3.ref());
    // expect upstream failure
    underTest.tell(createSubscription(4), probe4.ref());

    final String sid1 = probe1.expectMsgClass(SubscriptionCreated.class).getSubscriptionId();
    final String sid2 = probe2.expectMsgClass(SubscriptionCreated.class).getSubscriptionId();
    final String sid3 = probe3.expectMsgClass(SubscriptionCreated.class).getSubscriptionId();
    final String sid4 = probe4.expectMsgClass(SubscriptionCreated.class).getSubscriptionId();

    final List<?> sources = List.of(
            Source.single("t:1"),
            Source.single("t:2"),
            Source.from(List.of("t:3", "t:4")),
            new AskTimeoutException("mock error")
    );

    underTest.tell(request(sid1), probe1.ref());
    underTest.tell(request(sid2), probe2.ref());
    underTest.tell(request(sid3), probe3.ref());
    underTest.tell(request(sid4), probe4.ref());

    // there should be no upstream request until downstream requests.
    for (int i = 0; i < 4; ++i) {
        final StreamThings streamThings = conciergeForwarderProbe.expectMsgClass(StreamThings.class);
        final ActorRef sender = conciergeForwarderProbe.sender();
        final Object source = sources.get(getTag(streamThings) - 1);
        if (source instanceof Source) {
            ((Source<?, ?>) source).runWith(StreamRefs.sourceRef(), materializer)
                    .thenAccept(sourceRef -> sender.tell(sourceRef, ActorRef.noSender()));
        } else {
            sender.tell(source, ActorRef.noSender());
        }
    }

    probe1.expectMsg(hasNext(sid1, "t:1"));
    probe1.expectMsg(subscriptionComplete(sid1));

    probe2.expectMsg(hasNext(sid2, "t:2"));
    probe2.expectMsg(subscriptionComplete(sid2));

    probe3.expectMsg(hasNext(sid3, "t:3"));
    underTest.tell(CancelSubscription.of(sid3, DittoHeaders.empty()), probe3.ref());

    assertThat(probe4.expectMsgClass(SubscriptionFailed.class).getError())
            .isInstanceOf(GatewayInternalErrorException.class);

    CompletableFuture.allOf(
            CompletableFuture.runAsync(
                    () -> probe3.expectNoMessage(FiniteDuration.apply(250L, TimeUnit.MILLISECONDS))),
            CompletableFuture.runAsync(
                    () -> probe4.expectNoMessage(FiniteDuration.apply(250L, TimeUnit.MILLISECONDS)))
    ).join();
}
 

/**
 * Test delayed registration of task executor where the delay is introduced during connection from resource manager
 * to the registering task executor.
 */
@Test
public void testDelayedRegisterTaskExecutor() throws Exception {
	final Time fastTimeout = Time.milliseconds(1L);
	try {
		final OneShotLatch startConnection = new OneShotLatch();
		final OneShotLatch finishConnection = new OneShotLatch();

		// first registration is with blocking connection
		rpcService.setRpcGatewayFutureFunction(rpcGateway ->
			CompletableFuture.supplyAsync(
				() -> {
					startConnection.trigger();
					try {
						finishConnection.await();
					} catch (InterruptedException ignored) {}
					return rpcGateway;
				},
				TestingUtils.defaultExecutor()));

		CompletableFuture<RegistrationResponse> firstFuture =
			rmGateway.registerTaskExecutor(taskExecutorGateway.getAddress(), taskExecutorResourceID, dataPort, hardwareDescription, fastTimeout);
		try {
			firstFuture.get();
			fail("Should have failed because connection to taskmanager is delayed beyond timeout");
		} catch (Exception e) {
			assertThat(ExceptionUtils.stripExecutionException(e), instanceOf(AskTimeoutException.class));
		}

		startConnection.await();

		// second registration after timeout is with no delay, expecting it to be succeeded
		rpcService.resetRpcGatewayFutureFunction();
		CompletableFuture<RegistrationResponse> secondFuture =
			rmGateway.registerTaskExecutor(taskExecutorGateway.getAddress(), taskExecutorResourceID, dataPort, hardwareDescription, TIMEOUT);
		RegistrationResponse response = secondFuture.get();
		assertTrue(response instanceof TaskExecutorRegistrationSuccess);

		// on success, send slot report for taskmanager registration
		final SlotReport slotReport = new SlotReport(new SlotStatus(new SlotID(taskExecutorResourceID, 0), ResourceProfile.UNKNOWN));
		rmGateway.sendSlotReport(taskExecutorResourceID,
			((TaskExecutorRegistrationSuccess) response).getRegistrationId(), slotReport, TIMEOUT).get();

		// let the remaining part of the first registration proceed
		finishConnection.trigger();
		Thread.sleep(1L);

		// verify that the latest registration is valid not being unregistered by the delayed one
		final TaskManagerInfo taskManagerInfo = rmGateway.requestTaskManagerInfo(
			taskExecutorResourceID,
			TIMEOUT).get();
		assertThat(taskManagerInfo.getResourceId(), equalTo(taskExecutorResourceID));
		assertThat(taskManagerInfo.getNumberSlots(), equalTo(1));
	} finally {
		rpcService.resetRpcGatewayFutureFunction();
	}
}
 

/**
 * Test delayed registration of task executor where the delay is introduced during connection from resource manager
 * to the registering task executor.
 */
@Test
public void testDelayedRegisterTaskExecutor() throws Exception {
	final Time fastTimeout = Time.milliseconds(1L);
	try {
		final OneShotLatch startConnection = new OneShotLatch();
		final OneShotLatch finishConnection = new OneShotLatch();

		// first registration is with blocking connection
		rpcService.setRpcGatewayFutureFunction(rpcGateway ->
			CompletableFuture.supplyAsync(
				() -> {
					startConnection.trigger();
					try {
						finishConnection.await();
					} catch (InterruptedException ignored) {}
					return rpcGateway;
				},
				TestingUtils.defaultExecutor()));

		CompletableFuture<RegistrationResponse> firstFuture =
			rmGateway.registerTaskExecutor(taskExecutorGateway.getAddress(), taskExecutorResourceID, dataPort, hardwareDescription, fastTimeout);
		try {
			firstFuture.get();
			fail("Should have failed because connection to taskmanager is delayed beyond timeout");
		} catch (Exception e) {
			assertThat(ExceptionUtils.stripExecutionException(e), instanceOf(AskTimeoutException.class));
		}

		startConnection.await();

		// second registration after timeout is with no delay, expecting it to be succeeded
		rpcService.resetRpcGatewayFutureFunction();
		CompletableFuture<RegistrationResponse> secondFuture =
			rmGateway.registerTaskExecutor(taskExecutorGateway.getAddress(), taskExecutorResourceID, dataPort, hardwareDescription, TIMEOUT);
		RegistrationResponse response = secondFuture.get();
		assertTrue(response instanceof TaskExecutorRegistrationSuccess);

		// on success, send slot report for taskmanager registration
		final SlotReport slotReport = new SlotReport(new SlotStatus(new SlotID(taskExecutorResourceID, 0), ResourceProfile.UNKNOWN));
		rmGateway.sendSlotReport(taskExecutorResourceID,
			((TaskExecutorRegistrationSuccess) response).getRegistrationId(), slotReport, TIMEOUT).get();

		// let the remaining part of the first registration proceed
		finishConnection.trigger();
		Thread.sleep(1L);

		// verify that the latest registration is valid not being unregistered by the delayed one
		final TaskManagerInfo taskManagerInfo = rmGateway.requestTaskManagerInfo(
			taskExecutorResourceID,
			TIMEOUT).get();
		assertThat(taskManagerInfo.getResourceId(), equalTo(taskExecutorResourceID));
		assertThat(taskManagerInfo.getNumberSlots(), equalTo(1));
	} finally {
		rpcService.resetRpcGatewayFutureFunction();
	}
}
 

/**
 * Check whether response or error from a future is {@code AskTimeoutException}.
 *
 * @param response response from a future.
 * @param error error thrown in a future.
 * @return whether either is {@code AskTimeoutException}.
 */
protected static boolean isAskTimeoutException(final Object response, @Nullable final Throwable error) {
    return error instanceof AskTimeoutException || response instanceof AskTimeoutException;
}