[Fix] Retry on too many auth requests (#355)

## Changes
<!-- Summary of your changes that are easy to understand -->
This PR addresses an issue encountered during authentication, where the
server returns a 429 status code when fetching OIDC endpoints from a
well-known location, particularly during frequent authentication
requests. The update introduces a retry mechanism within the auth flow
to mitigate the impact of 429 errors.

Modified `ApiClient` object construction to allow unauthenticated calls
by including authentication function injection, where within the OIDC
workflow act as no-op actions. `DatabricksConfig` would now instantiate
an `ApiClient` instead of a standard `CommonsHttpClient` in these cases.
errors during a complete authentication workflow, especially under high
concurrent request conditions.

## Tests
<!-- How is this tested? -->
* Added unit tests to verify the retry functionality upon encountering a
429 response.
* Included a manual test to simulate real-world scenarios, ensuring the
retry mechanism effectively handles 429

---------

Signed-off-by: Omer Lachish <[email protected]>
Co-authored-by: Omer Lachish <[email protected]>

Author: rauchy

databricks-sdk-java/src/main/java/com/databricks/sdk/core/ApiClient.java

@@ -21,6 +21,7 @@
 import java.time.ZonedDateTime;
 import java.time.format.DateTimeFormatter;
 import java.util.*;
+import java.util.function.Function;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -29,55 +30,125 @@
  * guessing
  */
 public class ApiClient {
+  public static class Builder {
+    private Timer timer;
+    private Function<Void, Map<String, String>> authenticateFunc;
+    private Function<Void, String> getHostFunc;
+    private Function<Void, String> getAuthTypeFunc;
+    private Integer debugTruncateBytes;
+    private HttpClient httpClient;
+    private String accountId;
+    private RetryStrategyPicker retryStrategyPicker;
+    private boolean isDebugHeaders;
+
+    public Builder withDatabricksConfig(DatabricksConfig config) {
+      this.authenticateFunc = v -> config.authenticate();
+      this.getHostFunc = v -> config.getHost();
+      this.getAuthTypeFunc = v -> config.getAuthType();
+      this.httpClient = config.getHttpClient();
+      this.debugTruncateBytes = config.getDebugTruncateBytes();
+      this.accountId = config.getAccountId();
+      this.retryStrategyPicker = new RequestBasedRetryStrategyPicker(config.getHost());
+      this.isDebugHeaders = config.isDebugHeaders();
+
+      return this;
+    }
+
+    public Builder withTimer(Timer timer) {
+      this.timer = timer;
+      return this;
+    }
+
+    public Builder withAuthenticateFunc(Function<Void, Map<String, String>> authenticateFunc) {
+      this.authenticateFunc = authenticateFunc;
+      return this;
+    }
+
+    public Builder withGetHostFunc(Function<Void, String> getHostFunc) {
+      this.getHostFunc = getHostFunc;
+      return this;
+    }
+
+    public Builder withGetAuthTypeFunc(Function<Void, String> getAuthTypeFunc) {
+      this.getAuthTypeFunc = getAuthTypeFunc;
+      return this;
+    }
+
+    public Builder withHttpClient(HttpClient httpClient) {
+      this.httpClient = httpClient;
+      return this;
+    }
+
+    public Builder withRetryStrategyPicker(RetryStrategyPicker retryStrategyPicker) {
+      this.retryStrategyPicker = retryStrategyPicker;
+      return this;
+    }
+
+    public ApiClient build() {
+      return new ApiClient(this);
+    }
+  }
+
   private static final Logger LOG = LoggerFactory.getLogger(ApiClient.class);
 
   private final int maxAttempts;
 
   private final ObjectMapper mapper;
 
-  private final DatabricksConfig config;
-
   private final Random random;
 
   private final HttpClient httpClient;
   private final BodyLogger bodyLogger;
   private final RetryStrategyPicker retryStrategyPicker;
   private final Timer timer;
+  private final Function<Void, Map<String, String>> authenticateFunc;
+  private final Function<Void, String> getHostFunc;
+  private final Function<Void, String> getAuthTypeFunc;
+  private final String accountId;
+  private final boolean isDebugHeaders;
   private static final String RETRY_AFTER_HEADER = "retry-after";
 
   public ApiClient() {
     this(ConfigLoader.getDefault());
   }
 
   public String configuredAccountID() {
-    return config.getAccountId();
+    return accountId;
   }
 
   public ApiClient(DatabricksConfig config) {
     this(config, new SystemTimer());
   }
 
   public ApiClient(DatabricksConfig config, Timer timer) {
-    this.config = config;
-    config.resolve();
-
-    Integer rateLimit = config.getRateLimit();
-    if (rateLimit == null) {
-      rateLimit = 15;
-    }
-
-    Integer debugTruncateBytes = config.getDebugTruncateBytes();
+    this(new Builder().withDatabricksConfig(config.resolve()).withTimer(timer));
+  }
+
+  private ApiClient(Builder builder) {
+    this.timer = builder.timer != null ? builder.timer : new SystemTimer();
+    this.authenticateFunc =
+        builder.authenticateFunc != null
+            ? builder.authenticateFunc
+            : v -> new HashMap<String, String>();
+    this.getHostFunc = builder.getHostFunc != null ? builder.getHostFunc : v -> "";
+    this.getAuthTypeFunc = builder.getAuthTypeFunc != null ? builder.getAuthTypeFunc : v -> "";
+    this.httpClient = builder.httpClient;
+    this.accountId = builder.accountId;
+    this.retryStrategyPicker =
+        builder.retryStrategyPicker != null
+            ? builder.retryStrategyPicker
+            : new RequestBasedRetryStrategyPicker(this.getHostFunc.apply(null));
+    this.isDebugHeaders = builder.isDebugHeaders;
+
+    Integer debugTruncateBytes = builder.debugTruncateBytes;
     if (debugTruncateBytes == null) {
       debugTruncateBytes = 96;
     }
 
     maxAttempts = 4;
     mapper = SerDeUtils.createMapper();
     random = new Random();
-    httpClient = config.getHttpClient();
     bodyLogger = new BodyLogger(mapper, 1024, debugTruncateBytes);
-    retryStrategyPicker = new RequestBasedRetryStrategyPicker(this.config);
-    this.timer = timer;
   }
 
   private static <I> void setQuery(Request in, I entity) {
@@ -203,7 +274,7 @@ private <I> Request prepareBaseRequest(String method, String path, I in)
       InputStream body = (InputStream) in;
       return new Request(method, path, body);
     } else {
-      String body = serialize(in);
+      String body = (in instanceof String) ? (String) in : serialize(in);
       return new Request(method, path, body);
     }
   }
@@ -245,15 +316,19 @@ private Response executeInner(Request in, String path) {
       Response out = null;
 
       // Authenticate the request. Failures should not be retried.
-      in.withHeaders(config.authenticate());
+      in.withHeaders(authenticateFunc.apply(null));
 
       // Prepend host to URL only after config.authenticate().
       // This call may configure the host (e.g. in case of notebook native auth).
-      in.withUrl(config.getHost() + path);
+      in.withUrl(getHostFunc.apply(null) + path);
 
       // Set User-Agent with auth type info, which is available only
       // after the first invocation to config.authenticate()
-      String userAgent = String.format("%s auth/%s", UserAgent.asString(), config.getAuthType());
+      String userAgent = UserAgent.asString();
+      String authType = getAuthTypeFunc.apply(null);
+      if (authType != "") {
+        userAgent += String.format(" auth/%s", authType);
+      }
       in.withHeader("User-Agent", userAgent);
 
       // Make the request, catching any exceptions, as we may want to retry.
@@ -347,9 +422,9 @@ private String makeLogRecord(Request in, Response out) {
     StringBuilder sb = new StringBuilder();
     sb.append("> ");
     sb.append(in.getRequestLine());
-    if (config.isDebugHeaders()) {
+    if (this.isDebugHeaders) {
       sb.append("\n * Host: ");
-      sb.append(config.getHost());
+      sb.append(this.getHostFunc.apply(null));
       in.getHeaders()
           .forEach((header, value) -> sb.append(String.format("\n * %s: %s", header, value)));
     }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

@@ -581,12 +581,15 @@ private OpenIDConnectEndpoints fetchDefaultOidcEndpoints() throws IOException {
       return new OpenIDConnectEndpoints(prefix + "/v1/token", prefix + "/v1/authorize");
     }
 
-    String oidcEndpoint = getHost() + "/oidc/.well-known/oauth-authorization-server";
-    Response resp = getHttpClient().execute(new Request("GET", oidcEndpoint));
-    if (resp.getStatusCode() != 200) {
-      return null;
-    }
-    return new ObjectMapper().readValue(resp.getBody(), OpenIDConnectEndpoints.class);
+    ApiClient apiClient =
+        new ApiClient.Builder()
+            .withHttpClient(getHttpClient())
+            .withGetHostFunc(v -> getHost())
+            .build();
+    return apiClient.GET(
+        "/oidc/.well-known/oauth-authorization-server",
+        OpenIDConnectEndpoints.class,
+        new HashMap<>());
   }
 
   @Override

databricks-sdk-java/src/main/java/com/databricks/sdk/core/http/FormRequest.java

@@ -8,15 +8,15 @@ public FormRequest(String url, Map<String, String> form) {
   }
 
   public FormRequest(String method, String url, Map<String, String> form) {
-    super(method, url, mapToQuery(wrapValuesInList(form)));
+    super(method, url, wrapValuesInList(form));
     withHeader("Content-Type", "application/x-www-form-urlencoded");
   }
 
-  static Map<String, List<String>> wrapValuesInList(Map<String, String> map) {
+  public static String wrapValuesInList(Map<String, String> map) {
     Map<String, List<String>> m = new LinkedHashMap<>();
     for (Map.Entry<String, String> entry : map.entrySet()) {
       m.put(entry.getKey(), Collections.singletonList(entry.getValue()));
     }
-    return m;
+    return mapToQuery(m);
   }
 }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OpenIDConnectEndpoints.java

@@ -11,10 +11,14 @@
  */
 @JsonIgnoreProperties(ignoreUnknown = true)
 public class OpenIDConnectEndpoints {
+  @JsonProperty("token_endpoint")
   private String tokenEndpoint;
 
+  @JsonProperty("authorization_endpoint")
   private String authorizationEndpoint;
 
+  public OpenIDConnectEndpoints() {}
+
   public OpenIDConnectEndpoints(
       @JsonProperty("token_endpoint") String tokenEndpoint,
       @JsonProperty("authorization_endpoint") String authorizationEndpoint)

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/RefreshableTokenSource.java

@@ -1,11 +1,9 @@
 package com.databricks.sdk.core.oauth;
 
+import com.databricks.sdk.core.ApiClient;
 import com.databricks.sdk.core.DatabricksException;
 import com.databricks.sdk.core.http.FormRequest;
 import com.databricks.sdk.core.http.HttpClient;
-import com.databricks.sdk.core.http.Response;
-import com.fasterxml.jackson.databind.ObjectMapper;
-import java.io.IOException;
 import java.time.LocalDateTime;
 import java.time.temporal.ChronoUnit;
 import java.util.Base64;
@@ -62,17 +60,19 @@ protected static Token retrieveToken(
         headers.put(HttpHeaders.AUTHORIZATION, authHeaderValue);
         break;
     }
-    FormRequest req = new FormRequest(tokenUrl, params);
-    req.withHeaders(headers);
+    headers.put("Content-Type", "application/x-www-form-urlencoded");
     try {
-      Response rawResp = hc.execute(req);
-      OAuthResponse resp = new ObjectMapper().readValue(rawResp.getBody(), OAuthResponse.class);
+      ApiClient apiClient = new ApiClient.Builder().withHttpClient(hc).build();
+
+      OAuthResponse resp =
+          apiClient.POST(
+              tokenUrl, FormRequest.wrapValuesInList(params), OAuthResponse.class, headers);
       if (resp.getErrorCode() != null) {
         throw new IllegalArgumentException(resp.getErrorCode() + ": " + resp.getErrorSummary());
       }
       LocalDateTime expiry = LocalDateTime.now().plus(resp.getExpiresIn(), ChronoUnit.SECONDS);
       return new Token(resp.getAccessToken(), resp.getTokenType(), resp.getRefreshToken(), expiry);
-    } catch (IOException e) {
+    } catch (Exception e) {
       throw new DatabricksException("Failed to refresh credentials: " + e.getMessage(), e);
     }
   }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/retry/RequestBasedRetryStrategyPicker.java

@@ -1,6 +1,5 @@
 package com.databricks.sdk.core.retry;
 
-import com.databricks.sdk.core.DatabricksConfig;
 import com.databricks.sdk.core.http.Request;
 import java.util.AbstractMap;
 import java.util.Arrays;
@@ -37,15 +36,14 @@ public class RequestBasedRetryStrategyPicker implements RetryStrategyPicker {
   private static final IdempotentRequestRetryStrategy IDEMPOTENT_RETRY_STRATEGY =
       new IdempotentRequestRetryStrategy();
 
-  public RequestBasedRetryStrategyPicker(DatabricksConfig config) {
+  public RequestBasedRetryStrategyPicker(String host) {
     this.idempotentRequestsPattern =
         IDEMPOTENT_REQUESTS.stream()
             .map(
                 request ->
                     new AbstractMap.SimpleEntry<>(
                         request.getMethod(),
-                        Pattern.compile(
-                            config.getHost() + request.getUrl(), Pattern.CASE_INSENSITIVE)))
+                        Pattern.compile(host + request.getUrl(), Pattern.CASE_INSENSITIVE)))
             .collect(Collectors.toList());
   }
 

databricks-sdk-java/src/test/java/com/databricks/sdk/benchmark/DatabricksAuthLoadTest.java

@@ -0,0 +1,73 @@
+package com.databricks.sdk.benchmark;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import java.util.concurrent.*;
+
+/*
+ This test executes the authentication workflow 200 times concurrently and verifies that all 200 runs complete successfully.
+ It is designed to address a previously observed issue where multiple SDK operations needed to authenticate, causing the OIDC endpoints to rate-limit the requests.
+ Now that these endpoints are configured to retry upon receiving a 429 error, the test runs successfully.
+ However, since this test generates a large number of requests, it should be run manually rather than being included in CI processes.
+*/
+/*
+public class DatabricksAuthLoadTest implements GitHubUtils, ConfigResolving {
+
+  @Test
+  @Disabled
+  public void testConcurrentConfigBasicAuthAttrs() throws Exception {
+    int numThreads = 200;
+    ExecutorService executorService = Executors.newFixedThreadPool(numThreads);
+    List<Future<Boolean>> futures = new ArrayList<>();
+    int successCount = 0;
+    int failureCount = 0;
+
+    Callable<Boolean> task =
+        () -> {
+          try {
+            DatabricksConfig config =
+                new DatabricksConfig()
+                    .setHost("https://dbc-bb03964f-3f59.cloud.databricks.com")
+                    .setClientId("<<REDACTED>>")
+                    .setClientSecret("<<REDACTED>>");
+
+            config.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build());
+            config.authenticate();
+
+            assertEquals("oauth-m2m", config.getAuthType());
+
+            return true;
+          } catch (Exception e) {
+            System.err.println(
+                "DatabricksException occurred in thread " + Thread.currentThread().getName());
+            e.printStackTrace();
+            return false;
+          }
+        };
+
+    for (int i = 0; i < numThreads; i++) {
+      futures.add(executorService.submit(task));
+    }
+
+    executorService.shutdown();
+    if (!executorService.awaitTermination(60, TimeUnit.SECONDS)) {
+      executorService.shutdownNow();
+    }
+
+    for (Future<Boolean> future : futures) {
+      if (future.get()) {
+        successCount++;
+      } else {
+        failureCount++;
+      }
+    }
+
+    // Log the results
+    System.out.println("Number of successful threads: " + successCount);
+    System.out.println("Number of failed threads: " + failureCount);
+
+    // Optionally, you can assert that there were no failures
+    assertEquals(0, failureCount);
+  }
+}
+*/