Merge pull request #1322 from netikras/master

Exec.exec returning Future and hiding Process as an implementation detail

Author: k8s-ci-robot

kubernetes/src/main/java/io/kubernetes/client/custom/IOTrio.java

@@ -0,0 +1,84 @@
+/*
+Copyright 2020 The Kubernetes Authors.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+package io.kubernetes.client.custom;
+
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.function.BiConsumer;
+
+/**
+ * A collection of all the 3 main pipes used in stdio: STDIN, STDOUT and STDERR. As working with the
+ * named pipes is usually carried out asynchronously, this collection also provides means to
+ * initiate and handle the close() event. A close() initiator calls the {@link #close(int, long)}
+ * method expressing its intent to close the communication channel. Handlers are notified of this
+ * intent and its up to the handlers to decide what's to be done next. Calling {@link #close(int,
+ * long)} does not close the streams or do anything else besides notifying the handlers.
+ */
+public class IOTrio {
+  private InputStream stdout;
+  private InputStream stderr;
+  private OutputStream stdin;
+  private final Collection<BiConsumer<Integer, Long>> closeHandlers;
+
+  public IOTrio() {
+    closeHandlers = new ArrayList<>();
+  }
+
+  public InputStream getStdout() {
+    return stdout;
+  }
+
+  public void setStdout(InputStream stdout) {
+    this.stdout = stdout;
+  }
+
+  public InputStream getStderr() {
+    return stderr;
+  }
+
+  public void setStderr(InputStream stderr) {
+    this.stderr = stderr;
+  }
+
+  public OutputStream getStdin() {
+    return stdin;
+  }
+
+  public void setStdin(OutputStream stdin) {
+    this.stdin = stdin;
+  }
+
+  /**
+   * Capture the CLOSE intent and handle it accordingly.
+   *
+   * @param handler the handler that's invoked when someone intends to close this communication.
+   *     Multiple handlers can be registered
+   */
+  public void onClose(BiConsumer<Integer, Long> handler) {
+    closeHandlers.add(handler);
+  }
+
+  /**
+   * Express an intent to close this communication. This intent will be relayed to all the
+   * registered handlers and it's up to them what to do with it.
+   *
+   * @param code proposed exit code
+   * @param timeout time in milliseconds given to terminate this communication. Negative timeout
+   *     means no timeout (i.e. wait for as long as it takes). 0 means "stop it now".
+   */
+  public void close(int code, long timeout) {
+    closeHandlers.forEach(handler -> handler.accept(code, timeout));
+  }
+}

util/src/main/java/io/kubernetes/client/Exec.java

@@ -16,6 +16,7 @@
 
 import com.google.common.io.CharStreams;
 import com.google.gson.reflect.TypeToken;
+import io.kubernetes.client.custom.IOTrio;
 import io.kubernetes.client.openapi.ApiClient;
 import io.kubernetes.client.openapi.ApiException;
 import io.kubernetes.client.openapi.Configuration;
@@ -34,11 +35,17 @@
 import java.io.UnsupportedEncodingException;
 import java.lang.reflect.Type;
 import java.net.URLEncoder;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
 import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -191,6 +198,129 @@ public Process exec(
         .execute();
   }
 
