diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java index 8419f49e1290..a8b00689906e 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java @@ -1101,6 +1101,8 @@ private void closeImpl() throws SQLException { BigQueryJdbcMdc.clear(); BigQueryJdbcRootLogger.closeConnectionHandler(this.connectionId); BigQueryJdbcOpenTelemetry.unregisterConnection(this.connectionId); + BigQueryJdbcOpenTelemetry.releaseSdk(this.openTelemetry); + this.openTelemetry = null; } if (exceptionToThrow != null) { throw exceptionToThrow; @@ -1211,13 +1213,21 @@ private OpenTelemetry getOpenTelemetryInstance() { } private String resolveEffectiveCredentials() { - String creds = this.gcpTelemetryCredentials; + if (this.gcpTelemetryCredentials != null) { + return this.gcpTelemetryCredentials; + } + String authTypeStr = this.authProperties.get(BigQueryJdbcUrlUtility.OAUTH_TYPE_PROPERTY_NAME); - if (creds == null - && BigQueryJdbcOAuthUtility.AuthType.GOOGLE_SERVICE_ACCOUNT.name().equals(authTypeStr)) { - return this.authProperties.get(BigQueryJdbcUrlUtility.OAUTH_PVT_KEY_PROPERTY_NAME); + if (!BigQueryJdbcOAuthUtility.AuthType.GOOGLE_SERVICE_ACCOUNT.name().equals(authTypeStr)) { + return null; } - return creds; + + String pvtKey = this.authProperties.get(BigQueryJdbcUrlUtility.OAUTH_PVT_KEY_PROPERTY_NAME); + if (pvtKey != null) { + return pvtKey; + } + + return this.authProperties.get(BigQueryJdbcUrlUtility.OAUTH_PVT_KEY_PATH_PROPERTY_NAME); } private void validateTraceConfiguration(boolean isTraceEnabled, String effectiveCredentials) { diff --git a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java index 8544946af550..cee8fcb5b654 100644 --- a/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java +++ b/java-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcOpenTelemetry.java @@ -17,6 +17,7 @@ package com.google.cloud.bigquery.jdbc; import com.google.auth.Credentials; +import com.google.auth.oauth2.GoogleCredentials; import com.google.cloud.bigquery.exception.BigQueryJdbcRuntimeException; import com.google.cloud.logging.Logging; import com.google.cloud.logging.LoggingOptions; @@ -44,6 +45,8 @@ import java.util.Objects; import java.util.concurrent.Callable; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.logging.Handler; import java.util.logging.Logger; @@ -104,7 +107,16 @@ public int hashCode() { } } - private static final ConcurrentHashMap sdkCache = + private static final class CachedSdk { + final OpenTelemetrySdk sdk; + final AtomicInteger refCount = new AtomicInteger(1); + + CachedSdk(OpenTelemetrySdk sdk) { + this.sdk = sdk; + } + } + + private static final ConcurrentHashMap sdkCache = new ConcurrentHashMap<>(); static class TelemetryConfig { @@ -127,20 +139,6 @@ private BigQueryJdbcOpenTelemetry() {} static { ensureGlobalHandlerAttached(); - Runtime.getRuntime() - .addShutdownHook( - new Thread( - () -> { - for (OpenTelemetrySdk sdk : sdkCache.values()) { - try { - sdk.close(); - } catch (Exception e) { - // Ignore failures during shutdown to ensure all SDKs are attempted to be - // closed. Logging is avoided here because the logging system might have - // already been shut down by the JVM. - } - } - })); } public static synchronized void ensureGlobalHandlerAttached() { @@ -207,6 +205,44 @@ public static Logging createLoggingClient( } } + public static void releaseSdk(OpenTelemetry openTelemetry) { + if (openTelemetry == null + || openTelemetry == GlobalOpenTelemetry.get() + || openTelemetry == OpenTelemetry.noop()) { + return; + } + + AtomicBoolean shouldClose = new AtomicBoolean(false); + for (Map.Entry entry : sdkCache.entrySet()) { + if (entry.getValue().sdk == openTelemetry) { + sdkCache.computeIfPresent( + entry.getKey(), + (k, cachedSdk) -> { + if (cachedSdk.sdk != openTelemetry) { + return cachedSdk; + } + if (cachedSdk.refCount.decrementAndGet() <= 0) { + shouldClose.set(true); + return null; + } + return cachedSdk; + }); + break; + } + } + + if (shouldClose.get() && openTelemetry instanceof OpenTelemetrySdk) { + try { + ((OpenTelemetrySdk) openTelemetry).close(); + } catch (Exception e) { + // Swallow exceptions so that telemetry shutdown failures (e.g. flush timeouts) + // do not propagate and disrupt the core BigQueryConnection.close() logic. + // Throwing here could prevent the core connection from cleaning up correctly. + LOG.warning("Failed to close OpenTelemetry SDK: %s", e.getMessage()); + } + } + } + private static Credentials resolveCredentialsFromString(String credsString) { Map authProperties = new java.util.HashMap<>(); authProperties.put( @@ -295,70 +331,81 @@ public static OpenTelemetry getOpenTelemetry( gcpTelemetryProjectId, getCredentialsIdentifier(gcpTelemetryCredentials), enableGcpTraceExporter); - return sdkCache.computeIfAbsent( - key, - k -> { - Map props = new HashMap<>(); - - if (enableGcpTraceExporter) { - props.put(OTEL_TRACES_EXPORTER, EXPORTER_OTLP); - props.put(OTEL_EXPORTER_OTLP_ENDPOINT, OTLP_ENDPOINT_VALUE); - } else { - props.put(OTEL_TRACES_EXPORTER, EXPORTER_NONE); - } - - // Logs are handled directly via GCP logging - props.put(OTEL_LOGS_EXPORTER, EXPORTER_NONE); - // Metrics are deferred to a future phase - props.put(OTEL_METRICS_EXPORTER, EXPORTER_NONE); - - if (gcpTelemetryProjectId != null) { - props.put(GOOGLE_CLOUD_PROJECT, gcpTelemetryProjectId); - } - - // Set safe, generous default limits on attribute value lengths (32KB) to protect - // customers from GCP Cloud Trace 64KB span ingestion failures when logging massive - // exception stack traces or database schema metadata. - // Respect any existing user configuration overrides. - if (!props.containsKey(OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT)) { - props.put(OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT, DEFAULT_ATTRIBUTE_LENGTH_LIMIT); - } - if (!props.containsKey(OTEL_ATTRIBUTE_VALUE_LENGTH_LIMIT)) { - props.put(OTEL_ATTRIBUTE_VALUE_LENGTH_LIMIT, DEFAULT_ATTRIBUTE_LENGTH_LIMIT); - } - - AutoConfiguredOpenTelemetrySdk autoConfigured = - AutoConfiguredOpenTelemetrySdk.builder() - .addPropertiesSupplier(() -> props) - .addSpanExporterCustomizer( - (spanExporter, configProperties) -> { - if (gcpTelemetryCredentials == null) { - return spanExporter; - } - try { - Credentials credentials = - resolveCredentialsFromString(gcpTelemetryCredentials); - if (spanExporter instanceof OtlpHttpSpanExporter) { - return ((OtlpHttpSpanExporter) spanExporter) - .toBuilder().setHeaders(() -> getAuthHeaders(credentials)).build(); - } - if (spanExporter instanceof OtlpGrpcSpanExporter) { - return ((OtlpGrpcSpanExporter) spanExporter) - .toBuilder().setHeaders(() -> getAuthHeaders(credentials)).build(); - } - } catch (Exception e) { - LOG.warning( - e, - "Failed to resolve telemetry credentials. Telemetry will be exported using default OpenTelemetry configuration (custom authentication headers will not be injected)."); - } - return spanExporter; - }) - .build(); - - OpenTelemetrySdk sdk = autoConfigured.getOpenTelemetrySdk(); - - return sdk; - }); + return sdkCache.compute( + key, + (k, cachedSdk) -> { + if (cachedSdk != null) { + cachedSdk.refCount.incrementAndGet(); + return cachedSdk; + } + + Map props = new HashMap<>(); + + if (enableGcpTraceExporter) { + props.put(OTEL_TRACES_EXPORTER, EXPORTER_OTLP); + props.put(OTEL_EXPORTER_OTLP_ENDPOINT, OTLP_ENDPOINT_VALUE); + } else { + props.put(OTEL_TRACES_EXPORTER, EXPORTER_NONE); + } + + // Logs are handled directly via GCP logging + props.put(OTEL_LOGS_EXPORTER, EXPORTER_NONE); + // Metrics are deferred to a future phase + props.put(OTEL_METRICS_EXPORTER, EXPORTER_NONE); + + if (gcpTelemetryProjectId != null) { + props.put(GOOGLE_CLOUD_PROJECT, gcpTelemetryProjectId); + } + + // Set safe, generous default limits on attribute value lengths (32KB) to protect + // customers from GCP Cloud Trace 64KB span ingestion failures when logging massive + // exception stack traces or database schema metadata. + // Respect any existing user configuration overrides. + if (!props.containsKey(OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT)) { + props.put(OTEL_SPAN_ATTRIBUTE_VALUE_LENGTH_LIMIT, DEFAULT_ATTRIBUTE_LENGTH_LIMIT); + } + if (!props.containsKey(OTEL_ATTRIBUTE_VALUE_LENGTH_LIMIT)) { + props.put(OTEL_ATTRIBUTE_VALUE_LENGTH_LIMIT, DEFAULT_ATTRIBUTE_LENGTH_LIMIT); + } + + AutoConfiguredOpenTelemetrySdk autoConfigured = + AutoConfiguredOpenTelemetrySdk.builder() + .addPropertiesSupplier(() -> props) + .addSpanExporterCustomizer( + (spanExporter, configProperties) -> { + try { + Credentials credentials; + if (gcpTelemetryCredentials != null) { + credentials = resolveCredentialsFromString(gcpTelemetryCredentials); + } else { + credentials = GoogleCredentials.getApplicationDefault(); + } + if (spanExporter instanceof OtlpHttpSpanExporter) { + return ((OtlpHttpSpanExporter) spanExporter) + .toBuilder() + .setHeaders(() -> getAuthHeaders(credentials)) + .build(); + } + if (spanExporter instanceof OtlpGrpcSpanExporter) { + return ((OtlpGrpcSpanExporter) spanExporter) + .toBuilder() + .setHeaders(() -> getAuthHeaders(credentials)) + .build(); + } + } catch (Exception e) { + LOG.warning( + e, + "Failed to resolve telemetry credentials. Telemetry will be exported using default OpenTelemetry configuration (custom authentication headers will not be injected)."); + } + return spanExporter; + }) + .build(); + + OpenTelemetrySdk sdk = autoConfigured.getOpenTelemetrySdk(); + + return new CachedSdk(sdk); + }) + .sdk; } /** Gets a Tracer for the JDBC driver instrumentation scope. */ diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITOpenTelemetryTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITOpenTelemetryTest.java index 792ae5071841..9baf187b4d38 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITOpenTelemetryTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITOpenTelemetryTest.java @@ -306,8 +306,8 @@ private Trace verifyAndFetchTrace(String traceId) throws Exception { private T pollWithRetry(java.util.concurrent.Callable task) throws InterruptedException { int attempts = 0; - int maxAttempts = 30; // 30 attempts * 500ms = 15 seconds max delay - long delayMs = 500; // 500ms linear polling + int maxAttempts = 20; // 20 attempts * 3000ms = 60 seconds max delay + long delayMs = 3000; // 3000ms linear polling while (attempts < maxAttempts) { attempts++; @@ -321,6 +321,7 @@ private T pollWithRetry(java.util.concurrent.Callable task) throws Interr throw new RuntimeException("Test execution interrupted", e); } catch (Exception e) { // Ignore exceptions during remote lookup and retry + e.printStackTrace(); } if (attempts < maxAttempts) { Thread.sleep(delayMs);