Recover Azure App Service agents after overlapping restarts

communication/src/main/java/datadog/communication/ddagent/ExternalAgentLauncher.java

@@ -1,15 +1,27 @@
 package datadog.communication.ddagent;
 
+import static datadog.trace.util.AgentThreadFactory.AgentThread.PROCESS_SUPERVISOR;
 import static datadog.trace.util.ProcessSupervisor.ALWAYS_READY;
 import static datadog.trace.util.ProcessSupervisor.Health.HEALTHY;
 import static datadog.trace.util.ProcessSupervisor.Health.NEVER_CHECKED;
 import static datadog.trace.util.ProcessSupervisor.Health.READY_TO_START;
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
 import datadog.environment.OperatingSystem;
 import datadog.trace.api.Config;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.context.TraceScope;
+import datadog.trace.util.AgentThreadFactory;
 import datadog.trace.util.ProcessSupervisor;
 import java.io.Closeable;
 import java.io.File;
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.io.RandomAccessFile;
+import java.nio.charset.StandardCharsets;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.TimeoutException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -32,7 +44,7 @@ public ExternalAgentLauncher(Config config) {
 
         traceProcessSupervisor =
             new ProcessSupervisor(
-                "trace-agent", traceProcessBuilder, healthCheck(config.getAgentNamedPipe()));
+                "trace-agent", traceProcessBuilder, answersOnPipe(config.getAgentNamedPipe()));
       } else {
         log.warn("Trace agent path not set. Will not start trace agent process");
       }
@@ -45,7 +57,7 @@ public ExternalAgentLauncher(Config config) {
 
         dogStatsDProcessSupervisor =
             new ProcessSupervisor(
-                "dogstatsd", dogStatsDProcessBuilder, healthCheck(config.getDogStatsDNamedPipe()));
+                "dogstatsd", dogStatsDProcessBuilder, pipePresent(config.getDogStatsDNamedPipe()));
       } else {
         log.warn("DogStatsD path not set. Will not start DogStatsD process");
       }
@@ -63,18 +75,40 @@ public void close() {
     }
   }
 
