fix: retry aborted errors for writeAtLeastOnce (#2627)

* fix: retry aborted errors for writeAtLeastOnce The `writeAtLeastOnce` method could fail with an Aborted error if Cloud Spanner would abort the transaction during the single Commit RPC invocation that this method executes. This can for example happen if a schema change is executed while this method is being called. Fixes #2626 * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
googleapis · Sep 25, 2023 · 2addb19 · 2addb19
1 parent 9964bd5
commit 2addb19
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 3 deletions.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java
@@ -182,11 +182,11 @@ public CommitResponse writeAtLeastOnceWithOptions(
       }
       requestBuilder.setRequestOptions(requestOptionsBuilder.build());
     }
+    CommitRequest request = requestBuilder.build();
     Span span = tracer.spanBuilder(SpannerImpl.COMMIT).startSpan();
     try (Scope s = tracer.withSpan(span)) {
-      com.google.spanner.v1.CommitResponse response =
-          spanner.getRpc().commit(requestBuilder.build(), this.options);
-      return new CommitResponse(response);
+      return SpannerRetryHelper.runTxWithRetriesOnAborted(
+          () -> new CommitResponse(spanner.getRpc().commit(request, this.options)));
     } catch (RuntimeException e) {
       TraceUtil.setWithFailure(span, e);
       throw e;

diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java
@@ -525,6 +525,44 @@ public void testWrite() {
     assertEquals(Priority.PRIORITY_UNSPECIFIED, commit.getRequestOptions().getPriority());
   }
 
+  @Test
+  public void testWriteAborted() {
+    DatabaseClient client =
+        spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+    // Force the Commit RPC to return Aborted the first time it is called. The exception is cleared
+    // after the first call, so the retry should succeed.
+    mockSpanner.setCommitExecutionTime(
+        SimulatedExecutionTime.ofException(
+            mockSpanner.createAbortedException(ByteString.copyFromUtf8("test"))));
+    Timestamp timestamp =
+        client.write(
+            Collections.singletonList(
+                Mutation.newInsertBuilder("FOO").set("ID").to(1L).set("NAME").to("Bar").build()));
+    assertNotNull(timestamp);
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(2, commitRequests.size());
+  }
+
+  @Test
+  public void testWriteAtLeastOnceAborted() {
+    DatabaseClient client =
+        spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
+    // Force the Commit RPC to return Aborted the first time it is called. The exception is cleared
+    // after the first call, so the retry should succeed.
+    mockSpanner.setCommitExecutionTime(
+        SimulatedExecutionTime.ofException(
+            mockSpanner.createAbortedException(ByteString.copyFromUtf8("test"))));
+    Timestamp timestamp =
+        client.writeAtLeastOnce(
+            Collections.singletonList(
+                Mutation.newInsertBuilder("FOO").set("ID").to(1L).set("NAME").to("Bar").build()));
+    assertNotNull(timestamp);
+
+    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
+    assertEquals(2, commitRequests.size());
+  }
+
   @Test
   public void testWriteWithOptions() {
     DatabaseClient client =