Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Updated BigQueryRetryAlgorithm so that it can retry on RateLimit Errors using RegEx #1499

Merged
merged 5 commits into from Aug 11, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
Expand Up @@ -19,4 +19,9 @@
public class BigQueryErrorMessages {
public static final String RATE_LIMIT_EXCEEDED_MSG =
"Exceeded rate limits:"; // Error Message for RateLimitExceeded Error
public static final String JOB_RATE_LIMIT_EXCEEDED_MSG = "Job exceeded rate limits:";

public class RetryRegExPatterns {
public static final String RATE_LIMIT_EXCEEDED_REGEX = ".*exceed.*rate.*limit.*";
}
}
Expand Up @@ -240,7 +240,9 @@ public Page<FieldValueList> getNextPage() {
private static final BigQueryRetryConfig DEFAULT_RETRY_CONFIG =
BigQueryRetryConfig.newBuilder()
.retryOnMessage(BigQueryErrorMessages.RATE_LIMIT_EXCEEDED_MSG)
.build(); // retry config with Error Message for RateLimitExceeded Error
.retryOnMessage(BigQueryErrorMessages.JOB_RATE_LIMIT_EXCEEDED_MSG)
.retryOnRegEx(BigQueryErrorMessages.RetryRegExPatterns.RATE_LIMIT_EXCEEDED_REGEX)
.build(); // retry config with Error Messages and RegEx for RateLimitExceeded Error

BigQueryImpl(BigQueryOptions options) {
super(options);
Expand Down
Expand Up @@ -27,6 +27,7 @@
import com.google.api.gax.retrying.TimedRetryAlgorithmWithContext;
import java.util.Iterator;
import java.util.concurrent.CancellationException;
import java.util.regex.Pattern;

public class BigQueryRetryAlgorithm<ResponseT> extends RetryAlgorithm<ResponseT> {
private final BigQueryRetryConfig bigQueryRetryConfig;
Expand Down Expand Up @@ -69,17 +70,35 @@ private boolean shouldRetryBasedOnBigQueryRetryConfig(
*/
String errorDesc;
if (previousThrowable != null && (errorDesc = previousThrowable.getMessage()) != null) {
errorDesc = errorDesc.toLowerCase(); // for case insensitive comparison
for (Iterator<String> retriableMessages =
bigQueryRetryConfig.getRetriableErrorMessages().iterator();
retriableMessages.hasNext(); ) {
if (errorDesc.contains(retriableMessages.next())) { // Error message should be retried
if (errorDesc.contains(
retriableMessages
.next()
.toLowerCase())) { // Error message should be retried, implementing cases
// insensitive match
return true;
}
}
// Check if there's a regex which matches the error message. This avoids too many regex
// matches which is expensive
for (Iterator<String> retriableRegExes = bigQueryRetryConfig.getRetriableRegExes().iterator();
retriableRegExes.hasNext(); ) {
if (matchRegEx(retriableRegExes.next(), errorDesc)) {
return true;
}
}
}
return false;
}

public static boolean matchRegEx(
String retriableRegEx, String errorDesc) { // cases insensitive match regex matching
return Pattern.matches(retriableRegEx.toLowerCase(), errorDesc.toLowerCase());
}

/*Duplicating this method as it can not be inherited from the RetryAlgorithm due to the default access modifier*/
boolean shouldRetryBasedOnResult(
RetryingContext context, Throwable previousThrowable, ResponseT previousResponse) {
Expand Down
Expand Up @@ -21,18 +21,25 @@

public class BigQueryRetryConfig {
private final ImmutableSet<String> retriableErrorMessages;
private final ImmutableSet<String> retriableRegExes;

private BigQueryRetryConfig(Builder builder) {
retriableErrorMessages = builder.retriableErrorMessages.build();
retriableRegExes = builder.retriableRegExes.build();
}

public ImmutableSet<String> getRetriableErrorMessages() {
return retriableErrorMessages;
}

public ImmutableSet<String> getRetriableRegExes() {
return retriableRegExes;
}

// BigQueryRetryConfig builder
public static class Builder {
private final ImmutableSet.Builder<String> retriableErrorMessages = ImmutableSet.builder();
private final ImmutableSet.Builder<String> retriableRegExes = ImmutableSet.builder();

private Builder() {}

Expand All @@ -43,6 +50,13 @@ public final Builder retryOnMessage(String... errorMessages) {
return this;
}

public final Builder retryOnRegEx(String... regExPatterns) {
for (String regExPattern : regExPatterns) {
retriableRegExes.add(checkNotNull(regExPattern));
}
return this;
}

public BigQueryRetryConfig build() {
return new BigQueryRetryConfig(this);
}
Expand Down
Expand Up @@ -508,11 +508,12 @@ public class BigQueryImplTest {
.setEtag(ETAG)
.setVersion(1)
.build();

private BigQueryOptions options;
private BigQueryRpcFactory rpcFactoryMock;
private BigQueryRpc bigqueryRpcMock;
private BigQuery bigquery;
private static final String RATE_LIMIT_ERROR_MSG =
"Job exceeded rate limits: Your table exceeded quota for table update operations. For more information, see https://cloud.google.com/bigquery/docs/troubleshoot-quotas";

@Captor private ArgumentCaptor<Map<BigQueryRpc.Option, Object>> capturedOptions;
@Captor private ArgumentCaptor<com.google.api.services.bigquery.model.Job> jobCapture;
Expand Down Expand Up @@ -2439,9 +2440,7 @@ public void testFastQueryRateLimitIdempotency() throws Exception {
.thenThrow(new BigQueryException(504, "Gateway Timeout"))
.thenThrow(
new BigQueryException(
400,
BigQueryErrorMessages
.RATE_LIMIT_EXCEEDED_MSG)) // retrial on based on RATE_LIMIT_EXCEEDED_MSG
400, RATE_LIMIT_ERROR_MSG)) // retrial on based on RATE_LIMIT_EXCEEDED_MSG
.thenReturn(responsePb);

bigquery =
Expand Down Expand Up @@ -2470,6 +2469,30 @@ public void testFastQueryRateLimitIdempotency() throws Exception {
verify(bigqueryRpcMock, times(6)).queryRpc(eq(PROJECT), requestPbCapture.capture());
}

@Test
public void testRateLimitRegEx() throws Exception {
String msg2 =
"Job eceeded rate limits: Your table exceeded quota for table update operations. For more information, see https://cloud.google.com/bigquery/docs/troubleshoot-quotas";
String msg3 = "exceeded rate exceeded quota for table update";
String msg4 = "exceeded rate limits";
assertTrue(
BigQueryRetryAlgorithm.matchRegEx(
BigQueryErrorMessages.RetryRegExPatterns.RATE_LIMIT_EXCEEDED_REGEX,
RATE_LIMIT_ERROR_MSG));
assertFalse(
BigQueryRetryAlgorithm.matchRegEx(
BigQueryErrorMessages.RetryRegExPatterns.RATE_LIMIT_EXCEEDED_REGEX,
msg2.toLowerCase()));
assertFalse(
BigQueryRetryAlgorithm.matchRegEx(
BigQueryErrorMessages.RetryRegExPatterns.RATE_LIMIT_EXCEEDED_REGEX,
msg3.toLowerCase()));
assertTrue(
BigQueryRetryAlgorithm.matchRegEx(
BigQueryErrorMessages.RetryRegExPatterns.RATE_LIMIT_EXCEEDED_REGEX,
msg4.toLowerCase()));
}

@Test
public void testFastQueryDDLShouldRetry() throws Exception {
com.google.api.services.bigquery.model.QueryResponse responsePb =
Expand Down