TIKA-4817: reclaim a shared-mode worker abandoned by a client disconnect (#3015)

Author: tballison

tika-pipes/tika-pipes-core/src/main/java/org/apache/tika/pipes/core/PerClientServerManager.java

@@ -85,12 +85,15 @@ private static int forkHeapPercentage(int numClients) {
     private final Path tikaConfigPath;
     private final int clientId;
 
-    private Process process;
-    private ServerSocket serverSocket;
-    private Path tmpDir;
-    private int port = -1;
+    private volatile Process process;
+    private volatile ServerSocket serverSocket;
+    private volatile Path tmpDir;
+    private volatile int port = -1;
     private long filesProcessed = 0;
-    private boolean pendingRestart = false;
+    private volatile boolean pendingRestart = false;
+    // Set once by shutdown()/close(); guards a request thread from starting a fresh
+    // process after the manager has been torn down (which would leak the child).
+    private volatile boolean closed = false;
 
     public PerClientServerManager(PipesConfig pipesConfig, Path tikaConfigPath, int clientId) {
         this.pipesConfig = pipesConfig;
@@ -227,7 +230,10 @@ public int handleCrashAndGetExitCode() {
     }
 
     @Override
-    public void ensureRunning() throws IOException, InterruptedException, TimeoutException, ServerInitializationException {
+    public synchronized void ensureRunning() throws IOException, InterruptedException, TimeoutException, ServerInitializationException {
+        if (closed) {
+            throw new IllegalStateException("PerClientServerManager is closed");
+        }
         // Check if server is running AND not marked for restart
         if (isRunning() && !pendingRestart) {
             return;
@@ -237,25 +243,32 @@ public void ensureRunning() throws IOException, InterruptedException, TimeoutExc
 
     @Override
     public Socket connect(int socketTimeoutMillis) throws IOException, ServerInitializationException {
-        if (serverSocket == null) {
+        // Capture the socket up front: shutdown() may null the field concurrently, but this
+        // request keeps using (and detects the close on) the instance it started with.
+        ServerSocket ss = serverSocket;
+        if (ss == null) {
             throw new IllegalStateException("Server not started. Call ensureRunning() first.");
         }
 
         // Accept incoming connection from the server process
-        serverSocket.setSoTimeout(1000); // 1 second timeout for each poll
+        ss.setSoTimeout(1000); // 1 second timeout for each poll
         long startTime = System.currentTimeMillis();
 
         while (true) {
             try {
-                Socket socket = serverSocket.accept();
+                Socket socket = ss.accept();
                 socket.setSoTimeout(socketTimeoutMillis);
                 socket.setTcpNoDelay(true);
                 LOG.debug("clientId={}: accepted connection from server", clientId);
                 return socket;
             } catch (SocketTimeoutException e) {
-                // Check if the process died before connecting
-                if (!process.isAlive()) {
-                    int exitValue = process.exitValue();
+                // Check if the process died (or the manager was shut down) before connecting.
+                Process p = process;
+                if (p == null) {
+                    throw new IOException("Server manager was shut down while connecting");
+                }
+                if (!p.isAlive()) {
+                    int exitValue = p.exitValue();
                     LOG.error("clientId={}: Process exited with code {} before connecting to socket",
                             clientId, exitValue);
                     ServerProcessIO.surfaceCrashDiagnostics(LOG, "clientId=" + clientId, tmpDir);
@@ -281,10 +294,14 @@ public Socket connect(int socketTimeoutMillis) throws IOException, ServerInitial
         }
     }
 
-    private void startServer() throws IOException, InterruptedException, TimeoutException, ServerInitializationException {
-        // Clean up any previous server
+    private synchronized void startServer() throws IOException, InterruptedException, TimeoutException, ServerInitializationException {
+        if (closed) {
+            throw new IllegalStateException("PerClientServerManager is closed");
+        }
+        // Clean up any previous server (restart) -- teardown, not shutdown, so we do not
+        // mark the manager closed.
         if (process != null || serverSocket != null || tmpDir != null) {
-            shutdown();
+            teardown();
         }
 
         // Create new server socket to get a free port
@@ -348,8 +365,18 @@ private void startServer() throws IOException, InterruptedException, TimeoutExce
     }
 
     @Override
-    public void shutdown() throws InterruptedException {
-        LOG.debug("clientId={}: shutting down server", clientId);
+    public synchronized void shutdown() throws InterruptedException {
+        closed = true;
+        teardown();
+    }
+
+    /**
+     * Tears down the current server process, socket, and temp dir without marking the manager
+     * closed -- shared by the final {@link #shutdown()} and by {@link #startServer()} on restart.
+     * Callers hold the monitor.
+     */
+    private void teardown() throws InterruptedException {
+        LOG.debug("clientId={}: tearing down server", clientId);
 
         if (serverSocket != null) {
             try {

tika-pipes/tika-pipes-core/src/main/java/org/apache/tika/pipes/core/server/ConnectionHandler.java

@@ -250,13 +250,27 @@ private void loopUntilDone(FetchEmitTuple fetchEmitTuple, ParseContext mergedCon
         long totalTaskTimeoutMillis = limits.getTotalTaskTimeoutMillis();
         long heartbeatCounter = 1;
         boolean wroteIntermediateResult = false;
+        // If the client disconnects mid-parse we stop writing to the dead socket but keep
+        // polling and enforcing the timeouts below, so an abandoned worker that will not stop
+        // still trips checkTotalTimeout/checkProgressTimeout -> System.exit -> the shared JVM
+        // recycles it. Otherwise (per-JVM shared mode has no per-request process to kill) a
+        // runaway parse would spin forever with its heap pinned.
+        boolean clientGone = false;
 
         while (running) {
             // Check for intermediate result
             if (!wroteIntermediateResult) {
                 Metadata intermediate = intermediateResult.poll(100, TimeUnit.MILLISECONDS);
                 if (intermediate != null) {
-                    protocolIO.writeIntermediate(intermediate);
+                    if (!clientGone) {
+                        try {
+                            protocolIO.writeIntermediate(intermediate);
+                        } catch (IOException e) {
+                            clientGone = true;
+                            LOG.debug("handlerId={}: client gone (writing intermediate); keeping the "
+                                    + "worker under its timeout so a runaway parse is reclaimed", handlerId);
+                        }
+                    }
                     countDownLatch.countDown();
                     wroteIntermediateResult = true;
                 }
@@ -285,15 +299,27 @@ private void loopUntilDone(FetchEmitTuple fetchEmitTuple, ParseContext mergedCon
                 }
                 LOG.debug("handlerId={}: finished task id={} status={}", handlerId,
                         fetchEmitTuple.getId(), pipesResult.status());
-                protocolIO.writeFinished(pipesResult);
+                if (!clientGone) {
+                    try {
+                        protocolIO.writeFinished(pipesResult);
+                    } catch (IOException e) {
+                        LOG.debug("handlerId={}: client gone before final result could be sent", handlerId);
+                    }
+                }
                 return;
             }
 
             // Send fire-and-forget heartbeat
             long elapsed = (System.nanoTime() - startNanos) / 1_000_000L;
-            if (elapsed > heartbeatCounter * heartbeatIntervalMillis) {
+            if (!clientGone && elapsed > heartbeatCounter * heartbeatIntervalMillis) {
                 LOG.trace("handlerId={}: still processing, counter={}", handlerId, heartbeatCounter);
-                PipesMessage.working().write(output);
+                try {
+                    PipesMessage.working().write(output);
+                } catch (IOException e) {
+                    clientGone = true;
+                    LOG.debug("handlerId={}: client gone (heartbeat); keeping the worker under its "
+                            + "timeout so a runaway parse is reclaimed", handlerId);
+                }
                 heartbeatCounter++;
             }