content: package io.activej.dataflow.stream;
  
  import io.activej.dataflow.DataflowClient;
  import io.activej.dataflow.DataflowServer;
  import io.activej.dataflow.collector.ICollector;
  import io.activej.dataflow.collector.MergeCollector;
  import io.activej.dataflow.dataset.Dataset;
  import io.activej.dataflow.graph.DataflowGraph;
  import io.activej.dataflow.graph.Partition;
  import io.activej.dataflow.inject.DatasetIdModule;
  import io.activej.dataflow.node.StreamSorterStorageFactory;
  import io.activej.datastream.consumer.ToListStreamConsumer;
  import io.activej.datastream.processor.reducer.ReducerToAccumulator;
  import io.activej.datastream.processor.reducer.ReducerToResult;
  import io.activej.datastream.supplier.StreamSupplier;
  import io.activej.inject.Injector;
  import io.activej.inject.Key;
  import io.activej.inject.module.Module;
  import io.activej.inject.module.ModuleBuilder;
  import io.activej.reactor.Reactor;
  import io.activej.serializer.annotations.Deserialize;
  import io.activej.serializer.annotations.Serialize;
  import io.activej.serializer.stream.StreamCodec;
  import io.activej.serializer.stream.StreamCodecs;
  import io.activej.test.rules.ByteBufRule;
  import io.activej.test.rules.ClassBuilderConstantsRule;
  import io.activej.test.rules.EventloopRule;
  import org.junit.*;
  
  import java.net.InetSocketAddress;
  import java.util.*;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.function.Function;
  
  import static io.activej.dataflow.codec.SubtypeImpl.subtype;
  import static io.activej.dataflow.dataset.Datasets.*;
  import static io.activej.dataflow.graph.StreamSchemas.simple;
  import static io.activej.dataflow.helper.MergeStubStreamSorterStorage.FACTORY_STUB;
  import static io.activej.dataflow.inject.DatasetIdImpl.datasetId;
  import static io.activej.dataflow.stream.DataflowTest.*;
  import static io.activej.promise.TestUtils.await;
  import static io.activej.test.TestUtils.assertCompleteFn;
  import static java.util.Comparator.naturalOrder;
  import static org.junit.Assert.assertEquals;
  
  public class MapReduceTest {
  
  	@ClassRule
  	public static final EventloopRule eventloopRule = new EventloopRule();
  
  	@ClassRule
  	public static final ByteBufRule byteBufRule = new ByteBufRule();
  
  	@Rule
  	public final ClassBuilderConstantsRule classBuilderConstantsRule = new ClassBuilderConstantsRule();
  
  	private ExecutorService executor;
  	private ExecutorService sortingExecutor;
  
  	@Before
  	public void setUp() {
  		executor = Executors.newSingleThreadExecutor();
  		sortingExecutor = Executors.newSingleThreadExecutor();
  	}
  
  	@After
  	public void tearDown() {
  		executor.shutdownNow();
  		sortingExecutor.shutdownNow();
  	}
  
  	public static class StringCount {
  		@Serialize
  		public final String s;
  		@Serialize
  		public int count;
  
  		public StringCount(@Deserialize("s") String s, @Deserialize("count") int count) {
  			this.s = s;
  			this.count = count;
  		}
  
  		@Override
  		public boolean equals(Object o) {
  			if (this == o) return true;
  			if (!(o instanceof StringCount that)) return false;
  			return s.equals(that.s) && count == that.count;
  		}
  
  		@Override
  		public int hashCode() {
  			return Objects.hash(s, count);
  		}
  
  		@Override
  		public String toString() {
  			return "StringCount{s='" + s + '\'' + ", count=" + count + '}';
  		}
  	}
  
  	@Test
  	public void test() throws Exception {
  
  		InetSocketAddress address1 = getFreeListenAddress();
  		InetSocketAddress address2 = getFreeListenAddress();
  
  		Module common = createCommon(List.of(new Partition(address1), new Partition(address2)))
  			.install(createSerializersModule())
  			.bind(StreamSorterStorageFactory.class).toInstance(FACTORY_STUB)
  			.build();
  
  		Module serverCommon = createCommonServer(common, executor, sortingExecutor);
  		Module serverModule1 = ModuleBuilder.create()
  			.install(serverCommon)
  			.install(DatasetIdModule.create())
  			.bind(Integer.class, "dataflowPort").toInstance(address1.getPort())
  			.bind(datasetId("items")).toInstance(List.of(
  				"dog",
  				"cat",
  				"horse",
  				"cat"))
  			.build();
  
  		Module serverModule2 = ModuleBuilder.create()
  			.install(serverCommon)
  			.install(DatasetIdModule.create())
  			.bind(Integer.class, "dataflowPort").toInstance(address2.getPort())
  			.bind(datasetId("items")).toInstance(List.of(
  				"dog",
  				"cat"))
  			.build();
  
  		DataflowServer server1 = Injector.of(serverModule1).getInstance(DataflowServer.class);
  		DataflowServer server2 = Injector.of(serverModule2).getInstance(DataflowServer.class);
  
  		server1.listen();
  		server2.listen();
  
  		Injector clientInjector = Injector.of(createCommonClient(common));
  		DataflowClient client = clientInjector.getInstance(DataflowClient.class);
  		DataflowGraph graph = clientInjector.getInstance(DataflowGraph.class);
  
  		Dataset<String> items = datasetOfId("items", simple(String.class));
  		Dataset<StringCount> mappedItems = map(items, new StringMapFunction(), simple(StringCount.class));
  		Dataset<StringCount> reducedItems = sortReduceRepartitionReduce(mappedItems,
  			new StringReducer(), String.class, new StringKeyFunction(), Comparator.naturalOrder());
  		ICollector<StringCount> collector = MergeCollector.create(Reactor.getCurrentReactor(), reducedItems, client, new StringKeyFunction(), naturalOrder());
  		StreamSupplier<StringCount> resultSupplier = collector.compile(graph);
  		ToListStreamConsumer<StringCount> resultConsumer = ToListStreamConsumer.create();
  
  		resultSupplier.streamTo(resultConsumer).whenComplete(assertCompleteFn());
  
  		await(graph.execute()
  			.whenComplete(assertCompleteFn($ -> {
  				server1.close();
  				server2.close();
  			})));
  
  		System.out.println(resultConsumer.getList());
  
  		assertEquals(Set.of(
  			new StringCount("cat", 3),
  			new StringCount("dog", 2),
  			new StringCount("horse", 1)), new HashSet<>(resultConsumer.getList()));
  	}
  
  	public static class StringReducer extends ReducerToAccumulator<String, StringCount, StringCount> {
  		@Override
  		public StringCount createAccumulator(String key) {
  			return new StringCount(key, 0);
  		}
  
  		@Override
  		public StringCount accumulate(StringCount accumulator, StringCount value) {
  			accumulator.count += value.count;
  			return accumulator;
  		}
  
  		@Override
  		public StringCount combine(StringCount accumulator, StringCount anotherAccumulator) {
  			accumulator.count += anotherAccumulator.count;
  			return accumulator;
  		}
  	}
  
  	public static class StringMapFunction implements Function<String, StringCount> {
  		@Override
  		public StringCount apply(String s) {
  			return new StringCount(s, 1);
  		}
  	}
  
  	public static class StringKeyFunction implements Function<StringCount, String> {
  		@Override
  		public String apply(StringCount stringCount) {
  			return stringCount.s;
  		}
  	}
  
  	@SuppressWarnings("rawtypes")
  	private static Module createSerializersModule() {
  		return ModuleBuilder.create()
  			.bind(new Key<StreamCodec<StringKeyFunction>>(subtype(0)) {}).toInstance(StreamCodecs.singleton(new StringKeyFunction()))
  			.bind(new Key<StreamCodec<StringMapFunction>>(subtype(1)) {}).toInstance(StreamCodecs.singleton(new StringMapFunction()))
  			.bind(new Key<StreamCodec<Comparator<?>>>() {}).toInstance(StreamCodecs.singleton(Comparator.naturalOrder()))
  			.bind(new Key<StreamCodec<ReducerToResult>>() {}).toInstance(StreamCodecs.singleton(new StringReducer()))
  			.build();
  	}
  }
  