org.apache.flink.runtime.rpc.RpcGateway Java Examples

@Override
public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
	Class<?> declaringClass = method.getDeclaringClass();

	Object result;

	if (declaringClass.equals(AkkaBasedEndpoint.class) ||
		declaringClass.equals(Object.class) ||
		declaringClass.equals(RpcGateway.class) ||
		declaringClass.equals(StartStoppable.class) ||
		declaringClass.equals(MainThreadExecutable.class) ||
		declaringClass.equals(RpcServer.class)) {
		result = method.invoke(this, args);
	} else if (declaringClass.equals(FencedRpcGateway.class)) {
		throw new UnsupportedOperationException("AkkaInvocationHandler does not support the call FencedRpcGateway#" +
			method.getName() + ". This indicates that you retrieved a FencedRpcGateway without specifying a " +
			"fencing token. Please use RpcService#connect(RpcService, F, Time) with F being the fencing token to " +
			"retrieve a properly FencedRpcGateway.");
	} else {
		result = invokeRpc(method, args);
	}

	return result;
}
 

@Override
public <C extends RpcGateway> CompletableFuture<C> connect(
		final String address,
		final Class<C> clazz) {

	return connectInternal(
		address,
		clazz,
		(ActorRef actorRef) -> {
			Tuple2<String, String> addressHostname = extractAddressHostname(actorRef);

			return new AkkaInvocationHandler(
				addressHostname.f0,
				addressHostname.f1,
				actorRef,
				configuration.getTimeout(),
				configuration.getMaximumFramesize(),
				null);
		});
}
 

@Override
public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
	Class<?> declaringClass = method.getDeclaringClass();

	Object result;

	if (declaringClass.equals(AkkaBasedEndpoint.class) ||
		declaringClass.equals(Object.class) ||
		declaringClass.equals(RpcGateway.class) ||
		declaringClass.equals(StartStoppable.class) ||
		declaringClass.equals(MainThreadExecutable.class) ||
		declaringClass.equals(RpcServer.class)) {
		result = method.invoke(this, args);
	} else if (declaringClass.equals(FencedRpcGateway.class)) {
		throw new UnsupportedOperationException("AkkaInvocationHandler does not support the call FencedRpcGateway#" +
			method.getName() + ". This indicates that you retrieved a FencedRpcGateway without specifying a " +
			"fencing token. Please use RpcService#connect(RpcService, F, Time) with F being the fencing token to " +
			"retrieve a properly FencedRpcGateway.");
	} else {
		result = invokeRpc(method, args);
	}

	return result;
}
 

@Override
public <C extends RpcGateway> CompletableFuture<C> connect(
		final String address,
		final Class<C> clazz) {

	return connectInternal(
		address,
		clazz,
		(ActorRef actorRef) -> {
			Tuple2<String, String> addressHostname = extractAddressHostname(actorRef);

			return new AkkaInvocationHandler(
				addressHostname.f0,
				addressHostname.f1,
				actorRef,
				configuration.getTimeout(),
				configuration.getMaximumFramesize(),
				null);
		});
}
 

@Override
public <C extends RpcGateway> CompletableFuture<C> connect(
		final String address,
		final Class<C> clazz) {

	return connectInternal(
		address,
		clazz,
		(ActorRef actorRef) -> {
			Tuple2<String, String> addressHostname = extractAddressHostname(actorRef);

			return new AkkaInvocationHandler(
				addressHostname.f0,
				addressHostname.f1,
				actorRef,
				configuration.getTimeout(),
				configuration.getMaximumFramesize(),
				null,
				captureAskCallstacks);
		});
}
 

@Test
public void resolvesRunningAkkaRpcActor() throws Exception {
	final String endpointName = "foobar";

	try (RpcEndpoint simpleRpcEndpoint1 = createRpcEndpointWithRandomNameSuffix(endpointName);
		RpcEndpoint simpleRpcEndpoint2 = createRpcEndpointWithRandomNameSuffix(endpointName)) {

		simpleRpcEndpoint1.closeAsync().join();

		final String wildcardName = AkkaRpcServiceUtils.createWildcardName(endpointName);
		final String wildcardAddress = AkkaRpcServiceUtils.getLocalRpcUrl(wildcardName);
		final RpcGateway rpcGateway = akkaRpcService.connect(wildcardAddress, RpcGateway.class).join();

		assertThat(rpcGateway.getAddress(), is(equalTo(simpleRpcEndpoint2.getAddress())));
	}
}
 

