gemini-code-assist[bot] commented on code in PR #39255:
URL: https://github.com/apache/beam/pull/39255#discussion_r3548191042
##########
sdks/java/harness/src/main/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClient.java:
##########
@@ -124,27 +143,88 @@ public StreamObserver<Elements> getOutboundObserver(
return getMultiplexer(apiServiceDescriptor,
dataStreamId).getOutboundObserver();
}
+ @Override
+ public void retainDataStream(String dataStreamId) {
+ if (dataStreamId == null || dataStreamId.isEmpty()) {
+ // The default data stream is kept open for the lifetime of the client.
+ return;
+ }
+ synchronized (dataStreamLifecycleLock) {
+ dataStreamRefCounts.merge(dataStreamId, 1, Integer::sum);
+ }
+ }
+
+ @Override
+ public void releaseDataStream(String dataStreamId) {
+ if (dataStreamId == null || dataStreamId.isEmpty()) {
+ // The default data stream is kept open for the lifetime of the client.
+ return;
+ }
+ List<BeamFnDataGrpcMultiplexer> multiplexersToClose = new ArrayList<>();
+ synchronized (dataStreamLifecycleLock) {
+ Integer refCount = dataStreamRefCounts.get(dataStreamId);
+ if (refCount == null) {
+ LOG.warn("Released data stream {} which was not retained.",
dataStreamId);
+ return;
+ }
+ if (refCount > 1) {
+ dataStreamRefCounts.put(dataStreamId, refCount - 1);
+ return;
+ }
+ dataStreamRefCounts.remove(dataStreamId);
+ Iterator<Map.Entry<MultiplexerKey, BeamFnDataGrpcMultiplexer>> iterator =
+ multiplexerCache.entrySet().iterator();
+ while (iterator.hasNext()) {
+ Map.Entry<MultiplexerKey, BeamFnDataGrpcMultiplexer> entry =
iterator.next();
+ if (dataStreamId.equals(entry.getKey().dataStreamId)) {
+ multiplexersToClose.add(entry.getValue());
+ iterator.remove();
+ }
+ }
+ }
+ // Close outside of the lock as closing terminates the underlying gRPC
stream and may block.
+ for (BeamFnDataGrpcMultiplexer multiplexer : multiplexersToClose) {
+ LOG.debug("Closing multiplexer for released data stream {}",
dataStreamId);
+ try {
+ multiplexer.close();
+ } catch (Exception e) {
+ LOG.warn("Failed to close multiplexer for data stream {}",
dataStreamId, e);
+ }
+ }
+ }
+
private BeamFnDataGrpcMultiplexer getMultiplexer(
Endpoints.ApiServiceDescriptor apiServiceDescriptor, String
dataStreamId) {
MultiplexerKey key = new MultiplexerKey(apiServiceDescriptor,
dataStreamId);
- return multiplexerCache.computeIfAbsent(
- key,
- k -> {
- OutboundObserverFactory.BasicFactory<Elements, Elements>
baseOutboundObserverFactory =
- inboundObserver -> {
- BeamFnDataGrpc.BeamFnDataStub stub =
-
BeamFnDataGrpc.newStub(channelFactory.apply(apiServiceDescriptor));
- if (dataStreamId != null && !dataStreamId.isEmpty()) {
- Metadata headers = new Metadata();
- headers.put(
- Metadata.Key.of("data_stream_id",
Metadata.ASCII_STRING_MARSHALLER),
- dataStreamId);
- stub =
stub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(headers));
- }
- return stub.data(inboundObserver);
- };
- return new BeamFnDataGrpcMultiplexer(
- apiServiceDescriptor, outboundObserverFactory,
baseOutboundObserverFactory);
- });
+ BeamFnDataGrpcMultiplexer existingMultiplexer = multiplexerCache.get(key);
+ if (existingMultiplexer != null) {
+ return existingMultiplexer;
+ }
+ // Create under the lifecycle lock so that a named data stream being
concurrently closed by
+ // releaseDataStream is not observed in a partially removed state. Callers
are expected to
+ // retain named data streams for the duration of their usage which
prevents the returned
+ // multiplexer from being closed while in use.
+ synchronized (dataStreamLifecycleLock) {
+ return multiplexerCache.computeIfAbsent(
+ key,
+ k -> {
+ OutboundObserverFactory.BasicFactory<Elements, Elements>
baseOutboundObserverFactory =
+ inboundObserver -> {
+ BeamFnDataGrpc.BeamFnDataStub stub =
+
BeamFnDataGrpc.newStub(channelFactory.apply(apiServiceDescriptor));
+ if (dataStreamId != null && !dataStreamId.isEmpty()) {
+ Metadata headers = new Metadata();
+ headers.put(
+ Metadata.Key.of("data_stream_id",
Metadata.ASCII_STRING_MARSHALLER),
+ dataStreamId);
+ stub =
+
stub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(headers));
+ }
+ return stub.data(inboundObserver);
+ };
+ return new BeamFnDataGrpcMultiplexer(
+ apiServiceDescriptor, outboundObserverFactory,
baseOutboundObserverFactory);
+ });
+ }
Review Comment:
Since we are already inside the `synchronized (dataStreamLifecycleLock)`
block, using `computeIfAbsent` on `multiplexerCache` is redundant. We can
simplify this by performing a simple `get` and `put` check inside the
synchronized block, which also avoids allocating the lambda/closure on every
cache miss.
```java
synchronized (dataStreamLifecycleLock) {
BeamFnDataGrpcMultiplexer multiplexer = multiplexerCache.get(key);
if (multiplexer == null) {
OutboundObserverFactory.BasicFactory<Elements, Elements>
baseOutboundObserverFactory =
inboundObserver -> {
BeamFnDataGrpc.BeamFnDataStub stub =
BeamFnDataGrpc.newStub(channelFactory.apply(apiServiceDescriptor));
if (dataStreamId != null && !dataStreamId.isEmpty()) {
Metadata headers = new Metadata();
headers.put(
Metadata.Key.of("data_stream_id",
Metadata.ASCII_STRING_MARSHALLER),
dataStreamId);
stub =
stub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(headers));
}
return stub.data(inboundObserver);
};
multiplexer = new BeamFnDataGrpcMultiplexer(
apiServiceDescriptor, outboundObserverFactory,
baseOutboundObserverFactory);
multiplexerCache.put(key, multiplexer);
}
return multiplexer;
}
```
##########
sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClientTest.java:
##########
@@ -364,6 +366,116 @@ public StreamObserver<BeamFnApi.Elements> data(
}
}
+ @Test
+ public void testNamedDataStreamClosedWhenNoLongerRetained() throws Exception
{
+ String dataStreamId = "streamA";
+ AtomicInteger connectionCount = new AtomicInteger();
+ CountDownLatch firstConnection = new CountDownLatch(1);
+ CountDownLatch secondConnection = new CountDownLatch(2);
+ CountDownLatch streamTerminated = new CountDownLatch(1);
+
+ Endpoints.ApiServiceDescriptor apiServiceDescriptor =
+ Endpoints.ApiServiceDescriptor.newBuilder()
+ .setUrl(this.getClass().getName() + "-" + UUID.randomUUID())
+ .build();
+ Server server =
+ InProcessServerBuilder.forName(apiServiceDescriptor.getUrl())
+ .addService(
+ new BeamFnDataGrpc.BeamFnDataImplBase() {
+ @Override
+ public StreamObserver<BeamFnApi.Elements> data(
+ StreamObserver<BeamFnApi.Elements> outboundObserver) {
+ connectionCount.incrementAndGet();
+ firstConnection.countDown();
+ secondConnection.countDown();
+ return TestStreams.<BeamFnApi.Elements>withOnNext(elements
-> {})
+ .withOnError(streamTerminated::countDown)
+ .withOnCompleted(streamTerminated::countDown)
+ .build();
+ }
+ })
+ .build();
+ server.start();
+ try {
+ ManagedChannel channel =
+
InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
+
+ BeamFnDataGrpcClient clientFactory =
+ new BeamFnDataGrpcClient(
+ (Endpoints.ApiServiceDescriptor descriptor) -> channel,
+ OutboundObserverFactory.trivial());
+
+ // Two bundles concurrently use the same named data stream.
+ clientFactory.retainDataStream(dataStreamId);
+ clientFactory.retainDataStream(dataStreamId);
+ clientFactory.getOutboundObserver(apiServiceDescriptor, dataStreamId);
+ assertTrue(firstConnection.await(5, TimeUnit.SECONDS));
+ assertEquals(1, connectionCount.get());
+
+ // Releasing one of the two usages should not close the stream.
+ clientFactory.releaseDataStream(dataStreamId);
+ assertFalse(streamTerminated.await(100, TimeUnit.MILLISECONDS));
+
+ // Releasing the last usage should close the stream.
+ clientFactory.releaseDataStream(dataStreamId);
+ assertTrue(streamTerminated.await(5, TimeUnit.SECONDS));
+
+ // A subsequent usage of the same named data stream establishes a new
stream.
+ clientFactory.retainDataStream(dataStreamId);
+ clientFactory.getOutboundObserver(apiServiceDescriptor, dataStreamId);
+ assertTrue(secondConnection.await(5, TimeUnit.SECONDS));
+ assertEquals(2, connectionCount.get());
+ clientFactory.releaseDataStream(dataStreamId);
+ } finally {
+ server.shutdownNow();
+ }
Review Comment:
The `ManagedChannel` created for the test is not shut down in the `finally`
block, which can leak threads and other resources in the test JVM. It should be
shut down along with the server.
```java
ManagedChannel channel =
InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
try {
BeamFnDataGrpcClient clientFactory =
new BeamFnDataGrpcClient(
(Endpoints.ApiServiceDescriptor descriptor) -> channel,
OutboundObserverFactory.trivial());
// Two bundles concurrently use the same named data stream.
clientFactory.retainDataStream(dataStreamId);
clientFactory.retainDataStream(dataStreamId);
clientFactory.getOutboundObserver(apiServiceDescriptor, dataStreamId);
assertTrue(firstConnection.await(5, TimeUnit.SECONDS));
assertEquals(1, connectionCount.get());
// Releasing one of the two usages should not close the stream.
clientFactory.releaseDataStream(dataStreamId);
assertFalse(streamTerminated.await(100, TimeUnit.MILLISECONDS));
// Releasing the last usage should close the stream.
clientFactory.releaseDataStream(dataStreamId);
assertTrue(streamTerminated.await(5, TimeUnit.SECONDS));
// A subsequent usage of the same named data stream establishes a new
stream.
clientFactory.retainDataStream(dataStreamId);
clientFactory.getOutboundObserver(apiServiceDescriptor, dataStreamId);
assertTrue(secondConnection.await(5, TimeUnit.SECONDS));
assertEquals(2, connectionCount.get());
clientFactory.releaseDataStream(dataStreamId);
} finally {
channel.shutdownNow();
server.shutdownNow();
}
```
##########
sdks/java/harness/src/test/java/org/apache/beam/fn/harness/data/BeamFnDataGrpcClientTest.java:
##########
@@ -364,6 +366,116 @@ public StreamObserver<BeamFnApi.Elements> data(
}
}
+ @Test
+ public void testNamedDataStreamClosedWhenNoLongerRetained() throws Exception
{
+ String dataStreamId = "streamA";
+ AtomicInteger connectionCount = new AtomicInteger();
+ CountDownLatch firstConnection = new CountDownLatch(1);
+ CountDownLatch secondConnection = new CountDownLatch(2);
+ CountDownLatch streamTerminated = new CountDownLatch(1);
+
+ Endpoints.ApiServiceDescriptor apiServiceDescriptor =
+ Endpoints.ApiServiceDescriptor.newBuilder()
+ .setUrl(this.getClass().getName() + "-" + UUID.randomUUID())
+ .build();
+ Server server =
+ InProcessServerBuilder.forName(apiServiceDescriptor.getUrl())
+ .addService(
+ new BeamFnDataGrpc.BeamFnDataImplBase() {
+ @Override
+ public StreamObserver<BeamFnApi.Elements> data(
+ StreamObserver<BeamFnApi.Elements> outboundObserver) {
+ connectionCount.incrementAndGet();
+ firstConnection.countDown();
+ secondConnection.countDown();
+ return TestStreams.<BeamFnApi.Elements>withOnNext(elements
-> {})
+ .withOnError(streamTerminated::countDown)
+ .withOnCompleted(streamTerminated::countDown)
+ .build();
+ }
+ })
+ .build();
+ server.start();
+ try {
+ ManagedChannel channel =
+
InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
+
+ BeamFnDataGrpcClient clientFactory =
+ new BeamFnDataGrpcClient(
+ (Endpoints.ApiServiceDescriptor descriptor) -> channel,
+ OutboundObserverFactory.trivial());
+
+ // Two bundles concurrently use the same named data stream.
+ clientFactory.retainDataStream(dataStreamId);
+ clientFactory.retainDataStream(dataStreamId);
+ clientFactory.getOutboundObserver(apiServiceDescriptor, dataStreamId);
+ assertTrue(firstConnection.await(5, TimeUnit.SECONDS));
+ assertEquals(1, connectionCount.get());
+
+ // Releasing one of the two usages should not close the stream.
+ clientFactory.releaseDataStream(dataStreamId);
+ assertFalse(streamTerminated.await(100, TimeUnit.MILLISECONDS));
+
+ // Releasing the last usage should close the stream.
+ clientFactory.releaseDataStream(dataStreamId);
+ assertTrue(streamTerminated.await(5, TimeUnit.SECONDS));
+
+ // A subsequent usage of the same named data stream establishes a new
stream.
+ clientFactory.retainDataStream(dataStreamId);
+ clientFactory.getOutboundObserver(apiServiceDescriptor, dataStreamId);
+ assertTrue(secondConnection.await(5, TimeUnit.SECONDS));
+ assertEquals(2, connectionCount.get());
+ clientFactory.releaseDataStream(dataStreamId);
+ } finally {
+ server.shutdownNow();
+ }
+ }
+
+ @Test
+ public void testDefaultDataStreamIsNotClosedOnRelease() throws Exception {
+ CountDownLatch streamTerminated = new CountDownLatch(1);
+
+ Endpoints.ApiServiceDescriptor apiServiceDescriptor =
+ Endpoints.ApiServiceDescriptor.newBuilder()
+ .setUrl(this.getClass().getName() + "-" + UUID.randomUUID())
+ .build();
+ Server server =
+ InProcessServerBuilder.forName(apiServiceDescriptor.getUrl())
+ .addService(
+ new BeamFnDataGrpc.BeamFnDataImplBase() {
+ @Override
+ public StreamObserver<BeamFnApi.Elements> data(
+ StreamObserver<BeamFnApi.Elements> outboundObserver) {
+ return TestStreams.<BeamFnApi.Elements>withOnNext(elements
-> {})
+ .withOnError(streamTerminated::countDown)
+ .withOnCompleted(streamTerminated::countDown)
+ .build();
+ }
+ })
+ .build();
+ server.start();
+ try {
+ ManagedChannel channel =
+
InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
+
+ BeamFnDataGrpcClient clientFactory =
+ new BeamFnDataGrpcClient(
+ (Endpoints.ApiServiceDescriptor descriptor) -> channel,
+ OutboundObserverFactory.trivial());
+
+ clientFactory.retainDataStream("");
+ StreamObserver<BeamFnApi.Elements> outboundObserver =
+ clientFactory.getOutboundObserver(apiServiceDescriptor, "");
+ clientFactory.releaseDataStream("");
+
+ // The default data stream is kept open and reused even after release.
+ assertFalse(streamTerminated.await(100, TimeUnit.MILLISECONDS));
+ assertSame(outboundObserver,
clientFactory.getOutboundObserver(apiServiceDescriptor, ""));
+ } finally {
+ server.shutdownNow();
+ }
Review Comment:
The `ManagedChannel` created for the test is not shut down in the `finally`
block, which can leak threads and other resources in the test JVM. It should be
shut down along with the server.
```suggestion
ManagedChannel channel =
InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
try {
BeamFnDataGrpcClient clientFactory =
new BeamFnDataGrpcClient(
(Endpoints.ApiServiceDescriptor descriptor) -> channel,
OutboundObserverFactory.trivial());
clientFactory.retainDataStream("");
StreamObserver<BeamFnApi.Elements> outboundObserver =
clientFactory.getOutboundObserver(apiServiceDescriptor, "");
clientFactory.releaseDataStream("");
// The default data stream is kept open and reused even after release.
assertFalse(streamTerminated.await(100, TimeUnit.MILLISECONDS));
assertSame(outboundObserver,
clientFactory.getOutboundObserver(apiServiceDescriptor, ""));
} finally {
channel.shutdownNow();
server.shutdownNow();
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
