HADOOP-18406: Adds alignment context to call path for creating RPC proxy with multiple connections per user.

Fixes #4748

Signed-off-by: Owen O'Malley <oomalley@linkedin.com>
This commit is contained in:
Simba Dzinamarira 2022-08-17 09:33:33 -04:00 committed by Owen O'Malley
parent c37f01d95b
commit 4890ba5052
No known key found for this signature in database
GPG Key ID: D19EB09DAD1C5877
7 changed files with 49 additions and 20 deletions

View File

@ -80,9 +80,9 @@ public class ProtobufRpcEngine implements RpcEngine {
@Override @Override
@SuppressWarnings("unchecked") @SuppressWarnings("unchecked")
public <T> ProtocolProxy<T> getProxy(Class<T> protocol, long clientVersion, public <T> ProtocolProxy<T> getProxy(Class<T> protocol, long clientVersion,
ConnectionId connId, Configuration conf, SocketFactory factory) ConnectionId connId, Configuration conf, SocketFactory factory,
throws IOException { AlignmentContext alignmentContext) throws IOException {
final Invoker invoker = new Invoker(protocol, connId, conf, factory); final Invoker invoker = new Invoker(protocol, connId, conf, factory, alignmentContext);
return new ProtocolProxy<T>(protocol, (T) Proxy.newProxyInstance( return new ProtocolProxy<T>(protocol, (T) Proxy.newProxyInstance(
protocol.getClassLoader(), new Class[] {protocol}, invoker), false); protocol.getClassLoader(), new Class[] {protocol}, invoker), false);
} }
@ -126,7 +126,7 @@ public class ProtobufRpcEngine implements RpcEngine {
return new ProtocolProxy<ProtocolMetaInfoPB>(protocol, return new ProtocolProxy<ProtocolMetaInfoPB>(protocol,
(ProtocolMetaInfoPB) Proxy.newProxyInstance(protocol.getClassLoader(), (ProtocolMetaInfoPB) Proxy.newProxyInstance(protocol.getClassLoader(),
new Class[] { protocol }, new Invoker(protocol, connId, conf, new Class[] { protocol }, new Invoker(protocol, connId, conf,
factory)), false); factory, null)), false);
} }
protected static class Invoker implements RpcInvocationHandler { protected static class Invoker implements RpcInvocationHandler {
@ -147,9 +147,8 @@ public class ProtobufRpcEngine implements RpcEngine {
throws IOException { throws IOException {
this(protocol, Client.ConnectionId.getConnectionId( this(protocol, Client.ConnectionId.getConnectionId(
addr, protocol, ticket, rpcTimeout, connectionRetryPolicy, conf), addr, protocol, ticket, rpcTimeout, connectionRetryPolicy, conf),
conf, factory); conf, factory, alignmentContext);
this.fallbackToSimpleAuth = fallbackToSimpleAuth; this.fallbackToSimpleAuth = fallbackToSimpleAuth;
this.alignmentContext = alignmentContext;
} }
/** /**
@ -158,14 +157,16 @@ public class ProtobufRpcEngine implements RpcEngine {
* @param connId input connId. * @param connId input connId.
* @param conf input Configuration. * @param conf input Configuration.
* @param factory input factory. * @param factory input factory.
* @param alignmentContext Alignment context
*/ */
protected Invoker(Class<?> protocol, Client.ConnectionId connId, protected Invoker(Class<?> protocol, Client.ConnectionId connId,
Configuration conf, SocketFactory factory) { Configuration conf, SocketFactory factory, AlignmentContext alignmentContext) {
this.remoteId = connId; this.remoteId = connId;
this.client = CLIENTS.getClient(conf, factory, RpcWritable.Buffer.class); this.client = CLIENTS.getClient(conf, factory, RpcWritable.Buffer.class);
this.protocolName = RPC.getProtocolName(protocol); this.protocolName = RPC.getProtocolName(protocol);
this.clientProtocolVersion = RPC this.clientProtocolVersion = RPC
.getProtocolVersion(protocol); .getProtocolVersion(protocol);
this.alignmentContext = alignmentContext;
} }
private RequestHeaderProto constructRpcRequestHeader(Method method) { private RequestHeaderProto constructRpcRequestHeader(Method method) {

View File

@ -103,9 +103,9 @@ public class ProtobufRpcEngine2 implements RpcEngine {
@Override @Override
@SuppressWarnings("unchecked") @SuppressWarnings("unchecked")
public <T> ProtocolProxy<T> getProxy(Class<T> protocol, long clientVersion, public <T> ProtocolProxy<T> getProxy(Class<T> protocol, long clientVersion,
ConnectionId connId, Configuration conf, SocketFactory factory) ConnectionId connId, Configuration conf, SocketFactory factory,
throws IOException { AlignmentContext alignmentContext) throws IOException {
final Invoker invoker = new Invoker(protocol, connId, conf, factory); final Invoker invoker = new Invoker(protocol, connId, conf, factory, alignmentContext);
return new ProtocolProxy<T>(protocol, (T) Proxy.newProxyInstance( return new ProtocolProxy<T>(protocol, (T) Proxy.newProxyInstance(
protocol.getClassLoader(), new Class[] {protocol}, invoker), false); protocol.getClassLoader(), new Class[] {protocol}, invoker), false);
} }
@ -133,7 +133,7 @@ public class ProtobufRpcEngine2 implements RpcEngine {
return new ProtocolProxy<ProtocolMetaInfoPB>(protocol, return new ProtocolProxy<ProtocolMetaInfoPB>(protocol,
(ProtocolMetaInfoPB) Proxy.newProxyInstance(protocol.getClassLoader(), (ProtocolMetaInfoPB) Proxy.newProxyInstance(protocol.getClassLoader(),
new Class[]{protocol}, new Invoker(protocol, connId, conf, new Class[]{protocol}, new Invoker(protocol, connId, conf,
factory)), false); factory, null)), false);
} }
protected static class Invoker implements RpcInvocationHandler { protected static class Invoker implements RpcInvocationHandler {
@ -154,9 +154,8 @@ public class ProtobufRpcEngine2 implements RpcEngine {
throws IOException { throws IOException {
this(protocol, Client.ConnectionId.getConnectionId( this(protocol, Client.ConnectionId.getConnectionId(
addr, protocol, ticket, rpcTimeout, connectionRetryPolicy, conf), addr, protocol, ticket, rpcTimeout, connectionRetryPolicy, conf),
conf, factory); conf, factory, alignmentContext);
this.fallbackToSimpleAuth = fallbackToSimpleAuth; this.fallbackToSimpleAuth = fallbackToSimpleAuth;
this.alignmentContext = alignmentContext;
} }
/** /**
@ -166,14 +165,16 @@ public class ProtobufRpcEngine2 implements RpcEngine {
* @param connId input connId. * @param connId input connId.
* @param conf input Configuration. * @param conf input Configuration.
* @param factory input factory. * @param factory input factory.
* @param alignmentContext Alignment context
*/ */
protected Invoker(Class<?> protocol, Client.ConnectionId connId, protected Invoker(Class<?> protocol, Client.ConnectionId connId,
Configuration conf, SocketFactory factory) { Configuration conf, SocketFactory factory, AlignmentContext alignmentContext) {
this.remoteId = connId; this.remoteId = connId;
this.client = CLIENTS.getClient(conf, factory, RpcWritable.Buffer.class); this.client = CLIENTS.getClient(conf, factory, RpcWritable.Buffer.class);
this.protocolName = RPC.getProtocolName(protocol); this.protocolName = RPC.getProtocolName(protocol);
this.clientProtocolVersion = RPC this.clientProtocolVersion = RPC
.getProtocolVersion(protocol); .getProtocolVersion(protocol);
this.alignmentContext = alignmentContext;
} }
private RequestHeaderProto constructRpcRequestHeader(Method method) { private RequestHeaderProto constructRpcRequestHeader(Method method) {

View File

@ -558,11 +558,32 @@ public class RPC {
public static <T> ProtocolProxy<T> getProtocolProxy(Class<T> protocol, public static <T> ProtocolProxy<T> getProtocolProxy(Class<T> protocol,
long clientVersion, ConnectionId connId, Configuration conf, long clientVersion, ConnectionId connId, Configuration conf,
SocketFactory factory) throws IOException { SocketFactory factory) throws IOException {
return getProtocolProxy(protocol, clientVersion, connId, conf,
factory, null);
}
/**
* Get a protocol proxy that contains a proxy connection to a remote server
* and a set of methods that are supported by the server.
*
* @param <T> Generics Type T
* @param protocol protocol class
* @param clientVersion client's version
* @param connId client connection identifier
* @param conf configuration
* @param factory socket factory
* @param alignmentContext StateID alignment context
* @return the protocol proxy
* @throws IOException if the far end through a RemoteException
*/
public static <T> ProtocolProxy<T> getProtocolProxy(Class<T> protocol,
long clientVersion, ConnectionId connId, Configuration conf,
SocketFactory factory, AlignmentContext alignmentContext) throws IOException {
if (UserGroupInformation.isSecurityEnabled()) { if (UserGroupInformation.isSecurityEnabled()) {
SaslRpcServer.init(conf); SaslRpcServer.init(conf);
} }
return getProtocolEngine(protocol, conf).getProxy( return getProtocolEngine(protocol, conf).getProxy(
protocol, clientVersion, connId, conf, factory); protocol, clientVersion, connId, conf, factory, alignmentContext);
} }
/** /**

View File

@ -66,11 +66,13 @@ public interface RpcEngine {
* @param connId input ConnectionId. * @param connId input ConnectionId.
* @param conf input Configuration. * @param conf input Configuration.
* @param factory input factory. * @param factory input factory.
* @param alignmentContext Alignment context
* @throws IOException raised on errors performing I/O. * @throws IOException raised on errors performing I/O.
* @return ProtocolProxy. * @return ProtocolProxy.
*/ */
<T> ProtocolProxy<T> getProxy(Class<T> protocol, long clientVersion, <T> ProtocolProxy<T> getProxy(Class<T> protocol, long clientVersion,
Client.ConnectionId connId, Configuration conf, SocketFactory factory) Client.ConnectionId connId, Configuration conf, SocketFactory factory,
AlignmentContext alignmentContext)
throws IOException; throws IOException;
/** /**

View File

@ -315,16 +315,18 @@ public class WritableRpcEngine implements RpcEngine {
* @param connId input ConnectionId. * @param connId input ConnectionId.
* @param conf input Configuration. * @param conf input Configuration.
* @param factory input factory. * @param factory input factory.
* @param alignmentContext Alignment context
* @throws IOException raised on errors performing I/O. * @throws IOException raised on errors performing I/O.
* @return ProtocolProxy. * @return ProtocolProxy.
*/ */
@Override @Override
public <T> ProtocolProxy<T> getProxy(Class<T> protocol, long clientVersion, public <T> ProtocolProxy<T> getProxy(Class<T> protocol, long clientVersion,
Client.ConnectionId connId, Configuration conf, SocketFactory factory) Client.ConnectionId connId, Configuration conf, SocketFactory factory,
AlignmentContext alignmentContext)
throws IOException { throws IOException {
return getProxy(protocol, clientVersion, connId.getAddress(), return getProxy(protocol, clientVersion, connId.getAddress(),
connId.getTicket(), conf, factory, connId.getRpcTimeout(), connId.getTicket(), conf, factory, connId.getRpcTimeout(),
connId.getRetryPolicy(), null, null); connId.getRetryPolicy(), null, alignmentContext);
} }
/** /**

View File

@ -293,7 +293,8 @@ public class TestRPC extends TestRpcBase {
@Override @Override
public <T> ProtocolProxy<T> getProxy(Class<T> protocol, long clientVersion, public <T> ProtocolProxy<T> getProxy(Class<T> protocol, long clientVersion,
ConnectionId connId, Configuration conf, SocketFactory factory) ConnectionId connId, Configuration conf, SocketFactory factory,
AlignmentContext alignmentContext)
throws IOException { throws IOException {
throw new UnsupportedOperationException("This proxy is not supported"); throw new UnsupportedOperationException("This proxy is not supported");
} }

View File

@ -189,7 +189,8 @@ public class TestRpcBase {
0, 0,
connId, connId,
clientConf, clientConf,
NetUtils.getDefaultSocketFactory(clientConf)).getProxy(); NetUtils.getDefaultSocketFactory(clientConf),
null).getProxy();
} catch (IOException e) { } catch (IOException e) {
throw new ServiceException(e); throw new ServiceException(e);
} }