Merge pull request #493 from eddavisson/prev-txn

stickfigure · web-flow · commit 7fbb5718428f · 2024-04-19T10:04:10.000-07:00
Specify previous_transaction when retrying a txn
diff --git a/src/main/java/com/googlecode/objectify/cache/CachingAsyncDatastore.java b/src/main/java/com/googlecode/objectify/cache/CachingAsyncDatastore.java
@@ -3,6 +3,7 @@
 import com.google.cloud.datastore.Entity;
 import com.google.cloud.datastore.Key;
 import com.google.cloud.datastore.ReadOption;
+import com.google.protobuf.ByteString;
 import com.googlecode.objectify.cache.EntityMemcache.Bucket;
 import com.googlecode.objectify.impl.AsyncDatastore;
 import com.googlecode.objectify.impl.AsyncTransaction;
@@ -62,6 +63,11 @@ public AsyncTransaction newTransaction(final Runnable afterCommit) {
 		return new CachingAsyncTransaction(raw.newTransaction(afterCommit), memcache);
 	}
 
+	@Override
+	public AsyncTransaction newTransaction(final Runnable afterCommit, ByteString prevTxnHandle) {
+		return new CachingAsyncTransaction(raw.newTransaction(afterCommit, prevTxnHandle), memcache);
+	}
+
 	@Override
 	public Future<Map<Key, Entity>> get(final Collection<Key> keys, final ReadOption... options) {
 		final Map<Key, Bucket> soFar = this.memcache.getAll(keys);
diff --git a/src/main/java/com/googlecode/objectify/impl/AsyncDatastore.java b/src/main/java/com/googlecode/objectify/impl/AsyncDatastore.java
@@ -1,5 +1,7 @@
 package com.googlecode.objectify.impl;
 
+import com.google.protobuf.ByteString;
+
 /**
  * The new datastore SDK has a neat structure of interfaces and implementations (transaction, datastorereader, etc)
  * but doesn't currently support async operations. We need to shim in a Future-based API so that we can seamlessly
@@ -10,4 +12,8 @@ public interface AsyncDatastore extends AsyncDatastoreReaderWriter {
 	/**
 	 */
 	AsyncTransaction newTransaction(Runnable afterCommit);
+
+	/**
+	*/
+	AsyncTransaction newTransaction(Runnable afterCommit, ByteString prevTxnHandle);
 }
diff --git a/src/main/java/com/googlecode/objectify/impl/AsyncDatastoreImpl.java b/src/main/java/com/googlecode/objectify/impl/AsyncDatastoreImpl.java
@@ -1,6 +1,8 @@
 package com.googlecode.objectify.impl;
 
 import com.google.cloud.datastore.Datastore;
+import com.google.datastore.v1.TransactionOptions;
+import com.google.protobuf.ByteString;
 
 /** */
 public class AsyncDatastoreImpl extends AsyncDatastoreReaderWriterImpl implements AsyncDatastore {
@@ -15,6 +17,15 @@ public AsyncDatastoreImpl(final Datastore raw) {
 
 	@Override
 	public AsyncTransaction newTransaction(final Runnable afterCommit) {
-		return new AsyncTransactionImpl(datastore.newTransaction(), afterCommit);
+		return newTransaction(afterCommit, null);
+	}
+
+	@Override
+	public AsyncTransaction newTransaction(final Runnable afterCommit, final ByteString prevTxnHandle) {
+		TransactionOptions.Builder txnOptions = TransactionOptions.newBuilder();
+		if (prevTxnHandle != null) {
+			txnOptions.getReadWriteBuilder().setPreviousTransaction(prevTxnHandle);
+		}
+		return new AsyncTransactionImpl(datastore.newTransaction(txnOptions.build()), afterCommit);
 	}
 }
diff --git a/src/main/java/com/googlecode/objectify/impl/TransactorNo.java b/src/main/java/com/googlecode/objectify/impl/TransactorNo.java
@@ -2,10 +2,14 @@
 
 import com.google.cloud.datastore.DatastoreException;
 import com.google.common.base.Preconditions;
+import com.google.protobuf.ByteString;
 import com.google.rpc.Code;
 import com.googlecode.objectify.ObjectifyFactory;
 import com.googlecode.objectify.TxnType;
 import com.googlecode.objectify.Work;
+
+import java.util.concurrent.atomic.AtomicReference;
+
 import lombok.extern.slf4j.Slf4j;
 
 /**
@@ -90,9 +94,10 @@ public <R> R transactNew(final ObjectifyImpl parent, int limitTries, final Work<
 		Preconditions.checkArgument(limitTries >= 1);
 		final int ORIGINAL_TRIES = limitTries;
 
+		AtomicReference<ByteString> prevTxnHandle = new AtomicReference<>();
 		while (true) {
 			try {
-				return transactOnce(parent, work);
+				return transactOnce(parent, work, prevTxnHandle);
 			} catch (DatastoreException ex) {
 
 				// This doesn't work because the SDK considers all transactions to be non-retryable. Objectify has always
@@ -131,8 +136,10 @@ public AsyncDatastoreReaderWriter asyncDatastore(final ObjectifyImpl ofy) {
 	/**
 	 * One attempt at executing a transaction
 	 */
-	private <R> R transactOnce(final ObjectifyImpl parent, final Work<R> work) {
-		final ObjectifyImpl txnOfy = parent.factory().open(parent.getOptions(), new TransactorYes(parent.factory(), parent.getOptions().isCache(), this));
+	private <R> R transactOnce(final ObjectifyImpl parent, final Work<R> work, final AtomicReference<ByteString> prevTxnHandle) {
+		final ObjectifyImpl txnOfy = parent.factory().open(parent.getOptions(), new TransactorYes(parent.factory(), parent.getOptions().isCache(), this,
+			prevTxnHandle.get()));
+		prevTxnHandle.set(txnOfy.getTransaction().getTransactionHandle());
 
 		boolean committedSuccessfully = false;
 		try {
diff --git a/src/main/java/com/googlecode/objectify/impl/TransactorYes.java b/src/main/java/com/googlecode/objectify/impl/TransactorYes.java
@@ -1,5 +1,6 @@
 package com.googlecode.objectify.impl;
 
+import com.google.protobuf.ByteString;
 import com.googlecode.objectify.ObjectifyFactory;
 import com.googlecode.objectify.TxnType;
 import com.googlecode.objectify.Work;
@@ -20,10 +21,11 @@ class TransactorYes extends Transactor
 
 	/**
 	 */
-	TransactorYes(final ObjectifyFactory factory, final boolean cache, final TransactorNo parentTransactor) {
+	TransactorYes(final ObjectifyFactory factory, final boolean cache, final TransactorNo parentTransactor,
+	    ByteString prevTxnHandle) {
 		super(factory);
 
-		this.transaction = factory.asyncDatastore(cache).newTransaction(this::committed);
+		this.transaction = factory.asyncDatastore(cache).newTransaction(this::committed, prevTxnHandle);
 		this.parentTransactor = parentTransactor;
 	}
 
