fix: reuse clientId for invalidated databases

Author: olavloite

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

@@ -1704,10 +1704,9 @@ private void handleCreateSessionsFailure(SpannerException e, int count) {
           break;
         }
       }
-      this.resourceNotFoundException =
-          MoreObjects.firstNonNull(
-              this.resourceNotFoundException,
-              isDatabaseOrInstanceNotFound(e) ? (ResourceNotFoundException) e : null);
+      if (isDatabaseOrInstanceNotFound(e)) {
+        setResourceNotFoundException((ResourceNotFoundException) e);
+      }
     }
   }
 
@@ -1734,9 +1733,7 @@ private void handlePrepareSessionFailure(
             decrementPendingClosures(1);
           }
           allSessions.remove(session);
-          this.resourceNotFoundException =
-              MoreObjects.firstNonNull(
-                  this.resourceNotFoundException, (ResourceNotFoundException) e);
+          setResourceNotFoundException((ResourceNotFoundException) e);
         } else {
           releaseSession(session, Position.FIRST);
         }
@@ -1749,6 +1746,10 @@ private void handlePrepareSessionFailure(
     }
   }
 
+  void setResourceNotFoundException(ResourceNotFoundException e) {
+    this.resourceNotFoundException = MoreObjects.firstNonNull(this.resourceNotFoundException, e);
+  }
+
   private void decrementPendingClosures(int count) {
     pendingClosure -= count;
     if (pendingClosure == 0) {

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

@@ -162,16 +162,20 @@ public InstanceAdminClient getInstanceAdminClient() {
   public DatabaseClient getDatabaseClient(DatabaseId db) {
     synchronized (this) {
       Preconditions.checkState(!spannerIsClosed, "Cloud Spanner client has been closed");
+      String clientId = null;
       if (dbClients.containsKey(db) && !dbClients.get(db).pool.isValid()) {
         // Move the invalidated client to a separate list, so we can close it together with the
         // other database clients when the Spanner instance is closed.
         invalidatedDbClients.add(dbClients.get(db));
+        clientId = dbClients.get(db).clientId;
         dbClients.remove(db);
       }
       if (dbClients.containsKey(db)) {
         return dbClients.get(db);
       } else {
-        String clientId = nextDatabaseClientId(db);
+        if (clientId == null) {
+          clientId = nextDatabaseClientId(db);
+        }
         List<LabelValue> labelValues =
             ImmutableList.of(
                 LabelValue.create(clientId),

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

@@ -27,6 +27,7 @@
 import com.google.cloud.NoCredentials;
 import com.google.cloud.ServiceRpc;
 import com.google.cloud.grpc.GrpcTransportOptions;
+import com.google.cloud.spanner.SpannerException.DoNotConstructDirectly;
 import com.google.cloud.spanner.spi.v1.SpannerRpc;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
 import java.util.Collections;
@@ -187,7 +188,7 @@ public void testSpannerClosed() throws InterruptedException {
   }
 
   @Test
-  public void testClientId() {
+  public void testClientId() throws Exception {
     // Create a unique database id to be sure it has not yet been used in the lifetime of this JVM.
     String dbName =
         String.format("projects/p1/instances/i1/databases/%s", UUID.randomUUID().toString());
@@ -212,6 +213,13 @@ public void testClientId() {
     DatabaseClientImpl databaseClient2 = (DatabaseClientImpl) impl.getDatabaseClient(db2);
     assertThat(databaseClient2.clientId).isEqualTo("client-1");
 
+    // Getting a new database client for an invalidated database should use the same client id.
+    databaseClient.pool.setResourceNotFoundException(
+        new DatabaseNotFoundException(DoNotConstructDirectly.ALLOWED, "not found", null, null));
+    DatabaseClientImpl revalidated = (DatabaseClientImpl) impl.getDatabaseClient(db);
+    assertThat(revalidated).isNotSameInstanceAs(databaseClient);
+    assertThat(revalidated.clientId).isEqualTo(databaseClient.clientId);
+
     // Create a new Spanner instance. This will generate new database clients with new ids.
     try (Spanner spanner =
         SpannerOptions.newBuilder()
@@ -222,8 +230,7 @@ public void testClientId() {
 
       // Get a database client for the same database as the first database. As this goes through a
       // different Spanner instance with potentially different options, it will get a different
-      // client
-      // id.
+      // client id.
       DatabaseClientImpl databaseClient3 = (DatabaseClientImpl) spanner.getDatabaseClient(db);
       assertThat(databaseClient3.clientId).isEqualTo("client-2");
     }