Skip to content

Commit

Permalink
[FLINK-29377][rpc] Reusable timeout extraction
Browse files Browse the repository at this point in the history
  • Loading branch information
zentol committed Sep 22, 2022
1 parent b5cd9f3 commit b9e3dfe
Show file tree
Hide file tree
Showing 2 changed files with 90 additions and 59 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -18,15 +18,14 @@

package org.apache.flink.runtime.rpc.akka;

import org.apache.flink.api.common.time.Time;
import org.apache.flink.configuration.AkkaOptions;
import org.apache.flink.runtime.concurrent.akka.AkkaFutureUtils;
import org.apache.flink.runtime.rpc.FencedRpcGateway;
import org.apache.flink.runtime.rpc.Local;
import org.apache.flink.runtime.rpc.MainThreadExecutable;
import org.apache.flink.runtime.rpc.RpcGateway;
import org.apache.flink.runtime.rpc.RpcGatewayUtils;
import org.apache.flink.runtime.rpc.RpcServer;
import org.apache.flink.runtime.rpc.RpcTimeout;
import org.apache.flink.runtime.rpc.StartStoppable;
import org.apache.flink.runtime.rpc.exceptions.RecipientUnreachableException;
import org.apache.flink.runtime.rpc.exceptions.RpcException;
Expand All @@ -37,7 +36,6 @@
import org.apache.flink.runtime.rpc.messages.RunAsync;
import org.apache.flink.util.ExceptionUtils;
import org.apache.flink.util.Preconditions;
import org.apache.flink.util.TimeUtils;

import akka.actor.ActorRef;
import akka.pattern.Patterns;
Expand Down Expand Up @@ -220,7 +218,8 @@ private Object invokeRpc(Method method, Object[] args) throws Exception {
Class<?>[] parameterTypes = method.getParameterTypes();
final boolean isLocalRpcInvocation = method.getAnnotation(Local.class) != null;
Annotation[][] parameterAnnotations = method.getParameterAnnotations();
Duration futureTimeout = extractRpcTimeout(parameterAnnotations, args, timeout);
Duration futureTimeout =
RpcGatewayUtils.extractRpcTimeout(parameterAnnotations, args, timeout);

final RpcInvocation rpcInvocation =
createRpcInvocationMessage(
Expand Down Expand Up @@ -320,61 +319,6 @@ private RpcInvocation createRpcInvocationMessage(
// Helper methods
// ------------------------------------------------------------------------

/**
* Extracts the {@link RpcTimeout} annotated rpc timeout value from the list of given method
* arguments. If no {@link RpcTimeout} annotated parameter could be found, then the default
* timeout is returned.
*
* @param parameterAnnotations Parameter annotations
* @param args Array of arguments
* @param defaultTimeout Default timeout to return if no {@link RpcTimeout} annotated parameter
* has been found
* @return Timeout extracted from the array of arguments or the default timeout
*/
private static Duration extractRpcTimeout(
Annotation[][] parameterAnnotations, Object[] args, Duration defaultTimeout) {
if (args != null) {
Preconditions.checkArgument(parameterAnnotations.length == args.length);

for (int i = 0; i < parameterAnnotations.length; i++) {
if (isRpcTimeout(parameterAnnotations[i])) {
if (args[i] instanceof Time) {
return TimeUtils.toDuration((Time) args[i]);
} else if (args[i] instanceof Duration) {
return (Duration) args[i];
} else {
throw new RuntimeException(
"The rpc timeout parameter must be of type "
+ Time.class.getName()
+ " or "
+ Duration.class.getName()
+ ". The type "
+ args[i].getClass().getName()
+ " is not supported.");
}
}
}
}

return defaultTimeout;
}

/**
* Checks whether any of the annotations is of type {@link RpcTimeout}.
*
* @param annotations Array of annotations
* @return True if {@link RpcTimeout} was found; otherwise false
*/
private static boolean isRpcTimeout(Annotation[] annotations) {
for (Annotation annotation : annotations) {
if (annotation.annotationType().equals(RpcTimeout.class)) {
return true;
}
}

return false;
}

/**
* Sends the message to the RPC endpoint.
*
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package org.apache.flink.runtime.rpc;

import org.apache.flink.api.common.time.Time;
import org.apache.flink.util.Preconditions;
import org.apache.flink.util.TimeUtils;

import java.lang.annotation.Annotation;
import java.time.Duration;

/** Utils for {@link RpcGateway} implementations. */
public class RpcGatewayUtils {

/**
* Extracts the {@link RpcTimeout} annotated rpc timeout value from the list of given method
* arguments. If no {@link RpcTimeout} annotated parameter could be found, then the default
* timeout is returned.
*
* @param parameterAnnotations Parameter annotations
* @param args Array of arguments
* @param defaultTimeout Default timeout to return if no {@link RpcTimeout} annotated parameter
* has been found
* @return Timeout extracted from the array of arguments or the default timeout
*/
public static Duration extractRpcTimeout(
Annotation[][] parameterAnnotations, Object[] args, Duration defaultTimeout) {
if (args != null) {
Preconditions.checkArgument(parameterAnnotations.length == args.length);

for (int i = 0; i < parameterAnnotations.length; i++) {
if (isRpcTimeout(parameterAnnotations[i])) {
if (args[i] instanceof Time) {
return TimeUtils.toDuration((Time) args[i]);
} else if (args[i] instanceof Duration) {
return (Duration) args[i];
} else {
throw new RuntimeException(
"The rpc timeout parameter must be of type "
+ Time.class.getName()
+ " or "
+ Duration.class.getName()
+ ". The type "
+ args[i].getClass().getName()
+ " is not supported.");
}
}
}
}

return defaultTimeout;
}

/**
* Checks whether any of the annotations is of type {@link RpcTimeout}.
*
* @param annotations Array of annotations
* @return True if {@link RpcTimeout} was found; otherwise false
*/
private static boolean isRpcTimeout(Annotation[] annotations) {
for (Annotation annotation : annotations) {
if (annotation.annotationType().equals(RpcTimeout.class)) {
return true;
}
}

return false;
}

private RpcGatewayUtils() {}
}

0 comments on commit b9e3dfe

Please sign in to comment.