@Override
public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
	Class<?> declaringClass = method.getDeclaringClass();

	Object result;

	if (declaringClass.equals(AkkaBasedEndpoint.class) ||
		declaringClass.equals(Object.class) ||
		declaringClass.equals(RpcGateway.class) ||
		declaringClass.equals(StartStoppable.class) ||
		declaringClass.equals(MainThreadExecutable.class) ||
		declaringClass.equals(RpcServer.class)) {
		result = method.invoke(this, args);
	} else if (declaringClass.equals(FencedRpcGateway.class)) {
		throw new UnsupportedOperationException("AkkaInvocationHandler does not support the call FencedRpcGateway#" +
			method.getName() + ". This indicates that you retrieved a FencedRpcGateway without specifying a " +
			"fencing token. Please use RpcService#connect(RpcService, F, Time) with F being the fencing token to " +
			"retrieve a properly FencedRpcGateway.");
	} else {
		result = invokeRpc(method, args);
	}

	return result;
}
 

@Override
protected CompletableFuture<RpcGateway> createGateway(CompletableFuture<Tuple2<String, UUID>> leaderFuture) {
	CompletableFuture<RpcGateway> result;

	if (retrievalAttempt < 2) {
		result = FutureUtils.completedExceptionally(new FlinkException("Could not resolve the leader gateway."));
	} else {
		result = CompletableFuture.completedFuture(rpcGateway);
	}

	retrievalAttempt++;

	return result;
}
 

/**
 * Tests that the gateway retrieval is retried in case of a failure.
 */
@Test
public void testGatewayRetrievalFailures() throws Exception {
	final String address = "localhost";
	final UUID leaderId = UUID.randomUUID();

	RpcGateway rpcGateway = mock(RpcGateway.class);

	TestingLeaderGatewayRetriever leaderGatewayRetriever = new TestingLeaderGatewayRetriever(rpcGateway);
	SettableLeaderRetrievalService settableLeaderRetrievalService = new SettableLeaderRetrievalService();

	settableLeaderRetrievalService.start(leaderGatewayRetriever);

	CompletableFuture<RpcGateway> gatewayFuture = leaderGatewayRetriever.getFuture();

	// this triggers the first gateway retrieval attempt
	settableLeaderRetrievalService.notifyListener(address, leaderId);

	// check that the first future has been failed
	try {
		gatewayFuture.get();

		fail("The first future should have been failed.");
	} catch (ExecutionException ignored) {
		// that's what we expect
	}

	// the second attempt should fail as well
	assertFalse((leaderGatewayRetriever.getNow().isPresent()));

	// the third attempt should succeed
	assertEquals(rpcGateway, leaderGatewayRetriever.getNow().get());
}
 

@Override
protected CompletableFuture<RpcGateway> createGateway(CompletableFuture<Tuple2<String, UUID>> leaderFuture) {
	CompletableFuture<RpcGateway> result;

	if (retrievalAttempt < 2) {
		result = FutureUtils.completedExceptionally(new FlinkException("Could not resolve the leader gateway."));
	} else {
		result = CompletableFuture.completedFuture(rpcGateway);
	}

	retrievalAttempt++;

	return result;
}
 

/**
 * Tests that the gateway retrieval is retried in case of a failure.
 */
@Test
public void testGatewayRetrievalFailures() throws Exception {
	final String address = "localhost";
	final UUID leaderId = UUID.randomUUID();

	RpcGateway rpcGateway = mock(RpcGateway.class);

	TestingLeaderGatewayRetriever leaderGatewayRetriever = new TestingLeaderGatewayRetriever(rpcGateway);
	SettableLeaderRetrievalService settableLeaderRetrievalService = new SettableLeaderRetrievalService();

	settableLeaderRetrievalService.start(leaderGatewayRetriever);

	CompletableFuture<RpcGateway> gatewayFuture = leaderGatewayRetriever.getFuture();

	// this triggers the first gateway retrieval attempt
	settableLeaderRetrievalService.notifyListener(address, leaderId);

	// check that the first future has been failed
	try {
		gatewayFuture.get();

		fail("The first future should have been failed.");
	} catch (ExecutionException ignored) {
		// that's what we expect
	}

	// the second attempt should fail as well
	assertFalse((leaderGatewayRetriever.getNow().isPresent()));

	// the third attempt should succeed
	assertEquals(rpcGateway, leaderGatewayRetriever.getNow().get());
}
 

