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");
  }