From 7e49fc0c4fe912872e5f3cc67ee21b90597a90ed Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Wed, 23 Sep 2020 12:07:04 -0400 Subject: [PATCH 1/2] feat: Allow user-agents to be specified by both internal headers and user headers This is primarily meant to address the https://github.com/googleapis/java-bigtable/pull/404#pullrequestreview-480972135. Where both the hbase adapter and the core client needs to set the UserAgent. To reconcile the conflict this PR takes grpc's approach and concatenates the values. The reason this needs to happen now is that for DirectPath we need to ship clients that send a UserAgent, however Bigtable needs to keep stats to differentiate client usage. --- .../com/google/api/gax/rpc/ClientContext.java | 45 ++++++----- .../google/api/gax/rpc/ClientContextTest.java | 80 +++++++++++++++++++ 2 files changed, 107 insertions(+), 18 deletions(-) diff --git a/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java b/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java index 1dca243d9..ac2335235 100644 --- a/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java +++ b/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java @@ -42,10 +42,13 @@ import com.google.auto.value.AutoValue; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Sets; import java.io.IOException; import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.Executor; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; @@ -222,27 +225,33 @@ public static ClientContext create(StubSettings settings) throws IOException { * Project Id. */ private static Map getHeadersFromSettings(StubSettings settings) { - ImmutableMap.Builder headersBuilder = ImmutableMap.builder(); - if (settings.getQuotaProjectId() != null) { - headersBuilder.put(QUOTA_PROJECT_ID_HEADER_KEY, settings.getQuotaProjectId()); - for (Map.Entry entry : settings.getHeaderProvider().getHeaders().entrySet()) { - if (entry.getKey().equals(QUOTA_PROJECT_ID_HEADER_KEY)) { - continue; - } - headersBuilder.put(entry); + // Resolve conflicts when merging headers from multiple sources + Map userHeaders = settings.getHeaderProvider().getHeaders(); + Map internalHeaders = settings.getInternalHeaderProvider().getHeaders(); + Map conflictResolution = new HashMap<>(); + + Set conflicts = Sets.intersection(userHeaders.keySet(), internalHeaders.keySet()); + for (String key : conflicts) { + if ("user-agent".equals(key)) { + conflictResolution.put(key, userHeaders.get(key) + " " + internalHeaders.get(key)); + continue; } - for (Map.Entry entry : - settings.getInternalHeaderProvider().getHeaders().entrySet()) { - if (entry.getKey().equals(QUOTA_PROJECT_ID_HEADER_KEY)) { - continue; - } - headersBuilder.put(entry); + // Backwards compat: quota project id can conflict if its overriden in settings + if (QUOTA_PROJECT_ID_HEADER_KEY.equals(key) && settings.getQuotaProjectId() != null) { + continue; } - } else { - headersBuilder.putAll(settings.getHeaderProvider().getHeaders()); - headersBuilder.putAll(settings.getInternalHeaderProvider().getHeaders()); + throw new IllegalArgumentException("Header provider can't override the header: " + key); + } + if (settings.getQuotaProjectId() != null) { + conflictResolution.put(QUOTA_PROJECT_ID_HEADER_KEY, settings.getQuotaProjectId()); } - return headersBuilder.build(); + + Map effectiveHeaders = new HashMap<>(); + effectiveHeaders.putAll(internalHeaders); + effectiveHeaders.putAll(userHeaders); + effectiveHeaders.putAll(conflictResolution); + + return ImmutableMap.copyOf(effectiveHeaders); } @AutoValue.Builder diff --git a/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java b/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java index 0f7305ca2..71b03bcc8 100644 --- a/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java +++ b/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java @@ -36,6 +36,7 @@ import com.google.api.gax.core.CredentialsProvider; import com.google.api.gax.core.ExecutorProvider; import com.google.api.gax.core.FixedCredentialsProvider; +import com.google.api.gax.core.FixedExecutorProvider; import com.google.api.gax.rpc.testing.FakeChannel; import com.google.api.gax.rpc.testing.FakeClientSettings; import com.google.api.gax.rpc.testing.FakeTransportChannel; @@ -193,6 +194,14 @@ public TransportChannelProvider withCredentials(Credentials credentials) { } } + private class FakeChannelWithHeaders extends FakeChannel { + private final Map headers; + + FakeChannelWithHeaders(Map headers) { + this.headers = headers; + } + } + @Test public void testNoAutoCloseContextNeedsNoExecutor() throws Exception { runTest(false, false, false, false); @@ -526,4 +535,75 @@ public Credentials getCredentials() throws IOException { ClientContext clientContext = ClientContext.create(builder.build()); assertThat(clientContext.getCredentials().getRequestMetadata(null)).isEqualTo(metaData); } + + @Test + public void testUserAgentInternalOnly() throws Exception { + TransportChannelProvider transportChannelProvider = + new FakeTransportProvider( + FakeTransportChannel.create(new FakeChannel()), null, true, null, null); + + ClientSettings.Builder builder = + new FakeClientSettings.Builder() + .setExecutorProvider( + FixedExecutorProvider.create(Mockito.mock(ScheduledExecutorService.class))) + .setTransportChannelProvider(transportChannelProvider) + .setCredentialsProvider( + FixedCredentialsProvider.create(Mockito.mock(GoogleCredentials.class))); + + builder.setInternalHeaderProvider(FixedHeaderProvider.create("user-agent", "internal-agent")); + + ClientContext clientContext = ClientContext.create(builder.build()); + FakeTransportChannel transportChannel = + (FakeTransportChannel) clientContext.getTransportChannel(); + + assertThat(transportChannel.getHeaders()).containsEntry("user-agent", "internal-agent"); + } + + @Test + public void testUserAgentExternalOnly() throws Exception { + TransportChannelProvider transportChannelProvider = + new FakeTransportProvider( + FakeTransportChannel.create(new FakeChannel()), null, true, null, null); + + ClientSettings.Builder builder = + new FakeClientSettings.Builder() + .setExecutorProvider( + FixedExecutorProvider.create(Mockito.mock(ScheduledExecutorService.class))) + .setTransportChannelProvider(transportChannelProvider) + .setCredentialsProvider( + FixedCredentialsProvider.create(Mockito.mock(GoogleCredentials.class))); + + builder.setHeaderProvider(FixedHeaderProvider.create("user-agent", "user-supplied-agent")); + + ClientContext clientContext = ClientContext.create(builder.build()); + FakeTransportChannel transportChannel = + (FakeTransportChannel) clientContext.getTransportChannel(); + + assertThat(transportChannel.getHeaders()).containsEntry("user-agent", "user-supplied-agent"); + } + + @Test + public void testUserAgentConcat() throws Exception { + TransportChannelProvider transportChannelProvider = + new FakeTransportProvider( + FakeTransportChannel.create(new FakeChannel()), null, true, null, null); + + ClientSettings.Builder builder = + new FakeClientSettings.Builder() + .setExecutorProvider( + FixedExecutorProvider.create(Mockito.mock(ScheduledExecutorService.class))) + .setTransportChannelProvider(transportChannelProvider) + .setCredentialsProvider( + FixedCredentialsProvider.create(Mockito.mock(GoogleCredentials.class))); + + builder.setHeaderProvider(FixedHeaderProvider.create("user-agent", "user-supplied-agent")); + builder.setInternalHeaderProvider(FixedHeaderProvider.create("user-agent", "internal-agent")); + + ClientContext clientContext = ClientContext.create(builder.build()); + FakeTransportChannel transportChannel = + (FakeTransportChannel) clientContext.getTransportChannel(); + + assertThat(transportChannel.getHeaders()) + .containsEntry("user-agent", "user-supplied-agent internal-agent"); + } } From 429a5803590e0d186f6f8f95a47ad7c04233ba77 Mon Sep 17 00:00:00 2001 From: Igor Bernstein Date: Wed, 23 Sep 2020 12:09:28 -0400 Subject: [PATCH 2/2] remove stale code --- .../java/com/google/api/gax/rpc/ClientContextTest.java | 8 -------- 1 file changed, 8 deletions(-) diff --git a/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java b/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java index 71b03bcc8..58a007644 100644 --- a/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java +++ b/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java @@ -194,14 +194,6 @@ public TransportChannelProvider withCredentials(Credentials credentials) { } } - private class FakeChannelWithHeaders extends FakeChannel { - private final Map headers; - - FakeChannelWithHeaders(Map headers) { - this.headers = headers; - } - } - @Test public void testNoAutoCloseContextNeedsNoExecutor() throws Exception { runTest(false, false, false, false);