Skip to content

Commit

Permalink
retry datastore transactions (#1242) (#1932)
Browse files Browse the repository at this point in the history
* retry datastore transactions (#1242)
  • Loading branch information
neozwu authored Apr 18, 2017
1 parent cc8dc45 commit b355f4d
Show file tree
Hide file tree
Showing 5 changed files with 188 additions and 5 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ public RetryResult afterEval(Exception exception, RetryResult retryResult) {
public RetryResult beforeEval(Exception exception) {
if (exception instanceof BaseServiceException) {
boolean retriable = ((BaseServiceException) exception).isRetryable();
return retriable ? Interceptor.RetryResult.RETRY : Interceptor.RetryResult.NO_RETRY;
return retriable ? Interceptor.RetryResult.RETRY : Interceptor.RetryResult.CONTINUE_EVALUATION;
}
return Interceptor.RetryResult.CONTINUE_EVALUATION;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,11 @@

package com.google.cloud.datastore;

import com.google.api.gax.core.RetrySettings;
import com.google.cloud.BaseService;
import com.google.cloud.ExceptionHandler;
import com.google.cloud.RetryHelper;
import com.google.cloud.RetryHelper.RetryHelperException;
import com.google.api.gax.core.RetrySettings;
import com.google.cloud.ServiceOptions;
import com.google.cloud.datastore.ReadOption.EventualConsistency;
import com.google.cloud.datastore.spi.v1.DatastoreRpc;
Expand All @@ -31,7 +32,6 @@
import com.google.common.collect.Sets;
import com.google.datastore.v1.ReadOptions.ReadConsistency;
import com.google.protobuf.ByteString;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
Expand All @@ -47,6 +47,8 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas

private final DatastoreRpc datastoreRpc;
private final RetrySettings retrySettings;
private static final ExceptionHandler TRANSACTION_EXCEPTION_HANDLER =
TransactionExceptionHandler.build();

DatastoreImpl(DatastoreOptions options) {
super(options);
Expand All @@ -66,8 +68,22 @@ public Transaction newTransaction() {
}

@Override
public <T> T runInTransaction(TransactionCallable<T> callable) {
return DatastoreHelper.runInTransaction(this, callable);
public <T> T runInTransaction(final TransactionCallable<T> callable) {
final DatastoreImpl self = this;
try {
return RetryHelper.runWithRetries(
new Callable<T>() {
@Override
public T call() throws DatastoreException {
return DatastoreHelper.runInTransaction(self, callable);
}
},
retrySettings,
TRANSACTION_EXCEPTION_HANDLER,
getOptions().getClock());
} catch (RetryHelperException e) {
throw DatastoreException.translateAndThrow(e);
}
}

@Override
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
/*
* Copyright 2017 Google Inc. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package com.google.cloud.datastore;

import com.google.cloud.ExceptionHandler;
import com.google.cloud.ExceptionHandler.Interceptor;

public class TransactionExceptionHandler {
public static final Interceptor TRANSACTION_EXCEPTION_HANDLER_INTERCEPTOR =
new Interceptor() {

private static final long serialVersionUID = -1240723093072535978L;

private static final int ABORTED_CODE = 10;

@Override
public RetryResult beforeEval(Exception exception) {
if (exception instanceof DatastoreException) {
DatastoreException e = getInnerException((DatastoreException) exception);
if (e.getCode() == ABORTED_CODE || e.getReason().equals("ABORTED")) {
return Interceptor.RetryResult.RETRY;
}
}
return Interceptor.RetryResult.CONTINUE_EVALUATION;
}

@Override
public RetryResult afterEval(Exception exception, RetryResult retryResult) {
return Interceptor.RetryResult.CONTINUE_EVALUATION;
}

private DatastoreException getInnerException(DatastoreException exception) {
while (exception.getCause() instanceof DatastoreException) {
exception = (DatastoreException) exception.getCause();
}
return exception;
}
};

public static ExceptionHandler build() {
return ExceptionHandler.newBuilder()
.abortOn(RuntimeException.class)
.addInterceptors(
DatastoreImpl.EXCEPTION_HANDLER_INTERCEPTOR, TRANSACTION_EXCEPTION_HANDLER_INTERCEPTOR)
.build();
}

/** Intentionally private empty constructor to disable instantiation of this class. */
private TransactionExceptionHandler() {}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
/*
* Copyright 2017 Google Inc. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package com.google.cloud.datastore;

import static junit.framework.TestCase.assertFalse;
import static junit.framework.TestCase.assertTrue;

import com.google.cloud.BaseServiceException;
import com.google.cloud.ExceptionHandler;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;

/** Tests for {@link TransactionExceptionHandler}. */
public class TransactionExceptionHandlerTest {

@Rule public ExpectedException thrown = ExpectedException.none();

@Test
public void testShouldTry() {
ExceptionHandler handler =
ExceptionHandler.newBuilder()
.abortOn(RuntimeException.class)
.addInterceptors(DatastoreImpl.EXCEPTION_HANDLER_INTERCEPTOR)
.build();
ExceptionHandler transactionHandler = TransactionExceptionHandler.build();

assertFalse(handler.accept(new DatastoreException(10, "", "ABORTED", false, null)));
assertFalse(handler.accept(new DatastoreException(10, "", "", false, null)));
assertFalse(handler.accept(new DatastoreException(0, "", "", false, null)));

assertTrue(transactionHandler.accept(new DatastoreException(10, "", "ABORTED", false, null)));
assertTrue(transactionHandler.accept(new DatastoreException(10, "", "", false, null)));
assertFalse(transactionHandler.accept(new DatastoreException(0, "", "", false, null)));

DatastoreException nestedDatastoreException =
new DatastoreException(
BaseServiceException.UNKNOWN_CODE,
"",
null,
new DatastoreException(10, "", "ABORTED", false, null));

assertTrue(transactionHandler.accept(nestedDatastoreException));
assertFalse(handler.accept(nestedDatastoreException));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import com.google.cloud.datastore.Datastore;
import com.google.cloud.datastore.DatastoreException;
import com.google.cloud.datastore.DatastoreOptions;
import com.google.cloud.datastore.DatastoreReaderWriter;
import com.google.cloud.datastore.DateTime;
import com.google.cloud.datastore.DateTimeValue;
import com.google.cloud.datastore.Entity;
Expand Down Expand Up @@ -729,4 +730,46 @@ public void testDelete() {
assertNull(keys.next());
assertFalse(keys.hasNext());
}

@Test
public void testRunInTransaction() {
Datastore.TransactionCallable<Integer> callable1 =
new Datastore.TransactionCallable<Integer>() {
private Integer attempts = 1;

public Integer run(DatastoreReaderWriter transaction) {
transaction.get(KEY1);
if (attempts < 2) {
++attempts;
throw new DatastoreException(10, "", "ABORTED", false, null);
}

return attempts;
}
};

int result = DATASTORE.runInTransaction(callable1);
assertEquals(result, 2);

Datastore.TransactionCallable<Integer> callable2 =
new Datastore.TransactionCallable<Integer>() {
private Integer attempts = 1;

public Integer run(DatastoreReaderWriter transaction) {
transaction.get(KEY1);
if (attempts < 2) {
++attempts;
throw new DatastoreException(4, "", "DEADLINE_EXCEEDED", false, null);
}
return attempts;
}
};

try {
DATASTORE.runInTransaction(callable2);
fail("Expecting a failure");
} catch (DatastoreException expected) {
assertEquals(((DatastoreException) expected.getCause()).getCode(), 4);
}
}
}

0 comments on commit b355f4d

Please sign in to comment.