test: unflake StatementTimeoutTest

Author: olavloite

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

@@ -2302,6 +2302,7 @@ public void waitForRequestsToContain(Class<? extends AbstractMessage> type, long
       throws InterruptedException, TimeoutException {
     Stopwatch watch = Stopwatch.createStarted();
     while (countRequestsOfType(type) == 0) {
+      //noinspection BusyWait
       Thread.sleep(1L);
       if (watch.elapsed(TimeUnit.MILLISECONDS) > timeoutMillis) {
         throw new TimeoutException(

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementTimeoutTest.java

@@ -36,7 +36,6 @@
 import com.google.cloud.spanner.connection.ITAbstractSpannerTest.ITConnection;
 import com.google.cloud.spanner.connection.StatementExecutor.StatementExecutorType;
 import com.google.common.base.Stopwatch;
-import com.google.common.collect.Collections2;
 import com.google.longrunning.Operation;
 import com.google.protobuf.AbstractMessage;
 import com.google.protobuf.Any;
@@ -47,6 +46,9 @@
 import com.google.spanner.v1.ExecuteSqlRequest;
 import io.grpc.Status;
 import java.time.Duration;
+import java.util.ArrayList;
+import java.util.ConcurrentModificationException;
+import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
@@ -617,20 +619,20 @@ static void waitForRequestsToContain(Class<? extends AbstractMessage> request) {
   private void waitForDdlRequestOnServer() {
     try {
       Stopwatch watch = Stopwatch.createStarted();
-      while (Collections2.filter(
-              mockDatabaseAdmin.getRequests(),
-              input -> input.getClass().equals(UpdateDatabaseDdlRequest.class))
-          .isEmpty()) {
+      while (watch.elapsed(TimeUnit.MILLISECONDS) < EXECUTION_TIME_SLOW_STATEMENT) {
+        try {
+          List<AbstractMessage> requests = new ArrayList<>(mockDatabaseAdmin.getRequests());
+          if (requests.stream().anyMatch(request -> request instanceof UpdateDatabaseDdlRequest)) {
+            break;
+          }
+        } catch (ConcurrentModificationException ignore) {
+          // Just ignore and retry.
+        }
         //noinspection BusyWait
         Thread.sleep(1L);
-        if (watch.elapsed(TimeUnit.MILLISECONDS) > EXECUTION_TIME_SLOW_STATEMENT) {
-          throw new TimeoutException("Timeout while waiting for DDL request");
-        }
       }
     } catch (InterruptedException e) {
       throw SpannerExceptionFactory.propagateInterrupt(e);
-    } catch (TimeoutException e) {
-      throw SpannerExceptionFactory.propagateTimeout(e);
     }
   }
 
@@ -1010,6 +1012,7 @@ public void testCancelDdlBatch() {
       } finally {
         executor.shutdownNow();
       }
+      connection.closeAsync();
     }
   }
 
@@ -1036,6 +1039,7 @@ public void testCancelDdlAutocommit() {
       } finally {
         executor.shutdownNow();
       }
+      connection.closeAsync();
     }
   }
 
@@ -1049,6 +1053,8 @@ public void testTimeoutExceptionDdlAutocommit() {
       SpannerException e =
           assertThrows(SpannerException.class, () -> connection.execute(Statement.of(SLOW_DDL)));
       assertEquals(ErrorCode.DEADLINE_EXCEEDED, e.getErrorCode());
+
+      connection.closeAsync();
     }
   }