Skip to content

Commit 7c4ce96

Browse files
igorbernstein2igorbernstein2
authored
chore: support the rest of RLS headers for priming (#9749)
NOTE: this depends on and includes #9748
1 parent da1b9e5 commit 7c4ce96

File tree

5 files changed

+280
-165
lines changed

5 files changed

+280
-165
lines changed

‎bigtable/bigtable-proxy/src/main/java/com/google/cloud/bigtable/examples/proxy/core/CallLabels.java‎

Lines changed: 72 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import com.google.auto.value.AutoValue;
2020
import com.google.bigtable.v2.PingAndWarmRequest;
2121
import com.google.bigtable.v2.PingAndWarmRequest.Builder;
22+
import com.google.common.annotations.VisibleForTesting;
2223
import com.google.common.collect.ImmutableMap;
2324
import io.grpc.Metadata;
2425
import io.grpc.Metadata.Key;
@@ -29,6 +30,8 @@
2930
import java.util.Map;
3031
import java.util.Map.Entry;
3132
import java.util.Optional;
33+
import org.slf4j.Logger;
34+
import org.slf4j.LoggerFactory;
3235

3336
/**
3437
* A value class to encapsulate call identity.
@@ -43,9 +46,18 @@
4346
*/
4447
@AutoValue
4548
public abstract class CallLabels {
46-
public static final Key<String> REQUEST_PARAMS =
49+
private static final Logger LOG = LoggerFactory.getLogger(CallLabels.class);
50+
51+
// All RLS headers
52+
static final Key<String> REQUEST_PARAMS =
4753
Key.of("x-goog-request-params", Metadata.ASCII_STRING_MARSHALLER);
48-
private static final Key<String> API_CLIENT =
54+
static final Key<String> LEGACY_RESOURCE_PREFIX =
55+
Key.of("google-cloud-resource-prefix", Metadata.ASCII_STRING_MARSHALLER);
56+
static final Key<String> ROUTING_COOKIE =
57+
Key.of("x-goog-cbt-cookie-routing", Metadata.ASCII_STRING_MARSHALLER);
58+
static final Key<String> FEATURE_FLAGS =
59+
Key.of("bigtable-features", Metadata.ASCII_STRING_MARSHALLER);
60+
static final Key<String> API_CLIENT =
4961
Key.of("x-goog-api-client", Metadata.ASCII_STRING_MARSHALLER);
5062

5163
enum ResourceNameType {
@@ -74,36 +86,56 @@ static ResourceName create(ResourceNameType type, String value) {
7486
}
7587
}
7688

77-
public abstract Optional<String> getApiClient();
89+
public abstract String getMethodName();
7890

79-
public abstract Optional<String> getResourceName();
91+
abstract Optional<String> getRequestParams();
8092

81-
public abstract Optional<String> getAppProfileId();
93+
abstract Optional<String> getLegacyResourcePrefix();
8294

83-
public abstract String getMethodName();
95+
abstract Optional<String> getRoutingCookie();
96+
97+
abstract Optional<String> getEncodedFeatures();
98+
99+
public abstract Optional<String> getApiClient();
84100

85101
public static CallLabels create(MethodDescriptor<?, ?> method, Metadata headers) {
86102
Optional<String> apiClient = Optional.ofNullable(headers.get(API_CLIENT));
87103

88-
String requestParams = Optional.ofNullable(headers.get(REQUEST_PARAMS)).orElse("");
89-
String[] encodedKvPairs = requestParams.split("&");
90-
Optional<String> resourceName = extractResourceName(encodedKvPairs).map(ResourceName::getValue);
91-
Optional<String> appProfile = extractAppProfileId(encodedKvPairs);
104+
Optional<String> requestParams = Optional.ofNullable(headers.get(REQUEST_PARAMS));
105+
Optional<String> legacyResourcePrefix =
106+
Optional.ofNullable(headers.get(LEGACY_RESOURCE_PREFIX));
107+
Optional<String> routingCookie = Optional.ofNullable(headers.get(ROUTING_COOKIE));
108+
Optional<String> encodedFeatures = Optional.ofNullable(headers.get(FEATURE_FLAGS));
92109

93-
return create(method, apiClient, resourceName, appProfile);
110+
return create(
111+
method, requestParams, legacyResourcePrefix, routingCookie, encodedFeatures, apiClient);
94112
}
95113

114+
@VisibleForTesting
96115
public static CallLabels create(
97116
MethodDescriptor<?, ?> method,
98-
Optional<String> apiClient,
99-
Optional<String> resourceName,
100-
Optional<String> appProfile) {
117+
Optional<String> requestParams,
118+
Optional<String> legacyResourcePrefix,
119+
Optional<String> routingCookie,
120+
Optional<String> encodedFeatures,
121+
Optional<String> apiClient) {
101122

102123
return new AutoValue_CallLabels(
103-
apiClient, resourceName, appProfile, method.getFullMethodName());
124+
method.getFullMethodName(),
125+
requestParams,
126+
legacyResourcePrefix,
127+
routingCookie,
128+
encodedFeatures,
129+
apiClient);
104130
}
105131

106-
private static Optional<ResourceName> extractResourceName(String[] encodedKvPairs) {
132+
public Optional<String> extractResourceName() throws ParsingException {
133+
if (getRequestParams().isEmpty()) {
134+
return getLegacyResourcePrefix();
135+
}
136+
137+
String requestParams = getRequestParams().orElse("");
138+
String[] encodedKvPairs = requestParams.split("&");
107139
Optional<ResourceName> resourceName = Optional.empty();
108140

109141
for (String encodedKv : encodedKvPairs) {
@@ -132,10 +164,10 @@ private static Optional<ResourceName> extractResourceName(String[] encodedKvPair
132164

133165
resourceName = Optional.of(ResourceName.create(newType.get(), decodedValue));
134166
}
135-
return resourceName;
167+
return resourceName.map(ResourceName::getValue);
136168
}
137169

138-
private static Optional<ResourceNameType> findType(String encodedKey) {
170+
private static Optional<ResourceNameType> findType(String encodedKey) throws ParsingException {
139171
String decodedKey = percentDecode(encodedKey);
140172

141173
for (ResourceNameType type : ResourceNameType.values()) {
@@ -146,8 +178,10 @@ private static Optional<ResourceNameType> findType(String encodedKey) {
146178
return Optional.empty();
147179
}
148180

149-
private static Optional<String> extractAppProfileId(String[] encodedKvPairs) {
150-
for (String encodedPair : encodedKvPairs) {
181+
public Optional<String> extractAppProfileId() throws ParsingException {
182+
String requestParams = getRequestParams().orElse("");
183+
184+
for (String encodedPair : requestParams.split("&")) {
151185
if (!encodedPair.startsWith("app_profile_id=")) {
152186
continue;
153187
}
@@ -158,8 +192,12 @@ private static Optional<String> extractAppProfileId(String[] encodedKvPairs) {
158192
return Optional.empty();
159193
}
160194

161-
private static String percentDecode(String s) {
162-
return URLDecoder.decode(s, StandardCharsets.UTF_8);
195+
private static String percentDecode(String s) throws ParsingException {
196+
try {
197+
return URLDecoder.decode(s, StandardCharsets.UTF_8);
198+
} catch (RuntimeException e) {
199+
throw new ParsingException("Failed to url decode " + s, e);
200+
}
163201
}
164202

165203
@AutoValue
@@ -171,7 +209,9 @@ public abstract static class PrimingKey {
171209
abstract Optional<String> getAppProfileId();
172210

173211
public static Optional<PrimingKey> from(CallLabels labels) throws ParsingException {
174-
Optional<String> resourceName = labels.getResourceName();
212+
final ImmutableMap.Builder<String, String> md = ImmutableMap.builder();
213+
214+
Optional<String> resourceName = labels.extractResourceName();
175215
if (resourceName.isEmpty()) {
176216
return Optional.empty();
177217
}
@@ -188,12 +228,18 @@ public static Optional<PrimingKey> from(CallLabels labels) throws ParsingExcepti
188228
.append("name=")
189229
.append(URLEncoder.encode(instanceName, StandardCharsets.UTF_8));
190230

191-
Optional<String> appProfileId = labels.getAppProfileId();
231+
Optional<String> appProfileId = labels.extractAppProfileId();
192232
appProfileId.ifPresent(val -> reqParams.append("&app_profile_id=").append(val));
193-
194-
ImmutableMap.Builder<String, String> md = ImmutableMap.builder();
195233
md.put(REQUEST_PARAMS.name(), reqParams.toString());
196234

235+
labels
236+
.getLegacyResourcePrefix()
237+
.ifPresent(ignored -> md.put(LEGACY_RESOURCE_PREFIX.name(), instanceName));
238+
239+
labels.getRoutingCookie().ifPresent(c -> md.put(ROUTING_COOKIE.name(), c));
240+
241+
labels.getEncodedFeatures().ifPresent(c -> md.put(FEATURE_FLAGS.name(), c));
242+
197243
labels.getApiClient().ifPresent(c -> md.put(API_CLIENT.name(), c));
198244

199245
return Optional.of(

‎bigtable/bigtable-proxy/src/main/java/com/google/cloud/bigtable/examples/proxy/metrics/MetricsImpl.java‎

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,13 @@
1919
import com.google.auth.Credentials;
2020
import com.google.auto.value.AutoValue;
2121
import com.google.cloud.bigtable.examples.proxy.core.CallLabels;
22+
import com.google.cloud.bigtable.examples.proxy.core.CallLabels.ParsingException;
2223
import com.google.cloud.opentelemetry.metric.GoogleCloudMetricExporter;
2324
import com.google.cloud.opentelemetry.metric.MetricConfiguration;
2425
import io.grpc.Status;
2526
import io.opentelemetry.api.common.AttributeKey;
2627
import io.opentelemetry.api.common.Attributes;
28+
import io.opentelemetry.api.common.AttributesBuilder;
2729
import io.opentelemetry.api.metrics.DoubleHistogram;
2830
import io.opentelemetry.api.metrics.LongCounter;
2931
import io.opentelemetry.api.metrics.LongHistogram;
@@ -41,8 +43,12 @@
4143
import java.time.Duration;
4244
import java.util.concurrent.atomic.AtomicInteger;
4345
import java.util.function.Supplier;
46+
import org.slf4j.Logger;
47+
import org.slf4j.LoggerFactory;
4448

4549
public class MetricsImpl implements Closeable, Metrics {
50+
private static final Logger LOG = LoggerFactory.getLogger(MetricsImpl.class);
51+
4652
public static final InstrumentationScopeInfo INSTRUMENTATION_SCOPE_INFO =
4753
InstrumentationScopeInfo.create("bigtable-proxy");
4854

@@ -186,13 +192,30 @@ public void close() throws IOException {
186192

187193
@Override
188194
public MetricsAttributesImpl createAttributes(CallLabels callLabels) {
189-
return new AutoValue_MetricsImpl_MetricsAttributesImpl(
195+
AttributesBuilder attrs =
190196
Attributes.builder()
191-
.put(MetricsImpl.API_CLIENT_KEY, callLabels.getApiClient().orElse("<missing>"))
192-
.put(MetricsImpl.RESOURCE_KEY, callLabels.getResourceName().orElse("<missing>"))
193-
.put(MetricsImpl.APP_PROFILE_KEY, callLabels.getAppProfileId().orElse("<missing>"))
194-
.put(MetricsImpl.METHOD_KEY, callLabels.getMethodName())
195-
.build());
197+
.put(METHOD_KEY, callLabels.getMethodName())
198+
.put(API_CLIENT_KEY, callLabels.getApiClient().orElse("<missing>"));
199+
200+
String resourceValue;
201+
try {
202+
resourceValue = callLabels.extractResourceName().orElse("<missing>");
203+
} catch (ParsingException e) {
204+
LOG.atWarn().log("Failed to extract resource from callLabels: {}", callLabels, e);
205+
resourceValue = "<error>";
206+
}
207+
attrs.put(MetricsImpl.RESOURCE_KEY, resourceValue);
208+
209+
String appProfile;
210+
try {
211+
appProfile = callLabels.extractAppProfileId().orElse("<missing>");
212+
} catch (ParsingException e) {
213+
LOG.atWarn().log("Failed to extract app profile from callLabels: {}", callLabels, e);
214+
appProfile = "<error>";
215+
}
216+
attrs.put(MetricsImpl.APP_PROFILE_KEY, appProfile);
217+
218+
return new AutoValue_MetricsImpl_MetricsAttributesImpl(attrs.build());
196219
}
197220

198221
private static SdkMeterProvider createMeterProvider(Credentials credentials, String projectId) {

0 commit comments

Comments
 (0)