This is an automated email from the ASF dual-hosted git repository.
jbonofre pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-java.git
The following commit(s) were added to refs/heads/main by this push:
new 5a6314d1b GH-384: Wait for the server executor in FlightServer.close()
(#1333)
5a6314d1b is described below
commit 5a6314d1b1ed4bebeb2614275e4e5672c84bead1
Author: JB Onofré <[email protected]>
AuthorDate: Tue Oct 6 07:21:15 2026 +0200
GH-384: Wait for the server executor in FlightServer.close() (#1333)
## What's Changed
`TestDoExchange.tearDown` is flaky (most recently in
`testServerCancelLeak` on the macOS job) with:
```
java.lang.IllegalStateException: Memory was leaked by query. Memory leaked:
(8192)
Allocator(ROOT) 0/0/16384/2147483647 (res/actual/peak/limit)
```
This is a shutdown race rather than a real leak: the message reports
8192 leaked bytes, but the allocator state printed right after already
shows 0 bytes allocated.
gRPC reports the server as terminated once its transports are closed,
without waiting for the calls still running on the executor.
`FlightServer.close()` returned at that point, while such a call could
still be parsing an incoming message (which allocates a buffer,
discarded right after because the stream is already closed). Closing the
allocator in that window reports a leak.
`FlightServer.close()` now also waits for the executor it created,
within the existing 6 seconds budget, and logs a warning if the executor
does not terminate in time. Nothing changes for user-supplied executors
(Flight does not own them), nor for `shutdown()` / `awaitTermination()`.
Tests:
- `TestServerOptions.closeWaitsForDefaultExecutor` is a new test that
fails without the fix.
- `TestDoExchange.testClientClose` leaked the allocator on purpose to
work around the same race (`TODO(ARROW-9586)`). This workaround is
removed.
I reproduced the flakiness locally by looping the tests in parallel JVMs
(to get CPU contention):
| Scenario | Before | After |
|---|---|---|
| `testServerCancelLeak` | 61 failures / 12,000 runs | 0 / 12,000 |
| `testClientClose` without the workaround | 1,250 failures / 7,200 runs
| 0 / 7,200 |
Closes #384.
---
.../java/org/apache/arrow/flight/FlightServer.java | 32 ++++++++++++--------
.../org/apache/arrow/flight/TestDoExchange.java | 9 ------
.../org/apache/arrow/flight/TestServerOptions.java | 34 ++++++++++++++++++++++
3 files changed, 54 insertions(+), 21 deletions(-)
diff --git
a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightServer.java
b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightServer.java
index ac761457f..c003b541b 100644
--- a/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightServer.java
+++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightServer.java
@@ -145,25 +145,33 @@ public class FlightServer implements AutoCloseable {
/** Shutdown the server, waits for up to 6 seconds for successful shutdown
before returning. */
@Override
public void close() throws InterruptedException {
+ final long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(6);
shutdown();
final boolean terminated = awaitTermination(3000, TimeUnit.MILLISECONDS);
if (terminated) {
logger.debug("Server was terminated within 3s");
- return;
- }
-
- // get more aggressive in termination.
- server.shutdownNow();
+ } else {
+ // get more aggressive in termination.
+ server.shutdownNow();
+
+ int count = 0;
+ while (!server.isTerminated() && count < 30) {
+ count++;
+ logger.debug("Waiting for termination");
+ Thread.sleep(100);
+ }
- int count = 0;
- while (!server.isTerminated() && count < 30) {
- count++;
- logger.debug("Waiting for termination");
- Thread.sleep(100);
+ if (!server.isTerminated()) {
+ logger.warn("Couldn't shutdown server, resources likely will be
leaked.");
+ }
}
- if (!server.isTerminated()) {
- logger.warn("Couldn't shutdown server, resources likely will be
leaked.");
+ // gRPC reports the server as terminated once its transports are closed,
without waiting for
+ // the calls still running on the executor. Those calls may still allocate
or hold buffers, so
+ // wait for them too; otherwise closing the allocator right after the
server can report a leak.
+ if (grpcExecutor != null
+ && !grpcExecutor.awaitTermination(deadline - System.nanoTime(),
TimeUnit.NANOSECONDS)) {
+ logger.warn("Couldn't shutdown server executor, resources likely will be
leaked.");
}
}
diff --git
a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestDoExchange.java
b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestDoExchange.java
index 1cd6e95e3..8967c92b9 100644
---
a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestDoExchange.java
+++
b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestDoExchange.java
@@ -424,15 +424,6 @@ public class TestDoExchange {
client.doExchange(FlightDescriptor.command(EXCHANGE_DO_GET))) {
assertEquals(Producer.SCHEMA, stream.getReader().getSchema());
}
- // Intentionally leak the allocator in this test. gRPC has a bug where it
does not wait for all
- // calls to complete
- // when shutting down the server, so this test will fail otherwise because
it closes the
- // allocator while the
- // server-side call still has memory allocated.
- // TODO(ARROW-9586): fix this once we track outstanding RPCs outside of
gRPC.
- // https://stackoverflow.com/questions/46716024/
- allocator = null;
- client = null;
}
/** Test closing with Metadata can't lead to error. */
diff --git
a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestServerOptions.java
b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestServerOptions.java
index c39ac922c..b18e6cfef 100644
---
a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestServerOptions.java
+++
b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestServerOptions.java
@@ -86,6 +86,40 @@ public class TestServerOptions {
assertTrue(executor.isShutdown());
}
+ /**
+ * Make sure that closing the server waits for the calls still running on
the default executor, as
+ * they may still be using the allocator.
+ */
+ @Test
+ public void closeWaitsForDefaultExecutor() throws Exception {
+ final AtomicBoolean callFinished = new AtomicBoolean();
+ final FlightProducer producer =
+ new NoOpFlightProducer() {
+ @Override
+ public void doAction(
+ CallContext context, Action action, StreamListener<Result>
listener) {
+ // End the RPC, so that gRPC can shut down, but keep the executor
thread busy
+ listener.onCompleted();
+ try {
+ Thread.sleep(200);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return;
+ }
+ callFinished.set(true);
+ }
+ };
+ try (BufferAllocator a = new RootAllocator(Long.MAX_VALUE);
+ FlightServer server =
+ FlightServer.builder(a, forGrpcInsecure(LOCALHOST, 0),
producer).build().start()) {
+ try (FlightClient client = FlightClient.builder(a,
server.getLocation()).build()) {
+ client.doAction(new Action("")).forEachRemaining(result -> {});
+ }
+ server.close();
+ assertTrue(callFinished.get());
+ }
+ }
+
/** Make sure that if the user provides an executor to gRPC, then Flight
does not close it. */
@Test
public void suppliedExecutorNotClosed() throws Exception {