/**
 * Tests that the gateway retrieval is retried in case of a failure.
 */
@Test
public void testGatewayRetrievalFailures() throws Exception {
	final String address = "localhost";
	final UUID leaderId = UUID.randomUUID();

	RpcGateway rpcGateway = mock(RpcGateway.class);

	TestingLeaderGatewayRetriever leaderGatewayRetriever = new TestingLeaderGatewayRetriever(rpcGateway);
	SettableLeaderRetrievalService settableLeaderRetrievalService = new SettableLeaderRetrievalService();

	settableLeaderRetrievalService.start(leaderGatewayRetriever);

	CompletableFuture<RpcGateway> gatewayFuture = leaderGatewayRetriever.getFuture();

	// this triggers the first gateway retrieval attempt
	settableLeaderRetrievalService.notifyListener(address, leaderId);

	// check that the first future has been failed
	try {
		gatewayFuture.get();

		fail("The first future should have been failed.");
	} catch (ExecutionException ignored) {
		// that's what we expect
	}

	// the second attempt should fail as well
	assertFalse((leaderGatewayRetriever.getNow().isPresent()));

	// the third attempt should succeed
	assertEquals(rpcGateway, leaderGatewayRetriever.getNow().get());
}
 

@Override
protected CompletableFuture<RpcGateway> createGateway(CompletableFuture<Tuple2<String, UUID>> leaderFuture) {
	CompletableFuture<RpcGateway> result;

	if (retrievalAttempt < 2) {
		result = FutureUtils.completedExceptionally(new FlinkException("Could not resolve the leader gateway."));
	} else {
		result = CompletableFuture.completedFuture(rpcGateway);
	}

	retrievalAttempt++;

	return result;
}
 

private <C extends RpcEndpoint & RpcGateway> SupervisorActor.ActorRegistration registerAkkaRpcActor(C rpcEndpoint) {
	final Class<? extends AbstractActor> akkaRpcActorType;

	if (rpcEndpoint instanceof FencedRpcEndpoint) {
		akkaRpcActorType = FencedAkkaRpcActor.class;
	} else {
		akkaRpcActorType = AkkaRpcActor.class;
	}

	synchronized (lock) {
		checkState(!stopped, "RpcService is stopped");

		final SupervisorActor.StartAkkaRpcActorResponse startAkkaRpcActorResponse = SupervisorActor.startAkkaRpcActor(
			supervisor.getActor(),
			actorTerminationFuture -> Props.create(
				akkaRpcActorType,
				rpcEndpoint,
				actorTerminationFuture,
				getVersion(),
				configuration.getMaximumFramesize()),
			rpcEndpoint.getEndpointId());

		final SupervisorActor.ActorRegistration actorRegistration = startAkkaRpcActorResponse.orElseThrow(cause -> new AkkaRpcRuntimeException(
			String.format("Could not create the %s for %s.",
				AkkaRpcActor.class.getSimpleName(),
				rpcEndpoint.getEndpointId()),
			cause));

		actors.put(actorRegistration.getActorRef(), rpcEndpoint);

		return actorRegistration;
	}
}
 

private TestingLeaderGatewayRetriever(RpcGateway rpcGateway) {
	this.rpcGateway = rpcGateway;
}
 

private TestingLeaderGatewayRetriever(RpcGateway rpcGateway) {
	this.rpcGateway = rpcGateway;
}
 

