fix(dedup): bind the control socket before the app starts

.github/workflows/java-agent.yml

@@ -38,5 +38,8 @@ jobs:
       - name: Build
         run: mvn -B -DskipTests clean verify
 
+      - name: Unit tests
+        run: mvn -B -pl keploy-sdk -am test
+
       - name: Smoke test Java agent
         run: ./scripts/smoke-javaagent.sh

README.md

@@ -25,6 +25,8 @@ Keploy Enterprise drives dynamic dedup per testcase.
 
 Coverage is collected at per-testcase granularity, not process granularity.
 
+The agent binds `/tmp/coverage_control.sock` before the application's `main` runs, so the socket is listening before the application can answer a health check. It answers `ACK` to each command it took, and `IGNORED` to an `END` that no `START` for that testcase came before: nothing is published for it, and Enterprise keeps that testcase out of the comparison. On the TCP transport used in Kubernetes the same `END` is answered with `IGNORED <testcase>` and then `ACK`.
+
 ## How to Use
 
 ### 1. Download the Keploy Java Agent

keploy-sdk/src/main/java/io/keploy/dedup/KeployDedupAgent.java

@@ -84,12 +84,19 @@ static String encodeBytecodeFrame(String buildTag, String manifestJson, byte[] z
         return "MANIFEST " + b64Tag + " " + b64Manifest;
     }
 
-    private static final String CONTROL_SOCKET_PATH = "/tmp/coverage_control.sock";
+    // The unix control socket's path. Not final only so tests can bind one of
+    // their own; the agent never changes it.
+    static volatile String controlSocketPath = "/tmp/coverage_control.sock";
     private static final String DATA_SOCKET_PATH = "/tmp/coverage_data.sock";
     private static final String DEFAULT_JACOCO_HOST = "127.0.0.1";
     private static final int DEFAULT_JACOCO_PORT = 36320;
     private static final int SOCKET_TIMEOUT_MILLIS = 3000;
     private static final int SOCKET_BACKLOG = 50;
+    // The answer to an END not taken (no START for that test came first, as
+    // after a restart): nothing was published for it, so Keploy must not use
+    // the test's coverage. On the unix socket it is the whole answer; on TCP
+    // it is a frame naming the test, and ACK still follows (CoverageTcpClient).
+    static final String IGNORED_REPLY = "IGNORED";
 
     private static final AtomicBoolean STARTED = new AtomicBoolean(false);
     private static final AtomicBoolean SHUTDOWN_HOOK_REGISTERED = new AtomicBoolean(false);
