grpc
/
grpc-java
mirror of https://github.com/grpc/grpc-java.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Use generics for LoadBalancer to avoid ClientTransport exposure 

TransportManager.makeTransport() was added to remove the only other
reference to ClientTransport outside of core and the transports.
This commit is contained in:
Eric Anderson 2016-01-23 10:07:13 -08:00
parent bf42913c23
commit 1488010cb1
14 changed files with 136 additions and 137 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
core/src/main/java/io/grpc/LoadBalancer.java
View File

@ -33,8 +33,6 @@ package io.grpc;
import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.ListenableFuture;
import io.grpc.internal.ClientTransport;
import java.util.List; import java.util.List;
import javax.annotation.Nullable; import javax.annotation.Nullable;
@ -46,12 +44,14 @@ import javax.annotation.concurrent.ThreadSafe;
* *
* <p>Note to implementations: all methods are expected to return quickly. Any work that may block * <p>Note to implementations: all methods are expected to return quickly. Any work that may block
* should be done asynchronously. * should be done asynchronously.
*
* @param T the transport type to balance
*/ */
// TODO(zhangkun83): since it's also used for non-loadbalancing cases like pick-first, // TODO(zhangkun83): since it's also used for non-loadbalancing cases like pick-first,
// "RequestRouter" might be a better name. // "RequestRouter" might be a better name.
@ExperimentalApi @ExperimentalApi
@ThreadSafe @ThreadSafe
public abstract class LoadBalancer { public abstract class LoadBalancer<T> {
/** /**
* Pick a transport that Channel will use for next RPC. * Pick a transport that Channel will use for next RPC.
* *
@ -61,8 +61,7 @@ public abstract class LoadBalancer {
* *
* @param requestKey for affinity-based routing * @param requestKey for affinity-based routing
*/ */
public abstract ListenableFuture<ClientTransport> pickTransport( public abstract ListenableFuture<T> pickTransport(@Nullable RequestKey requestKey);
@Nullable RequestKey requestKey);
/** /**
* Shuts down this {@code LoadBalancer}. * Shuts down this {@code LoadBalancer}.
@ -86,13 +85,12 @@ public abstract class LoadBalancer {
/** /**
* Called when a transport is fully connected and ready to accept traffic. * Called when a transport is fully connected and ready to accept traffic.
*/ */
public void transportReady(EquivalentAddressGroup addressGroup, ClientTransport transport) { } public void transportReady(EquivalentAddressGroup addressGroup, T transport) { }
/** /**
* Called when a transport is shutting down. * Called when a transport is shutting down.
*/ */
public void transportShutdown( public void transportShutdown(EquivalentAddressGroup addressGroup, T transport, Status s) { }
EquivalentAddressGroup addressGroup, ClientTransport transport, Status s) { }
public abstract static class Factory { public abstract static class Factory {
/** /**
@ -102,6 +100,6 @@ public abstract class LoadBalancer {
* @param tm the interface where an {@code LoadBalancer} implementation gets connected * @param tm the interface where an {@code LoadBalancer} implementation gets connected
* transports from * transports from
*/ */
public abstract LoadBalancer newLoadBalancer(String serviceName, TransportManager tm); public abstract <T> LoadBalancer<T> newLoadBalancer(String serviceName, TransportManager<T> tm);
} }
} }

24
core/src/main/java/io/grpc/SimpleLoadBalancerFactory.java
View File

