Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
feat: Updated BigQueryRetryAlgorithm so that it can retry on RateLi…
…mit Errors using RegEx (#1499)

* Modified RATE_LIMIT_ERROR_MSG and added testRateLimitRegEx for testing regex matching

* Added JOB_RATE_LIMIT_EXCEEDED_MSG and RATE_LIMIT_EXCEEDED_REGEX pattern

* Implemented retryOnRegEx to include retry on RegEx patterns

* Updated DEFAULT_RETRY_CONFIG to include retry on JOB_RATE_LIMIT_EXCEEDED_MSG and RATE_LIMIT_EXCEEDED_REGEX

* Modified shouldRetryBasedOnBigQueryRetryConfig method to try RegEx matching if all the contains checks fail
  • Loading branch information
prash-mi committed Aug 11, 2021
1 parent b705052 commit ec68c11
Show file tree
Hide file tree
Showing 5 changed files with 69 additions and 6 deletions.
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

0 comments on commit ec68c11

Please sign in to comment.