diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index ece6d862f87b..39c8fe265653 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -89,6 +89,8 @@ static class MultiplexedSessionTransaction extends SessionImpl { private boolean done; + private boolean channelHintReleased; + MultiplexedSessionTransaction( MultiplexedSessionDatabaseClient client, ISpan span, @@ -123,19 +125,18 @@ void onReadDone() { if (this.singleUse && getActiveTransaction() != null) { getActiveTransaction().close(); setActive(null); - if (this.singleUseChannelHint != NO_CHANNEL_HINT) { - this.client.channelUsage.clear(this.singleUseChannelHint); - } - this.client.numCurrentSingleUseTransactions.decrementAndGet(); + releaseSingleUseChannelHint(); } } @Override public CommitResponse writeAtLeastOnceWithOptions( Iterable mutations, TransactionOption... options) throws SpannerException { - CommitResponse response = super.writeAtLeastOnceWithOptions(mutations, options); - onTransactionDone(); - return response; + try { + return super.writeAtLeastOnceWithOptions(mutations, options); + } finally { + onTransactionDone(); + } } @Override @@ -150,6 +151,30 @@ void onTransactionDone() { if (markedDone) { client.numSessionsReleased.incrementAndGet(); } + if (this.singleUse) { + // Mutation-only operations (writeAtLeastOnceWithOptions, batchWriteAtLeastOnce) never + // return a ResultSet and thus never call onReadDone(), so the channel hint must also be + // released here. + releaseSingleUseChannelHint(); + } + } + + /** + * Releases the channel hint that was reserved for this single-use transaction, so a subsequent + * single-use transaction can use the channel. This method is idempotent; the hint is released + * by whichever of {@link #onReadDone()} and {@link #onTransactionDone()} is called first. + */ + private void releaseSingleUseChannelHint() { + boolean release = false; + synchronized (this) { + if (!this.channelHintReleased) { + this.channelHintReleased = true; + release = true; + } + } + if (release) { + this.client.releaseSingleUseChannelHint(this.singleUseChannelHint); + } } @Override @@ -456,7 +481,11 @@ private DelayedMultiplexedSessionTransaction createDelayedMultiplexSessionTransa } private int getSingleUseChannelHint() { + // The counter tracks the number of single-use transactions that currently hold a channel hint, + // so it is only incremented when a hint is actually reserved, and any transient increment is + // rolled back before returning NO_CHANNEL_HINT. if (this.numCurrentSingleUseTransactions.incrementAndGet() > this.numChannels) { + this.numCurrentSingleUseTransactions.decrementAndGet(); return NO_CHANNEL_HINT; } synchronized (this.channelUsage) { @@ -466,6 +495,7 @@ private int getSingleUseChannelHint() { // This then means that all channels have already been assigned to single-use transactions, // and that we should not use a specific channel, but rather pick a random one. if (channel == this.numChannels) { + this.numCurrentSingleUseTransactions.decrementAndGet(); return NO_CHANNEL_HINT; } this.channelUsage.set(channel); @@ -473,6 +503,18 @@ private int getSingleUseChannelHint() { } } + private void releaseSingleUseChannelHint(int channelHint) { + if (channelHint == NO_CHANNEL_HINT) { + return; + } + // Take the same monitor as getSingleUseChannelHint(); the BitSet is not thread-safe and is + // shared by all clients of the same SpannerImpl. + synchronized (this.channelUsage) { + this.channelUsage.clear(channelHint); + } + this.numCurrentSingleUseTransactions.decrementAndGet(); + } + private final AbstractLazyInitializer dialectSupplier = new AbstractLazyInitializer() { @Override diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java index c4cd8ba611f0..2d98f8cb2230 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClientMockServerTest.java @@ -535,6 +535,126 @@ public void testWriteAtLeastOnce() { assertEquals(1L, client.multiplexedSessionDatabaseClient.getNumSessionsReleased().get()); } + @Test + public void testWriteAtLeastOnceReleasesSingleUseChannelHint() throws Exception { + // Channel hints are only used when the gRPC-GCP extension is disabled. + try (Spanner testSpanner = createSpannerWithoutGrpcGcp()) { + DatabaseClientImpl client = + (DatabaseClientImpl) testSpanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + int numChannels = testSpanner.getOptions().getNumChannels(); + + // Execute more mutation-only transactions than there are channels. Each transaction should + // release its channel hint when it is done, so the channel-spreading optimization keeps + // working for subsequent single-use transactions. + for (int i = 0; i < numChannels + 1; i++) { + assertNotNull(MockSpannerTestActions.writeAtLeastOnceInsertMutation(client)); + } + assertEquals(0, getChannelUsage(client).cardinality()); + assertEquals(0, getNumCurrentSingleUseTransactions(client).get()); + + // A subsequent single-use read-only transaction should still get a channel hint. The hint is + // held from the creation of the transaction until the result set has been consumed or closed. + try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) { + assertEquals(1, getChannelUsage(client).cardinality()); + assertEquals(1, getNumCurrentSingleUseTransactions(client).get()); + //noinspection StatementWithEmptyBody + while (resultSet.next()) {} + } + assertEquals(0, getChannelUsage(client).cardinality()); + assertEquals(0, getNumCurrentSingleUseTransactions(client).get()); + } + } + + @Test + public void testFailedWriteAtLeastOnceReleasesSingleUseChannelHint() throws Exception { + try (Spanner testSpanner = createSpannerWithoutGrpcGcp()) { + DatabaseClientImpl client = + (DatabaseClientImpl) testSpanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + // Make sure that the multiplexed session has been created, so that the write transaction + // below does not use the delayed path (which never reserves a channel hint). + awaitMultiplexedSessionCreation(client); + mockSpanner.setCommitExecutionTime( + SimulatedExecutionTime.ofException( + Status.FAILED_PRECONDITION.withDescription("test").asRuntimeException())); + + SpannerException exception = + assertThrows( + SpannerException.class, + () -> MockSpannerTestActions.writeAtLeastOnceInsertMutation(client)); + + assertEquals(ErrorCode.FAILED_PRECONDITION, exception.getErrorCode()); + assertEquals(0, getChannelUsage(client).cardinality()); + assertEquals(0, getNumCurrentSingleUseTransactions(client).get()); + } + } + + @Test + public void testBatchWriteAtLeastOnceReleasesSingleUseChannelHint() throws Exception { + try (Spanner testSpanner = createSpannerWithoutGrpcGcp()) { + DatabaseClientImpl client = + (DatabaseClientImpl) testSpanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + // Make sure that the multiplexed session has been created, so that the batch write below + // does not use the delayed path (which never reserves a channel hint). + awaitMultiplexedSessionCreation(client); + + ServerStream responseStream = + client.batchWriteAtLeastOnce( + ImmutableList.of( + MutationGroup.of( + Mutation.newInsertBuilder("FOO") + .set("ID") + .to(1L) + .set("NAME") + .to("Bar") + .build()))); + for (BatchWriteResponse response : responseStream) { + assertNotNull(response); + } + + assertEquals(0, getChannelUsage(client).cardinality()); + assertEquals(0, getNumCurrentSingleUseTransactions(client).get()); + } + } + + private void awaitMultiplexedSessionCreation(DatabaseClientImpl client) { + try (ResultSet resultSet = client.singleUse().executeQuery(STATEMENT)) { + //noinspection StatementWithEmptyBody + while (resultSet.next()) {} + } + } + + private Spanner createSpannerWithoutGrpcGcp() { + return SpannerOptions.newBuilder() + .setProjectId("test-project") + .setChannelProvider(channelProvider) + .setCredentials(NoCredentials.getInstance()) + .disableGrpcGcpExtension() + .setSessionPoolOption( + SessionPoolOptions.newBuilder() + .setUseMultiplexedSession(true) + .setUseMultiplexedSessionForRW(true) + .setUseMultiplexedSessionPartitionedOps(true) + .setFailOnSessionLeak() + .build()) + .build() + .getService(); + } + + private static BitSet getChannelUsage(DatabaseClientImpl client) throws Exception { + java.lang.reflect.Field field = + MultiplexedSessionDatabaseClient.class.getDeclaredField("channelUsage"); + field.setAccessible(true); + return (BitSet) field.get(client.multiplexedSessionDatabaseClient); + } + + private static AtomicInteger getNumCurrentSingleUseTransactions(DatabaseClientImpl client) + throws Exception { + java.lang.reflect.Field field = + MultiplexedSessionDatabaseClient.class.getDeclaredField("numCurrentSingleUseTransactions"); + field.setAccessible(true); + return (AtomicInteger) field.get(client.multiplexedSessionDatabaseClient); + } + @Test public void testWriteAtLeastOnceWithCommitStats() { DatabaseClientImpl client =