fix #4154: adding a callback for stream consumption

Author: shawkins

kubernetes-client-api/src/main/java/io/fabric8/kubernetes/client/StreamConsumer.java

@@ -0,0 +1,60 @@
+/**
+ * Copyright (C) 2015 Red Hat, Inc.
+ *
+ * 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.fabric8.kubernetes.client;
+
+import java.io.OutputStream;
+import java.io.PipedInputStream;
+import java.io.PipedOutputStream;
+import java.nio.ByteBuffer;
+import java.nio.channels.Channels;
+import java.nio.channels.WritableByteChannel;
+import java.util.concurrent.CompletionStage;
+
+public interface StreamConsumer {
+
+  public static StreamConsumer newStreamConsumer(OutputStream os) {
+    checkForPiped(os);
+    return newBlockingStreamConsumer(Channels.newChannel(os));
+  }
+
+  public static StreamConsumer newBlockingStreamConsumer(WritableByteChannel channel) {
+    return buffer -> {
+      int remaining = buffer.remaining();
+      if (channel.write(buffer) != remaining) {
+        throw new KubernetesClientException("Unsucessful blocking write");
+      }
+      return null;
+    };
+  }
+
+  public static void checkForPiped(Object object) {
+    if (object instanceof PipedOutputStream || object instanceof PipedInputStream) {
+      throw new KubernetesClientException("Piped streams should not be used");
+    }
+  }
+
+  /**
+   * A callback for consuming a stream as a series of {@link ByteBuffer}s
+   *
+   * @param buffer
+   * @return a {@link CompletionStage} that is completed when the buffer has been fully consumed,
+   *         or null if it was already consumed
+   * @throws Exception
+   */
+  CompletionStage<?> consume(ByteBuffer buffer) throws Exception;
+
+}

kubernetes-client-api/src/main/java/io/fabric8/kubernetes/client/dsl/Loggable.java

@@ -16,6 +16,8 @@
 
 package io.fabric8.kubernetes.client.dsl;
 