-  private static ProcessSupervisor.HealthCheck healthCheck(String pipeName) {
+  /** Healthy while the agent's pipe is present and the agent answers a request on it. */
+  private static ProcessSupervisor.HealthCheck answersOnPipe(String pipeName) {
+    return healthCheck(pipeName, true);
+  }
+
+  /** Healthy while the agent's pipe is present. DogStatsD never replies on its pipe. */
+  private static ProcessSupervisor.HealthCheck pipePresent(String pipeName) {
+    return healthCheck(pipeName, false);
+  }
+
+  private static ProcessSupervisor.HealthCheck healthCheck(String pipeName, boolean checkResponse) {
     return null != pipeName && !pipeName.trim().isEmpty()
-        ? new NamedPipeHealthCheck(pipeName)
+        ? new NamedPipeHealthCheck(pipeName, checkResponse)
         : ALWAYS_READY;
   }
 
   static final class NamedPipeHealthCheck implements ProcessSupervisor.HealthCheck {
     private static final String NAMED_PIPE_PREFIX = "\\\\.\\pipe\\";
 
+    private static final byte[] INFO_REQUEST =
+        "GET /info HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n"
+            .getBytes(StandardCharsets.US_ASCII);
+    private static final int MAX_STATUS_LINE_LENGTH = 64;
+    private static final long RESPONSE_TIMEOUT_MILLIS = 2_000;
+
     private final File pipe;
+    private final boolean checkResponse;
 
-    NamedPipeHealthCheck(String pipeName) {
+    // A request that got no answer in time. Only the supervisor thread runs the health check, so
+    // this needs no synchronization.
+    private FutureTask<Boolean> timedOutInfoRequest;
+
+    NamedPipeHealthCheck(String pipeName, boolean checkResponse) {
+      this.checkResponse = checkResponse;
       if (pipeName.startsWith(NAMED_PIPE_PREFIX)) {
         this.pipe = new File(pipeName);
       } else {
@@ -91,7 +125,7 @@ public ProcessSupervisor.Health run(ProcessSupervisor.Health previousHealth)
 
         double delayMillis = 50;
         for (int retries = 0; retries < 7; retries++) {
-          if (!pipe.exists()) {
+          if (!pipeExists()) {
             return READY_TO_START; // no longer bound, start our own external process
           }
 
@@ -100,15 +134,117 @@ public ProcessSupervisor.Health run(ProcessSupervisor.Health previousHealth)
           delayMillis = delayMillis * 1.75;
         }
 
-        return HEALTHY; // use existing external process
+        return respondsOrReadyToStart(); // use existing external process if it responds
       }
 
       // otherwise just check that the pipe is still bound
-      if (pipe.exists()) {
-        return HEALTHY; // keep using external process
+      if (pipeExists()) {
+        return respondsOrReadyToStart(); // keep using external process if it responds
       } else {
         return READY_TO_START; // start our own process
       }
     }
+
+    private ProcessSupervisor.Health respondsOrReadyToStart() throws InterruptedException {
+      return !checkResponse || agentResponds() ? HEALTHY : READY_TO_START;
+    }
+
+    /**
+     * Sends a request on another thread and waits up to {@link #RESPONSE_TIMEOUT_MILLIS} for the
+     * answer. Reads from a pipe cannot time out, so a request to a hung agent stays blocked until
+     * the agent exits. Until then later checks fail without sending another request. An answer that
+     * arrives after the timeout still counts at the next check, so a slow agent is not restarted.
+     */
+    private boolean agentResponds() throws InterruptedException {
+      FutureTask<Boolean> lateRequest = timedOutInfoRequest;
+      if (lateRequest != null) {
+        if (!lateRequest.isDone()) {
+          return false; // still waiting for an answer
+        }
+        timedOutInfoRequest = null;
+        if (answered(lateRequest)) {
+          return true;
+        }
+      }
+      FutureTask<Boolean> infoRequest = new FutureTask<>(this::sendInfoRequest);
+      AgentThreadFactory.newAgentThread(PROCESS_SUPERVISOR, "-health-check", infoRequest, true)
+          .start();
+      try {
+        return infoRequest.get(RESPONSE_TIMEOUT_MILLIS, MILLISECONDS);
+      } catch (TimeoutException e) {
+        timedOutInfoRequest = infoRequest;
+        return false;
+      } catch (ExecutionException e) {
+        return false;
+      }
+    }
+
+    /** Returns the result of a finished request, without waiting. */
+    private static boolean answered(FutureTask<Boolean> finishedRequest)
+        throws InterruptedException {
+      try {
+        return finishedRequest.get();
+      } catch (ExecutionException e) {
+        return false;
+      }
+    }
+
+    /**
+     * Sends a minimal request over a pipe handle of its own, outside the tracer's HTTP clients.
+     *
+     * @return true if the agent answers with an HTTP status line, whatever the status
+     */
+    private boolean sendInfoRequest() throws IOException {
+      try (TraceScope ignored = AgentTracer.get().muteTracing()) {
+        RandomAccessFile file;
+        try {
+          file = new RandomAccessFile(pipe, "rw");
+        } catch (FileNotFoundException e) {
+          // a pipe with all instances busy cannot be asked, but an agent is serving it
+          return !isPipeNotFound(e);
+        }
+        try {
+          file.write(INFO_REQUEST);
+          StringBuilder statusLine = new StringBuilder();
+          for (int c = file.read();
+              c != -1 && c != '\r' && statusLine.length() < MAX_STATUS_LINE_LENGTH;
+              c = file.read()) {
+            statusLine.append((char) c);
+          }
+          return statusLine.toString().startsWith("HTTP/1.");
+        } finally {
+          file.close();
+        }
+      }
+    }
+
+    /**
+     * {@link File#exists()} can report false for a pipe whose instances are all busy, for example
+     * under load. In that case open the pipe once to tell busy from absent.
+     */
+    private boolean pipeExists() {
+      if (pipe.exists()) {
+        return true;
+      }
+      if (!OperatingSystem.isWindows()) {
+        return false;
+      }
+      try (RandomAccessFile ignored = new RandomAccessFile(pipe, "rw")) {
+        return true;
+      } catch (FileNotFoundException e) {
+        return !isPipeNotFound(e);
+      } catch (IOException e) {
+        return true; // opened, then failed to close
+      }
+    }
+
+    /**
+     * Only "cannot find the file" (ERROR_FILE_NOT_FOUND) means no process is serving the pipe.
+     * Errors such as "All pipe instances are busy" (ERROR_PIPE_BUSY) mean that one is.
+     */
+    private static boolean isPipeNotFound(FileNotFoundException e) {
+      String message = e.getMessage();
+      return message != null && message.contains("cannot find the file");
+    }
   }
 }

internal-api/internal-api-9/src/main/java/datadog/trace/util/JDK9ProcessPidFunction.java

@@ -0,0 +1,15 @@
+package datadog.trace.util;
+
+import java.util.function.Function;
+
+/** Use standard API to retrieve the PID of a child process on Java9+. */
+public final class JDK9ProcessPidFunction implements Function<Process, String> {
+  @Override
+  public String apply(Process process) {
+    try {
+      return Long.toString(process.pid());
+    } catch (Throwable e) {
+      return "";
+    }
+  }
+}

internal-api/src/main/java/datadog/trace/util/PidHelper.java

@@ -13,6 +13,7 @@
 import java.util.Collections;
 import java.util.Objects;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -38,6 +39,31 @@ public static long getPidAsLong() {
     return PID_AS_LONG;
   }
 
+  private static final Function<Process, String> CHILD_PID = findChildPidFunction();
+
+  /** Returns the PID of a child process, or an empty string when the JVM cannot provide it. */
+  public static String getPid(Process process) {
+    if (null == process || null == CHILD_PID) {
+      return "";
+    }
+    return CHILD_PID.apply(process);
+  }
+
+  @SuppressWarnings("unchecked")
+  private static Function<Process, String> findChildPidFunction() {
+    if (JavaVirtualMachine.isJavaVersionAtLeast(9)) {
+      try {
+        return (Function<Process, String>)
+            Class.forName("datadog.trace.util.JDK9ProcessPidFunction")
+                .getDeclaredConstructor()
+                .newInstance();
+      } catch (Throwable e) {
+        log.debug("JDK9ProcessPidFunction not available", e);
+      }
+    }
+    return null; // Java 8 has no public API for the PID of a child process
+  }
+
   @SuppressWarnings("unchecked")
   private static String findPid() {
     String pid = "";