From 0c495e7e72260d8eb746393d87bb32853154a575 Mon Sep 17 00:00:00 2001
From: Ehren Kret <ehren@signal.org>
Date: Wed, 29 Apr 2020 17:10:13 -0700
Subject: [PATCH] Workaround lack of internal retry on transaction rollback

The get endpoint for key fetching can fail if the transaction cannot
complete because of simultaneous modification. Clients currently
receive 500 from this and retry if it happens, but this test case runs
into it without retrying and then complains that not all the threads
completed successfully. This workaround adds some retry attempts.
---
 .../textsecuregcm/tests/storage/KeysTest.java  | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/KeysTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/KeysTest.java
index cd9076966..614445ab2 100644
--- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/KeysTest.java
+++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/KeysTest.java
@@ -3,16 +3,16 @@ package org.whispersystems.textsecuregcm.tests.storage;
 import com.opentable.db.postgres.embedded.LiquibasePreparer;
 import com.opentable.db.postgres.junit.EmbeddedPostgresRules;
 import com.opentable.db.postgres.junit.PreparedDbRule;
-import org.jdbi.v3.core.HandleCallback;
+import io.github.resilience4j.circuitbreaker.CircuitBreakerOpenException;
 import org.jdbi.v3.core.HandleConsumer;
 import org.jdbi.v3.core.Jdbi;
+import org.jdbi.v3.core.statement.UnableToExecuteStatementException;
 import org.jdbi.v3.core.transaction.SerializableTransactionRunner;
 import org.jdbi.v3.core.transaction.TransactionException;
 import org.jdbi.v3.core.transaction.TransactionIsolationLevel;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
-import org.postgresql.util.PSQLException;
 import org.whispersystems.textsecuregcm.configuration.CircuitBreakerConfiguration;
 import org.whispersystems.textsecuregcm.entities.PreKey;
 import org.whispersystems.textsecuregcm.storage.FaultTolerantDatabase;
@@ -25,7 +25,6 @@ import java.sql.SQLException;
 import java.util.LinkedList;
 import java.util.List;
 
-import io.github.resilience4j.circuitbreaker.CircuitBreakerOpenException;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doThrow;
@@ -246,7 +245,18 @@ public class KeysTest {
 
     for (int i=0;i<20;i++) {
       Thread thread = new Thread(() -> {
-        List<KeyRecord> results = keys.get("+14152222222");
+        List<KeyRecord> results = null;
+        final int MAX_RETRIES = 5;
+        for (int retryAttempt = 0; results == null && retryAttempt < MAX_RETRIES; ++retryAttempt) {
+          try {
+            results = keys.get("+14152222222");
+          } catch (UnableToExecuteStatementException e) {
+            if (retryAttempt == MAX_RETRIES - 1) {
+              throw e;
+            }
+          }
+        }
+        assertThat(results).isNotNull();
         assertThat(results.size()).isEqualTo(2);
       });
       thread.start();