+import io.fabric8.kubernetes.client.StreamConsumer;
+
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.io.PipedOutputStream;
@@ -70,7 +72,11 @@ public interface Loggable {
    * @param out {@link OutputStream} for storing logs
    * @return returns a Closeable interface for log watch
    */
-  LogWatch watchLog(OutputStream out);
+  default LogWatch watchLog(OutputStream out) {
+    return watchLog(StreamConsumer.newStreamConsumer(out), true);
+  }
+
+  LogWatch watchLog(StreamConsumer consumer, boolean blocking);
 
   /**
    * While waiting for Pod logs, how long shall we wait until a Pod

kubernetes-client-api/src/main/java/io/fabric8/kubernetes/client/dsl/TtyExecErrorable.java

@@ -15,6 +15,8 @@
  */
 package io.fabric8.kubernetes.client.dsl;
 
+import io.fabric8.kubernetes.client.StreamConsumer;
+
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.io.PipedOutputStream;
@@ -27,7 +29,11 @@ public interface TtyExecErrorable extends
    * <p>
    * In particular do no use a {@link PipedOutputStream} - use {@link #redirectingError()} instead
    */
-  TtyExecErrorChannelable writingError(OutputStream in);
+  default TtyExecErrorChannelable writingError(OutputStream in) {
+    return writingError(StreamConsumer.newStreamConsumer(in), true);
+  }
+
+  TtyExecErrorChannelable writingError(StreamConsumer consumer, boolean blocking);
 
   /**
    * If the {@link ExecWatch} should terminate when a stdErr message is received.

kubernetes-client-api/src/main/java/io/fabric8/kubernetes/client/dsl/TtyExecOutputErrorable.java

@@ -15,6 +15,8 @@
  */
 package io.fabric8.kubernetes.client.dsl;
 
+import io.fabric8.kubernetes.client.StreamConsumer;
+
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.io.PipedOutputStream;
@@ -27,7 +29,11 @@ public interface TtyExecOutputErrorable extends
    * <p>
    * In particular do no use a {@link PipedOutputStream} - use {@link #redirectingOutput()} instead
    */
-  TtyExecErrorable writingOutput(OutputStream in);
+  default TtyExecErrorable writingOutput(OutputStream in) {
+    return writingOutput(StreamConsumer.newStreamConsumer(in), true);
+  }
+
+  TtyExecErrorable writingOutput(StreamConsumer consumer, boolean blocking);
 
   /**
    * Will provide an {@link InputStream} via {@link ExecWatch#getOutput()}

kubernetes-client-api/src/main/java/io/fabric8/kubernetes/client/http/BufferUtil.java

@@ -76,14 +76,14 @@ public static ByteBuffer copy(ByteBuffer buffer) {
 
   /**
    * Very rudimentary method to check if the provided ByteBuffer contains text.
-   * 
+   *
    * @return true if the buffer contains text, false otherwise.
    */
   public static boolean isPlainText(ByteBuffer originalBuffer) {
     if (originalBuffer == null) {
       return false;
     }
-    final ByteBuffer buffer = copy(originalBuffer);
+    final ByteBuffer buffer = originalBuffer.asReadOnlyBuffer();
     final CharsetDecoder decoder = StandardCharsets.UTF_8.newDecoder();
     try {
       decoder.decode(buffer);

kubernetes-client-api/src/main/java/io/fabric8/kubernetes/client/http/WebSocket.java

@@ -42,6 +42,8 @@ default void onMessage(WebSocket webSocket, String text) {
     /**
      * Called once the full binary message has been built. {@link WebSocket#request()} must
      * be called to receive more messages.
+     *
+     * @param bytes which will not further used nor modified by the {@link HttpClient}
      */
     default void onMessage(WebSocket webSocket, ByteBuffer bytes) {
       webSocket.request();

kubernetes-client/src/main/java/io/fabric8/kubernetes/client/dsl/internal/ExecWebSocketListener.java

@@ -20,6 +20,7 @@
 import io.fabric8.kubernetes.api.model.Status;
 import io.fabric8.kubernetes.api.model.StatusCause;
 import io.fabric8.kubernetes.client.KubernetesClientException;
+import io.fabric8.kubernetes.client.StreamConsumer;
 import io.fabric8.kubernetes.client.dsl.ExecListener;
 import io.fabric8.kubernetes.client.dsl.ExecListener.Response;
 import io.fabric8.kubernetes.client.dsl.ExecWatch;
@@ -37,8 +38,6 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
-import java.nio.channels.Channels;
-import java.nio.channels.WritableByteChannel;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -51,6 +50,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
 
 /**
  * A {@link WebSocket.Listener} for exec operations.
@@ -106,7 +106,7 @@ public ListenerStream(String name) {
       this.name = name;
     }
 
-    private void handle(ByteBuffer byteString, WebSocket webSocket) throws IOException {
+    private void handle(ByteBuffer byteString, WebSocket webSocket) throws Exception {
       if (handler != null) {
         if (LOGGER.isDebugEnabled()) {
           LOGGER.debug("exec message received {} bytes on channel {}", byteString.remaining(), name);
@@ -178,37 +178,46 @@ private ListenerStream createStream(String name, StreamContext streamContext) {
     if (streamContext == null) {
       return stream;
     }
-    OutputStream os = streamContext.getOutputStream();
-    if (os == null) {
+    StreamConsumer consumer = streamContext.getConsumer();
+    if (consumer == null) {
       // redirecting
       stream.inputStream = new ExecWatchInputStream(() -> this.webSocketRef.get().request());
       this.exitCode.whenComplete(stream.inputStream::onExit);
       stream.handler = b -> stream.inputStream.consume(Arrays.asList(b));
     } else {
-      WritableByteChannel channel = Channels.newChannel(os);
-      stream.handler = b -> asyncWrite(channel, b);
+      stream.handler = b -> consume(consumer, b, streamContext.isBlocking() ? serialExecutor : Runnable::run,
+          this::postConsume);
     }
     return stream;
   }
 
-  private void asyncWrite(WritableByteChannel channel, ByteBuffer b) {
-    CompletableFuture.runAsync(() -> {
+  public static void consume(StreamConsumer consumer, ByteBuffer bytes, Executor executor, Consumer<Throwable> postConsume) {
+    CompletableFuture.supplyAsync(() -> {
       try {
-        channel.write(b);
-      } catch (IOException e) {
+        return consumer.consume(bytes);
+      } catch (Exception e) {
         throw KubernetesClientException.launderThrowable(e);
       }
-    }, serialExecutor).whenComplete((v, t) -> {
-      webSocketRef.get().request();
-      if (t != null) {
-        if (closed.get()) {
-          LOGGER.debug("Stream write failed after close", t);
-        } else {
-          // This could happen if the user simply closes their stream prior to completion
-          LOGGER.warn("Stream write failed", t);
-        }
+    }, executor)
+        .whenComplete((cs, t) -> {
+          if (cs != null) {
+            cs.whenComplete((v, t1) -> postConsume.accept(t1));
+          } else {
+            postConsume.accept(t);
+          }
+        });
+  }
+
+  private void postConsume(Throwable t) {
+    webSocketRef.get().request();
+    if (t != null) {
+      if (closed.get()) {
+        LOGGER.debug("Stream write failed after close", t);
+      } else {
+        // This could happen if the user simply closes their stream prior to completion
+        LOGGER.warn("Stream write failed", t);
       }
-    });
+    }
   }
 
   @Override
@@ -299,6 +308,7 @@ public void onError(WebSocket webSocket, Throwable t, boolean connectionError) {
   @Override
   public void onMessage(WebSocket webSocket, String text) {
     LOGGER.debug("Exec Web Socket: onMessage(String)");
+    // this is unexpected and will likely just result in an exception
     onMessage(webSocket, ByteBuffer.wrap(text.getBytes(StandardCharsets.UTF_8)));
   }
 
@@ -337,7 +347,7 @@ public void onMessage(WebSocket webSocket, ByteBuffer bytes) {
         default:
           throw new IOException("Unknown stream ID " + streamID);
       }
-    } catch (IOException e) {
+    } catch (Exception e) {
       throw KubernetesClientException.launderThrowable(e);
     } finally {
       if (close) {

kubernetes-client/src/main/java/io/fabric8/kubernetes/client/dsl/internal/LogWatchCallback.java

@@ -17,17 +17,17 @@
 
 import io.fabric8.kubernetes.client.KubernetesClientException;
 import io.fabric8.kubernetes.client.RequestConfigBuilder;
+import io.fabric8.kubernetes.client.StreamConsumer;
 import io.fabric8.kubernetes.client.dsl.LogWatch;
 import io.fabric8.kubernetes.client.http.AsyncBody;
+import io.fabric8.kubernetes.client.http.BufferUtil;
 import io.fabric8.kubernetes.client.http.HttpClient;
 import io.fabric8.kubernetes.client.http.HttpRequest;
 import io.fabric8.kubernetes.client.utils.internal.SerialExecutor;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.io.IOException;
 import java.io.InputStream;
-import java.io.OutputStream;
 import java.net.URL;
 import java.nio.ByteBuffer;
 import java.nio.channels.Channels;
@@ -41,20 +41,18 @@ public class LogWatchCallback implements LogWatch, AutoCloseable {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(LogWatchCallback.class);
 
-  private final OutputStream out;
-  private WritableByteChannel outChannel;
+  private final StreamConsumer consumer;
+  private final boolean blocking;
   private final OperationContext context;
   private volatile InputStream output;
 
   private final AtomicBoolean closed = new AtomicBoolean(false);
   private final CompletableFuture<AsyncBody> asyncBody = new CompletableFuture<>();
   private final SerialExecutor serialExecutor;
 
-  public LogWatchCallback(OutputStream out, OperationContext context) {
-    this.out = out;
-    if (out != null) {
-      outChannel = Channels.newChannel(out);
-    }
+  public LogWatchCallback(StreamConsumer consumer, boolean blocking, OperationContext context) {
+    this.consumer = consumer;
+    this.blocking = blocking;
     this.context = context;
     this.serialExecutor = new SerialExecutor(context.getExecutor());
   }
@@ -79,7 +77,7 @@ public LogWatchCallback callAndWait(HttpClient client, URL url) {
             .withRequestTimeout(0).build())
         .readTimeout(0, TimeUnit.MILLISECONDS).build();
 
-    if (out == null) {
+    if (consumer == null) {
       // we can pass the input stream directly to the consumer
       clone.sendAsync(request, InputStream.class).whenComplete((r, e) -> {
         if (e != null) {
@@ -91,45 +89,40 @@ public LogWatchCallback callAndWait(HttpClient client, URL url) {
       }).join();
     } else {
       // we need to write the bytes to the given output
-      // we don't know if the write will be blocking, so hand it off to another thread
-      clone.consumeBytes(request, (buffers, a) -> CompletableFuture.runAsync(() -> {
-        for (ByteBuffer byteBuffer : buffers) {
-          try {
-            outChannel.write(byteBuffer);
-          } catch (IOException e1) {
-            throw KubernetesClientException.launderThrowable(e1);
-          }
-        }
-      }, serialExecutor).whenComplete((v, t) -> {
-        if (t != null) {
-          a.cancel();
-          onFailure(t);
-        } else if (!closed.get()) {
-          a.consume();
-        } else {
-          a.cancel();
-        }
-      })).whenComplete((a, e) -> {
-        if (e != null) {
-          onFailure(e);
-        }
-        if (a != null) {
-          asyncBody.complete(a.body());
-          a.body().consume();
-          a.body().done().whenComplete((v, t) -> CompletableFuture.runAsync(() -> {
-            if (t != null) {
-              onFailure(t);
-            } else {
-              cleanUp();
+      clone.consumeBytes(request, (buffers, a) -> ExecWebSocketListener.consume(consumer,
+          ByteBuffer.wrap(BufferUtil.toArray(buffers)), blocking ? serialExecutor : Runnable::run, (t) -> postConsume(a, t)))
+          .whenComplete((a, e) -> {
+            if (e != null) {
+              onFailure(e);
             }
-          }, serialExecutor));
-        }
-      });
+            if (a != null) {
+              asyncBody.complete(a.body());
+              a.body().consume();
+              a.body().done().whenComplete((v, t) -> CompletableFuture.runAsync(() -> {
+                if (t != null) {
+                  onFailure(t);
+                } else {
+                  cleanUp();
+                }
+              }, serialExecutor));
+            }
+          });
     }
 
     return this;
   }
 
+  private void postConsume(AsyncBody a, Throwable t) {
+    if (t != null) {
+      a.cancel();
+      onFailure(t);
+    } else if (!closed.get()) {
+      a.consume();
+    } else {
+      a.cancel();
+    }
+  }
+
   @Override
   public InputStream getOutput() {
     return output;