diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java index eb497dc284..0db55bcccb 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionPool.java +++ b/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) { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java index 66551d9abc..4e937459cf 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java +++ b/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 labelValues = ImmutableList.of( LabelValue.create(clientId), diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java index 30a5a8e68c..25f3381bd1 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerImplTest.java +++ b/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()