grpclb: expose the token metadata key. (#2944)

It is accessed by the server-side load-reporting code which is from a
different package.
This commit is contained in:
Kun Zhang 2017-04-25 15:17:26 -07:00 committed by GitHub
parent 49bde5494b
commit c1a2fb412a
3 changed files with 15 additions and 11 deletions

View File

@ -33,6 +33,7 @@ package io.grpc.grpclb;
import io.grpc.Attributes; import io.grpc.Attributes;
import io.grpc.ExperimentalApi; import io.grpc.ExperimentalApi;
import io.grpc.Metadata;
/** /**
* Constants for the GRPCLB load-balancer. * Constants for the GRPCLB load-balancer.
@ -61,5 +62,12 @@ public final class GrpclbConstants {
public static final Attributes.Key<String> ATTR_LB_ADDR_AUTHORITY = public static final Attributes.Key<String> ATTR_LB_ADDR_AUTHORITY =
Attributes.Key.of("io.grpc.grpclb.lbAddrAuthority"); Attributes.Key.of("io.grpc.grpclb.lbAddrAuthority");
/**
* The opaque token given by the remote balancer for each returned server address. The client
* will send this token with any requests sent to the associated server.
*/
public static final Metadata.Key<String> TOKEN_METADATA_KEY =
Metadata.Key.of("lb-token", Metadata.ASCII_STRING_MARSHALLER);
private GrpclbConstants() { } private GrpclbConstants() { }
} }

View File

@ -95,10 +95,6 @@ class GrpclbLoadBalancer extends LoadBalancer implements WithLogId {
private static final Attributes.Key<AtomicReference<ConnectivityStateInfo>> STATE_INFO = private static final Attributes.Key<AtomicReference<ConnectivityStateInfo>> STATE_INFO =
Attributes.Key.of("io.grpc.grpclb.GrpclbLoadBalancer.stateInfo"); Attributes.Key.of("io.grpc.grpclb.GrpclbLoadBalancer.stateInfo");
@VisibleForTesting
static final Metadata.Key<String> TOKEN_KEY =
Metadata.Key.of("lb-token", Metadata.ASCII_STRING_MARSHALLER);
@VisibleForTesting @VisibleForTesting
static final RoundRobinEntry DROP_ENTRY = static final RoundRobinEntry DROP_ENTRY =
new RoundRobinEntry(Status.UNAVAILABLE.withDescription("Drop requested by balancer")); new RoundRobinEntry(Status.UNAVAILABLE.withDescription("Drop requested by balancer"));
@ -548,8 +544,8 @@ class GrpclbLoadBalancer extends LoadBalancer implements WithLogId {
void updateHeaders(Metadata headers) { void updateHeaders(Metadata headers) {
if (token != null) { if (token != null) {
headers.discardAll(TOKEN_KEY); headers.discardAll(GrpclbConstants.TOKEN_METADATA_KEY);
headers.put(TOKEN_KEY, token); headers.put(GrpclbConstants.TOKEN_METADATA_KEY, token);
} }
} }

View File

@ -291,30 +291,30 @@ public class GrpclbLoadBalancerTest {
when(args1.getHeaders()).thenReturn(headers1); when(args1.getHeaders()).thenReturn(headers1);
assertSame(r1.result, picker.pickSubchannel(args1)); assertSame(r1.result, picker.pickSubchannel(args1));
verify(args1).getHeaders(); verify(args1).getHeaders();
assertFalse(headers1.containsKey(GrpclbLoadBalancer.TOKEN_KEY)); assertFalse(headers1.containsKey(GrpclbConstants.TOKEN_METADATA_KEY));
PickSubchannelArgs args2 = mock(PickSubchannelArgs.class); PickSubchannelArgs args2 = mock(PickSubchannelArgs.class);
Metadata headers2 = new Metadata(); Metadata headers2 = new Metadata();
// The existing token on the headers will be replaced // The existing token on the headers will be replaced
headers2.put(GrpclbLoadBalancer.TOKEN_KEY, "LBTOKEN__OLD"); headers2.put(GrpclbConstants.TOKEN_METADATA_KEY, "LBTOKEN__OLD");
when(args2.getHeaders()).thenReturn(headers2); when(args2.getHeaders()).thenReturn(headers2);
assertSame(r2.result, picker.pickSubchannel(args2)); assertSame(r2.result, picker.pickSubchannel(args2));
verify(args2).getHeaders(); verify(args2).getHeaders();
assertThat(headers2.getAll(GrpclbLoadBalancer.TOKEN_KEY)).containsExactly("LBTOKEN0001"); assertThat(headers2.getAll(GrpclbConstants.TOKEN_METADATA_KEY)).containsExactly("LBTOKEN0001");
PickSubchannelArgs args3 = mock(PickSubchannelArgs.class); PickSubchannelArgs args3 = mock(PickSubchannelArgs.class);
Metadata headers3 = new Metadata(); Metadata headers3 = new Metadata();
when(args3.getHeaders()).thenReturn(headers3); when(args3.getHeaders()).thenReturn(headers3);
assertSame(r3.result, picker.pickSubchannel(args3)); assertSame(r3.result, picker.pickSubchannel(args3));
verify(args3).getHeaders(); verify(args3).getHeaders();
assertThat(headers3.getAll(GrpclbLoadBalancer.TOKEN_KEY)).containsExactly("LBTOKEN0002"); assertThat(headers3.getAll(GrpclbConstants.TOKEN_METADATA_KEY)).containsExactly("LBTOKEN0002");
PickSubchannelArgs args4 = mock(PickSubchannelArgs.class); PickSubchannelArgs args4 = mock(PickSubchannelArgs.class);
Metadata headers4 = new Metadata(); Metadata headers4 = new Metadata();
when(args4.getHeaders()).thenReturn(headers4); when(args4.getHeaders()).thenReturn(headers4);
assertSame(r1.result, picker.pickSubchannel(args4)); assertSame(r1.result, picker.pickSubchannel(args4));
verify(args4).getHeaders(); verify(args4).getHeaders();
assertFalse(headers4.containsKey(GrpclbLoadBalancer.TOKEN_KEY)); assertFalse(headers4.containsKey(GrpclbConstants.TOKEN_METADATA_KEY));
verify(subchannel, never()).getAttributes(); verify(subchannel, never()).getAttributes();
} }