fix: reuse clientId for invalidated databases (#206)

Author: olavloite

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

@@ -1708,10 +1708,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);
+      }
     }
   }
 
@@ -1738,9 +1737,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);
         }
@@ -1753,6 +1750,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

@@ -185,16 +185,20 @@ public InstanceAdminClient getInstanceAdminClient() {
   public DatabaseClient getDatabaseClient(DatabaseId db) {
     synchronized (this) {
       checkClosed();
+      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.SpannerImpl.ClosedException;
 import com.google.cloud.spanner.spi.v1.SpannerRpc;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
@@ -190,7 +191,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());
@@ -215,6 +216,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()