@Override
public <C extends RpcEndpoint & RpcGateway> RpcServer startServer(C rpcEndpoint) {
	checkNotNull(rpcEndpoint, "rpc endpoint");

	final SupervisorActor.ActorRegistration actorRegistration = registerAkkaRpcActor(rpcEndpoint);
	final ActorRef actorRef = actorRegistration.getActorRef();
	final CompletableFuture<Void> actorTerminationFuture = actorRegistration.getTerminationFuture();

	LOG.info("Starting RPC endpoint for {} at {} .", rpcEndpoint.getClass().getName(), actorRef.path());

	final String akkaAddress = AkkaUtils.getAkkaURL(actorSystem, actorRef);
	final String hostname;
	Option<String> host = actorRef.path().address().host();
	if (host.isEmpty()) {
		hostname = "localhost";
	} else {
		hostname = host.get();
	}

	Set<Class<?>> implementedRpcGateways = new HashSet<>(RpcUtils.extractImplementedRpcGateways(rpcEndpoint.getClass()));

	implementedRpcGateways.add(RpcServer.class);
	implementedRpcGateways.add(AkkaBasedEndpoint.class);

	final InvocationHandler akkaInvocationHandler;

	if (rpcEndpoint instanceof FencedRpcEndpoint) {
		// a FencedRpcEndpoint needs a FencedAkkaInvocationHandler
		akkaInvocationHandler = new FencedAkkaInvocationHandler<>(
			akkaAddress,
			hostname,
			actorRef,
			configuration.getTimeout(),
			configuration.getMaximumFramesize(),
			actorTerminationFuture,
			((FencedRpcEndpoint<?>) rpcEndpoint)::getFencingToken,
			captureAskCallstacks);

		implementedRpcGateways.add(FencedMainThreadExecutable.class);
	} else {
		akkaInvocationHandler = new AkkaInvocationHandler(
			akkaAddress,
			hostname,
			actorRef,
			configuration.getTimeout(),
			configuration.getMaximumFramesize(),
			actorTerminationFuture,
			captureAskCallstacks);
	}

	// Rather than using the System ClassLoader directly, we derive the ClassLoader
	// from this class . That works better in cases where Flink runs embedded and all Flink
	// code is loaded dynamically (for example from an OSGI bundle) through a custom ClassLoader
	ClassLoader classLoader = getClass().getClassLoader();

	@SuppressWarnings("unchecked")
	RpcServer server = (RpcServer) Proxy.newProxyInstance(
		classLoader,
		implementedRpcGateways.toArray(new Class<?>[implementedRpcGateways.size()]),
		akkaInvocationHandler);

	return server;
}
 

@Override
public <C extends RpcEndpoint & RpcGateway> RpcServer startServer(C rpcEndpoint) {
	checkNotNull(rpcEndpoint, "rpc endpoint");

	CompletableFuture<Void> terminationFuture = new CompletableFuture<>();
	final Props akkaRpcActorProps;

	if (rpcEndpoint instanceof FencedRpcEndpoint) {
		akkaRpcActorProps = Props.create(
			FencedAkkaRpcActor.class,
			rpcEndpoint,
			terminationFuture,
			getVersion(),
			configuration.getMaximumFramesize());
	} else {
		akkaRpcActorProps = Props.create(
			AkkaRpcActor.class,
			rpcEndpoint,
			terminationFuture,
			getVersion(),
			configuration.getMaximumFramesize());
	}

	ActorRef actorRef;

	synchronized (lock) {
		checkState(!stopped, "RpcService is stopped");
		actorRef = actorSystem.actorOf(akkaRpcActorProps, rpcEndpoint.getEndpointId());
		actors.put(actorRef, rpcEndpoint);
	}

	LOG.info("Starting RPC endpoint for {} at {} .", rpcEndpoint.getClass().getName(), actorRef.path());

	final String akkaAddress = AkkaUtils.getAkkaURL(actorSystem, actorRef);
	final String hostname;
	Option<String> host = actorRef.path().address().host();
	if (host.isEmpty()) {
		hostname = "localhost";
	} else {
		hostname = host.get();
	}

	Set<Class<?>> implementedRpcGateways = new HashSet<>(RpcUtils.extractImplementedRpcGateways(rpcEndpoint.getClass()));

	implementedRpcGateways.add(RpcServer.class);
	implementedRpcGateways.add(AkkaBasedEndpoint.class);

	final InvocationHandler akkaInvocationHandler;

	if (rpcEndpoint instanceof FencedRpcEndpoint) {
		// a FencedRpcEndpoint needs a FencedAkkaInvocationHandler
		akkaInvocationHandler = new FencedAkkaInvocationHandler<>(
			akkaAddress,
			hostname,
			actorRef,
			configuration.getTimeout(),
			configuration.getMaximumFramesize(),
			terminationFuture,
			((FencedRpcEndpoint<?>) rpcEndpoint)::getFencingToken);

		implementedRpcGateways.add(FencedMainThreadExecutable.class);
	} else {
		akkaInvocationHandler = new AkkaInvocationHandler(
			akkaAddress,
			hostname,
			actorRef,
			configuration.getTimeout(),
			configuration.getMaximumFramesize(),
			terminationFuture);
	}

	// Rather than using the System ClassLoader directly, we derive the ClassLoader
	// from this class . That works better in cases where Flink runs embedded and all Flink
	// code is loaded dynamically (for example from an OSGI bundle) through a custom ClassLoader
	ClassLoader classLoader = getClass().getClassLoader();

	@SuppressWarnings("unchecked")
	RpcServer server = (RpcServer) Proxy.newProxyInstance(
		classLoader,
		implementedRpcGateways.toArray(new Class<?>[implementedRpcGateways.size()]),
		akkaInvocationHandler);

	return server;
}
 