@@ -170,7 +177,24 @@ public static boolean start() {
         } else {
             // local/docker: SDK owns the unix control socket and pushes coverage
             // back over the unix data socket on a pod-shared /tmp.
-            CommandServer server = new CommandServer(collector, new CoveragePublisher(new File(DATA_SOCKET_PATH)));
+            //
+            // Bind here, on the caller's thread, not on the accept thread:
+            // premain runs before the app's main, so the socket is listening
+            // before any app code can answer a health check, and Keploy's
+            // first START finds it. Bound on the accept thread, it could come
+            // up after the app answered and Keploy had fired its first tests,
+            // whose START and END then reached no one: their coverage was lost.
+            CommandServer server = new CommandServer(collector,
+                    new CoveragePublisher(new File(DATA_SOCKET_PATH)), new File(controlSocketPath));
+            try {
+                server.bind();
+            } catch (Throwable t) {
+                // Never let premain fail the app's start: run it without dedup.
+                server.close();
+                STARTED.set(false);
+                log(Level.SEVERE, "Java dedup control socket server is unavailable", t);
+                return false;
+            }
             worker = server;
             coverageWorker = server;
             threadName = "keploy-java-dedup-control";
@@ -332,26 +356,32 @@ private static final class CommandServer implements Runnable, Closeable {
 
         private final CoverageCollector collector;
         private final CoveragePublisher publisher;
+        private final File controlSocket;
         private final AtomicBoolean running = new AtomicBoolean(true);
         private final Object testCaseLock = new Object();
         private volatile AFUNIXServerSocket serverSocket;
         private String activeTestId = "";
 
-        CommandServer(CoverageCollector collector, CoveragePublisher publisher) {
+        CommandServer(CoverageCollector collector, CoveragePublisher publisher, File controlSocket) {
             this.collector = collector;
             this.publisher = publisher;
+            this.controlSocket = controlSocket;
         }
 
-        @Override
-        public void run() {
-            File controlSocket = new File(CONTROL_SOCKET_PATH);
+        // bind replaces a stale socket file and listens on the control socket.
+        // start() calls it before it returns, so a command sent once start()
+        // has returned waits in the backlog for run() instead of being refused.
+        // On a failure the caller closes this server, which closes the socket.
+        void bind() throws IOException {
             deleteSocketFile(controlSocket);
+            serverSocket = AFUNIXServerSocket.newInstance();
+            serverSocket.bind(AFUNIXSocketAddress.of(controlSocket), SOCKET_BACKLOG);
+            relaxSocketPermissions(controlSocket);
+        }
 
-            try (AFUNIXServerSocket localServer = AFUNIXServerSocket.newInstance()) {
-                localServer.bind(AFUNIXSocketAddress.of(controlSocket), SOCKET_BACKLOG);
-                relaxSocketPermissions(controlSocket);
-                serverSocket = localServer;
-
+        @Override
+        public void run() {
+            try (AFUNIXServerSocket localServer = serverSocket) {
                 while (running.get()) {
                     try {
                         handle(localServer.accept());
@@ -407,7 +437,9 @@ private void dispatch(CoverageCommand command, OutputStream outputStream) {
                                 "Ignoring mismatched END command. expected=" + activeTestId + ", actual="
                                         + command.testId,
                                 null);
-                        writeAck(outputStream);
+                        // Not ACK: this END publishes nothing, and an ACK told
+                        // Keploy it had, so it waited for a record that never came.
+                        writeReply(outputStream, IGNORED_REPLY);
                         return;
                     }
 
@@ -429,11 +461,15 @@ private void dispatch(CoverageCommand command, OutputStream outputStream) {
         }
 
         private void writeAck(OutputStream outputStream) {
+            writeReply(outputStream, "ACK");
+        }
+
+        private void writeReply(OutputStream outputStream, String reply) {
             try {
-                outputStream.write("ACK\n".getBytes(StandardCharsets.UTF_8));
+                outputStream.write((reply + "\n").getBytes(StandardCharsets.UTF_8));
                 outputStream.flush();
             } catch (IOException e) {
-                log(Level.FINE, "Java dedup control client disconnected before ACK was delivered", e);
+                log(Level.FINE, "Java dedup control client disconnected before " + reply + " was delivered", e);
             }
         }
 
@@ -459,6 +495,7 @@ public void close() {
      *   collector -> SDK : "START <id>" | "END <id>"
      *   SDK -> collector : "ACK"                        (after START reset)
      *                      "COV <compact-json>" + "ACK" (after END dump)
+     *                      "IGNORED <id>" + "ACK"         (after END with no START)
      * </pre>
      * The collector starts listening only when replay begins, so the connect loop
      * retries until it is reachable.
@@ -559,6 +596,11 @@ private void dispatch(CoverageCommand command, OutputStream out) throws IOExcept
                                 "Ignoring mismatched END command. expected=" + activeTestId + ", actual="
                                         + command.testId,
                                 null);
+                        // Say so before the ACK: with no COV the collector has no
+                        // other way to tell this test went unmeasured. ACK still
+                        // follows, as it ends the collector's wait, and a collector
+                        // that does not know this frame skips it.
+                        writeLine(out, IGNORED_REPLY + " " + command.testId);
                         writeLine(out, "ACK");
                         return;
                     }

keploy-sdk/src/test/java/io/keploy/dedup/ControlSocketTest.java

@@ -0,0 +1,100 @@
+package io.keploy.dedup;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStreamReader;
+import java.io.OutputStream;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.attribute.PosixFilePermission;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.newsclub.net.unix.AFUNIXSocket;
+import org.newsclub.net.unix.AFUNIXSocketAddress;
+
+/**
+ * The unix control socket Keploy sends START and END on. Keploy fires its first
+ * test as soon as the app answers its health check, so the socket must listen
+ * by the time start() (premain, before the app's main) returns; and an END the
+ * SDK does not take must not be answered as taken.
+ */
+public class ControlSocketTest {
+
+    private String savedPath;
+    private File dir;
+
+    @Before
+    public void setUp() throws IOException {
+        savedPath = KeployDedupAgent.controlSocketPath;
+        dir = Files.createTempDirectory("keploy-dedup-control").toFile();
+        KeployDedupAgent.controlSocketPath = new File(dir, "control.sock").getAbsolutePath();
+    }
+
+    @After
+    public void tearDown() {
+        KeployDedupAgent.stop();
+        KeployDedupAgent.controlSocketPath = savedPath;
+        File[] left = dir.listFiles();
+        if (left != null) {
+            for (File f : left) {
+                f.delete();
+            }
+        }
+        dir.delete();
+    }
+
+    @Test
+    public void controlSocketListensWhenStartReturns() throws Exception {
+        assertTrue(KeployDedupAgent.start());
+        File socket = new File(KeployDedupAgent.controlSocketPath);
+        // No wait and no retry: a command Keploy sends the moment start() has
+        // returned must find the socket listening.
+        assertTrue("the control socket is bound when start() returns", socket.exists());
+        try (AFUNIXSocket client = AFUNIXSocket.newInstance()) {
+            client.connect(AFUNIXSocketAddress.of(socket), 3000);
+        }
+        // Keploy's agent may run as another user (sudo, a container).
+        assertTrue("anyone may write to the control socket",
+                Files.getPosixFilePermissions(socket.toPath()).contains(PosixFilePermission.OTHERS_WRITE));
+    }
+
+    @Test
+    public void startReplacesAStaleSocketFile() throws Exception {
+        File socket = new File(KeployDedupAgent.controlSocketPath);
+        assertTrue(socket.createNewFile());
+        assertTrue(KeployDedupAgent.start());
+        assertEquals("IGNORED", send("END test-set-0/test-1"));
+    }
+
+    @Test
+    public void startReportsFailureWhenTheControlSocketCannotBind() {
+        KeployDedupAgent.controlSocketPath = new File(new File(dir, "missing"), "control.sock").getAbsolutePath();
+        assertFalse("start() with no control socket", KeployDedupAgent.start());
+        assertFalse(KeployDedupAgent.isStarted());
+    }
+
+    @Test
+    public void endWithNoStartIsAnsweredIgnoredNotAck() throws Exception {
+        assertTrue(KeployDedupAgent.start());
+        assertEquals("IGNORED", send("END test-set-0/test-47"));
+    }
+
+    private static String send(String command) throws IOException {
+        try (AFUNIXSocket client = AFUNIXSocket.newInstance()) {
+            client.connect(AFUNIXSocketAddress.of(new File(KeployDedupAgent.controlSocketPath)), 3000);
+            client.setSoTimeout(30000);
+            OutputStream out = client.getOutputStream();
+            out.write((command + "\n").getBytes(StandardCharsets.UTF_8));
+            out.flush();
+            BufferedReader in = new BufferedReader(
+                    new InputStreamReader(client.getInputStream(), StandardCharsets.UTF_8));
+            return in.readLine();
+        }
+    }
+}

keploy-sdk/src/test/java/io/keploy/dedup/CoverageTcpClientTest.java

@@ -0,0 +1,62 @@
+package io.keploy.dedup;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStreamReader;
+import java.io.OutputStream;
+import java.net.InetAddress;
+import java.net.ServerSocket;
+import java.net.Socket;
+import java.nio.charset.StandardCharsets;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+/**
+ * The TCP transport, with the test as the collector the SDK dials (k8s). An END
+ * the SDK does not take publishes no COV, so it must say which test went
+ * unmeasured, and must still answer ACK, which is what ends the collector's wait.
+ */
+public class CoverageTcpClientTest {
+
+    private static final String ENDPOINT_PROPERTY = "keploy.coverage.endpoint";
+
+    private ServerSocket collector;
+
+    @Before
+    public void setUp() throws IOException {
+        collector = new ServerSocket(0, 1, InetAddress.getLoopbackAddress());
+        collector.setSoTimeout(30000);
+        System.setProperty(ENDPOINT_PROPERTY, "127.0.0.1:" + collector.getLocalPort());
+    }
+
+    @After
+    public void tearDown() throws IOException {
+        KeployDedupAgent.stop();
+        System.clearProperty(ENDPOINT_PROPERTY);
+        collector.close();
+    }
+
+    @Test
+    public void endWithNoStartNamesTheTestThenAnswersAck() throws Exception {
+        assertTrue(KeployDedupAgent.start());
+        try (Socket sdk = collector.accept()) {
+            sdk.setSoTimeout(30000);
+            OutputStream out = sdk.getOutputStream();
+            BufferedReader in = new BufferedReader(
+                    new InputStreamReader(sdk.getInputStream(), StandardCharsets.UTF_8));
+
+            // Two, so that a missing ACK shows as the next test's frame in its place.
+            out.write("END test-set-0/test-47\nEND test-set-0/test-48\n".getBytes(StandardCharsets.UTF_8));
+            out.flush();
+
+            assertEquals("IGNORED test-set-0/test-47", in.readLine());
+            assertEquals("ACK", in.readLine());
+            assertEquals("IGNORED test-set-0/test-48", in.readLine());
+            assertEquals("ACK", in.readLine());
+        }
+    }
+}