[Fix] Retry on too many auth requests #355

rauchy · 2024-10-02T09:20:04Z

Changes

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

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

mgyucht · 2024-10-02T14:01:17Z

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

+  public ApiClient(
+      DatabricksConfig config,
+      Timer timer,
+      Function<Void, Map<String, String>> authenticateFunc,
+      Function<Void, String> getHostFunc,
+      Function<Void, String> getAuthTypeFunc) {


private ApiClient(Builder builder) {...}

mgyucht · 2024-10-02T14:01:39Z

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

+    this(
+        config,
+        timer,
+        v -> config.authenticate(),
+        v -> config.getHost(),
+        v -> config.getAuthType());


this(new Builder().withConfig(config.resolve()).withTimer(timer))

mgyucht · 2024-10-02T14:02:39Z

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

+    }
+
+    public ApiClient build() {
+      return new ApiClient(config, timer, authenticateFunc, getHostFunc, getAuthTypeFunc);


return new ApiClient(this);

mgyucht · 2024-10-02T14:10:12Z

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

+    this(
+        config,
+        timer,
+        v -> config.authenticate(),


config::authenticate

See #355 (comment)

mgyucht · 2024-10-02T14:16:42Z

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

+  static class Builder {
+    private DatabricksConfig config;
+    private Timer timer;
+    private Function<Void, Map<String, String>> authenticateFunc;


Can we use Supplier for these?

Suppliers look good for replacing the use of Function here, but some functions (authenticate as an example) are overloaded, and Supplier and the :: syntax can't deal with those. So if we choose to use Supplier here, some functions will still be stored under a Function variable, which kinda makes it even more complicated :/

mgyucht · 2024-10-02T14:17:56Z

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

-      String userAgent = String.format("%s auth/%s", UserAgent.asString(), config.getAuthType());
+      String authType = getAuthTypeFunc.apply(null);
+      String userAgent = UserAgent.asString();
+      if (authType != "") {


If getAuthTypeFunc is null, we can jst check if that func is null

mgyucht

Looking good. Few comments, otherwise we're on the right track.

mgyucht · 2024-10-03T11:33:24Z

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

+    private RetryStrategyPicker retryStrategyPicker;
+    private boolean isDebugHeaders;
+
+    public Builder withDatabricksConfig(DatabricksConfig config) {


Actually, this could even be private, since it is only used in the default constructor.

mgyucht · 2024-10-03T11:34:35Z

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

@@ -202,6 +267,8 @@ private <I> Request prepareBaseRequest(String method, String path, I in)
    } else if (InputStream.class.isAssignableFrom(in.getClass())) {
      InputStream body = (InputStream) in;
      return new Request(method, path, body);
+    } else if (in instanceof String) {


Suggestion: leave a comment that this assumes that the request is already serialized.

mgyucht · 2024-10-03T11:36:12Z

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

-      rateLimit = 15;
-    }
+  private ApiClient(Builder builder) {
+    this.timer = builder.timer;


Suggestion: let's use sane defaults when those values are not provided in the builder. Ideally ApiClient.Builder().build() would make a more or less working ApiClient. Here, we can use SystemTimer() if not specified.

mgyucht · 2024-10-03T11:36:32Z

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

-    return new ObjectMapper().readValue(resp.getBody(), OpenIDConnectEndpoints.class);
+    ApiClient apiClient =
+        new ApiClient.Builder()
+            .withDatabricksConfig(this)


Do we need to specify the config directly here?

mgyucht · 2024-10-03T11:36:48Z

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

+            .withTimer(new SystemTimer())
+            .withAuthenticateFunc(v -> new HashMap<String, String>())
+            .withGetHostFunc(v -> getHost())
+            .withGetAuthTypeFunc(v -> "")


This we should probably leave out? Otherwise it is not null according to your check.

mgyucht · 2024-10-03T11:38:33Z

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

+        new ApiClient.Builder()
+            .withDatabricksConfig(this)
+            .withTimer(new SystemTimer())
+            .withAuthenticateFunc(v -> new HashMap<String, String>())


If you follow my suggestion from above, we can leave this out.

mgyucht · 2024-10-03T11:39:42Z

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

+              .withGetAuthTypeFunc(v -> "")
+              .withGetHostFunc(v -> "")


ditto, these can be removed if you take my suggestions

mgyucht · 2024-10-03T11:40:08Z

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

+    this.getAuthTypeFunc = builder.getAuthTypeFunc;
+    this.httpClient = builder.httpClient;
+    this.accountId = builder.accountId;
+    this.retryStrategyPicker = builder.retryStrategyPicker;


Same here, can we use the default request-based retry strategy picker if none is provided?

mgyucht · 2024-10-03T11:41:18Z

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

@@ -0,0 +1,73 @@
+// These are auto-generated tests for Unified Authentication


This is not autogenerated?

mgyucht · 2024-10-03T11:42:01Z

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

+import java.util.concurrent.*;
+import org.junit.jupiter.api.Test;
+
+public class DatabricksAuthLoadTest implements GitHubUtils, ConfigResolving {


Let's leave a comment here about this test.

Also small suggestion: let's move this into a benchmark package. It'll be a new convention within Java SDK tests to keep things a bit more organized for now.

rauchy · 2024-10-03T12:51:56Z

Great comments @mgyucht, thanks for the thorough review.

The null checks make sense. I eventually went with a more "tell, don't ask" approach by providing no-ops in the private ApiClient constructor. This way we could avoid null paranoia.

mgyucht

LGTM. Make sure to run make fmt to fix the format check.

I will take a look at how Maven is running tests to see why this test is being executed. In the meantime, we should be able to use the Disabled annotation (https://junit.org/junit5/docs/5.2.0/api/org/junit/jupiter/api/Disabled.html) to disable the test explicitly.

Signed-off-by: Omer Lachish <rauchy@users.noreply.github.com>

…sponses when fetching OIDC endpoints Signed-off-by: Omer Lachish <rauchy@users.noreply.github.com>

Signed-off-by: Omer Lachish <rauchy@users.noreply.github.com>

…he time Signed-off-by: Omer Lachish <rauchy@users.noreply.github.com>

### Bug Fixes * Retry on too many auth requests ([#355](#355)).

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

### Bug Fixes * Retry on too many auth requests ([#355](#355)). Co-authored-by: Omer Lachish <rauchy@users.noreply.github.com>

rauchy requested a review from mgyucht October 2, 2024 09:20

rauchy force-pushed the rauchy/retry-on-too-many-requests branch from 5c68f5d to 7409ea2 Compare October 2, 2024 10:16

mgyucht reviewed Oct 2, 2024

View reviewed changes

mgyucht reviewed Oct 3, 2024

View reviewed changes

rauchy force-pushed the rauchy/retry-on-too-many-requests branch from e36b927 to 68ae507 Compare October 3, 2024 12:49

rauchy requested a review from mgyucht October 3, 2024 12:52

rauchy force-pushed the rauchy/retry-on-too-many-requests branch from c6c811c to 68ae507 Compare October 3, 2024 13:01

mgyucht approved these changes Oct 3, 2024

View reviewed changes

rauchy force-pushed the rauchy/retry-on-too-many-requests branch 9 times, most recently from a50b790 to 1731133 Compare October 3, 2024 15:12

rauchy added 11 commits October 3, 2024 17:19

allow FixtureServer to respond with a specific status code

f3656de

Signed-off-by: Omer Lachish <rauchy@users.noreply.github.com>

use ApiClient instead of CommonsHttpClient to allow retries on 429 re…

babd8d1

…sponses when fetching OIDC endpoints Signed-off-by: Omer Lachish <rauchy@users.noreply.github.com>

formatting

f9b1cf0

Signed-off-by: Omer Lachish <rauchy@users.noreply.github.com>

decouple DatabricksConfig from ApiClient

4a7eab6

Signed-off-by: Omer Lachish <rauchy@users.noreply.github.com>

retry on 429s coming in from /oidc/v1/token as well

d75013f

Signed-off-by: Omer Lachish <rauchy@users.noreply.github.com>

check for getAuthTypeFunc injection

bcf0554

Signed-off-by: Omer Lachish <rauchy@users.noreply.github.com>

add a manual auth load test

a863160

Signed-off-by: Omer Lachish <rauchy@users.noreply.github.com>

address some review comments

cb5df62

Signed-off-by: Omer Lachish <rauchy@users.noreply.github.com>

move DatabricksAuthLoadTest to benchmarks and document it

35a10f8

Signed-off-by: Omer Lachish <rauchy@users.noreply.github.com>

removed redundant import

ef1267f

Signed-off-by: Omer Lachish <rauchy@users.noreply.github.com>

disable the auth load benchmark test to prevent it from running all t…

86f3f8b

…he time Signed-off-by: Omer Lachish <rauchy@users.noreply.github.com>

rauchy force-pushed the rauchy/retry-on-too-many-requests branch from 1731133 to 86f3f8b Compare October 3, 2024 15:19

rauchy added this pull request to the merge queue Oct 3, 2024

Merged via the queue into main with commit ed4992d Oct 3, 2024
10 checks passed

rauchy deleted the rauchy/retry-on-too-many-requests branch October 3, 2024 15:30

rauchy added a commit that referenced this pull request Oct 3, 2024

[Release] Release v0.32.1

054f5c8

### Bug Fixes * Retry on too many auth requests ([#355](#355)).

rauchy mentioned this pull request Oct 3, 2024

[Release] Release v0.32.1 #356

Merged

rauchy added a commit that referenced this pull request Oct 3, 2024

[Release] Release v0.32.1

aa3fb52

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

github-merge-queue bot pushed a commit that referenced this pull request Oct 3, 2024

[Release] Release v0.32.1 (#356)

32211bf

### Bug Fixes * Retry on too many auth requests ([#355](#355)). Co-authored-by: Omer Lachish <rauchy@users.noreply.github.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Fix] Retry on too many auth requests #355

[Fix] Retry on too many auth requests #355

rauchy commented Oct 2, 2024

mgyucht Oct 2, 2024

rauchy Oct 3, 2024

mgyucht Oct 2, 2024

rauchy Oct 3, 2024

mgyucht Oct 2, 2024

rauchy Oct 3, 2024

mgyucht Oct 2, 2024

rauchy Oct 3, 2024

mgyucht Oct 2, 2024

rauchy Oct 3, 2024

mgyucht Oct 2, 2024

rauchy Oct 3, 2024

mgyucht left a comment

mgyucht Oct 3, 2024

mgyucht Oct 3, 2024

mgyucht Oct 3, 2024

mgyucht Oct 3, 2024

mgyucht Oct 3, 2024

mgyucht Oct 3, 2024

mgyucht Oct 3, 2024

mgyucht Oct 3, 2024

mgyucht Oct 3, 2024

rauchy Oct 3, 2024

mgyucht Oct 3, 2024

rauchy Oct 3, 2024

rauchy commented Oct 3, 2024

mgyucht left a comment

		@@ -0,0 +1,73 @@
		// These are auto-generated tests for Unified Authentication

[Fix] Retry on too many auth requests #355

[Fix] Retry on too many auth requests #355

Conversation

rauchy commented Oct 2, 2024

Changes

Tests

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mgyucht left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rauchy commented Oct 3, 2024

mgyucht left a comment

Choose a reason for hiding this comment