Add PartitonHub, #21880

* FixedSizePartitionHub in Artery
* expose consumer queue size

akka-docs/src/test/java/jdocs/stream/HubDocTest.java

@@ -11,17 +11,25 @@ import akka.japi.Pair;
 import akka.stream.ActorMaterializer;
 import akka.stream.KillSwitches;
 import akka.stream.Materializer;
+import akka.stream.ThrottleMode;
 import akka.stream.UniqueKillSwitch;
 import akka.stream.javadsl.*;
+import akka.stream.javadsl.PartitionHub.ConsumerInfo;
+
 import jdocs.AbstractJavaTest;
 import akka.testkit.javadsl.TestKit;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
+import scala.concurrent.duration.Duration;
 import scala.concurrent.duration.FiniteDuration;
 
+import java.util.List;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.TimeUnit;
+import java.util.function.BiFunction;
+import java.util.function.Supplier;
+import java.util.function.ToLongBiFunction;
 
 public class HubDocTest extends AbstractJavaTest {
 
@@ -137,4 +145,136 @@ public class HubDocTest extends AbstractJavaTest {
     killSwitch.shutdown();
     //#pub-sub-4
   }
+  
+  @Test
+  public void dynamicPartition() {
+    // Used to be able to clean up the running stream
+    ActorMaterializer materializer = ActorMaterializer.create(system);
+
+    //#partition-hub
+    // A simple producer that publishes a new "message-n" every second
+    Source<String, Cancellable> producer = Source.tick(
+      FiniteDuration.create(1, TimeUnit.SECONDS),
+      FiniteDuration.create(1, TimeUnit.SECONDS),
+      "message"
+    ).zipWith(Source.range(0, 100), (a, b) -> a + "-" + b);
+
+    // Attach a PartitionHub Sink to the producer. This will materialize to a
+    // corresponding Source.
+    // (We need to use toMat and Keep.right since by default the materialized
+    // value to the left is used)
+    RunnableGraph<Source<String, NotUsed>> runnableGraph =
+      producer.toMat(PartitionHub.of(
+          String.class, 
+          (size, elem) -> Math.abs(elem.hashCode()) % size,
+          2, 256), Keep.right());
+
+    // By running/materializing the producer, we get back a Source, which
+    // gives us access to the elements published by the producer.
+    Source<String, NotUsed> fromProducer = runnableGraph.run(materializer);
+
+    // Print out messages from the producer in two independent consumers
+    fromProducer.runForeach(msg -> System.out.println("consumer1: " + msg), materializer);
+    fromProducer.runForeach(msg -> System.out.println("consumer2: " + msg), materializer);
+    //#partition-hub
+    
+    // Cleanup
+    materializer.shutdown();
+  }
+  
+  //#partition-hub-stateful-function
+  // Using a class since variable must otherwise be final.
+  // New instance is created for each materialization of the PartitionHub.
+  static class RoundRobin<T> implements ToLongBiFunction<ConsumerInfo, T> {
+
+    private long i = -1;
+    
+    @Override
+    public long applyAsLong(ConsumerInfo info, T elem) {
+      i++;
+      return info.consumerIdByIdx((int) (i % info.size()));
+    }
+  }
+  //#partition-hub-stateful-function
+  
+  @Test
+  public void dynamicStatefulPartition() {
+    // Used to be able to clean up the running stream
+    ActorMaterializer materializer = ActorMaterializer.create(system);
+
+    //#partition-hub-stateful
+    // A simple producer that publishes a new "message-n" every second
+    Source<String, Cancellable> producer = Source.tick(
+      FiniteDuration.create(1, TimeUnit.SECONDS),
+      FiniteDuration.create(1, TimeUnit.SECONDS),
+      "message"
+    ).zipWith(Source.range(0, 100), (a, b) -> a + "-" + b);
+    
+    // Attach a PartitionHub Sink to the producer. This will materialize to a
+    // corresponding Source.
+    // (We need to use toMat and Keep.right since by default the materialized
+    // value to the left is used)
+    RunnableGraph<Source<String, NotUsed>> runnableGraph =
+      producer.toMat(
+        PartitionHub.ofStateful(
+          String.class,
+          () -> new RoundRobin<String>(),
+          2, 
+          256),
+        Keep.right());
+
+    // By running/materializing the producer, we get back a Source, which
+    // gives us access to the elements published by the producer.
+    Source<String, NotUsed> fromProducer = runnableGraph.run(materializer);
+
+    // Print out messages from the producer in two independent consumers
+    fromProducer.runForeach(msg -> System.out.println("consumer1: " + msg), materializer);
+    fromProducer.runForeach(msg -> System.out.println("consumer2: " + msg), materializer);
+    //#partition-hub-stateful
+    
+    // Cleanup
+    materializer.shutdown();
+  }
+  
+  @Test
+  public void dynamicFastestPartition() {
+    // Used to be able to clean up the running stream
+    ActorMaterializer materializer = ActorMaterializer.create(system);
+
+    //#partition-hub-fastest
+    Source<Integer, NotUsed> producer = Source.range(0, 100);
+
+    // ConsumerInfo.queueSize is the approximate number of buffered elements for a consumer.
+    // Note that this is a moving target since the elements are consumed concurrently.
+    RunnableGraph<Source<Integer, NotUsed>> runnableGraph =
+      producer.toMat(
+        PartitionHub.ofStateful(
+          Integer.class,
+          () -> (info, elem) -> {
+            final List<Object> ids = info.getConsumerIds();
+            int minValue = info.queueSize(0);
+            long fastest = info.consumerIdByIdx(0);
+            for (int i = 1; i < ids.size(); i++) {
+              int value = info.queueSize(i);
+              if (value < minValue) {
+                  minValue = value;
+                  fastest = info.consumerIdByIdx(i);
+              }
+            }
+            return fastest;
+          },
+          2, 
+          8),
+        Keep.right());
+
+    Source<Integer, NotUsed> fromProducer = runnableGraph.run(materializer);
+
+    fromProducer.runForeach(msg -> System.out.println("consumer1: " + msg), materializer);
+    fromProducer.throttle(10, Duration.create(100, TimeUnit.MILLISECONDS), 10, ThrottleMode.shaping())
+      .runForeach(msg -> System.out.println("consumer2: " + msg), materializer);
+    //#partition-hub-fastest
+
+    // Cleanup
+    materializer.shutdown();
+  }
 }