fix(spanner): pin inline read-write transactions to default endpoint

Author: rahul2393

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java

@@ -785,6 +785,7 @@ private void maybeTrackReadWriteBegin(TransactionSelector selector) {
       if (selector.getSelectorCase() == TransactionSelector.SelectorCase.BEGIN
           && !selector.getBegin().hasReadOnly()) {
         shouldRecordTransactionAffinity = true;
+        allowDefaultAffinity = true;
       }
     }
 

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java

@@ -930,6 +930,98 @@ public void readOnlyTransactionRoutesEachReadIndependently() throws Exception {
     assertThat(harness.defaultManagedChannel.callCount()).isEqualTo(2);
   }
 
+  @Test
+  public void readWriteInlinedBeginExecuteSqlOnDefaultRecordsDefaultAffinity() throws Exception {
+    TestHarness harness = createHarness();
+    ByteString transactionId = ByteString.copyFromUtf8("rw-inline-sql-default");
+
+    ClientCall<ExecuteSqlRequest, ResultSet> beginCall =
+        harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT);
+    beginCall.start(new CapturingListener<ResultSet>(), new Metadata());
+    beginCall.sendMessage(
+        ExecuteSqlRequest.newBuilder()
+            .setSession(SESSION)
+            .setTransaction(
+                TransactionSelector.newBuilder()
+                    .setBegin(
+                        TransactionOptions.newBuilder()
+                            .setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance())
+                            .build()))
+            .build());
+
+    @SuppressWarnings("unchecked")
+    RecordingClientCall<ExecuteSqlRequest, ResultSet> beginDelegate =
+        (RecordingClientCall<ExecuteSqlRequest, ResultSet>)
+            harness.defaultManagedChannel.latestCall();
+    beginDelegate.emitOnMessage(
+        ResultSet.newBuilder()
+            .setMetadata(
+                ResultSetMetadata.newBuilder()
+                    .setTransaction(Transaction.newBuilder().setId(transactionId)))
+            .build());
+
+    seedCache(harness, createTwoRangeCacheUpdate());
+
+    ClientCall<ExecuteSqlRequest, ResultSet> nextCall =
+        harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT);
+    nextCall.start(new CapturingListener<ResultSet>(), new Metadata());
+    nextCall.sendMessage(
+        ExecuteSqlRequest.newBuilder()
+            .setSession(SESSION)
+            .setTransaction(TransactionSelector.newBuilder().setId(transactionId))
+            .setRoutingHint(RoutingHint.newBuilder().setKey(bytes("b")).build())
+            .build());
+
+    assertThat(harness.defaultManagedChannel.callCount()).isEqualTo(3);
+    assertThat(harness.endpointCache.callCountForAddress("server-a:1234")).isEqualTo(0);
+  }
+
+  @Test
+  public void readWriteInlinedBeginReadOnDefaultRecordsDefaultAffinity() throws Exception {
+    TestHarness harness = createHarness();
+    ByteString transactionId = ByteString.copyFromUtf8("rw-inline-read-default");
+
+    ClientCall<ReadRequest, PartialResultSet> beginCall =
+        harness.channel.newCall(SpannerGrpc.getStreamingReadMethod(), CallOptions.DEFAULT);
+    beginCall.start(new CapturingListener<PartialResultSet>(), new Metadata());
+    beginCall.sendMessage(
+        ReadRequest.newBuilder()
+            .setSession(SESSION)
+            .setTransaction(
+                TransactionSelector.newBuilder()
+                    .setBegin(
+                        TransactionOptions.newBuilder()
+                            .setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance())
+                            .build()))
+            .build());
+
+    @SuppressWarnings("unchecked")
+    RecordingClientCall<ReadRequest, PartialResultSet> beginDelegate =
+        (RecordingClientCall<ReadRequest, PartialResultSet>)
+            harness.defaultManagedChannel.latestCall();
+    beginDelegate.emitOnMessage(
+        PartialResultSet.newBuilder()
+            .setMetadata(
+                ResultSetMetadata.newBuilder()
+                    .setTransaction(Transaction.newBuilder().setId(transactionId)))
+            .build());
+
+    seedCache(harness, createTwoRangeCacheUpdate());
+
+    ClientCall<ReadRequest, PartialResultSet> nextCall =
+        harness.channel.newCall(SpannerGrpc.getStreamingReadMethod(), CallOptions.DEFAULT);
+    nextCall.start(new CapturingListener<PartialResultSet>(), new Metadata());
+    nextCall.sendMessage(
+        ReadRequest.newBuilder()
+            .setSession(SESSION)
+            .setTransaction(TransactionSelector.newBuilder().setId(transactionId))
+            .setRoutingHint(RoutingHint.newBuilder().setKey(bytes("b")).build())
+            .build());
+
+    assertThat(harness.defaultManagedChannel.callCount()).isEqualTo(3);
+    assertThat(harness.endpointCache.callCountForAddress("server-a:1234")).isEqualTo(0);
+  }
+
   @Test
   public void readOnlyInlinedBeginExecuteSqlRoutesSubsequentRequestsIndependently()
       throws Exception {