fix(qwp): fix a memory-access fault and worker leak during shutdown

Author: kafka1991

core/src/main/java/io/questdb/client/cutlass/qwp/client/sf/cursor/SegmentManager.java

@@ -67,6 +67,7 @@ public final class SegmentManager implements QuietCloseable {
     public static final long DISK_FULL_LOG_THROTTLE_NANOS = 30_000_000_000L; // 30 s
     public static final long UNLIMITED_TOTAL_BYTES = Long.MAX_VALUE;
     private static final Logger LOG = LoggerFactory.getLogger(SegmentManager.class);
+    private static final long WORKER_JOIN_TIMEOUT_MILLIS = 5_000L;
 
     private final AtomicLong fileGeneration = new AtomicLong();
     private final Object lock = new Object();
@@ -107,6 +108,7 @@ public final class SegmentManager implements QuietCloseable {
     // {@link #lock}; the lock covers both paths so the counter stays
     // consistent across registration boundaries.
     private long totalBytes;
+    private long workerJoinTimeoutMillis = WORKER_JOIN_TIMEOUT_MILLIS;
     // volatile because wakeWorker() reads workerThread without holding the
     // monitor; the synchronized start()/close() pair handles the
     // start-vs-close ordering.
@@ -168,10 +170,28 @@ public synchronized void close() {
         Thread t = workerThread;
         if (t != null) {
             LockSupport.unpark(t);
+            // A pending interrupt on the caller makes Thread.join() throw at
+            // once; clear it so the join actually reaps the worker (which
+            // still owns segment files), then restore it for the rest of the
+            // interrupted-teardown protocol.
+            boolean interrupted = Thread.interrupted();
+            long deadlineNanos = System.nanoTime() + workerJoinTimeoutMillis * 1_000_000L;
             try {
-                t.join(5_000);
-            } catch (InterruptedException ignored) {
-                Thread.currentThread().interrupt();
+                while (t.isAlive()) {
+                    long remainingMillis = (deadlineNanos - System.nanoTime()) / 1_000_000L;
+                    if (remainingMillis <= 0) {
+                        break;
+                    }
+                    try {
+                        t.join(remainingMillis);
+                    } catch (InterruptedException ignored) {
+                        interrupted = true;
+                    }
+                }
+            } finally {
+                if (interrupted) {
+                    Thread.currentThread().interrupt();
+                }
             }
             if (t.isAlive()) {
                 LOG.warn("SegmentManager worker did not stop before close wait completed; "
@@ -287,6 +307,11 @@ public void setBeforeTrimSyncHook(Runnable hook) {
         this.beforeTrimSyncHook = hook;
     }
 
+    @TestOnly
+    public void setWorkerJoinTimeoutMillis(long millis) {
+        this.workerJoinTimeoutMillis = millis;
+    }
+
     public synchronized void start() {
         if (workerThread != null) {
             throw new IllegalStateException("already started");

core/src/test/java/io/questdb/client/test/cutlass/qwp/client/sf/cursor/SegmentManagerCloseRaceTest.java

@@ -170,10 +170,7 @@ public void testCloseDoesNotFreePathScratchWhenWorkerStillAlive() throws Excepti
                 Assert.assertTrue("precondition: path scratch should be allocated",
                         readPathScratchImpl(manager) != 0L);
 
-                // Exercise the same branch as a timed-out join without making
-                // the test sleep for 5 seconds: join() returns while the worker
-                // is still alive. close() must leave worker-owned native memory
-                // alone so the worker can resume safely.
+                manager.setWorkerJoinTimeoutMillis(50L);
                 Thread.currentThread().interrupt();
                 manager.close();
                 Assert.assertTrue("close should preserve interrupted status",
@@ -185,6 +182,7 @@ public void testCloseDoesNotFreePathScratchWhenWorkerStillAlive() throws Excepti
                         readPathScratchImpl(manager) != 0L);
 
                 releaseWorker.countDown();
+                manager.setWorkerJoinTimeoutMillis(TimeUnit.SECONDS.toMillis(60));
                 manager.close();
                 managerClosed = true;
                 Assert.assertNull("successful close should clear workerThread",
@@ -206,6 +204,85 @@ public void testCloseDoesNotFreePathScratchWhenWorkerStillAlive() throws Excepti
         });
     }
 
+    @Test(timeout = 15_000L)
+    public void testInterruptedCallerDoesNotAbandonReapableWorker() throws Exception {
+        TestUtils.assertMemoryLeak(() -> {
+            long segSize = MmapSegment.HEADER_SIZE + (MmapSegment.FRAME_HEADER_SIZE + 32);
+            String slot = tmpDir + "/interrupt-slot";
+            Assert.assertEquals(0, Files.mkdir(slot, Files.DIR_MODE_DEFAULT));
+            MmapSegment initial = MmapSegment.create(slot + "/sf-initial.sfa", 0L, segSize);
+            SegmentRing ring = new SegmentRing(initial, segSize);
+            SegmentManager manager = new SegmentManager(segSize, TimeUnit.SECONDS.toNanos(60));
+            CountDownLatch workerBlocked = new CountDownLatch(1);
+            CountDownLatch releaseWorker = new CountDownLatch(1);
+            CountDownLatch closeReturned = new CountDownLatch(1);
+            AtomicBoolean fired = new AtomicBoolean();
+            AtomicBoolean interruptPreserved = new AtomicBoolean();
+            AtomicReference<Throwable> hookErr = new AtomicReference<>();
+            AtomicReference<Throwable> closeErr = new AtomicReference<>();
+            boolean managerClosed = false;
+            try {
+                manager.register(ring, slot);
+                manager.setBeforeInstallSyncHook(() -> {
+                    if (!fired.compareAndSet(false, true)) return;
+                    workerBlocked.countDown();
+                    try {
+                        if (!releaseWorker.await(10, TimeUnit.SECONDS)) {
+                            hookErr.compareAndSet(null,
+                                    new AssertionError("timed out waiting for test to release worker"));
+                        }
+                    } catch (Throwable t) {
+                        hookErr.compareAndSet(null, t);
+                    }
+                });
+                manager.start();
+                Assert.assertTrue("worker did not reach install hook",
+                        workerBlocked.await(5, TimeUnit.SECONDS));
+
+                Thread closer = new Thread(() -> {
+                    Thread.currentThread().interrupt();
+                    try {
+                        manager.close();
+                    } catch (Throwable t) {
+                        closeErr.compareAndSet(null, t);
+                    } finally {
+                        interruptPreserved.set(Thread.currentThread().isInterrupted());
+                        closeReturned.countDown();
+                    }
+                }, "interrupted-closer");
+                closer.start();
+
+                Assert.assertFalse("interrupted close() abandoned a live worker instead of waiting",
+                        closeReturned.await(300, TimeUnit.MILLISECONDS));
+
+                releaseWorker.countDown();
+                Assert.assertTrue("close() never returned after the worker was released",
+                        closeReturned.await(10, TimeUnit.SECONDS));
+                closer.join(TimeUnit.SECONDS.toMillis(5));
+                managerClosed = readWorkerThread(manager) == null;
+
+                if (closeErr.get() != null) {
+                    throw new AssertionError("close() threw", closeErr.get());
+                }
+                Assert.assertNull("close() must reap the worker despite the pending interrupt",
+                        readWorkerThread(manager));
+                Assert.assertTrue("close() must restore the caller's interrupt status",
+                        interruptPreserved.get());
+                if (hookErr.get() != null) {
+                    throw new AssertionError("install hook failed", hookErr.get());
+                }
+            } finally {
+                manager.setBeforeInstallSyncHook(null);
+                releaseWorker.countDown();
+                if (!managerClosed) {
+                    Thread.interrupted();
+                    manager.close();
+                }
+                ring.close();
+            }
+        });
+    }
+
     private static void cleanupRecursively(String dir) {
         if (!Files.exists(dir)) return;
         long find = Files.findFirst(dir);