@ -36,7 +36,6 @@ import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.ListenableFuture;
import io.grpc.internal.BlankFutureProvider; import io.grpc.internal.BlankFutureProvider;
import io.grpc.internal.ClientTransport;
import java.net.SocketAddress; import java.net.SocketAddress;
import java.util.ArrayList; import java.util.ArrayList;
@ -63,29 +62,28 @@ public final class SimpleLoadBalancerFactory extends LoadBalancer.Factory {
} }
@Override @Override
public LoadBalancer newLoadBalancer(String serviceName, TransportManager tm) { public <T> LoadBalancer<T> newLoadBalancer(String serviceName, TransportManager<T> tm) {
return new SimpleLoadBalancer(tm); return new SimpleLoadBalancer<T>(tm);
} }
private static class SimpleLoadBalancer extends LoadBalancer { private static class SimpleLoadBalancer<T> extends LoadBalancer<T> {
private final Object lock = new Object(); private final Object lock = new Object();
@GuardedBy("lock") @GuardedBy("lock")
private EquivalentAddressGroup addresses; private EquivalentAddressGroup addresses;
@GuardedBy("lock") @GuardedBy("lock")
private final BlankFutureProvider<ClientTransport> pendingPicks = private final BlankFutureProvider<T> pendingPicks = new BlankFutureProvider<T>();
new BlankFutureProvider<ClientTransport>();
@GuardedBy("lock") @GuardedBy("lock")
private StatusException nameResolutionError; private StatusException nameResolutionError;
private final TransportManager tm; private final TransportManager<T> tm;
private SimpleLoadBalancer(TransportManager tm) { private SimpleLoadBalancer(TransportManager<T> tm) {
this.tm = tm; this.tm = tm;
} }
@Override @Override
public ListenableFuture<ClientTransport> pickTransport(@Nullable RequestKey requestKey) { public ListenableFuture<T> pickTransport(@Nullable RequestKey requestKey) {
EquivalentAddressGroup addressesCopy; EquivalentAddressGroup addressesCopy;
synchronized (lock) { synchronized (lock) {
addressesCopy = addresses; addressesCopy = addresses;
@ -102,7 +100,7 @@ public final class SimpleLoadBalancerFactory extends LoadBalancer.Factory {
@Override @Override
public void handleResolvedAddresses( public void handleResolvedAddresses(
List<ResolvedServerInfo> updatedServers, Attributes config) { List<ResolvedServerInfo> updatedServers, Attributes config) {
BlankFutureProvider.FulfillmentBatch<ClientTransport> pendingPicksFulfillmentBatch; BlankFutureProvider.FulfillmentBatch<T> pendingPicksFulfillmentBatch;
final EquivalentAddressGroup newAddresses; final EquivalentAddressGroup newAddresses;
synchronized (lock) { synchronized (lock) {
ArrayList<SocketAddress> newAddressList = ArrayList<SocketAddress> newAddressList =
@ -118,8 +116,8 @@ public final class SimpleLoadBalancerFactory extends LoadBalancer.Factory {
nameResolutionError = null; nameResolutionError = null;
pendingPicksFulfillmentBatch = pendingPicks.createFulfillmentBatch(); pendingPicksFulfillmentBatch = pendingPicks.createFulfillmentBatch();
} }
pendingPicksFulfillmentBatch.link(new Supplier<ListenableFuture<ClientTransport>>() { pendingPicksFulfillmentBatch.link(new Supplier<ListenableFuture<T>>() {
@Override public ListenableFuture<ClientTransport> get() { @Override public ListenableFuture<T> get() {
return tm.getTransport(newAddresses); return tm.getTransport(newAddresses);
} }
}); });
@ -127,7 +125,7 @@ public final class SimpleLoadBalancerFactory extends LoadBalancer.Factory {
@Override @Override
public void handleNameResolutionError(Status error) { public void handleNameResolutionError(Status error) {
BlankFutureProvider.FulfillmentBatch<ClientTransport> pendingPicksFulfillmentBatch; BlankFutureProvider.FulfillmentBatch<T> pendingPicksFulfillmentBatch;
StatusException statusException = StatusException statusException =
error.augmentDescription("Name resolution failed").asException(); error.augmentDescription("Name resolution failed").asException();
synchronized (lock) { synchronized (lock) {

12
core/src/main/java/io/grpc/TransportManager.java
View File

@ -33,15 +33,13 @@ package io.grpc;
import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.ListenableFuture;
import io.grpc.internal.ClientTransport;
import java.util.Collection; import java.util.Collection;
/** /**
* Manages transport life-cycles and provide ready-to-use transports. * Manages transport life-cycles and provide ready-to-use transports.
*/ */
@ExperimentalApi @ExperimentalApi
public abstract class TransportManager { public abstract class TransportManager<T> {
/** /**
* Advises this {@code TransportManager} to retain transports only to these servers, for warming * Advises this {@code TransportManager} to retain transports only to these servers, for warming
* up connections and discarding unused connections. * up connections and discarding unused connections.
@ -59,6 +57,10 @@ public abstract class TransportManager {
// TODO(zhangkun83): GrpcLoadBalancer will use this to get transport to connect to LB servers, // TODO(zhangkun83): GrpcLoadBalancer will use this to get transport to connect to LB servers,
// which would have a different authority than the primary servers. We need to figure out how to // which would have a different authority than the primary servers. We need to figure out how to
// do it. // do it.
public abstract ListenableFuture<ClientTransport> getTransport( public abstract ListenableFuture<T> getTransport(EquivalentAddressGroup addressGroup);
EquivalentAddressGroup addressGroup);
/**
* Returns a channel that uses {@code transport}; useful for issuing RPCs on a transport.
*/
public abstract Channel makeChannel(T transport);
} }

10
core/src/main/java/io/grpc/internal/ManagedChannelImpl.java
View File

@ -125,7 +125,7 @@ public final class ManagedChannelImpl extends ManagedChannel {
private final Channel interceptorChannel; private final Channel interceptorChannel;
private final NameResolver nameResolver; private final NameResolver nameResolver;
private final LoadBalancer loadBalancer; private final LoadBalancer<ClientTransport> loadBalancer;
/** /**
* Maps EquivalentAddressGroups to transports for that server. * Maps EquivalentAddressGroups to transports for that server.
@ -357,7 +357,7 @@ public final class ManagedChannelImpl extends ManagedChannel {
transportFactory.release(); transportFactory.release();
} }
private final TransportManager tm = new TransportManager() { private final TransportManager<ClientTransport> tm = new TransportManager<ClientTransport>() {
@Override @Override
public void updateRetainedTransports(Collection<EquivalentAddressGroup> addrs) { public void updateRetainedTransports(Collection<EquivalentAddressGroup> addrs) {
// TODO(zhangkun83): warm-up new servers and discard removed servers. // TODO(zhangkun83): warm-up new servers and discard removed servers.
@ -396,5 +396,11 @@ public final class ManagedChannelImpl extends ManagedChannel {
} }
return ts.obtainActiveTransport(); return ts.obtainActiveTransport();
} }
@Override
public Channel makeChannel(ClientTransport transport) {
return new SingleTransportChannel(
transport, executor, scheduledExecutor, authority());
}
}; };
} }

2
core/src/main/java/io/grpc/internal/SingleTransportChannel.java
View File

@ -50,7 +50,7 @@ import java.util.concurrent.ScheduledExecutorService;
/** /**
* A {@link Channel} that wraps a {@link ClientTransport}. * A {@link Channel} that wraps a {@link ClientTransport}.
*/ */
public final class SingleTransportChannel extends Channel { final class SingleTransportChannel extends Channel {
private final ClientTransport transport; private final ClientTransport transport;
private final Executor executor; private final Executor executor;

16
core/src/main/java/io/grpc/internal/TransportSet.java
View File

@ -105,7 +105,7 @@ final class TransportSet {
@GuardedBy("lock") @GuardedBy("lock")
private final Collection<ClientTransport> transports = new ArrayList<ClientTransport>(); private final Collection<ClientTransport> transports = new ArrayList<ClientTransport>();
private final LoadBalancer loadBalancer; private final LoadBalancer<ClientTransport> loadBalancer;
@GuardedBy("lock") @GuardedBy("lock")
private boolean shutdown; private boolean shutdown;
@ -117,17 +117,19 @@ final class TransportSet {
@Nullable @Nullable
private volatile UncancellableTransportFuture activeTransportFuture; private volatile UncancellableTransportFuture activeTransportFuture;
TransportSet(EquivalentAddressGroup addressGroup, String authority, LoadBalancer loadBalancer, TransportSet(EquivalentAddressGroup addressGroup, String authority,
BackoffPolicy.Provider backoffPolicyProvider, ClientTransportFactory transportFactory, LoadBalancer<ClientTransport> loadBalancer, BackoffPolicy.Provider backoffPolicyProvider,
ScheduledExecutorService scheduledExecutor, Callback callback) { ClientTransportFactory transportFactory, ScheduledExecutorService scheduledExecutor,
Callback callback) {
this(addressGroup, authority, loadBalancer, backoffPolicyProvider, transportFactory, this(addressGroup, authority, loadBalancer, backoffPolicyProvider, transportFactory,
scheduledExecutor, callback, Stopwatch.createUnstarted()); scheduledExecutor, callback, Stopwatch.createUnstarted());
} }
@VisibleForTesting @VisibleForTesting
TransportSet(EquivalentAddressGroup addressGroup, String authority, LoadBalancer loadBalancer, TransportSet(EquivalentAddressGroup addressGroup, String authority,
BackoffPolicy.Provider backoffPolicyProvider, ClientTransportFactory transportFactory, LoadBalancer<ClientTransport> loadBalancer, BackoffPolicy.Provider backoffPolicyProvider,
ScheduledExecutorService scheduledExecutor, Callback callback, Stopwatch backoffWatch) { ClientTransportFactory transportFactory, ScheduledExecutorService scheduledExecutor,
Callback callback, Stopwatch backoffWatch) {
this.addressGroup = Preconditions.checkNotNull(addressGroup, "addressGroup"); this.addressGroup = Preconditions.checkNotNull(addressGroup, "addressGroup");
this.authority = authority; this.authority = authority;
this.loadBalancer = loadBalancer; this.loadBalancer = loadBalancer;

22
core/src/test/java/io/grpc/SimpleLoadBalancerTest.java
View File

@ -37,7 +37,6 @@ import static org.junit.Assert.assertNotSame;
import static org.junit.Assert.assertSame; import static org.junit.Assert.assertSame;
import static org.mockito.Matchers.any; import static org.mockito.Matchers.any;
import static org.mockito.Matchers.eq; import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never; import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times; import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verify;
@ -47,8 +46,6 @@ import static org.mockito.Mockito.when;
import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture; import com.google.common.util.concurrent.SettableFuture;
import io.grpc.internal.ClientTransport;
import org.junit.Before; import org.junit.Before;
import org.junit.Test; import org.junit.Test;
import org.junit.runner.RunWith; import org.junit.runner.RunWith;
@ -62,13 +59,13 @@ import java.util.ArrayList;
/** Unit test for {@link SimpleLoadBalancerFactory}. */ /** Unit test for {@link SimpleLoadBalancerFactory}. */
@RunWith(JUnit4.class) @RunWith(JUnit4.class)
public class SimpleLoadBalancerTest { public class SimpleLoadBalancerTest {
private LoadBalancer loadBalancer; private LoadBalancer<Transport> loadBalancer;
private ArrayList<ResolvedServerInfo> servers; private ArrayList<ResolvedServerInfo> servers;
private EquivalentAddressGroup addressGroup; private EquivalentAddressGroup addressGroup;
@Mock @Mock
private TransportManager mockTransportManager; private TransportManager<Transport> mockTransportManager;
@Before @Before
public void setUp() { public void setUp() {
@ -87,12 +84,12 @@ public class SimpleLoadBalancerTest {
@Test @Test
public void pickBeforeResolved() throws Exception { public void pickBeforeResolved() throws Exception {
ClientTransport mockTransport = mock(ClientTransport.class); Transport mockTransport = new Transport();
SettableFuture<ClientTransport> sourceFuture = SettableFuture.create(); SettableFuture<Transport> sourceFuture = SettableFuture.create();
when(mockTransportManager.getTransport(eq(addressGroup))) when(mockTransportManager.getTransport(eq(addressGroup)))
.thenReturn(sourceFuture); .thenReturn(sourceFuture);
ListenableFuture<ClientTransport> f1 = loadBalancer.pickTransport(null); ListenableFuture<Transport> f1 = loadBalancer.pickTransport(null);
ListenableFuture<ClientTransport> f2 = loadBalancer.pickTransport(null); ListenableFuture<Transport> f2 = loadBalancer.pickTransport(null);
assertNotNull(f1); assertNotNull(f1);
assertNotNull(f2); assertNotNull(f2);
assertNotSame(f1, f2); assertNotSame(f1, f2);
@ -113,12 +110,12 @@ public class SimpleLoadBalancerTest {
@Test @Test
public void pickAfterResolved() throws Exception { public void pickAfterResolved() throws Exception {
ClientTransport mockTransport = mock(ClientTransport.class); Transport mockTransport = new Transport();
SettableFuture<ClientTransport> sourceFuture = SettableFuture.create(); SettableFuture<Transport> sourceFuture = SettableFuture.create();
when(mockTransportManager.getTransport(eq(addressGroup))) when(mockTransportManager.getTransport(eq(addressGroup)))
.thenReturn(sourceFuture); .thenReturn(sourceFuture);
loadBalancer.handleResolvedAddresses(servers, Attributes.EMPTY); loadBalancer.handleResolvedAddresses(servers, Attributes.EMPTY);
ListenableFuture<ClientTransport> f = loadBalancer.pickTransport(null); ListenableFuture<Transport> f = loadBalancer.pickTransport(null);
assertSame(sourceFuture, f); assertSame(sourceFuture, f);
assertFalse(f.isDone()); assertFalse(f.isDone());
sourceFuture.set(mockTransport); sourceFuture.set(mockTransport);
@ -139,4 +136,5 @@ public class SimpleLoadBalancerTest {
} }
} }
private static class Transport {}
} }

6
core/src/test/java/io/grpc/internal/ManagedChannelImplTest.java
View File

@ -529,15 +529,15 @@ public class ManagedChannelImplTest {
private class SpyingLoadBalancerFactory extends LoadBalancer.Factory { private class SpyingLoadBalancerFactory extends LoadBalancer.Factory {
private final LoadBalancer.Factory delegate; private final LoadBalancer.Factory delegate;
private final List<LoadBalancer> balancers = new ArrayList<LoadBalancer>(); private final List<LoadBalancer<?>> balancers = new ArrayList<LoadBalancer<?>>();
private SpyingLoadBalancerFactory(LoadBalancer.Factory delegate) { private SpyingLoadBalancerFactory(LoadBalancer.Factory delegate) {
this.delegate = delegate; this.delegate = delegate;
} }
@Override @Override
public LoadBalancer newLoadBalancer(String serviceName, TransportManager tm) { public <T> LoadBalancer<T> newLoadBalancer(String serviceName, TransportManager<T> tm) {
LoadBalancer lb = spy(delegate.newLoadBalancer(serviceName, tm)); LoadBalancer<T> lb = spy(delegate.newLoadBalancer(serviceName, tm));
balancers.add(lb); balancers.add(lb);
return lb; return lb;
} }

14
core/src/test/java/io/grpc/internal/ManagedChannelImplTransportManagerTest.java
View File

@ -62,6 +62,7 @@ import org.junit.Test;
import org.junit.runner.RunWith; import org.junit.runner.RunWith;
import org.junit.runners.JUnit4; import org.junit.runners.JUnit4;
import org.mockito.ArgumentCaptor; import org.mockito.ArgumentCaptor;
import org.mockito.Matchers;
import org.mockito.Mock; import org.mockito.Mock;
import org.mockito.MockitoAnnotations; import org.mockito.MockitoAnnotations;
import org.mockito.invocation.InvocationOnMock; import org.mockito.invocation.InvocationOnMock;
@ -114,15 +115,18 @@ public class ManagedChannelImplTransportManagerTest {
@Mock private BackoffPolicy.Provider mockBackoffPolicyProvider; @Mock private BackoffPolicy.Provider mockBackoffPolicyProvider;
@Mock private BackoffPolicy mockBackoffPolicy; @Mock private BackoffPolicy mockBackoffPolicy;
private TransportManager tm; private TransportManager<ClientTransport> tm;
@Before @Before
public void setUp() { public void setUp() {
MockitoAnnotations.initMocks(this); MockitoAnnotations.initMocks(this);
when(mockBackoffPolicyProvider.get()).thenReturn(mockBackoffPolicy); when(mockBackoffPolicyProvider.get()).thenReturn(mockBackoffPolicy);
when(mockLoadBalancerFactory.newLoadBalancer(anyString(), any(TransportManager.class))) @SuppressWarnings("unchecked")
.thenReturn(mock(LoadBalancer.class)); LoadBalancer<ClientTransport> loadBalancer = mock(LoadBalancer.class);
when(mockLoadBalancerFactory
.newLoadBalancer(anyString(), Matchers.<TransportManager<ClientTransport>>any()))
.thenReturn(loadBalancer);
channel = new ManagedChannelImpl("fake://target", mockBackoffPolicyProvider, channel = new ManagedChannelImpl("fake://target", mockBackoffPolicyProvider,
nameResolverFactory, Attributes.EMPTY, mockLoadBalancerFactory, nameResolverFactory, Attributes.EMPTY, mockLoadBalancerFactory,
@ -130,7 +134,8 @@ public class ManagedChannelImplTransportManagerTest {
CompressorRegistry.getDefaultInstance(), executor, null, CompressorRegistry.getDefaultInstance(), executor, null,
Collections.<ClientInterceptor>emptyList()); Collections.<ClientInterceptor>emptyList());
ArgumentCaptor<TransportManager> tmCaptor = ArgumentCaptor.forClass(TransportManager.class); ArgumentCaptor<TransportManager<ClientTransport>> tmCaptor
= ArgumentCaptor.forClass(null);
verify(mockLoadBalancerFactory).newLoadBalancer(anyString(), tmCaptor.capture()); verify(mockLoadBalancerFactory).newLoadBalancer(anyString(), tmCaptor.capture());
tm = tmCaptor.getValue(); tm = tmCaptor.getValue();
} }
@ -259,5 +264,4 @@ public class ManagedChannelImplTransportManagerTest {
verify(mockBackoffPolicyProvider, times(backoffReset)).get(); verify(mockBackoffPolicyProvider, times(backoffReset)).get();
verify(mockBackoffPolicy, times(++backoffConsulted)).nextBackoffMillis(); verify(mockBackoffPolicy, times(++backoffConsulted)).nextBackoffMillis();
} }
} }

2
core/src/test/java/io/grpc/internal/TransportSetTest.java
View File

@ -70,7 +70,7 @@ public class TransportSetTest {
private FakeClock fakeClock; private FakeClock fakeClock;
@Mock private LoadBalancer mockLoadBalancer; @Mock private LoadBalancer<ClientTransport> mockLoadBalancer;
@Mock private BackoffPolicy mockBackoffPolicy1; @Mock private BackoffPolicy mockBackoffPolicy1;
@Mock private BackoffPolicy mockBackoffPolicy2; @Mock private BackoffPolicy mockBackoffPolicy2;
@Mock private BackoffPolicy mockBackoffPolicy3; @Mock private BackoffPolicy mockBackoffPolicy3;

62
grpclb/src/main/java/io/grpc/grpclb/GrpclbLoadBalancer.java
View File

@ -49,10 +49,8 @@ import io.grpc.StatusException;
import io.grpc.TransportManager; import io.grpc.TransportManager;
import io.grpc.internal.BlankFutureProvider; import io.grpc.internal.BlankFutureProvider;
import io.grpc.internal.BlankFutureProvider.FulfillmentBatch; import io.grpc.internal.BlankFutureProvider.FulfillmentBatch;
import io.grpc.internal.ClientTransport;
import io.grpc.internal.GrpcUtil; import io.grpc.internal.GrpcUtil;
import io.grpc.internal.SharedResourceHolder; import io.grpc.internal.SharedResourceHolder;
import io.grpc.internal.SingleTransportChannel;
import io.grpc.stub.StreamObserver; import io.grpc.stub.StreamObserver;
import java.net.InetSocketAddress; import java.net.InetSocketAddress;
@ -62,7 +60,6 @@ import java.util.HashMap;
import java.util.HashSet; import java.util.HashSet;
import java.util.List; import java.util.List;
import java.util.concurrent.ExecutorService; import java.util.concurrent.ExecutorService;
import java.util.concurrent.ScheduledExecutorService;
import java.util.logging.Logger; import java.util.logging.Logger;
import javax.annotation.Nullable; import javax.annotation.Nullable;
@ -71,17 +68,16 @@ import javax.annotation.concurrent.GuardedBy;
/** /**
* A {@link LoadBalancer} that uses the GRPCLB protocol. * A {@link LoadBalancer} that uses the GRPCLB protocol.
*/ */
class GrpclbLoadBalancer extends LoadBalancer { class GrpclbLoadBalancer<T> extends LoadBalancer<T> {
private static final Logger logger = Logger.getLogger(GrpclbLoadBalancer.class.getName()); private static final Logger logger = Logger.getLogger(GrpclbLoadBalancer.class.getName());
private final Object lock = new Object(); private final Object lock = new Object();
private final String serviceName; private final String serviceName;
private final TransportManager tm; private final TransportManager<T> tm;
// General states // General states
@GuardedBy("lock") @GuardedBy("lock")
private final BlankFutureProvider<ClientTransport> pendingPicks = private final BlankFutureProvider<T> pendingPicks = new BlankFutureProvider<T>();
new BlankFutureProvider<ClientTransport>();
@GuardedBy("lock") @GuardedBy("lock")
private Throwable lastError; private Throwable lastError;
@ -92,9 +88,9 @@ class GrpclbLoadBalancer extends LoadBalancer {
@GuardedBy("lock") @GuardedBy("lock")
private EquivalentAddressGroup lbAddresses; private EquivalentAddressGroup lbAddresses;
@GuardedBy("lock") @GuardedBy("lock")
private ClientTransport lbTransport; private T lbTransport;
@GuardedBy("lock") @GuardedBy("lock")
private ListenableFuture<ClientTransport> directTransport; private ListenableFuture<T> directTransport;
@GuardedBy("lock") @GuardedBy("lock")
private StreamObserver<LoadBalanceResponse> lbResponseObserver; private StreamObserver<LoadBalanceResponse> lbResponseObserver;
@GuardedBy("lock") @GuardedBy("lock")
@ -105,16 +101,14 @@ class GrpclbLoadBalancer extends LoadBalancer {
private HashMap<SocketAddress, ResolvedServerInfo> servers; private HashMap<SocketAddress, ResolvedServerInfo> servers;
@GuardedBy("lock") @GuardedBy("lock")
@VisibleForTesting @VisibleForTesting
private RoundRobinServerList roundRobinServerList; private RoundRobinServerList<T> roundRobinServerList;
private ExecutorService executor; private ExecutorService executor;
private ScheduledExecutorService deadlineCancellationExecutor;
GrpclbLoadBalancer(String serviceName, TransportManager tm) { GrpclbLoadBalancer(String serviceName, TransportManager<T> tm) {
this.serviceName = serviceName; this.serviceName = serviceName;
this.tm = tm; this.tm = tm;
executor = SharedResourceHolder.get(GrpcUtil.SHARED_CHANNEL_EXECUTOR); executor = SharedResourceHolder.get(GrpcUtil.SHARED_CHANNEL_EXECUTOR);
deadlineCancellationExecutor = SharedResourceHolder.get(GrpcUtil.TIMER_SERVICE);
} }
@VisibleForTesting @VisibleForTesting
@ -125,15 +119,15 @@ class GrpclbLoadBalancer extends LoadBalancer {
} }
@VisibleForTesting @VisibleForTesting
RoundRobinServerList getRoundRobinServerList() { RoundRobinServerList<T> getRoundRobinServerList() {
synchronized (lock) { synchronized (lock) {
return roundRobinServerList; return roundRobinServerList;
} }
} }
@Override @Override
public ListenableFuture<ClientTransport> pickTransport(@Nullable RequestKey requestKey) { public ListenableFuture<T> pickTransport(@Nullable RequestKey requestKey) {
RoundRobinServerList serverListCopy; RoundRobinServerList<T> serverListCopy;
synchronized (lock) { synchronized (lock) {
Preconditions.checkState(!closed, "already closed"); Preconditions.checkState(!closed, "already closed");
if (directTransport != null) { if (directTransport != null) {
@ -178,11 +172,11 @@ class GrpclbLoadBalancer extends LoadBalancer {
Preconditions.checkNotNull(lbAddresses, "lbAddresses"); Preconditions.checkNotNull(lbAddresses, "lbAddresses");
// TODO(zhangkun83): LB servers may use an authority different from the service's. // TODO(zhangkun83): LB servers may use an authority different from the service's.
// getTransport() will need to add an argument for the authority. // getTransport() will need to add an argument for the authority.
ListenableFuture<ClientTransport> transportFuture = tm.getTransport(lbAddresses); ListenableFuture<T> transportFuture = tm.getTransport(lbAddresses);
Futures.addCallback( Futures.addCallback(
Preconditions.checkNotNull(transportFuture), Preconditions.checkNotNull(transportFuture),
new FutureCallback<ClientTransport>() { new FutureCallback<T>() {
@Override public void onSuccess(ClientTransport transport) { @Override public void onSuccess(T transport) {
synchronized (lock) { synchronized (lock) {
if (closed) { if (closed) {
return; return;
@ -221,9 +215,8 @@ class GrpclbLoadBalancer extends LoadBalancer {
@VisibleForTesting // to be mocked in tests @VisibleForTesting // to be mocked in tests
@GuardedBy("lock") @GuardedBy("lock")
void sendLbRequest(ClientTransport transport, LoadBalanceRequest request) { void sendLbRequest(T transport, LoadBalanceRequest request) {
Channel channel = new SingleTransportChannel(transport, executor, Channel channel = tm.makeChannel(transport);
deadlineCancellationExecutor, serviceName);
LoadBalancerGrpc.LoadBalancerStub stub = LoadBalancerGrpc.newStub(channel); LoadBalancerGrpc.LoadBalancerStub stub = LoadBalancerGrpc.newStub(channel);
lbRequestWriter = stub.balanceLoad(lbResponseObserver); lbRequestWriter = stub.balanceLoad(lbResponseObserver);
lbRequestWriter.onNext(request); lbRequestWriter.onNext(request);
@ -245,14 +238,11 @@ class GrpclbLoadBalancer extends LoadBalancer {
lbRequestWriter.onCompleted(); lbRequestWriter.onCompleted();
} }
executor = SharedResourceHolder.release(GrpcUtil.SHARED_CHANNEL_EXECUTOR, executor); executor = SharedResourceHolder.release(GrpcUtil.SHARED_CHANNEL_EXECUTOR, executor);
deadlineCancellationExecutor = SharedResourceHolder.release(
GrpcUtil.TIMER_SERVICE, deadlineCancellationExecutor);
} }
} }
@Override @Override
public void transportShutdown( public void transportShutdown(EquivalentAddressGroup addressGroup, T transport, Status status) {
EquivalentAddressGroup addressGroup, ClientTransport transport, Status status) {
handleError(status.augmentDescription("Transport to LB server closed")); handleError(status.augmentDescription("Transport to LB server closed"));
synchronized (lock) { synchronized (lock) {
if (transport == lbTransport) { if (transport == lbTransport) {
@ -262,7 +252,7 @@ class GrpclbLoadBalancer extends LoadBalancer {
} }
private void handleError(Status error) { private void handleError(Status error) {
FulfillmentBatch<ClientTransport> pendingPicksFulfillmentBatch; FulfillmentBatch<T> pendingPicksFulfillmentBatch;
StatusException statusException = error.asException(); StatusException statusException = error.asException();
synchronized (lock) { synchronized (lock) {
lastError = statusException; lastError = statusException;
@ -291,7 +281,7 @@ class GrpclbLoadBalancer extends LoadBalancer {
logger.info("Got a LB response: " + response); logger.info("Got a LB response: " + response);
InitialLoadBalanceResponse initialResponse = response.getInitialResponse(); InitialLoadBalanceResponse initialResponse = response.getInitialResponse();
// TODO(zhangkun83): make use of initialResponse // TODO(zhangkun83): make use of initialResponse
RoundRobinServerList.Builder listBuilder = new RoundRobinServerList.Builder(tm); RoundRobinServerList.Builder<T> listBuilder = new RoundRobinServerList.Builder<T>(tm);
ServerList serverList = response.getServerList(); ServerList serverList = response.getServerList();
HashMap<SocketAddress, ResolvedServerInfo> newServerMap = HashMap<SocketAddress, ResolvedServerInfo> newServerMap =
new HashMap<SocketAddress, ResolvedServerInfo>(); new HashMap<SocketAddress, ResolvedServerInfo>();
@ -310,13 +300,13 @@ class GrpclbLoadBalancer extends LoadBalancer {
} }
} }
} }
final RoundRobinServerList newRoundRobinServerList = listBuilder.build(); final RoundRobinServerList<T> newRoundRobinServerList = listBuilder.build();
if (newRoundRobinServerList.size() == 0) { if (newRoundRobinServerList.size() == 0) {
// initialResponse and serverList are under a oneof group. If initialResponse is set, // initialResponse and serverList are under a oneof group. If initialResponse is set,
// serverList will be empty. // serverList will be empty.
return; return;
} }
FulfillmentBatch<ClientTransport> pendingPicksFulfillmentBatch; FulfillmentBatch<T> pendingPicksFulfillmentBatch;
synchronized (lock) { synchronized (lock) {
if (lbResponseObserver != this) { if (lbResponseObserver != this) {
// Make sure I am still the current stream. // Make sure I am still the current stream.
@ -328,9 +318,9 @@ class GrpclbLoadBalancer extends LoadBalancer {
} }
updateRetainedTransports(); updateRetainedTransports();
pendingPicksFulfillmentBatch.link( pendingPicksFulfillmentBatch.link(
new Supplier<ListenableFuture<ClientTransport>>() { new Supplier<ListenableFuture<T>>() {
@Override @Override
public ListenableFuture<ClientTransport> get() { public ListenableFuture<T> get() {
return newRoundRobinServerList.getTransportForNextServer(); return newRoundRobinServerList.getTransportForNextServer();
} }
}); });
@ -348,8 +338,8 @@ class GrpclbLoadBalancer extends LoadBalancer {
private void onStreamClosed(Status status) { private void onStreamClosed(Status status) {
if (status.getCode() == Status.Code.UNIMPLEMENTED) { if (status.getCode() == Status.Code.UNIMPLEMENTED) {
FulfillmentBatch<ClientTransport> pendingPicksFulfillmentBatch; FulfillmentBatch<T> pendingPicksFulfillmentBatch;
final ListenableFuture<ClientTransport> transportFuture; final ListenableFuture<T> transportFuture;
// This LB transport doesn't seem to be an actual LB server, if the LB address comes // This LB transport doesn't seem to be an actual LB server, if the LB address comes
// directly from NameResolver, just use it to serve normal RPCs. // directly from NameResolver, just use it to serve normal RPCs.
// TODO(zhangkun83): check if lbAddresses are from NameResolver after we start getting // TODO(zhangkun83): check if lbAddresses are from NameResolver after we start getting
@ -362,9 +352,9 @@ class GrpclbLoadBalancer extends LoadBalancer {
pendingPicksFulfillmentBatch = pendingPicks.createFulfillmentBatch(); pendingPicksFulfillmentBatch = pendingPicks.createFulfillmentBatch();
} }
pendingPicksFulfillmentBatch.link( pendingPicksFulfillmentBatch.link(
new Supplier<ListenableFuture<ClientTransport>>() { new Supplier<ListenableFuture<T>>() {
@Override @Override
public ListenableFuture<ClientTransport> get() { public ListenableFuture<T> get() {
return transportFuture; return transportFuture;
} }
}); });

4
grpclb/src/main/java/io/grpc/grpclb/GrpclbLoadBalancerFactory.java
View File

@ -56,7 +56,7 @@ public class GrpclbLoadBalancerFactory extends LoadBalancer.Factory {
} }
@Override @Override
public LoadBalancer newLoadBalancer(String serviceName, TransportManager tm) { public <T> LoadBalancer<T> newLoadBalancer(String serviceName, TransportManager<T> tm) {
return new GrpclbLoadBalancer(serviceName, tm); return new GrpclbLoadBalancer<T>(serviceName, tm);
} }
} }

19
grpclb/src/main/java/io/grpc/grpclb/RoundRobinServerList.java
View File

@ -39,7 +39,6 @@ import com.google.common.util.concurrent.ListenableFuture;
import io.grpc.EquivalentAddressGroup; import io.grpc.EquivalentAddressGroup;
import io.grpc.TransportManager; import io.grpc.TransportManager;
import io.grpc.internal.ClientTransport;
import java.net.InetSocketAddress; import java.net.InetSocketAddress;
import java.util.Iterator; import java.util.Iterator;
@ -55,18 +54,18 @@ import javax.annotation.concurrent.ThreadSafe;
// TODO(zhangkun83): possibly move it to io.grpc.internal, as it can also be used by the round-robin // TODO(zhangkun83): possibly move it to io.grpc.internal, as it can also be used by the round-robin
// LoadBalancer. // LoadBalancer.
@ThreadSafe @ThreadSafe
class RoundRobinServerList { class RoundRobinServerList<T> {
private final TransportManager tm; private final TransportManager<T> tm;
private final List<EquivalentAddressGroup> list; private final List<EquivalentAddressGroup> list;
private final Iterator<EquivalentAddressGroup> cyclingIter; private final Iterator<EquivalentAddressGroup> cyclingIter;
private RoundRobinServerList(TransportManager tm, List<EquivalentAddressGroup> list) { private RoundRobinServerList(TransportManager<T> tm, List<EquivalentAddressGroup> list) {
this.tm = tm; this.tm = tm;
this.list = list; this.list = list;
this.cyclingIter = Iterables.cycle(list).iterator(); this.cyclingIter = Iterables.cycle(list).iterator();
} }
ListenableFuture<ClientTransport> getTransportForNextServer() { ListenableFuture<T> getTransportForNextServer() {
EquivalentAddressGroup currentServer; EquivalentAddressGroup currentServer;
synchronized (cyclingIter) { synchronized (cyclingIter) {
// TODO(zhangkun83): receive transportShutdown and transportReady events, then skip addresses // TODO(zhangkun83): receive transportShutdown and transportReady events, then skip addresses
@ -92,12 +91,12 @@ class RoundRobinServerList {
} }
@NotThreadSafe @NotThreadSafe
static class Builder { static class Builder<T> {
private final ImmutableList.Builder<EquivalentAddressGroup> listBuilder = private final ImmutableList.Builder<EquivalentAddressGroup> listBuilder =
ImmutableList.builder(); ImmutableList.builder();
private final TransportManager tm; private final TransportManager<T> tm;
Builder(TransportManager tm) { Builder(TransportManager<T> tm) {
this.tm = tm; this.tm = tm;
} }
@ -108,8 +107,8 @@ class RoundRobinServerList {
listBuilder.add(new EquivalentAddressGroup(addr)); listBuilder.add(new EquivalentAddressGroup(addr));
} }
RoundRobinServerList build() { RoundRobinServerList<T> build() {
return new RoundRobinServerList(tm, listBuilder.build()); return new RoundRobinServerList<T>(tm, listBuilder.build());
} }
} }
} }

64
grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java
View File

@ -55,7 +55,6 @@ import io.grpc.EquivalentAddressGroup;
import io.grpc.ResolvedServerInfo; import io.grpc.ResolvedServerInfo;
import io.grpc.Status; import io.grpc.Status;
import io.grpc.TransportManager; import io.grpc.TransportManager;
import io.grpc.internal.ClientTransport;
import org.junit.Test; import org.junit.Test;
import org.junit.runner.RunWith; import org.junit.runner.RunWith;
@ -77,7 +76,8 @@ import java.util.concurrent.TimeUnit;
public class GrpclbLoadBalancerTest { public class GrpclbLoadBalancerTest {
private static final String serviceName = "testlbservice"; private static final String serviceName = "testlbservice";
private final TransportManager mockTransportManager = mock(TransportManager.class); @SuppressWarnings("unchecked")
private final TransportManager<Transport> mockTransportManager = mock(TransportManager.class);
// The test subject // The test subject
private TestGrpclbLoadBalancer loadBalancer = new TestGrpclbLoadBalancer(); private TestGrpclbLoadBalancer loadBalancer = new TestGrpclbLoadBalancer();
@ -85,28 +85,28 @@ public class GrpclbLoadBalancerTest {
// Current addresses of the LB server // Current addresses of the LB server
private EquivalentAddressGroup lbAddressGroup; private EquivalentAddressGroup lbAddressGroup;
// The future of the currently requested transport for an LB server // The future of the currently requested transport for an LB server
private SettableFuture<ClientTransport> lbTransportFuture; private SettableFuture<Transport> lbTransportFuture;
@Test @Test
public void balancing() throws Exception { public void balancing() throws Exception {
List<ResolvedServerInfo> servers = createResolvedServerInfoList(4000, 4001); List<ResolvedServerInfo> servers = createResolvedServerInfoList(4000, 4001);
// Set up mocks // Set up mocks
List<ClientTransport> transports = new ArrayList<ClientTransport>(servers.size()); List<Transport> transports = new ArrayList<Transport>(servers.size());
List<SettableFuture<ClientTransport>> transportFutures = List<SettableFuture<Transport>> transportFutures =
new ArrayList<SettableFuture<ClientTransport>>(servers.size()); new ArrayList<SettableFuture<Transport>>(servers.size());
for (ResolvedServerInfo server : servers) { for (ResolvedServerInfo server : servers) {
transports.add( transports.add(
mock(ClientTransport.class, withSettings().name("Transport for " + server.toString()))); mock(Transport.class, withSettings().name("Transport for " + server.toString())));
SettableFuture<ClientTransport> future = SettableFuture.create(); SettableFuture<Transport> future = SettableFuture.create();
transportFutures.add(future); transportFutures.add(future);
when(mockTransportManager.getTransport(eq(new EquivalentAddressGroup(server.getAddress())))) when(mockTransportManager.getTransport(eq(new EquivalentAddressGroup(server.getAddress()))))
.thenReturn(future); .thenReturn(future);
} }
ListenableFuture<ClientTransport> pick0; ListenableFuture<Transport> pick0;
ListenableFuture<ClientTransport> pick1; ListenableFuture<Transport> pick1;
// Pick before name resolved // Pick before name resolved
pick0 = loadBalancer.pickTransport(null); pick0 = loadBalancer.pickTransport(null);
@ -118,7 +118,7 @@ public class GrpclbLoadBalancerTest {
pick1 = loadBalancer.pickTransport(null); pick1 = loadBalancer.pickTransport(null);
// Make the transport for LB server ready // Make the transport for LB server ready
ClientTransport lbTransport = mock(ClientTransport.class); Transport lbTransport = new Transport();
lbTransportFuture.set(lbTransport); lbTransportFuture.set(lbTransport);
// An LB request is sent // An LB request is sent
SendLbRequestArgs sentLbRequest = loadBalancer.sentLbRequests.poll(1000, TimeUnit.SECONDS); SendLbRequestArgs sentLbRequest = loadBalancer.sentLbRequests.poll(1000, TimeUnit.SECONDS);
@ -173,7 +173,7 @@ public class GrpclbLoadBalancerTest {
simulateLbAddressResolved(30001); simulateLbAddressResolved(30001);
// Make the transport for LB server ready // Make the transport for LB server ready
ClientTransport lbTransport = mock(ClientTransport.class); Transport lbTransport = new Transport();
lbTransportFuture.set(lbTransport); lbTransportFuture.set(lbTransport);
// An LB request is sent // An LB request is sent
@ -207,7 +207,7 @@ public class GrpclbLoadBalancerTest {
verify(mockTransportManager).getTransport(eq(lbAddressGroup)); verify(mockTransportManager).getTransport(eq(lbAddressGroup));
// Make the transport for LB server ready // Make the transport for LB server ready
ClientTransport lbTransport = mock(ClientTransport.class); Transport lbTransport = new Transport();
lbTransportFuture.set(lbTransport); lbTransportFuture.set(lbTransport);
// An LB request is sent // An LB request is sent
@ -221,7 +221,7 @@ public class GrpclbLoadBalancerTest {
assertNotEquals(lbAddress1, lbAddress2); assertNotEquals(lbAddress1, lbAddress2);
verify(mockTransportManager).updateRetainedTransports(eq(Collections.singleton(lbAddress2))); verify(mockTransportManager).updateRetainedTransports(eq(Collections.singleton(lbAddress2)));
verify(mockTransportManager).getTransport(eq(lbAddressGroup)); verify(mockTransportManager).getTransport(eq(lbAddressGroup));
lbTransport = mock(ClientTransport.class); lbTransport = new Transport();
lbTransportFuture.set(lbTransport); lbTransportFuture.set(lbTransport);
// Another LB request is sent // Another LB request is sent
@ -243,7 +243,7 @@ public class GrpclbLoadBalancerTest {
simulateLbAddressResolved(30001); simulateLbAddressResolved(30001);
// Make the transport for LB server ready // Make the transport for LB server ready
lbTransportFuture.set(mock(ClientTransport.class)); lbTransportFuture.set(new Transport());
// An LB request is sent // An LB request is sent
assertNotNull(loadBalancer.sentLbRequests.poll(1000, TimeUnit.SECONDS)); assertNotNull(loadBalancer.sentLbRequests.poll(1000, TimeUnit.SECONDS));
@ -271,10 +271,10 @@ public class GrpclbLoadBalancerTest {
simulateLbAddressResolved(30001); simulateLbAddressResolved(30001);
// First pick, will be pending // First pick, will be pending
ListenableFuture<ClientTransport> pick = loadBalancer.pickTransport(null); ListenableFuture<Transport> pick = loadBalancer.pickTransport(null);
// Make the transport for LB server ready // Make the transport for LB server ready
ClientTransport lbTransport = mock(ClientTransport.class); Transport lbTransport = new Transport();
lbTransportFuture.set(lbTransport); lbTransportFuture.set(lbTransport);
// An LB request is sent // An LB request is sent
@ -304,10 +304,10 @@ public class GrpclbLoadBalancerTest {
simulateLbAddressResolved(30001); simulateLbAddressResolved(30001);
// First pick, will be pending // First pick, will be pending
ListenableFuture<ClientTransport> pick = loadBalancer.pickTransport(null); ListenableFuture<Transport> pick = loadBalancer.pickTransport(null);
// Make the transport for LB server ready // Make the transport for LB server ready
ClientTransport lbTransport = mock(ClientTransport.class); Transport lbTransport = new Transport();
lbTransportFuture.set(lbTransport); lbTransportFuture.set(lbTransport);
// An LB request is sent // An LB request is sent
@ -356,7 +356,7 @@ public class GrpclbLoadBalancerTest {
simulateLbAddressResolved(30001); simulateLbAddressResolved(30001);
// Make the transport for LB server ready // Make the transport for LB server ready
ClientTransport lbTransport = mock(ClientTransport.class); Transport lbTransport = new Transport();
lbTransportFuture.set(lbTransport); lbTransportFuture.set(lbTransport);
// An LB request is sent // An LB request is sent
@ -378,7 +378,7 @@ public class GrpclbLoadBalancerTest {
verify(mockTransportManager, times(2)).getTransport(eq(lbAddressGroup)); verify(mockTransportManager, times(2)).getTransport(eq(lbAddressGroup));
// Make the new transport ready // Make the new transport ready
lbTransportFuture.set(mock(ClientTransport.class)); lbTransportFuture.set(new Transport());
// Another LB request is sent // Another LB request is sent
assertNotNull(loadBalancer.sentLbRequests.poll(1000, TimeUnit.SECONDS)); assertNotNull(loadBalancer.sentLbRequests.poll(1000, TimeUnit.SECONDS));
@ -389,10 +389,10 @@ public class GrpclbLoadBalancerTest {
simulateLbAddressResolved(30001); simulateLbAddressResolved(30001);
// First pick, will be pending // First pick, will be pending
ListenableFuture<ClientTransport> pick = loadBalancer.pickTransport(null); ListenableFuture<Transport> pick = loadBalancer.pickTransport(null);
// Make the transport for LB server ready // Make the transport for LB server ready
ClientTransport lbTransport = mock(ClientTransport.class); Transport lbTransport = new Transport();
lbTransportFuture.set(lbTransport); lbTransportFuture.set(lbTransport);
// An LB request is sent // An LB request is sent
@ -417,13 +417,13 @@ public class GrpclbLoadBalancerTest {
} }
@Test public void nameResolutionFailed() throws Exception { @Test public void nameResolutionFailed() throws Exception {
ListenableFuture<ClientTransport> pick0 = loadBalancer.pickTransport(null); ListenableFuture<Transport> pick0 = loadBalancer.pickTransport(null);
assertFalse(pick0.isDone()); assertFalse(pick0.isDone());
loadBalancer.handleNameResolutionError(Status.UNAVAILABLE); loadBalancer.handleNameResolutionError(Status.UNAVAILABLE);
assertTrue(pick0.isDone()); assertTrue(pick0.isDone());
ListenableFuture<ClientTransport> pick1 = loadBalancer.pickTransport(null); ListenableFuture<Transport> pick1 = loadBalancer.pickTransport(null);
assertTrue(pick1.isDone()); assertTrue(pick1.isDone());
assertFutureFailedWithError(pick0, Status.Code.UNAVAILABLE, "Name resolution failed"); assertFutureFailedWithError(pick0, Status.Code.UNAVAILABLE, "Name resolution failed");
assertFutureFailedWithError(pick1, Status.Code.UNAVAILABLE, "Name resolution failed"); assertFutureFailedWithError(pick1, Status.Code.UNAVAILABLE, "Name resolution failed");
@ -434,7 +434,7 @@ public class GrpclbLoadBalancerTest {
simulateLbAddressResolved(30001); simulateLbAddressResolved(30001);
// Make the transport for LB server ready // Make the transport for LB server ready
ClientTransport lbTransport = mock(ClientTransport.class); Transport lbTransport = new Transport();
lbTransportFuture.set(lbTransport); lbTransportFuture.set(lbTransport);
// An LB request is sent // An LB request is sent
@ -469,7 +469,7 @@ public class GrpclbLoadBalancerTest {
ResolvedServerInfo lbServerInfo = new ResolvedServerInfo( ResolvedServerInfo lbServerInfo = new ResolvedServerInfo(
new InetSocketAddress("127.0.0.1", lbPort), Attributes.EMPTY); new InetSocketAddress("127.0.0.1", lbPort), Attributes.EMPTY);
lbAddressGroup = buildAddressGroup(lbServerInfo); lbAddressGroup = buildAddressGroup(lbServerInfo);
ClientTransport lbTransport = mock(ClientTransport.class); Transport lbTransport = new Transport();
lbTransportFuture = SettableFuture.create(); lbTransportFuture = SettableFuture.create();
when(mockTransportManager.getTransport(eq(lbAddressGroup))).thenReturn(lbTransportFuture); when(mockTransportManager.getTransport(eq(lbAddressGroup))).thenReturn(lbTransportFuture);
loadBalancer.handleResolvedAddresses(Collections.singletonList(lbServerInfo), Attributes.EMPTY); loadBalancer.handleResolvedAddresses(Collections.singletonList(lbServerInfo), Attributes.EMPTY);
@ -490,7 +490,7 @@ public class GrpclbLoadBalancerTest {
* A slightly modified {@link GrpclbLoadBalancerTest} that saves LB requests in a queue instead of * A slightly modified {@link GrpclbLoadBalancerTest} that saves LB requests in a queue instead of
* sending them out. * sending them out.
*/ */
private class TestGrpclbLoadBalancer extends GrpclbLoadBalancer { private class TestGrpclbLoadBalancer extends GrpclbLoadBalancer<Transport> {
final LinkedBlockingQueue<SendLbRequestArgs> sentLbRequests = final LinkedBlockingQueue<SendLbRequestArgs> sentLbRequests =
new LinkedBlockingQueue<SendLbRequestArgs>(); new LinkedBlockingQueue<SendLbRequestArgs>();
@ -498,16 +498,16 @@ public class GrpclbLoadBalancerTest {
super(serviceName, mockTransportManager); super(serviceName, mockTransportManager);
} }
@Override void sendLbRequest(ClientTransport transport, LoadBalanceRequest request) { @Override void sendLbRequest(Transport transport, LoadBalanceRequest request) {
sentLbRequests.add(new SendLbRequestArgs(transport, request)); sentLbRequests.add(new SendLbRequestArgs(transport, request));
} }
} }
private static class SendLbRequestArgs { private static class SendLbRequestArgs {
final ClientTransport transport; final Transport transport;
final LoadBalanceRequest request; final LoadBalanceRequest request;
SendLbRequestArgs(ClientTransport transport, LoadBalanceRequest request) { SendLbRequestArgs(Transport transport, LoadBalanceRequest request) {
this.transport = transport; this.transport = transport;
this.request = request; this.request = request;
} }
@ -563,4 +563,6 @@ public class GrpclbLoadBalancerTest {
} }
} }
} }
public static class Transport {}
} }