Skip to content

Commit 33ce2f7

Browse files
lqiu96lqiu96
committed
chore: Address PR feedback
1 parent 51399b5 commit 33ce2f7

4 files changed

Lines changed: 172 additions & 61 deletions

File tree

‎java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/BuiltInDatastoreMetricsProvider.java‎

Lines changed: 38 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,14 @@
1717
package com.google.cloud.datastore.telemetry;
1818

1919
import com.google.auth.Credentials;
20+
import com.google.cloud.NoCredentials;
2021
import io.opentelemetry.api.OpenTelemetry;
2122
import io.opentelemetry.api.common.Attributes;
2223
import io.opentelemetry.api.common.AttributesBuilder;
2324
import io.opentelemetry.sdk.OpenTelemetrySdk;
2425
import io.opentelemetry.sdk.metrics.SdkMeterProvider;
2526
import io.opentelemetry.sdk.metrics.SdkMeterProviderBuilder;
2627
import io.opentelemetry.sdk.resources.Resource;
27-
import java.io.IOException;
2828
import java.lang.management.ManagementFactory;
2929
import java.lang.reflect.Method;
3030
import java.net.InetAddress;
@@ -56,7 +56,6 @@ public final class BuiltInDatastoreMetricsProvider {
5656
private static final Logger logger =
5757
Logger.getLogger(BuiltInDatastoreMetricsProvider.class.getName());
5858

59-
private static volatile String taskId;
6059
private static volatile String location;
6160
private static final String DEFAULT_LOCATION = "global";
6261

@@ -94,25 +93,34 @@ private Map<String, String> buildClientAttributes() {
9493
*/
9594
@Nullable
9695
public OpenTelemetry createOpenTelemetry(
97-
@Nonnull String projectId, @Nonnull String databaseId, @Nullable Credentials credentials) {
98-
try {
99-
SdkMeterProviderBuilder sdkMeterProviderBuilder = SdkMeterProvider.builder();
100-
// Register Datastore-specific views and the PeriodicMetricReader.
101-
DatastoreBuiltInMetricsView.registerBuiltinMetrics(
102-
DatastoreCloudMonitoringExporter.create(projectId, credentials), sdkMeterProviderBuilder);
103-
// Configure the monitored resource attributes for this specific client.
104-
sdkMeterProviderBuilder.setResource(
105-
Resource.create(createResourceAttributes(projectId, databaseId)));
106-
SdkMeterProvider sdkMeterProvider = sdkMeterProviderBuilder.build();
107-
return OpenTelemetrySdk.builder().setMeterProvider(sdkMeterProvider).build();
108-
} catch (IOException ex) {
96+
@Nonnull String projectId, @Nonnull String databaseId, @Nonnull Credentials credentials) {
97+
SdkMeterProviderBuilder sdkMeterProviderBuilder = SdkMeterProvider.builder();
98+
99+
// Generate unique client attributes (including unique taskId) for this specific client
100+
// instance.
101+
Map<String, String> clientAttributes = buildClientAttributes();
102+
103+
if (credentials instanceof NoCredentials) {
109104
logger.log(
110105
Level.WARNING,
111-
"Unable to create OpenTelemetry instance for client side metrics, will skip exporting"
112-
+ " built-in metrics",
113-
ex);
106+
"Built-in metrics exporting is disabled when using NoCredentials (emulator).");
107+
return null;
108+
}
109+
110+
DatastoreCloudMonitoringExporter exporter =
111+
DatastoreCloudMonitoringExporter.create(
112+
projectId, databaseId, credentials, clientAttributes);
113+
if (exporter == null) {
114114
return null;
115115
}
116+
117+
// Register Datastore-specific views and the PeriodicMetricReader.
118+
DatastoreBuiltInMetricsView.registerBuiltinMetrics(exporter, sdkMeterProviderBuilder);
119+
// Configure the monitored resource attributes for this specific client.
120+
sdkMeterProviderBuilder.setResource(
121+
Resource.create(createResourceAttributes(projectId, databaseId)));
122+
SdkMeterProvider sdkMeterProvider = sdkMeterProviderBuilder.build();
123+
return OpenTelemetrySdk.builder().setMeterProvider(sdkMeterProvider).build();
116124
}
117125

118126
/**
@@ -173,22 +181,23 @@ Map<String, String> getClientAttributes() {
173181
* <p>For Java 9 and later, the PID is obtained using the ProcessHandle API. For Java 8, the PID
174182
* is extracted from ManagementFactory.getRuntimeMXBean().getName().
175183
*
184+
* <p><b>Note</b>: This method generates a new value every time it is called to ensure that each
185+
* client instance gets a unique ID. It should be called sparingly (e.g., once per client
186+
* creation) to avoid performance overhead from UUID generation and hostname lookup.
187+
*
176188
* @return a unique identifier string.
177189
*/
178190
private static String getDefaultTaskValue() {
179-
if (taskId == null) {
180-
String identifier = UUID.randomUUID().toString();
181-
String pid = getProcessId();
182-
183-
try {
184-
String hostname = InetAddress.getLocalHost().getHostName();
185-
taskId = identifier + "@" + pid + "@" + hostname;
186-
} catch (UnknownHostException e) {
187-
logger.log(Level.INFO, "Unable to get the hostname.", e);
188-
taskId = identifier + "@" + pid + "@localhost";
189-
}
191+
String identifier = UUID.randomUUID().toString();
192+
String pid = getProcessId();
193+
194+
try {
195+
String hostname = InetAddress.getLocalHost().getHostName();
196+
return identifier + "@" + pid + "@" + hostname;
197+
} catch (UnknownHostException e) {
198+
logger.log(Level.CONFIG, "Unable to get the hostname.", e);
199+
return identifier + "@" + pid + "@localhost";
190200
}
191-
return taskId;
192201
}
193202

194203
private static String getProcessId() {

‎java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DatastoreCloudMonitoringExporter.java‎

Lines changed: 90 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,10 @@
1919
import com.google.api.core.ApiFuture;
2020
import com.google.api.core.ApiFutureCallback;
2121
import com.google.api.core.ApiFutures;
22-
import com.google.api.gax.core.CredentialsProvider;
2322
import com.google.api.gax.core.FixedCredentialsProvider;
24-
import com.google.api.gax.core.NoCredentialsProvider;
2523
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
2624
import com.google.api.gax.rpc.PermissionDeniedException;
2725
import com.google.auth.Credentials;
28-
import com.google.cloud.NoCredentials;
2926
import com.google.cloud.monitoring.v3.MetricServiceClient;
3027
import com.google.cloud.monitoring.v3.MetricServiceSettings;
3128
import com.google.common.annotations.VisibleForTesting;
@@ -45,7 +42,10 @@
4542
import java.util.ArrayList;
4643
import java.util.Collection;
4744
import java.util.List;
45+
import java.util.Map;
46+
import java.util.concurrent.ConcurrentHashMap;
4847
import java.util.concurrent.atomic.AtomicBoolean;
48+
import java.util.concurrent.atomic.AtomicInteger;
4949
import java.util.logging.Level;
5050
import java.util.logging.Logger;
5151
import javax.annotation.Nonnull;
@@ -71,7 +71,29 @@ class DatastoreCloudMonitoringExporter implements MetricExporter {
7171
private static final Logger logger =
7272
Logger.getLogger(DatastoreCloudMonitoringExporter.class.getName());
7373

74+
/**
75+
* Wrapper class to hold a {@link MetricServiceClient} and its reference count. This is used to
76+
* share the client across multiple exporter instances.
77+
*/
78+
static class CachedMetricsClient {
79+
final MetricServiceClient client;
80+
final AtomicInteger refCount = new AtomicInteger(0);
81+
82+
CachedMetricsClient(MetricServiceClient client) {
83+
this.client = client;
84+
}
85+
}
86+
87+
/**
88+
* Shared cache for {@link MetricServiceClient} instances, keyed by "projectId:databaseId". This
89+
* prevents creating a new gRPC client for every exporter instance, reducing resource usage.
90+
* Reference counting is used to safely shut down the client when no longer needed.
91+
*/
92+
static final ConcurrentHashMap<String, CachedMetricsClient> METRICS_CLIENT_CACHE =
93+
new ConcurrentHashMap<>();
94+
7495
private final MetricServiceClient client;
96+
private final Map<String, String> clientAttributes;
7597

7698
// This is the quota limit from Cloud Monitoring. More details in
7799
// https://cloud.google.com/monitoring/quotas#custom_metrics_quotas.
@@ -84,7 +106,8 @@ class DatastoreCloudMonitoringExporter implements MetricExporter {
84106
// Flag to prevent log spam of any export failures
85107
private final AtomicBoolean datastoreExportFailureLogged = new AtomicBoolean(false);
86108

87-
private final String datastoreProjectId;
109+
private final String projectId;
110+
private final String databaseId;
88111

89112
/**
90113
* Creates a new instance of the exporter.
@@ -102,10 +125,47 @@ class DatastoreCloudMonitoringExporter implements MetricExporter {
102125
* @param projectId the GCP project ID where metrics will be exported.
103126
* @param credentials the credentials used to authenticate with Cloud Monitoring.
104127
* @return a new {@link DatastoreCloudMonitoringExporter} instance.
105-
* @throws IOException if the {@link MetricServiceClient} fails to initialize.
106128
*/
129+
@Nullable
107130
static DatastoreCloudMonitoringExporter create(
108-
String projectId, @Nullable Credentials credentials) throws IOException {
131+
String projectId,
132+
String databaseId,
133+
Credentials credentials,
134+
Map<String, String> clientAttributes) {
135+
String key = projectId + ":" + databaseId;
136+
137+
// Use compute to acquire or create the client atomically with reference counting.
138+
// If creation fails, we log the error and return null so it's not added to the map.
139+
CachedMetricsClient cachedMetricsClient =
140+
METRICS_CLIENT_CACHE.compute(
141+
key,
142+
(k, v) -> {
143+
if (v == null) {
144+
try {
145+
v = new CachedMetricsClient(createMetricServiceClient(credentials));
146+
} catch (IOException e) {
147+
logger.log(
148+
Level.WARNING,
149+
"Failed to create MetricServiceClient for metrics export. Monitoring will be disabled.",
150+
e);
151+
return null; // Do not add to map
152+
}
153+
}
154+
v.refCount.incrementAndGet();
155+
return v;
156+
});
157+
158+
// If there is no client in the cache (creation failed), return null.
159+
if (cachedMetricsClient == null) {
160+
return null;
161+
}
162+
163+
return new DatastoreCloudMonitoringExporter(
164+
projectId, databaseId, cachedMetricsClient.client, clientAttributes);
165+
}
166+
167+
private static MetricServiceClient createMetricServiceClient(Credentials credentials)
168+
throws IOException {
109169
MetricServiceSettings.Builder settingsBuilder = MetricServiceSettings.newBuilder();
110170

111171
InstantiatingGrpcChannelProvider transportChannelProvider =
@@ -114,26 +174,25 @@ static DatastoreCloudMonitoringExporter create(
114174
.build();
115175
settingsBuilder.setTransportChannelProvider(transportChannelProvider);
116176

117-
CredentialsProvider credentialsProvider;
118-
if (credentials == null || credentials instanceof NoCredentials) {
119-
credentialsProvider = NoCredentialsProvider.create();
120-
} else {
121-
credentialsProvider = FixedCredentialsProvider.create(credentials);
122-
}
123-
settingsBuilder.setCredentialsProvider(credentialsProvider);
177+
settingsBuilder.setCredentialsProvider(FixedCredentialsProvider.create(credentials));
124178

125179
settingsBuilder
126180
.createTimeSeriesSettings()
127181
.setSimpleTimeoutNoRetriesDuration(Duration.ofMinutes(1));
128182

129-
return new DatastoreCloudMonitoringExporter(
130-
projectId, MetricServiceClient.create(settingsBuilder.build()));
183+
return MetricServiceClient.create(settingsBuilder.build());
131184
}
132185

133186
@VisibleForTesting
134-
DatastoreCloudMonitoringExporter(String projectId, MetricServiceClient client) {
187+
DatastoreCloudMonitoringExporter(
188+
String projectId,
189+
String databaseId,
190+
MetricServiceClient client,
191+
Map<String, String> clientAttributes) {
135192
this.client = client;
136-
this.datastoreProjectId = projectId;
193+
this.projectId = projectId;
194+
this.databaseId = databaseId;
195+
this.clientAttributes = clientAttributes;
137196
}
138197

139198
/**
@@ -159,8 +218,7 @@ public CompletableResultCode export(@Nonnull Collection<MetricData> collection)
159218
// Convert OTel MetricData to Cloud Monitoring TimeSeries.
160219
datastoreTimeSeries =
161220
DatastoreCloudMonitoringExporterUtils.convertToDatastoreTimeSeries(
162-
new ArrayList<>(collection),
163-
BuiltInDatastoreMetricsProvider.INSTANCE.getClientAttributes());
221+
new ArrayList<>(collection), clientAttributes);
164222
} catch (Throwable e) {
165223
logger.log(
166224
Level.WARNING,
@@ -169,7 +227,7 @@ public CompletableResultCode export(@Nonnull Collection<MetricData> collection)
169227
return CompletableResultCode.ofFailure();
170228
}
171229

172-
ProjectName projectName = ProjectName.of(datastoreProjectId);
230+
ProjectName projectName = ProjectName.of(projectId);
173231

174232
// Perform the actual network call to Cloud Monitoring.
175233
ApiFuture<List<Empty>> futureList = exportTimeSeriesInBatch(projectName, datastoreTimeSeries);
@@ -249,7 +307,18 @@ public CompletableResultCode shutdown() {
249307
}
250308
CompletableResultCode shutdownResult = new CompletableResultCode();
251309
try {
252-
client.shutdown();
310+
String key = projectId + ":" + databaseId;
311+
// Atomically decrement reference count and cleanup if zero.
312+
METRICS_CLIENT_CACHE.compute(
313+
key,
314+
(k, v) -> {
315+
if (v != null && v.refCount.decrementAndGet() == 0) {
316+
v.client.shutdown();
317+
return null; // Remove from map to prevent leaks
318+
}
319+
320+
return v;
321+
});
253322
shutdownResult.succeed();
254323
} catch (Throwable e) {
255324
logger.log(Level.WARNING, "failed to shutdown the monitoring client", e);

‎java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DatastoreCloudMonitoringExporterTest.java‎

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,12 +77,15 @@ public class DatastoreCloudMonitoringExporterTest {
7777
public void setUp() {
7878
mockMetricServiceStub = createMock(MetricServiceStub.class);
7979
fakeMetricServiceClient = new FakeMetricServiceClient(mockMetricServiceStub);
80-
exporter = new DatastoreCloudMonitoringExporter(PROJECT_ID, fakeMetricServiceClient);
8180

8281
Map<String, String> clientAttributes =
8382
BuiltInDatastoreMetricsProvider.INSTANCE.getClientAttributes();
8483
this.clientUid = clientAttributes.get(CLIENT_UID_KEY.getKey());
8584

85+
exporter =
86+
new DatastoreCloudMonitoringExporter(
87+
PROJECT_ID, DATABASE_ID, fakeMetricServiceClient, clientAttributes);
88+
8689
attributes =
8790
Attributes.builder()
8891
.put(DATABASE_ID_KEY, DATABASE_ID)
@@ -152,6 +155,41 @@ public void testExportingSumData() {
152155
verify(mockMetricServiceStub, mockCallable);
153156
}
154157

158+
@Test
159+
public void testClientCacheReferenceCounting() {
160+
MetricServiceClient mockClient = createMock(MetricServiceClient.class);
161+
expect(mockClient.isShutdown()).andReturn(false).anyTimes();
162+
mockClient.shutdown();
163+
EasyMock.expectLastCall(); // Expect shutdown when refCount reaches 0
164+
replay(mockClient);
165+
166+
String key = PROJECT_ID + ":" + DATABASE_ID;
167+
DatastoreCloudMonitoringExporter.CachedMetricsClient cachedMetricsClient =
168+
new DatastoreCloudMonitoringExporter.CachedMetricsClient(mockClient);
169+
cachedMetricsClient.refCount.set(2); // Simulate 2 references
170+
DatastoreCloudMonitoringExporter.METRICS_CLIENT_CACHE.put(key, cachedMetricsClient);
171+
172+
DatastoreCloudMonitoringExporter exporter1 =
173+
new DatastoreCloudMonitoringExporter(
174+
PROJECT_ID, DATABASE_ID, mockClient, Collections.emptyMap());
175+
176+
// First shutdown should decrement refCount to 1, but not close client
177+
exporter1.shutdown();
178+
assertThat(cachedMetricsClient.refCount.get()).isEqualTo(1);
179+
assertThat(DatastoreCloudMonitoringExporter.METRICS_CLIENT_CACHE.containsKey(key)).isTrue();
180+
181+
DatastoreCloudMonitoringExporter exporter2 =
182+
new DatastoreCloudMonitoringExporter(
183+
PROJECT_ID, DATABASE_ID, mockClient, Collections.emptyMap());
184+
185+
// Second shutdown should decrement refCount to 0, close client, and remove from cache
186+
exporter2.shutdown();
187+
assertThat(cachedMetricsClient.refCount.get()).isEqualTo(0);
188+
assertThat(DatastoreCloudMonitoringExporter.METRICS_CLIENT_CACHE.containsKey(key)).isFalse();
189+
190+
verify(mockClient);
191+
}
192+
155193
private static class FakeMetricServiceClient extends MetricServiceClient {
156194
protected FakeMetricServiceClient(MetricServiceStub stub) {
157195
super(stub);

‎java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DatastoreMetricsRecorderTest.java‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -70,25 +70,20 @@ public void tracingEnabledButMetricsDisabledAndBuiltInDisabled_returnsNoOp() {
7070
}
7171

7272
@Test
73-
public void defaultOptionsWithBuiltInMetricsEnabled_returnsOpenTelemetryRecorder() {
73+
public void defaultOptionsWithBuiltInMetricsEnabled_butNoCredentials_returnsNoOpRecorder() {
7474
// Explicitly enable built-in metrics export
7575
DatastoreOptions options =
76-
baseOptions()
76+
baseOptions() // Uses NoCredentials by default
7777
.setOpenTelemetryOptions(
7878
DatastoreOpenTelemetryOptions.newBuilder()
7979
.setExportBuiltinMetricsToGoogleCloudMonitoring(true)
8080
.build())
8181
.build();
8282
DatastoreMetricsRecorder recorder = DatastoreMetricsRecorder.getInstance(options);
8383

84-
// If we're not in an emulator, it should return a OpenTelemetryDatastoreMetricsRecorder
85-
// (Note: System.getenv() is hard to mock, but we assume it's not set in test environment)
86-
String emulatorHost = System.getenv(DatastoreOptions.LOCAL_HOST_ENV_VAR);
87-
if (emulatorHost == null || emulatorHost.isEmpty()) {
88-
assertThat(recorder).isInstanceOf(OpenTelemetryDatastoreMetricsRecorder.class);
89-
} else {
90-
assertThat(recorder).isInstanceOf(NoOpDatastoreMetricsRecorder.class);
91-
}
84+
// Since baseOptions() uses NoCredentials, it should return NoOpDatastoreMetricsRecorder
85+
// as we don't want to send metrics for local emulator logic.
86+
assertThat(recorder).isInstanceOf(NoOpDatastoreMetricsRecorder.class);
9287
}
9388

9489
@Test

0 commit comments

Comments
 (0)