SAMZA-2055: Async high level api

https://cwiki.apache.org/confluence/display/SAMZA/SEP-21%3A+Samza+Async+API+for+High+Level

Author: mynameborat <bharath.kumarasubramanian@gmail.com>

Reviewers: Jagadish <jagadish@apache.org>

Closes apache#905 from mynameborat/async-high-level-api

Author: mynameborat

samza-api/src/main/java/org/apache/samza/operators/MessageStream.java

@@ -22,7 +22,9 @@
 import java.util.ArrayList;
 import java.util.Collection;
 
+import java.util.concurrent.CompletionStage;
 import org.apache.samza.annotation.InterfaceStability;
+import org.apache.samza.operators.functions.AsyncFlatMapFunction;
 import org.apache.samza.operators.functions.FilterFunction;
 import org.apache.samza.operators.functions.FlatMapFunction;
 import org.apache.samza.operators.functions.JoinFunction;
@@ -67,6 +69,31 @@ public interface MessageStream<M> {
    */
   <OM> MessageStream<OM> flatMap(FlatMapFunction<? super M, ? extends OM> flatMapFn);
 
+  /**
+   * Applies the provided 1:n transformation asynchronously to this {@link MessageStream}. The asynchronous transformation
+   * is specified through {@link AsyncFlatMapFunction}. The results are emitted to the downstream operators upon the
+   * completion of the {@link CompletionStage} returned from the {@link AsyncFlatMapFunction}.
+   * <p>
+   * The operator can operate in two modes depending on <i>task.max.concurrency.</i>.
+   * <ul>
+   *   <li>
+   *     Serialized (task.max.concurrency=1) - In this mode, each invocation of the {@link AsyncFlatMapFunction} is guaranteed
+   *     to happen-before next invocation.
+   *   </li>
+   *   <li>
+   *     Parallel (task.max.concurrency&gt;1) - In this mode, multiple invocations can happen in parallel without happens-before guarantee
+   *     and the {@link AsyncFlatMapFunction} is required synchronize any shared state. The operator doesn't provide any ordering guarantees.
+   *     i.e The results corresponding to each invocation of this operator might not be emitted in the same order as invocations.
+   *     By extension, the operator chain that follows it also doesn't have any ordering guarantees.
+   *   </li>
+   * </ul>
+   *
+   * @param asyncFlatMapFn the async function to transform a message to zero or more messages
+   * @param <OM> the type of messages in the transformed {@link MessageStream}
+   * @return the transformed {@link MessageStream}
+   */
+  <OM> MessageStream<OM> flatMapAsync(AsyncFlatMapFunction<? super M, ? extends OM> asyncFlatMapFn);
+
   /**
    * Applies the provided function to messages in this {@link MessageStream} and returns the
    * filtered {@link MessageStream}.

samza-api/src/main/java/org/apache/samza/operators/functions/AsyncFlatMapFunction.java

@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.samza.operators.functions;
+
+import java.io.Serializable;
+import java.util.Collection;
+import java.util.concurrent.CompletionStage;
+import org.apache.samza.SamzaException;
+import org.apache.samza.annotation.InterfaceStability;
+
+
+/**
+ * Asynchronous variant of the {@link FlatMapFunction} used in tandem with {@link org.apache.samza.operators.MessageStream#flatMapAsync(AsyncFlatMapFunction)}
+ * to transform a collection of 0 or more messages.
+ * <p>
+ * Typically, {@link AsyncFlatMapFunction} is used for describing complex transformations that involve IO operations or remote calls.
+ * The following pseudo code demonstrates a sample implementation of {@link AsyncFlatMapFunction} that sends out an email
+ * and returns the status asynchronously.
+ * <pre> {@code
+ *    AsyncFlatMapFunction<Email, Status> asyncEmailSender = (Email message) -> {
+ *        ...
+ *
+ *        Request<Email> emailRequest = buildEmailRequest(message);
+ *        Future<EmailResponse> emailResponseFuture = emailClient.sendRequest(emailRequest); // send email asynchronously
+ *        ...
+ *
+ *        return new CompletableFuture<>(emailResponseFuture)
+ *             .thenApply(response -> fetchStatus(response);
+ *    }
+ * }
+ * </pre>
+ *
+ * <p>
+ *   The function needs to be thread safe in case of task.max.concurrency&gt;1. It also needs to coordinate any
+ *   shared state since happens-before is not guaranteed between the messages delivered to the function. Refer to
+ *   {@link org.apache.samza.operators.MessageStream#flatMapAsync(AsyncFlatMapFunction)} docs for more details on the modes
+ *   and guarantees.
+ *
+ * <p>
+ *   For each invocation, the {@link CompletionStage} returned by the function should be completed successfully/exceptionally
+ *   within task.callback.timeout.ms; failure to do so will result in {@link SamzaException} bringing down the application.
+ *
+ * @param <M>  type of the input message
+ * @param <OM>  type of the transformed messages
+ */
+@InterfaceStability.Unstable
+@FunctionalInterface
+public interface AsyncFlatMapFunction<M, OM> extends InitableFunction, ClosableFunction, Serializable {
+
+  /**
+   * Transforms the provided message into a collection of 0 or more messages.
+   *
+   * @param message  the input message to be transformed
+   * @return  a {@link CompletionStage} of a {@link Collection} of transformed messages
+   */
+  CompletionStage<Collection<OM>> apply(M message);
+}

samza-core/src/main/java/org/apache/samza/operators/MessageStreamImpl.java

@@ -25,12 +25,14 @@
 
 import org.apache.samza.SamzaException;
 import org.apache.samza.application.descriptors.StreamApplicationDescriptorImpl;
+import org.apache.samza.operators.functions.AsyncFlatMapFunction;
 import org.apache.samza.operators.functions.FilterFunction;
 import org.apache.samza.operators.functions.FlatMapFunction;
 import org.apache.samza.operators.functions.JoinFunction;
 import org.apache.samza.operators.functions.MapFunction;
 import org.apache.samza.operators.functions.SinkFunction;
 import org.apache.samza.operators.functions.StreamTableJoinFunction;
+import org.apache.samza.operators.spec.AsyncFlatMapOperatorSpec;
 import org.apache.samza.operators.spec.BroadcastOperatorSpec;
 import org.apache.samza.operators.spec.JoinOperatorSpec;
 import org.apache.samza.operators.spec.OperatorSpec;
@@ -102,6 +104,14 @@ public <TM> MessageStream<TM> flatMap(FlatMapFunction<? super M, ? extends TM> f
     return new MessageStreamImpl<>(this.streamAppDesc, op);
   }
 
+  @Override
+  public <OM> MessageStream<OM> flatMapAsync(AsyncFlatMapFunction<? super M, ? extends OM> flatMapFn) {
+    String opId = this.streamAppDesc.getNextOpId(OpCode.ASYNC_FLAT_MAP);
+    AsyncFlatMapOperatorSpec<M, OM> op = OperatorSpecs.createAsyncOperatorSpec(flatMapFn, opId);
+    this.operatorSpec.registerNextOperatorSpec(op);
+    return new MessageStreamImpl<>(this.streamAppDesc, op);
+  }
+
   @Override
   public void sink(SinkFunction<? super M> sinkFn) {
     String opId = this.streamAppDesc.getNextOpId(OpCode.SINK);

samza-core/src/main/java/org/apache/samza/operators/impl/AsyncFlatmapOperatorImpl.java

@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.samza.operators.impl;
+
+import java.util.Collection;
+import java.util.concurrent.CompletionStage;
+import org.apache.samza.context.Context;
+import org.apache.samza.operators.functions.AsyncFlatMapFunction;
+import org.apache.samza.operators.spec.AsyncFlatMapOperatorSpec;
+import org.apache.samza.operators.spec.OperatorSpec;
+import org.apache.samza.task.MessageCollector;
+import org.apache.samza.task.TaskCoordinator;
+
+
+public class AsyncFlatmapOperatorImpl<M, RM> extends OperatorImpl<M, RM> {
+  private final AsyncFlatMapOperatorSpec<M, RM> opSpec;
+  private final AsyncFlatMapFunction<M, RM> transformFn;
+
+  AsyncFlatmapOperatorImpl(AsyncFlatMapOperatorSpec<M, RM> opSpec) {
+    this.opSpec = opSpec;
+    this.transformFn = opSpec.getTransformFn();
+  }
+  @Override
+  protected void handleInit(Context context) {
+    this.transformFn.init(context);
+  }
+
+  @Override
+  protected CompletionStage<Collection<RM>> handleMessageAsync(M message, MessageCollector collector,
+      TaskCoordinator coordinator) {
+    return transformFn.apply(message);
+  }
+
+  @Override
+  protected void handleClose() {
+  }
+
+  @Override
+  protected OperatorSpec<M, RM> getOperatorSpec() {
+    return opSpec;
+  }
+}

samza-core/src/main/java/org/apache/samza/operators/impl/BroadcastOperatorImpl.java

@@ -19,6 +19,8 @@
 
 package org.apache.samza.operators.impl;
 
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
 import org.apache.samza.context.Context;
 import org.apache.samza.operators.spec.BroadcastOperatorSpec;
 import org.apache.samza.operators.spec.OperatorSpec;
@@ -50,9 +52,10 @@ protected void handleInit(Context context) {
   }
 
   @Override
-  protected Collection<Void> handleMessage(M message, MessageCollector collector, TaskCoordinator coordinator) {
+  protected CompletionStage<Collection<Void>> handleMessageAsync(M message, MessageCollector collector,
+      TaskCoordinator coordinator) {
     collector.send(new OutgoingMessageEnvelope(systemStream, 0, null, message));
-    return Collections.emptyList();
+    return CompletableFuture.completedFuture(Collections.emptyList());
   }
 
   @Override

samza-core/src/main/java/org/apache/samza/operators/impl/StreamOperatorImpl.java renamed to samza-core/src/main/java/org/apache/samza/operators/impl/FlatmapOperatorImpl.java

@@ -18,6 +18,8 @@
  */
 package org.apache.samza.operators.impl;
 
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
 import org.apache.samza.context.Context;
 import org.apache.samza.operators.functions.FlatMapFunction;
 import org.apache.samza.operators.spec.OperatorSpec;
@@ -34,12 +36,12 @@
  * @param <M>  the type of input message
  * @param <RM>  the type of result
  */
-class StreamOperatorImpl<M, RM> extends OperatorImpl<M, RM> {
+class FlatmapOperatorImpl<M, RM> extends OperatorImpl<M, RM> {
 
   private final StreamOperatorSpec<M, RM> streamOpSpec;
   private final FlatMapFunction<M, RM> transformFn;
 
-  StreamOperatorImpl(StreamOperatorSpec<M, RM> streamOpSpec) {
+  FlatmapOperatorImpl(StreamOperatorSpec<M, RM> streamOpSpec) {
     this.streamOpSpec = streamOpSpec;
     this.transformFn = streamOpSpec.getTransformFn();
   }
@@ -50,9 +52,9 @@ protected void handleInit(Context context) {
   }
 
   @Override
-  public Collection<RM> handleMessage(M message, MessageCollector collector,
+  protected CompletionStage<Collection<RM>> handleMessageAsync(M message, MessageCollector collector,
       TaskCoordinator coordinator) {
-    return this.transformFn.apply(message);
+    return CompletableFuture.completedFuture(this.transformFn.apply(message));
   }
 
   @Override

samza-core/src/main/java/org/apache/samza/operators/impl/InputOperatorImpl.java

@@ -18,6 +18,9 @@
  */
 package org.apache.samza.operators.impl;
 
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
 import org.apache.samza.context.Context;
 import org.apache.samza.operators.KV;
 import org.apache.samza.system.descriptors.InputTransformer;
@@ -47,18 +50,21 @@ protected void handleInit(Context context) {
   }
 
   @Override
-  public Collection<Object> handleMessage(IncomingMessageEnvelope ime, MessageCollector collector, TaskCoordinator coordinator) {
-    Object message;
+  protected CompletionStage<Collection<Object>> handleMessageAsync(IncomingMessageEnvelope message,
+      MessageCollector collector, TaskCoordinator coordinator) {
+    Object result;
     InputTransformer transformer = inputOpSpec.getTransformer();
     if (transformer != null) {
-      message = transformer.apply(ime);
+      result = transformer.apply(message);
     } else {
-      message = this.inputOpSpec.isKeyed() ? KV.of(ime.getKey(), ime.getMessage()) : ime.getMessage();
+      result = this.inputOpSpec.isKeyed() ? KV.of(message.getKey(), message.getMessage()) : message.getMessage();
     }
-    if (message != null) {
-      return Collections.singletonList(message);
-    }
-    return Collections.emptyList();
+
+    Collection<Object> output = Optional.ofNullable(result)
+        .map(Collections::singletonList)
+        .orElse(Collections.emptyList());
+
+    return CompletableFuture.completedFuture(output);
   }
 
   @Override