+  /**
+   * A convenience method. Executes a command remotely on a pod and monitors for events in that
+   * execution. The monitored events are: <br>
+   * - connection established (onOpen) <br>
+   * - connection closed (onClosed) <br>
+   * - execution error occurred (onError) <br>
+   * This method also allows to specify a MAX timeout for the execution and returns a future in
+   * order to monitor the execution flow. <br>
+   * onError and onClosed callbacks are invoked asynchronously, in a separate thread. <br>
+   *
+   * @param namespace a namespace the target pod "lives" in
+   * @param podName a name of the pod to exec the command on
+   * @param onOpen a callback invoked upon the connection established event.
+   * @param onClosed a callback invoked upon the process termination. Return code might not always
+   *     be there. N.B. this callback is invoked before the returned {@link Future} is completed.
+   * @param onError a callback to handle k8s errors (NOT the command errors/stderr!)
+   * @param timeoutMs timeout in milliseconds for the execution. I.e. the execution will take this
+   *     many ms or less. If the timeout command is running longer than the allowed timeout, the
+   *     command will be "asked" to terminate gracefully. If the command is still running after the
+   *     grace period, the sigkill will be issued. If null is passed, the timeout will not be used
+   *     and will wait for process to exit itself.
+   * @param tty whether you need tty to pipe the data. TTY mode will trim some binary data in order
+   *     to make it possible to show on screen (tty)
+   * @param command a tokenized command to run on the pod
+   * @return a {@link Future} promise representing this execution. Unless something goes south, the
+   *     promise will contain the process return exit code. If the timeoutMs is non-null and the
+   *     timeout expires before the process exits, promise will contain {@link Integer#MAX_VALUE}.
+   * @throws IOException
+   */
+  public Future<Integer> exec(
+      String namespace,
+      String podName,
+      Consumer<IOTrio> onOpen,
+      BiConsumer<Integer, IOTrio> onClosed,
+      BiConsumer<Throwable, IOTrio> onError,
+      Long timeoutMs,
+      boolean tty,
+      String... command)
+      throws IOException {
+    CompletableFuture<Integer> future = new CompletableFuture<>();
+    IOTrio io = new IOTrio();
+    String cmdStr = Arrays.toString(command);
+    BiConsumer<Throwable, IOTrio> errHandler =
+        (err, errIO) -> {
+          if (onError != null) {
+            onError.accept(err, errIO);
+          }
+        };
+    try {
+      Process process = exec(namespace, podName, command, null, true, tty);
+
+      io.onClose(
+          (code, timeout) -> {
+            process.destroy();
+            waitForProcessToExit(process, timeout, cmdStr, err -> errHandler.accept(err, io));
+            // processWaitingThread will handle the rest
+          });
+      io.setStdin(process.getOutputStream());
+      io.setStdout(process.getInputStream());
+      io.setStderr(process.getErrorStream());
+      runAsync(
+          "Process-Waiting-Thread-" + command[0] + "-" + podName,
+          () -> {
+            Supplier<Integer> returnCode = process::exitValue;
+            try {
+              log.debug("Waiting for process to close in {} ms: {}", timeoutMs, cmdStr);
+              boolean beforeTimout =
+                  waitForProcessToExit(
+                      process, timeoutMs, cmdStr, err -> errHandler.accept(err, io));
+              if (!beforeTimout) {
+                returnCode = () -> Integer.MAX_VALUE;
+              }
+            } catch (Exception e) {
+              errHandler.accept(e, io);
+            }
+            log.debug("process.onExit({}): {}", returnCode.get(), cmdStr);
+            if (onClosed != null) {
+              onClosed.accept(returnCode.get(), io);
+            }
+            future.complete(returnCode.get());
+          });
+      if (onOpen != null) {
+        onOpen.accept(io);
+      }
+    } catch (ApiException e) {
+      errHandler.accept(e, io);
+      future.completeExceptionally(e);
+    }
+    return future;
+  }
+
+  protected void runAsync(String taskName, Runnable task) {
+    Thread thread = new Thread(task);
+    thread.setName(taskName);
+    thread.start();
+  }
+
+  private boolean waitForProcessToExit(
+      Process process, Long timeoutMs, String cmdStr, Consumer<Exception> onError) {
+    boolean beforeTimeout = true;
+    if (timeoutMs != null && timeoutMs >= 0) {
+      boolean exited = false;
+      try {
+        exited = process.waitFor(timeoutMs, TimeUnit.MILLISECONDS);
+      } catch (InterruptedException e) {
+        onError.accept(e);
+      }
+      log.debug("Process closed={}: {}", exited, cmdStr);
+      if (!exited && process.isAlive()) {
+        beforeTimeout = false;
+        log.warn("Process timed out after {} ms. Shutting down: {}", timeoutMs, cmdStr);
+        process.destroy();
+      }
+    } else {
+      try {
+        process.waitFor();
+      } catch (InterruptedException e) {
+        onError.accept(e);
+      }
+    }
+    return beforeTimeout;
+  }
+
   public final class ExecutionBuilder {
     private final String namespace;
     private final String name;