private TestingLeaderGatewayRetriever(RpcGateway rpcGateway) {
	this.rpcGateway = rpcGateway;
}
 

@Override
public <C extends RpcEndpoint & RpcGateway> RpcServer startServer(C rpcEndpoint) {
	checkNotNull(rpcEndpoint, "rpc endpoint");

	CompletableFuture<Void> terminationFuture = new CompletableFuture<>();
	final Props akkaRpcActorProps;

	if (rpcEndpoint instanceof FencedRpcEndpoint) {
		akkaRpcActorProps = Props.create(
			FencedAkkaRpcActor.class,
			rpcEndpoint,
			terminationFuture,
			getVersion(),
			configuration.getMaximumFramesize());
	} else {
		akkaRpcActorProps = Props.create(
			AkkaRpcActor.class,
			rpcEndpoint,
			terminationFuture,
			getVersion(),
			configuration.getMaximumFramesize());
	}

	ActorRef actorRef;

	synchronized (lock) {
		checkState(!stopped, "RpcService is stopped");
		actorRef = actorSystem.actorOf(akkaRpcActorProps, rpcEndpoint.getEndpointId());
		actors.put(actorRef, rpcEndpoint);
	}

	LOG.info("Starting RPC endpoint for {} at {} .", rpcEndpoint.getClass().getName(), actorRef.path());

	final String akkaAddress = AkkaUtils.getAkkaURL(actorSystem, actorRef);
	final String hostname;
	Option<String> host = actorRef.path().address().host();
	if (host.isEmpty()) {
		hostname = "localhost";
	} else {
		hostname = host.get();
	}

	Set<Class<?>> implementedRpcGateways = new HashSet<>(RpcUtils.extractImplementedRpcGateways(rpcEndpoint.getClass()));

	implementedRpcGateways.add(RpcServer.class);
	implementedRpcGateways.add(AkkaBasedEndpoint.class);

	final InvocationHandler akkaInvocationHandler;

	if (rpcEndpoint instanceof FencedRpcEndpoint) {
		// a FencedRpcEndpoint needs a FencedAkkaInvocationHandler
		akkaInvocationHandler = new FencedAkkaInvocationHandler<>(
			akkaAddress,
			hostname,
			actorRef,
			configuration.getTimeout(),
			configuration.getMaximumFramesize(),
			terminationFuture,
			((FencedRpcEndpoint<?>) rpcEndpoint)::getFencingToken);

		implementedRpcGateways.add(FencedMainThreadExecutable.class);
	} else {
		akkaInvocationHandler = new AkkaInvocationHandler(
			akkaAddress,
			hostname,
			actorRef,
			configuration.getTimeout(),
			configuration.getMaximumFramesize(),
			terminationFuture);
	}

	// Rather than using the System ClassLoader directly, we derive the ClassLoader
	// from this class . That works better in cases where Flink runs embedded and all Flink
	// code is loaded dynamically (for example from an OSGI bundle) through a custom ClassLoader
	ClassLoader classLoader = getClass().getClassLoader();

	@SuppressWarnings("unchecked")
	RpcServer server = (RpcServer) Proxy.newProxyInstance(
		classLoader,
		implementedRpcGateways.toArray(new Class<?>[implementedRpcGateways.size()]),
		akkaInvocationHandler);

	return server;
}