Notify on exec process end and read exit code

Author: rjeberhard

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

@@ -18,6 +18,8 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.util.concurrent.TimeUnit;
+
 import org.apache.commons.lang3.StringUtils;
 
 public class Exec {
@@ -170,20 +172,57 @@ public Process exec(
       throws ApiException, IOException {
     String path = makePath(namespace, name, command, container, stdin, tty);
 
-    WebSocketStreamHandler handler = new WebSocketStreamHandler();
-    ExecProcess exec = new ExecProcess(handler);
+    ExecProcess exec = new ExecProcess();
+    WebSocketStreamHandler handler = exec.getHandler();
     WebSockets.stream(path, "GET", apiClient, handler);
 
     return exec;
   }
 
   private static class ExecProcess extends Process {
-    WebSocketStreamHandler streamHandler;
-    private int statusCode;
+    private final WebSocketStreamHandler streamHandler;
+    private volatile int statusCode;
 
-    public ExecProcess(WebSocketStreamHandler handler) throws IOException {
-      this.streamHandler = handler;
+    public ExecProcess() throws IOException {
       this.statusCode = -1;
+      this.streamHandler = new WebSocketStreamHandler() {
+        @Override
+        public void close() {
+          if (statusCode == -1) {
+            InputStream inputStream = getInputStream(3);
+            int exitCode = 0;
+            try {
+              int available = inputStream.available();
+              if (available > 0) {
+                byte[] b = new byte[available];
+                inputStream.read(b);
+                String result = new String(b, "UTF-8");
+                int idx = result.lastIndexOf(':');
+                if (idx > 0) {
+                  try {
+                    exitCode = Integer.parseInt(result.substring(idx + 1).trim());
+                  } catch (NumberFormatException nfe) {
+                    // no-op
+                  }
+                }
+              }
+            } catch (IOException e) {
+            }
+            
+            // notify of process completion
+            synchronized (ExecProcess.this) {
+              statusCode = exitCode;
+              ExecProcess.this.notifyAll();
+            }
+          }
+          super.close();
+        }
+        
+      };
+    }
+    
+    private WebSocketStreamHandler getHandler() {
+      return streamHandler;
     }
 
     @Override
@@ -208,11 +247,24 @@ public int waitFor() throws InterruptedException {
       }
       return statusCode;
     }
-
+    
+    @Override
+    public boolean waitFor(long timeout, TimeUnit unit)
+        throws InterruptedException {
+      synchronized (this) {
+        this.wait(TimeUnit.MILLISECONDS.convert(timeout, unit));
+      }
+      return !isAlive();
+    }
+    
+    @Override
     public int exitValue() {
+      if (statusCode == -1)
+        throw new IllegalThreadStateException();
       return statusCode;
     }
 
+    @Override
     public void destroy() {
       streamHandler.close();
     }