Add Circuit breaker implementation for Telemetry Push (#925)

## Description
<!-- Provide a brief summary of the changes made and the issue they aim
to address.-->
Added CircuitBreaker implementation using resilience4j library. Started
with hard coded values for CircuitBreaker config.

## Testing
<!-- Describe how the changes have been tested-->

Added thorough unit tests for states: open, closed, closed -> open,
closed -> open -> closed

## Additional Notes to the Reviewer
<!-- Share any additional context or insights that may help the reviewer
understand the changes better. This could include challenges faced,
limitations, or compromises made during the development process.
Also, mention any areas of the code that you would like the reviewer to
focus on specifically. -->

---------

Signed-off-by: Gopal Lal <[email protected]>

Author: gopalldb

NEXT_CHANGELOG.md

@@ -6,6 +6,7 @@
 - Added support for providing custom HTTP options: `HttpMaxConnectionsPerRoute` and `HttpConnectionRequestTimeout`.
 - Add V2 of chunk download using async http client with corresponding implementations of AbstractRemoteChunkProvider and 
 AbstractArrowResultChunk
+- Added CircuitBreaker support to handle transient failures in the Telemetry.
 
 ### Updated
 

pom.xml

@@ -69,6 +69,7 @@
     <async-httpclient.version>5.3.1</async-httpclient.version>
     <netty.version>4.2.0.Final</netty.version>
     <grpc.version>1.71.0</grpc.version>
+    <resilience4j.version>2.2.0</resilience4j.version>
   </properties>
   <dependencyManagement>
     <!-- Force safe version of commons-lang3 https://nvd.nist.gov/vuln/detail/CVE-2025-48924 -->
@@ -269,6 +270,16 @@
       <artifactId>httpcore5</artifactId>
       <version>${async-httpclient.version}</version>
     </dependency>
+    <dependency>
+      <groupId>io.github.resilience4j</groupId>
+      <artifactId>resilience4j-circuitbreaker</artifactId>
+      <version>${resilience4j.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>io.github.resilience4j</groupId>
+      <artifactId>resilience4j-core</artifactId>
+      <version>${resilience4j.version}</version>
+    </dependency>
   </dependencies>
 
   <build>

src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java

@@ -926,6 +926,11 @@ public int getChunkReadyTimeoutSeconds() {
     return Integer.parseInt(getParameter(DatabricksJdbcUrlParams.CHUNK_READY_TIMEOUT_SECONDS));
   }
 
+  @Override
+  public boolean isTelemetryCircuitBreakerEnabled() {
+    return getParameter(DatabricksJdbcUrlParams.TELEMETRY_CIRCUIT_BREAKER_ENABLED).equals("1");
+  }
+
   @Override
   public int getHttpMaxConnectionsPerRoute() {
     int maxConnectionsPerRoute = DEFAULT_MAX_HTTP_CONNECTIONS_PER_ROUTE;

src/main/java/com/databricks/jdbc/api/internal/IDatabricksConnectionContext.java

@@ -342,6 +342,9 @@ public interface IDatabricksConnectionContext {
   /** Returns the flush interval in milliseconds for telemetry */
   int getTelemetryFlushIntervalInMilliseconds();
 
+  /** Returns whether circuit breaker is enabled for telemetry */
+  boolean isTelemetryCircuitBreakerEnabled();
+
   /** Returns the maximum number of HTTP connections per route */
   int getHttpMaxConnectionsPerRoute();
 

src/main/java/com/databricks/jdbc/common/DatabricksJdbcUrlParams.java

@@ -146,6 +146,8 @@ public enum DatabricksJdbcUrlParams {
       "MaxVolumeOperationConcurrentPresignedRequests",
       "Maximum number of concurrent presigned requests",
       "50"),
+  TELEMETRY_CIRCUIT_BREAKER_ENABLED(
+      "TelemetryCircuitBreakerEnabled", "Enable circuit breaker for telemetry", "0"),
   HTTP_MAX_CONNECTIONS_PER_ROUTE(
       "HttpMaxConnectionsPerRoute", "Maximum connections per route for HTTP client", "1000"),
   HTTP_CONNECTION_REQUEST_TIMEOUT(

src/main/java/com/databricks/jdbc/telemetry/CircuitBreakerManager.java

@@ -0,0 +1,106 @@
+package com.databricks.jdbc.telemetry;
+
+import com.databricks.jdbc.exception.DatabricksHttpException;
+import com.databricks.jdbc.exception.DatabricksParsingException;
+import com.databricks.jdbc.log.JdbcLogger;
+import com.databricks.jdbc.log.JdbcLoggerFactory;
+import io.github.resilience4j.circuitbreaker.CircuitBreaker;
+import io.github.resilience4j.circuitbreaker.CircuitBreakerConfig;
+import java.net.ConnectException;
+import java.net.NoRouteToHostException;
+import java.net.SocketTimeoutException;
+import java.net.UnknownHostException;
+import java.time.Duration;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.RejectedExecutionException;
+import org.apache.arrow.util.VisibleForTesting;
+import org.apache.http.client.HttpResponseException;
+
+/**
+ * CircuitBreakerManager is a singleton that manages circuit breakers for different hosts. It
+ * initializes circuit breakers with a predefined configuration and provides methods to retrieve or
+ * reset them.
+ */
+public class CircuitBreakerManager {
+  private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(CircuitBreakerManager.class);
+
+  // Singleton instance of CircuitBreakerManager
+  private static final CircuitBreakerManager INSTANCE = new CircuitBreakerManager();
+
+  // Method to get the singleton instance
+  public static CircuitBreakerManager getInstance() {
+    return INSTANCE;
+  }
+
+  private final Map<String, CircuitBreaker> breakerPerHost = new ConcurrentHashMap<>();
+  private final CircuitBreakerConfig config;
+
+  // Private constructor to prevent instantiation
+  private CircuitBreakerManager() {
+    // Initialize config hard coded for now
+    config =
+        CircuitBreakerConfig.custom()
+            .failureRateThreshold(50) // Opens if 50%+ fail
+            .minimumNumberOfCalls(20) // Minimum sample size
+            .slidingWindowSize(30) // Keep recent 30 calls in window
+            .waitDurationInOpenState(Duration.ofSeconds(30)) // Cool-down before retrying
+            .permittedNumberOfCallsInHalfOpenState(3) // Retry with 3 test calls
+            .slidingWindowType(CircuitBreakerConfig.SlidingWindowType.COUNT_BASED)
+            // Exceptions that should be treated as failures
+            .recordExceptions(
+                // Server unavailability errors
+                ConnectException.class,
+                SocketTimeoutException.class,
+                NoRouteToHostException.class,
+                UnknownHostException.class,
+                // Resource exhausted errors
+                RejectedExecutionException.class,
+                // Memory errors also can lean to open state
+                OutOfMemoryError.class,
+                // HTTP errors that indicate server issues
+                HttpResponseException.class,
+                // Databricks specific exceptions
+                DatabricksHttpException.class)
+            .ignoreExceptions(
+                // Exceptions that can be ignored, the logic is that due to these exception,
+                // the execution will not even reach the server, so no point in opening the
+                // circuit breaker
+                DatabricksParsingException.class,
+                IllegalArgumentException.class,
+                NullPointerException.class)
+            .build();
+    LOGGER.debug("CircuitBreakerManager initialized");
+  }
+
+  /**
+   * Retrieves the CircuitBreaker for the specified host. If it does not exist, it creates a new one
+   * with the default configuration.
+   *
+   * @param host The host for which to retrieve the CircuitBreaker.
+   * @return The CircuitBreaker instance for the specified host.
+   */
+  public CircuitBreaker getCircuitBreaker(String host) {
+    return breakerPerHost.computeIfAbsent(
+        host,
+        h -> {
+          CircuitBreaker breaker = CircuitBreaker.of("telemetry-client-" + h, config);
+          breaker
+              .getEventPublisher()
+              .onStateTransition(
+                  event -> {
+                    LOGGER.debug(
+                        "CircuitBreaker for host [{}] transitioned from {} to {}",
+                        h,
+                        event.getStateTransition().getFromState(),
+                        event.getStateTransition().getToState());
+                  });
+          return breaker;
+        });
+  }
+
+  @VisibleForTesting
+  void resetCircuitBreaker(String host) {
+    breakerPerHost.remove(host);
+  }
+}

src/main/java/com/databricks/jdbc/telemetry/CircuitBreakerTelemetryPushClient.java

@@ -0,0 +1,75 @@
+package com.databricks.jdbc.telemetry;
+
+import com.databricks.jdbc.log.JdbcLogger;
+import com.databricks.jdbc.log.JdbcLoggerFactory;
+import com.databricks.jdbc.model.telemetry.TelemetryRequest;
+import io.github.resilience4j.circuitbreaker.CircuitBreaker;
+import io.github.resilience4j.circuitbreaker.CircuitBreakerConfig;
+import org.apache.arrow.util.VisibleForTesting;
+
+/**
+ * TelemetryClient wrapper that implements circuit breaker pattern using Resilience4j. This wrapper
+ * handles server unavailability and resource exhausted errors by temporarily stopping telemetry
+ * requests when the service is experiencing issues.
+ */
+public class CircuitBreakerTelemetryPushClient implements ITelemetryPushClient {
+
+  private static final JdbcLogger LOGGER =
+      JdbcLoggerFactory.getLogger(CircuitBreakerTelemetryPushClient.class);
+  private static final String TEST_CLIENT_NAME = "telemetry-client-test";
+
+  private final ITelemetryPushClient delegate;
+  private final CircuitBreaker circuitBreaker;
+  private final String host;
+
+  CircuitBreakerTelemetryPushClient(ITelemetryPushClient delegate, String host) {
+    this.delegate = delegate;
+    this.host = host;
+    this.circuitBreaker = CircuitBreakerManager.getInstance().getCircuitBreaker(host);
+  }
+
+  @VisibleForTesting
+  CircuitBreakerTelemetryPushClient(
+      ITelemetryPushClient delegate, String host, CircuitBreakerConfig config) {
+    this.delegate = delegate;
+    this.host = host;
+    this.circuitBreaker = CircuitBreaker.of(TEST_CLIENT_NAME, config);
+  }
+
+  @Override
+  public void pushEvent(TelemetryRequest request) {
+    try {
+      circuitBreaker.executeCallable(
+          () -> {
+            delegate.pushEvent(request);
+            return null;
+          });
+    } catch (Exception e) {
+      LOGGER.debug("Failed to export telemetry for host [{}]: {}", host, e.getMessage());
+
+      if (circuitBreaker.getState() == CircuitBreaker.State.OPEN) {
+        LOGGER.debug("CircuitBreaker for host [{}] is OPEN - dropping telemetry", host);
+      }
+    }
+  }
+
+  /**
+   * Get the current state of the circuit breaker.
+   *
+   * @return The current circuit breaker state
+   */
+  @VisibleForTesting
+  CircuitBreaker.State getCircuitBreakerState() {
+    return circuitBreaker.getState();
+  }
+
+  /**
+   * Get metrics about the circuit breaker.
+   *
+   * @return Circuit breaker metrics
+   */
+  @VisibleForTesting
+  CircuitBreaker.Metrics getCircuitBreakerMetrics() {
+    return circuitBreaker.getMetrics();
+  }
+}

src/main/java/com/databricks/jdbc/telemetry/ITelemetryClient.java

@@ -2,8 +2,22 @@
 
 import com.databricks.jdbc.model.telemetry.TelemetryFrontendLog;
 
+/**
+ * Interface for telemetry clients that handle the export of telemetry events. The implementation
+ * further handles batching and flushing of telemetry events to the backend service.
+ */
 public interface ITelemetryClient {
+
+  /**
+   * Exports a telemetry event to the backend service.
+   *
+   * @param event The telemetry event to be exported.
+   */
   void exportEvent(TelemetryFrontendLog event);
 
+  /**
+   * Closes the telemetry client, releasing any resources it holds. This method should be called on
+   * closure of connection or JVM shutdown.
+   */
   void close();
 }

src/main/java/com/databricks/jdbc/telemetry/ITelemetryPushClient.java

@@ -0,0 +1,20 @@
+package com.databricks.jdbc.telemetry;
+
+import com.databricks.jdbc.model.telemetry.TelemetryRequest;
+
+/**
+ * Interface for pushing telemetry events. Implementations should handle the actual transmission of
+ * telemetry data. This actually pushes the data to the telemetry service. The circuit breaker is
+ * engaged here to prevent overwhelming the service with requests during outages or high error
+ * rates.
+ */
+interface ITelemetryPushClient {
+
+  /**
+   * Pushes a telemetry request to the telemetry service.
+   *
+   * @param request The telemetry request to be sent.
+   * @throws Exception If there is an error while pushing the event.
+   */
+  void pushEvent(TelemetryRequest request) throws Exception;
+}

src/main/java/com/databricks/jdbc/telemetry/TelemetryClient.java

@@ -19,8 +19,8 @@ public class TelemetryClient implements ITelemetryClient {
   private final IDatabricksConnectionContext context;
   private final DatabricksConfig databricksConfig;
   private final int eventsBatchSize;
-  private final boolean isAuthEnabled;
   private final ExecutorService executorService;
+  private final ITelemetryPushClient telemetryPushClient;
   private final ScheduledExecutorService scheduledExecutorService;
   private List<TelemetryFrontendLog> eventsBatch;
   private volatile long lastFlushedTime;
@@ -46,29 +46,33 @@ public TelemetryClient(
       DatabricksConfig config) {
     this.eventsBatch = new LinkedList<>();
     this.eventsBatchSize = connectionContext.getTelemetryBatchSize();
-    this.isAuthEnabled = true;
     this.context = connectionContext;
     this.databricksConfig = config;
     this.executorService = executorService;
     this.scheduledExecutorService =
         Executors.newSingleThreadScheduledExecutor(createSchedulerThreadFactory());
     this.flushIntervalMillis = context.getTelemetryFlushIntervalInMilliseconds();
     this.lastFlushedTime = System.currentTimeMillis();
+    this.telemetryPushClient =
+        TelemetryClientFactory.getTelemetryPushClient(
+            true /* isAuthEnabled */, context, databricksConfig);
     schedulePeriodicFlush();
   }
 
   public TelemetryClient(
       IDatabricksConnectionContext connectionContext, ExecutorService executorService) {
     this.eventsBatch = new LinkedList<>();
     this.eventsBatchSize = connectionContext.getTelemetryBatchSize();
-    this.isAuthEnabled = false;
     this.context = connectionContext;
     this.databricksConfig = null;
     this.executorService = executorService;
     this.scheduledExecutorService =
         Executors.newSingleThreadScheduledExecutor(createSchedulerThreadFactory());
     this.flushIntervalMillis = context.getTelemetryFlushIntervalInMilliseconds();
     this.lastFlushedTime = System.currentTimeMillis();
+    this.telemetryPushClient =
+        TelemetryClientFactory.getTelemetryPushClient(
+            false /* isAuthEnabled */, context, null /* databricksConfig */);
     schedulePeriodicFlush();
   }
 
@@ -118,8 +122,7 @@ private void flush(boolean forceFlush) {
     synchronized (this) {
       if (!forceFlush ? isBatchFull() : !eventsBatch.isEmpty()) {
         List<TelemetryFrontendLog> logsToBeFlushed = eventsBatch;
-        executorService.submit(
-            new TelemetryPushTask(logsToBeFlushed, isAuthEnabled, context, databricksConfig));
+        executorService.submit(new TelemetryPushTask(logsToBeFlushed, telemetryPushClient));
         eventsBatch = new LinkedList<>();
       }
       lastFlushedTime = System.currentTimeMillis();