feat: add support for Optimistic Concurrency Control (#2142)

* feat: add support for Optimistic Concurrency Control

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
Co-authored-by: Rajat Bhatta <[email protected]>

Author: 3 people

google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java

@@ -87,6 +87,17 @@ public static TransactionOption commitStats() {
     return COMMIT_STATS_OPTION;
   }
 
+  /**
+   * Specifying this instructs the transaction to request Optimistic Lock from the backend. In this
+   * concurrency mode, operations during the execution phase, i.e., reads and queries, are performed
+   * without acquiring locks, and transactional consistency is ensured by running a validation
+   * process in the commit phase (when any needed locks are acquired). The validation process
+   * succeeds only if there are no conflicting committed transactions (that committed mutations to
+   * the read data at a commit timestamp after the read timestamp).
+   */
+  public static TransactionOption optimisticLock() {
+    return OPTIMISTIC_LOCK_OPTION;
+  }
   /**
    * Specifying this will cause the read to yield at most this many rows. This should be greater
    * than 0.
@@ -207,6 +218,16 @@ void appendToOptions(Options options) {
 
   static final CommitStatsOption COMMIT_STATS_OPTION = new CommitStatsOption();
 
+  /** Option to request Optimistic Concurrency Control for read/write transactions. */
+  static final class OptimisticLockOption extends InternalOption implements TransactionOption {
+    @Override
+    void appendToOptions(Options options) {
+      options.withOptimisticLock = true;
+    }
+  }
+
+  static final OptimisticLockOption OPTIMISTIC_LOCK_OPTION = new OptimisticLockOption();
+
   /** Option pertaining to flow control. */
   static final class FlowControlOption extends InternalOption implements ReadAndQueryOption {
     final int prefetchChunks;
@@ -299,6 +320,7 @@ void appendToOptions(Options options) {
   private String tag;
   private String etag;
   private Boolean validateOnly;
+  private Boolean withOptimisticLock;
 
   // Construction is via factory methods below.
   private Options() {}
@@ -387,6 +409,10 @@ Boolean validateOnly() {
     return validateOnly;
   }
 
+  Boolean withOptimisticLock() {
+    return withOptimisticLock;
+  }
+
   @Override
   public String toString() {
     StringBuilder b = new StringBuilder();
@@ -420,6 +446,9 @@ public String toString() {
     if (validateOnly != null) {
       b.append("validateOnly: ").append(validateOnly).append(' ');
     }
+    if (withOptimisticLock != null) {
+      b.append("withOptimisticLock: ").append(withOptimisticLock).append(' ');
+    }
     return b.toString();
   }
 
@@ -453,7 +482,8 @@ public boolean equals(Object o) {
         && Objects.equals(priority(), that.priority())
         && Objects.equals(tag(), that.tag())
         && Objects.equals(etag(), that.etag())
-        && Objects.equals(validateOnly(), that.validateOnly());
+        && Objects.equals(validateOnly(), that.validateOnly())
+        && Objects.equals(withOptimisticLock(), that.withOptimisticLock());
   }
 
   @Override
@@ -492,6 +522,9 @@ public int hashCode() {
     if (validateOnly != null) {
       result = 31 * result + validateOnly.hashCode();
     }
+    if (withOptimisticLock != null) {
+      result = 31 * result + withOptimisticLock.hashCode();
+    }
     return result;
   }
 

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java

@@ -67,6 +67,14 @@ static void throwIfTransactionsPending() {
     }
   }
 
+  static TransactionOptions createReadWriteTransactionOptions(Options options) {
+    TransactionOptions.ReadWrite.Builder readWrite = TransactionOptions.ReadWrite.newBuilder();
+    if (options.withOptimisticLock() == Boolean.TRUE) {
+      readWrite.setReadLockMode(TransactionOptions.ReadWrite.ReadLockMode.OPTIMISTIC);
+    }
+    return TransactionOptions.newBuilder().setReadWrite(readWrite).build();
+  }
+
   /**
    * Represents a transaction within a session. "Transaction" here is used in the general sense,
    * which covers standalone reads, standalone writes, single-use and multi-use read-only
@@ -299,14 +307,16 @@ ByteString beginTransaction() {
   }
 
   ApiFuture<ByteString> beginTransactionAsync() {
+    return beginTransactionAsync(Options.fromTransactionOptions());
+  }
+
+  ApiFuture<ByteString> beginTransactionAsync(Options transactionOptions) {
     final SettableApiFuture<ByteString> res = SettableApiFuture.create();
     final Span span = tracer.spanBuilder(SpannerImpl.BEGIN_TRANSACTION).startSpan();
     final BeginTransactionRequest request =
         BeginTransactionRequest.newBuilder()
             .setSession(name)
-            .setOptions(
-                TransactionOptions.newBuilder()
-                    .setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance()))
+            .setOptions(createReadWriteTransactionOptions(transactionOptions))
             .build();
     final ApiFuture<Transaction> requestFuture =
         spanner.getRpc().beginTransactionAsync(request, options);

google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java

@@ -255,7 +255,7 @@ ApiFuture<Void> ensureTxnAsync() {
 
     private void createTxnAsync(final SettableApiFuture<Void> res) {
       span.addAnnotation("Creating Transaction");
-      final ApiFuture<ByteString> fut = session.beginTransactionAsync();
+      final ApiFuture<ByteString> fut = session.beginTransactionAsync(options);
       fut.addListener(
           () -> {
             try {
@@ -493,9 +493,7 @@ TransactionSelector getTransactionSelector() {
           }
           if (tx == null) {
             return TransactionSelector.newBuilder()
-                .setBegin(
-                    TransactionOptions.newBuilder()
-                        .setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance()))
+                .setBegin(SessionImpl.createReadWriteTransactionOptions(options))
                 .build();
           } else {
             // Wait for the transaction to come available. The tx.get() call will fail with an

google-cloud-spanner/src/test/java/com/google/cloud/spanner/OptionsTest.java

@@ -561,4 +561,24 @@ public void transactionEquality() {
     o3 = Options.fromTransactionOptions(Options.tag("app=spanner,env=stage"));
     assertThat(o2.equals(o3)).isFalse();
   }
+
+  @Test
+  public void optimisticLockEquality() {
+    Options option1 = Options.fromTransactionOptions(Options.optimisticLock());
+    Options option2 = Options.fromTransactionOptions(Options.optimisticLock());
+    Options option3 = Options.fromReadOptions();
+
+    assertEquals(option1, option2);
+    assertNotEquals(option1, option3);
+  }
+
+  @Test
+  public void optimisticLockHashCode() {
+    Options option1 = Options.fromTransactionOptions(Options.optimisticLock());
+    Options option2 = Options.fromTransactionOptions(Options.optimisticLock());
+    Options option3 = Options.fromReadOptions();
+
+    assertEquals(option1.hashCode(), option2.hashCode());
+    assertNotEquals(option1.hashCode(), option3.hashCode());
+  }
 }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/ReadWriteTransactionWithInlineBeginTest.java

@@ -23,25 +23,33 @@
 import com.google.api.gax.grpc.testing.LocalChannelProvider;
 import com.google.cloud.NoCredentials;
 import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.common.collect.Iterables;
 import com.google.protobuf.AbstractMessage;
 import com.google.protobuf.ListValue;
 import com.google.spanner.v1.BeginTransactionRequest;
+import com.google.spanner.v1.ExecuteSqlRequest;
+import com.google.spanner.v1.ReadRequest;
 import com.google.spanner.v1.ResultSetMetadata;
 import com.google.spanner.v1.StructType;
 import com.google.spanner.v1.StructType.Field;
+import com.google.spanner.v1.TransactionOptions;
+import com.google.spanner.v1.TransactionOptions.ReadWrite;
 import com.google.spanner.v1.TypeCode;
 import io.grpc.Server;
 import io.grpc.Status;
 import io.grpc.inprocess.InProcessServerBuilder;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
 import org.junit.After;
 import org.junit.AfterClass;
 import org.junit.Before;
@@ -85,6 +93,10 @@ public class ReadWriteTransactionWithInlineBeginTest {
                   .build())
           .setMetadata(SELECT1_METADATA)
           .build();
+  private static final TransactionOptions OPTIMISTIC_LOCK_OPTIONS =
+      TransactionOptions.newBuilder()
+          .setReadWrite(ReadWrite.newBuilder().setReadLockMode(ReadWrite.ReadLockMode.OPTIMISTIC))
+          .build();
   private Spanner spanner;
   private DatabaseClient client;
 
@@ -102,6 +114,9 @@ public static void startStaticServer() throws IOException {
         StatementResult.exception(
             INVALID_SELECT_STATEMENT,
             Status.INVALID_ARGUMENT.withDescription("invalid statement").asRuntimeException()));
+    mockSpanner.putStatementResult(
+        StatementResult.read(
+            "FOO", KeySet.all(), Collections.singletonList("ID"), SELECT1_RESULTSET));
 
     String uniqueName = InProcessServerBuilder.generateName();
     server =
@@ -389,6 +404,76 @@ public void failedBatchUpdateAndThenUpdate() {
     assertThat(countTransactionsStarted()).isEqualTo(2);
   }
 
+  @Test
+  public void executeSqlWithOptimisticConcurrencyControl() {
+    client
+        .readWriteTransaction(Options.optimisticLock())
+        .run(
+            transaction -> {
+              try (ResultSet rs = transaction.executeQuery(SELECT1)) {
+                while (rs.next()) {
+                  assertEquals(rs.getLong(0), 1);
+                }
+              }
+              return null;
+            });
+    Collection<AbstractMessage> requests =
+        mockSpanner.getRequests().stream()
+            .filter(msg -> msg.getClass().equals(ExecuteSqlRequest.class))
+            .collect(Collectors.toList());
+    assertEquals(requests.size(), 1);
+    ExecuteSqlRequest request = (ExecuteSqlRequest) Iterables.getOnlyElement(requests);
+    assertEquals(request.getTransaction().getBegin(), OPTIMISTIC_LOCK_OPTIONS);
+  }
+
+  @Test
+  public void readWithOptimisticConcurrencyControl() {
+    client
+        .readWriteTransaction(Options.optimisticLock())
+        .run(
+            transaction -> {
+              try (ResultSet rs =
+                  transaction.read("FOO", KeySet.all(), Collections.singletonList("ID"))) {
+                while (rs.next()) {
+                  assertEquals(rs.getLong(0), 1);
+                }
+              }
+              return null;
+            });
+    Collection<AbstractMessage> requests =
+        mockSpanner.getRequests().stream()
+            .filter(msg -> msg.getClass().equals(ReadRequest.class))
+            .collect(Collectors.toList());
+    assertEquals(requests.size(), 1);
+    ReadRequest request = (ReadRequest) Iterables.getOnlyElement(requests);
+    assertThat(request.getTransaction().getBegin()).isEqualTo(OPTIMISTIC_LOCK_OPTIONS);
+  }
+
+  @Test
+  public void beginTransactionWithOptimisticConcurrencyControl() {
+    client
+        .readWriteTransaction(Options.optimisticLock())
+        .run(
+            transaction -> {
+              // Instead of adding the BeginTransaction option to the next statement, the client
+              // library will force a complete retry of the entire transaction, and use an explicit
+              // BeginTransaction RPC invocation for that transaction in order to include the failed
+              // statement in the transaction as well.
+              try (ResultSet rs = transaction.executeQuery(INVALID_SELECT_STATEMENT)) {
+                SpannerException e = assertThrows(SpannerException.class, () -> rs.next());
+                assertEquals(ErrorCode.INVALID_ARGUMENT, e.getErrorCode());
+              }
+              return transaction.executeUpdate(UPDATE_STATEMENT);
+            });
+    Collection<AbstractMessage> requests =
+        mockSpanner.getRequests().stream()
+            .filter(msg -> msg.getClass().equals(BeginTransactionRequest.class))
+            .collect(Collectors.toList());
+    assertEquals(requests.size(), 1);
+    BeginTransactionRequest request = (BeginTransactionRequest) Iterables.getOnlyElement(requests);
+    assertEquals(request.getOptions(), OPTIMISTIC_LOCK_OPTIONS);
+  }
+
   @Test
   public void failedQueryAndThenUpdate() {
     Long updateCount =

google-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java

@@ -530,6 +530,21 @@ public void readWriteTransactionRead() throws InterruptedException {
                 }));
   }
 
+  @Test
+  public void readWriteTransactionReadWithOptimisticLock() throws InterruptedException {
+    TransactionRunner runner = client.readWriteTransaction(Options.optimisticLock());
+    assertThrowsSessionNotFoundIfShouldFail(
+        () ->
+            runner.run(
+                transaction -> {
+                  try (ResultSet rs =
+                      transaction.read("FOO", KeySet.all(), Collections.singletonList("BAR"))) {
+                    while (rs.next()) {}
+                  }
+                  return null;
+                }));
+  }
+
   @Test
   public void readWriteTransactionReadUsingIndex() throws InterruptedException {
     TransactionRunner runner = client.readWriteTransaction();

google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java

@@ -785,7 +785,7 @@ public void testSessionNotFoundReadWriteTransaction() {
           .thenReturn(ApiFutures.immediateFuture(Empty.getDefaultInstance()));
       when(closedSession.newTransaction(Options.fromTransactionOptions()))
           .thenReturn(closedTransactionContext);
-      when(closedSession.beginTransactionAsync()).thenThrow(sessionNotFound);
+      when(closedSession.beginTransactionAsync(any())).thenThrow(sessionNotFound);
       TransactionRunnerImpl closedTransactionRunner = new TransactionRunnerImpl(closedSession);
       closedTransactionRunner.setSpan(mock(Span.class));
       when(closedSession.readWriteTransaction()).thenReturn(closedTransactionRunner);
@@ -798,7 +798,7 @@ public void testSessionNotFoundReadWriteTransaction() {
       final TransactionContextImpl openTransactionContext = mock(TransactionContextImpl.class);
       when(openSession.newTransaction(Options.fromTransactionOptions()))
           .thenReturn(openTransactionContext);
-      when(openSession.beginTransactionAsync())
+      when(openSession.beginTransactionAsync(any()))
           .thenReturn(ApiFutures.immediateFuture(ByteString.copyFromUtf8("open-txn")));
       TransactionRunnerImpl openTransactionRunner = new TransactionRunnerImpl(openSession);
       openTransactionRunner.setSpan(mock(Span.class));