akka.stream.javadsl.Flow Java Examples

private SourceQueueWithComplete<T> getSourceQueue(final ActorMaterializer materializer) {
    // Log stream completion and failure at level ERROR because the stream is supposed to survive forever.
    final Attributes streamLogLevels =
            Attributes.logLevels(Attributes.logLevelDebug(), Attributes.logLevelError(),
                    Attributes.logLevelError());

    return Source.<T>queue(getBufferSize(), OverflowStrategy.dropNew())
            .map(this::incrementDequeueCounter)
            .log("graph-actor-stream-1-dequeued", logger)
            .withAttributes(streamLogLevels)
            .via(Flow.fromFunction(this::beforeProcessMessage))
            .log("graph-actor-stream-2-preprocessed", logger)
            .withAttributes(streamLogLevels)
            .via(processMessageFlow())
            .log("graph-actor-stream-3-processed", logger)
            .withAttributes(streamLogLevels)
            .to(processedMessageSink())
            .run(materializer);
}
 

/**
 * Create a flow that delays its elements with exponential backoff and reset after no element passes through the
 * stream for a fixed period of time. Stream elements represent failures and are logged as errors.
 * <ul>
 * <li>Emits when: upstream emits.</li>
 * <li>Completes when: upstream completes.</li>
 * <li>Cancels when: never - this is a part of a feedback loop.</li>
 * <li>Fails when: {@code maxRestarts} elements pass through the stream.</li>
 * </ul>
 *
 * @param minBackoff Minimum backoff duration.
 * @param maxBackoff Maximum backoff duration.
 * @param maxRestarts Maximum number of tolerated failures. Tolerate no failure if 0. Tolerate arbitrarily many
 * failures if negative.
 * @param recovery The period after which backoff is reset to {@code minBackoff} if no failure occurred.
 * @param <E> Type of stream elements.
 * @return A delayed flow.
 */
private static <E> Flow<E, E, NotUsed> backoff(final Duration minBackoff, final Duration maxBackoff,
        final int maxRestarts, final Duration recovery) {
    final Flow<E, E, NotUsed> neverCancelFlowWithErrorLogging = Flow.<E>create()
            .log("resume-source-errors-flow")
            .withAttributes(logLevels(logLevelInfo(), logLevelDebug(), logLevelInfo()))
            .via(new NeverCancelFlow<>());

    final Flow<E, E, NotUsed> upstream = maxRestarts < 0
            ? neverCancelFlowWithErrorLogging
            : maxRestarts == 0
            ? neverCancelFlowWithErrorLogging.flatMapConcat(ResumeSource::failWithLimitsReached)
            : neverCancelFlowWithErrorLogging.limit(maxRestarts);

    return upstream.statefulMapConcat(() -> new StatefulBackoffFunction<>(minBackoff, maxBackoff, recovery))
            .flatMapConcat(pair ->
                    Source.single(pair.first())
                            .delay(pair.second(), OverflowStrategy.backpressure())
            );
}
 

private Source<SudoRetrieveThingResponse, NotUsed> sudoRetrieveThing(final ThingId thingId) {
    final SudoRetrieveThing command =
            SudoRetrieveThing.withOriginalSchemaVersion(thingId, DittoHeaders.empty());
    final CompletionStage<Source<SudoRetrieveThingResponse, NotUsed>> responseFuture =
            // using default thread-pool for asking Things shard region
            Patterns.ask(thingsShardRegion, command, thingsTimeout)
                    .handle((response, error) -> {
                        if (response instanceof SudoRetrieveThingResponse) {
                            return Source.single((SudoRetrieveThingResponse) response);
                        } else {
                            if (error != null) {
                                log.error("Failed " + command, error);
                            } else if (!(response instanceof ThingNotAccessibleException)) {
                                log.error("Unexpected response for <{}>: <{}>", command, response);
                            }
                            return Source.empty();
                        }
                    });

    return Source.fromSourceCompletionStage(responseFuture)
            .viaMat(Flow.create(), Keep.none());
}
 

private Source<Entry<Enforcer>, NotUsed> readCachedEnforcer(final Metadata metadata,
        final EntityIdWithResourceType policyId, final int iteration) {

    final Source<Entry<Enforcer>, ?> lazySource = Source.lazily(() -> {
        final CompletionStage<Source<Entry<Enforcer>, NotUsed>> enforcerFuture = policyEnforcerCache.get(policyId)
                .thenApply(optionalEnforcerEntry -> {
                    if (shouldReloadCache(optionalEnforcerEntry.orElse(null), metadata, iteration)) {
                        // invalid entry; invalidate and retry after delay
                        policyEnforcerCache.invalidate(policyId);
                        return readCachedEnforcer(metadata, policyId, iteration + 1)
                                .initialDelay(cacheRetryDelay);
                    } else {
                        return optionalEnforcerEntry.map(Source::single)
                                .orElse(ENFORCER_NONEXISTENT);
                    }
                })
                .exceptionally(error -> {
                    log.error("Failed to read policyEnforcerCache", error);
                    return ENFORCER_NONEXISTENT;
                });

        return Source.fromSourceCompletionStage(enforcerFuture);
    });

    return lazySource.viaMat(Flow.create(), Keep.none());
}
 

/**
 * Create a flow from seeds to either elements or failure with reason and the previous N elements that passed
 * through the stream before the failure. For maximum information, even elements emitted 2 or more failures ago
 * are kept in the look-behind queue.
 * <ul>
 * <li>Emits when: stream created by {@code resume} emits or fails.</li>
 * <li>Completes when: a stream created by {@code resume} completes.</li>
 * <li>Cancels when: downstream cancels.</li>
 * <li>Fails when: never.</li>
 * </ul>
 *
 * @param resume Creator of a stream of elements from a resumption seed.
 * @param lookBehind How many elements to keep in memory to create the next seed on failure.
 * @param <S> Type of seeds.
 * @param <E> Type of elements.
 * @return A never-failing flow.
 */
private static <S, E> Flow<S, Either<FailureWithLookBehind<E>, E>, NotUsed> resumeWithFailuresAppended(
        final Function<S, Source<E, ?>> resume, final int lookBehind,
        final Function<Throwable, Optional<Throwable>> resumeOrMapError) {

    return Flow.<S>create()
            .flatMapConcat(seed -> resume.apply(seed)
                    .<Envelope<E>>map(Element::new)
                    .concat(Source.single(new EndOfStream<>()))
                    .recoverWithRetries(1,
                            new PFBuilder<Throwable, Graph<SourceShape<Envelope<E>>, NotUsed>>()
                                    .matchAny(error -> resumeOrMapError.apply(error)
                                            .<Source<Envelope<E>, NotUsed>>map(Source::failed)
                                            .orElseGet(() -> Source.single(new Error<>(error)))
                                    )
                                    .build())
            )
            .via(new EndStreamOnEOS<>())
            .buffer(1, OverflowStrategy.backpressure())
            .statefulMapConcat(() -> new StatefulLookBehindFunction<>(lookBehind));
}
 

@SuppressWarnings("unchecked") // java 8 can't handle graph DSL types
private static <A, B, C, D> Graph<FlowShape<A, C>, NotUsed> assembleFlows(
        final Flow<A, B, NotUsed> stage1Flow,
        final Flow<B, C, NotUsed> stage2Flow,
        final Flow<B, D, NotUsed> sideChannelFlow,
        final Flow<Pair<C, D>, C, NotUsed> resultProcessorFlow) {

    return GraphDSL.create(builder -> {
        final FlowShape<A, B> stage1 = builder.add(stage1Flow);
        final FlowShape<B, C> stage2 = builder.add(stage2Flow);
        final FlowShape<B, D> sideChannel = builder.add(sideChannelFlow);
        final FlowShape<Pair<C, D>, C> resultProcessor = builder.add(resultProcessorFlow);

        final UniformFanOutShape<B, B> broadcast = builder.add(Broadcast.create(2));
        final FanInShape2<C, D, Pair<C, D>> zip = builder.add(Zip.create());

        builder.from(stage1.out()).toInlet(broadcast.in());
        builder.from(broadcast.out(0)).toInlet(stage2.in());
        builder.from(broadcast.out(1)).toInlet(sideChannel.in());
        builder.from(stage2.out()).toInlet(zip.in0());
        builder.from(sideChannel.out()).toInlet(zip.in1());
        builder.from(zip.out()).toInlet(resultProcessor.in());

        return FlowShape.of(stage1.in(), resultProcessor.out());
    });
}
 

private void newBinding() {
    requestQueue = new LinkedBlockingQueue<>();
    responseQueue = new LinkedBlockingQueue<>();

    final Flow<HttpRequest, HttpResponse, NotUsed> handler =
            Flow.fromGraph(KillSwitches.<HttpRequest>single())
                    .mapAsync(1, request -> {
                        requestQueue.offer(request);
                        return responseQueue.take();
                    })
                    .mapMaterializedValue(killSwitch -> {
                        Objects.requireNonNull(killSwitchTrigger.peek())
                                .thenAccept(_void -> killSwitch.shutdown());
                        return NotUsed.getInstance();
                    });
    binding = Http.get(actorSystem).bindAndHandle(handler, ConnectHttp.toHost("127.0.0.1", 0), mat)
            .toCompletableFuture()
            .join();
}
 

@Before
public void createActorSystem() {
    // create actor system with deactivated hostname blacklist to connect to localhost
    actorSystem = ActorSystem.create(getClass().getSimpleName(),
            TestConstants.CONFIG.withValue("ditto.connectivity.connection.http-push.blacklisted-hostnames",
                    ConfigValueFactory.fromAnyRef("")));
    mat = ActorMaterializer.create(actorSystem);
    requestQueue = new LinkedBlockingQueue<>();
    responseQueue = new LinkedBlockingQueue<>();
    handler = Flow.fromFunction(request -> {
        requestQueue.offer(request);
        return responseQueue.take();
    });
    binding = Http.get(actorSystem)
            .bindAndHandle(handler, ConnectHttp.toHost("127.0.0.1", 0), mat)
            .toCompletableFuture()
            .join();
    connection = getHttpConnectionBuilderToLocalBinding(false, binding.localAddress().getPort()).build();
}
 

@SuppressWarnings("unused")
private DispatcherActor(final ActorRef enforcerActor,
        final ActorRef pubSubMediator,
        final Flow<ImmutableDispatch, ImmutableDispatch, NotUsed> handler) {

    super(WithDittoHeaders.class);

    enforcementConfig = DittoConciergeConfig.of(
            DefaultScopedConfig.dittoScoped(getContext().getSystem().settings().config())
    ).getEnforcementConfig();

    this.handler = handler;
    final Props props = ThingsAggregatorActor.props(enforcerActor);
    thingsAggregatorActor = getContext().actorOf(props, ThingsAggregatorActor.ACTOR_NAME);

    initActor(getSelf(), pubSubMediator);
}
 

private Flow<Message,Message,?> flow(ActorRef shardRegion) {
    return Flow.<Message>create()
        .map(msg -> {
            if (msg.isText()) {
                log.warn("Ignoring unexpected text-kind web socket message {}", msg);
                return Option.<Query.EventEnvelope>none();
            } else {
                return Option.<Query.EventEnvelope>some(Query.EventEnvelope.parseFrom(msg.asBinaryMessage().getStrictData().toArray()));
            }
        })
        .filter(o -> o.isDefined())
        .map(o -> o.get())
        .mapAsync(maxInFlight, e -> ask(shardRegion, e, timeout))
        .map(resp -> (Long) resp)
        .map(l -> BinaryMessage.create(ByteString.fromArray(EventsPersisted.newBuilder().setOffset(l).build().toByteArray())));
}
 

@Override
public Flow<EventEnvelope,Long,?> uploadFlow() {
    ClientConnectionSettings settings = ClientConnectionSettings.create(system.settings().config());
    
    return Flow.<EventEnvelope>create()
        .map(e -> (Message) BinaryMessage.create(serialize(e)))
        .via(Http.get(system).webSocketClientFlow(WebSocketRequest.create(uri), connectionContext, Optional.empty(), settings, system.log()))
        .map(msg -> {
            if (msg.isText()) {
                log.warn("Ignoring unexpected text-type WS message {}", msg);
                return 0l;
            } else {
                EventsPersisted applied = EventsPersisted.parseFrom(
                    msg.asBinaryMessage().getStrictData().iterator().asInputStream());
                return applied.hasOffset() ? applied.getOffset() : 0l;
            }})
        .filter(l -> l > 0);
}
 

private Flow<Message, String, NotUsed> getStrictifyFlow(final HttpRequest request,
        final CharSequence correlationId) {
    return Flow.<Message>create()
            .via(Flow.fromFunction(msg -> {
                IN_COUNTER.increment();
                return msg;
            }))
            .filter(Message::isText)
            .map(Message::asTextMessage)
            .map(textMsg -> {
                if (textMsg.isStrict()) {
                    return Source.single(textMsg.getStrictText());
                } else {
                    return textMsg.getStreamedText();
                }
            })
            .flatMapConcat(textMsg -> textMsg.fold("", (str1, str2) -> str1 + str2))
            .via(incomingMessageSniffer.toAsyncFlow(request))
            .via(Flow.fromFunction(result -> {
                LOGGER.withCorrelationId(correlationId).debug("Received incoming WebSocket message: {}", result);
                return result;
            }))
            .withAttributes(Attributes.createLogLevels(Logging.DebugLevel(), Logging.DebugLevel(),
                    Logging.WarningLevel()));

}
 

@SuppressWarnings("unchecked")
private static <T> Flow<DittoRuntimeException, Message, NotUsed> joinOutgoingFlows(
        final Source<T, NotUsed> eventAndResponseSource,
        final Flow<DittoRuntimeException, T, NotUsed> errorFlow,
        final Flow<T, Message, NotUsed> messageFlow) {

    return Flow.fromGraph(GraphDSL.create3(eventAndResponseSource, errorFlow, messageFlow,
            (notUsed1, notUsed2, notUsed3) -> notUsed1,
            (builder, eventsAndResponses, errors, messages) -> {
                final UniformFanInShape<T, T> merge = builder.add(Merge.create(2, true));

                builder.from(eventsAndResponses).toFanIn(merge);
                builder.from(errors).toFanIn(merge);
                builder.from(merge.out()).toInlet(messages.in());

                return FlowShape.of(errors.in(), messages.out());
            }));
}
 

/**
 * Returns a flow that buffers up to [maximumBatchSize] JSON events and writing them together.
 * 
 * Buffering is only done if upstream is faster than downstream. If there is demand from downstream,
 * also slower batches will be written.
 */
public static Flow<JSONEvent,ByteString,NotUsed> flow(int maximumBatchSize) {
    return Flow.of(JSONEvent.class)
        .batch(maximumBatchSize, event -> {
            List<JSONEvent> l = new ArrayList<>();
            l.add(event);
            return l;
        }, (list, event) -> {
            list.add(event);
            return list;
        }).via(new JacksonWriter());
}
 

/**
 * Returns a flow that buffers up to [maximumBatchSize] XML events and writing them together.
 * 
 * Buffering is only done if upstream is faster than downstream. If there is demand from downstream,
 * also slower batches will be written.
 */
public static Flow<XMLEvent,ByteString,NotUsed> flow(int maximumBatchSize) {
    return Flow.of(XMLEvent.class)
        .batch(maximumBatchSize, event -> {
            List<XMLEvent> l = new ArrayList<>();
            l.add(event);
            return l;
        }, (list, event) -> {
            list.add(event);
            return list;
        }).via(new StaxWriter());
}
 

@Override
protected void setupMocks(final TestProbe probe) {
    connectionFactory = mock(KafkaConnectionFactory.class);
    when(connectionFactory.newFlow())
            .thenReturn(
                    Flow.fromFunction(envelope -> {
                        final ProducerMessage.Message<String, String, Object> message =
                                (ProducerMessage.Message<String, String, Object>) envelope;
                        received.add(message);
                        return createResult(message);
                    }));
}
 

/**
 * Creates a flow that feeds each event into [selector], and:
 * - If the selector emits a result, that result, and further results that also match, are fed through [target], and then merged into the stream.
 * - If the selector emits {@link ReadProtocol#none()}, the event is passed downstream unchanged
 * - If the selector emits a failure, the stream is failed.
 */
public static <E> Flow<E, E, NotUsed> filter(ReadProtocol<E,E> selector, Graph<FlowShape<E,E>, ?> target) {
    return Flow
        .<E>create()
        .via(insertMarkers(selector))
        .splitWhen(Marker.class::isInstance)
        .prefixAndTail(1)
        .flatMapConcat(t -> {
            // This is safe because all of the "tail" in "prefixAndTail" above comes from the original stream of E.
            Source<E, NotUsed> source = uncheckedCast(t.second());
            return isSelected(t.first().get(0)) ? source.via(target) : source;
        })
        .concatSubstreams();
}
 

private Flow<ByteString, ByteString, CompletionStage<HttpResponse>> createFlow(HttpMethod method, Uri uri, Option<ContentType> contentType, Predicate<HttpResponse> isSuccess, HttpHeader... headers) {
    Sink<ByteString, Publisher<ByteString>> in = Sink.asPublisher(AsPublisher.WITH_FANOUT); // akka internally recreates this twice, on some errors...
    Source<ByteString, Subscriber<ByteString>> out = Source.asSubscriber();
    
    return Flow.fromSinkAndSourceMat(in, out, Keep.both()).mapMaterializedValue(pair -> {
        RequestEntity entity;
        if (contentType.isDefined()) {
            Source<ByteString, NotUsed> inReader = Source.fromPublisher(pair.first());
            entity = HttpEntities.createChunked(contentType.get(), inReader);
        } else {
            entity = HttpEntities.EMPTY;
        }
        HttpRequest rq = HttpRequest.create().withMethod(method).withUri(uri).addHeaders(Arrays.asList(headers)).withEntity(entity);
        
        return http.singleRequest(rq).thenApply(resp -> {
            if (isSuccess.test(resp)) {
                resp.entity().getDataBytes()
                    .runWith(Sink.fromSubscriber(pair.second()), materializer);
            } else {
                log.info("Http responded error: {} for request {}", resp, rq);
                resp.discardEntityBytes(materializer);
                pair.second().onError(new IllegalStateException("Unsuccessful HTTP response: " + resp + " for " + rq));
            }
            return resp;
        }).exceptionally(x -> {
            Throwable cause = (x instanceof CompletionException) ? x.getCause() : x;
            if (!(cause instanceof IllegalStateException)) {
                log.info("Could not make http request " + rq, cause);
            }
            pair.second().onError(cause);
            throw (cause instanceof RuntimeException) ? (RuntimeException) x : new RuntimeException(cause);
        });
    });
}
 

private CompletionStage<F.Either<Result, Flow<JsonNode, JsonNode, ?>>>
  createActorFlow2(Http.RequestHeader request) {
    return CompletableFuture.completedFuture(
      request.session()
      .getOptional("username")
      .map(username ->
        F.Either.<Result, Flow<JsonNode, JsonNode, ?>>Right(
          createFlowForActor()))
      .orElseGet(() -> F.Either.Left(forbidden())));
}
 

private Flow<Integer, Double, NotUsed> computeAverage() {
    return Flow.of(Integer.class).grouped(2).mapAsyncUnordered(8, integers ->
            CompletableFuture.supplyAsync(() -> integers
                    .stream()
                    .mapToDouble(v -> v)
                    .average()
                    .orElse(-1.0)));
}
 

private Sink<EventEnvelope,NotUsed> filteredDataCenterSink() {
    log.debug("filteredDataCenterSink()");
    return Flow.<EventEnvelope>create()
        .mapAsync(parallelism, e -> {
            return visibilityRepo.isVisibleTo(dataCenter, e.persistenceId()).thenApply(v -> {
                log.debug("Visibility of {}: {}", e, v);
                return Tuple.of(e,v);});
        })
        .filter(t -> t._2)
        .map(t -> t._1)
        .via(dataCenter.uploadFlow())
        .map(EventDelivered::new)
        .to(Sink.actorRef(self(), new Failure(new IllegalStateException("Remote datacenter closed connection"))));
}
 

/**
 * Create a flow that splits 1 outbound signal into many as follows.
 * <ol>
 * <li>
 *   Targets with matching filtered topics without extra fields are grouped into 1 outbound signal, followed by
 * </li>
 * <li>one outbound signal for each target with a matching filtered topic with extra fields.</li>
 * </ol>
 * The matching filtered topic is attached in the latter case.
 * Consequently, for each outbound signal leaving this flow, if it has a filtered topic attached,
 * then it has 1 unique target with a matching topic with extra fields.
 * This satisfies the precondition of {@code this#enrichAndFilterSignal}.
 *
 * @return the flow.
 */
private static Flow<OutboundSignalWithId, Pair<OutboundSignalWithId, FilteredTopic>, NotUsed> splitByTargetExtraFieldsFlow() {
    return Flow.<OutboundSignalWithId>create()
            .mapConcat(outboundSignal -> {
                final Pair<List<Target>, List<Pair<Target, FilteredTopic>>> splitTargets =
                        splitTargetsByExtraFields(outboundSignal);

                final boolean shouldSendSignalWithoutExtraFields =
                        !splitTargets.first().isEmpty() ||
                                isTwinCommandResponseWithReplyTarget(outboundSignal.getSource()) ||
                                outboundSignal.getTargets().isEmpty(); // no target - this is an error response
                final Stream<Pair<OutboundSignalWithId, FilteredTopic>> outboundSignalWithoutExtraFields =
                        shouldSendSignalWithoutExtraFields
                                ? Stream.of(Pair.create(outboundSignal.setTargets(splitTargets.first()), null))
                                : Stream.empty();

                final Stream<Pair<OutboundSignalWithId, FilteredTopic>> outboundSignalWithExtraFields =
                        splitTargets.second().stream()
                                .map(targetAndSelector -> Pair.create(
                                        outboundSignal.setTargets(
                                                Collections.singletonList(targetAndSelector.first())),
                                        targetAndSelector.second()));

                return Stream.concat(outboundSignalWithoutExtraFields, outboundSignalWithExtraFields)
                        .collect(Collectors.toList());
            });
}
 

private static <T> Flow<Pair<T, StartedTimer>, T, NotUsed> createStopTimerFlow() {
    return Flow.fromFunction(pair -> {
        try {
            pair.second().stop();
        } catch (final IllegalStateException e) {
            // it is okay if the timer stopped already; simply return the result.
        }
        return pair.first();
    });
}
 

/**
 * Create a new flow through the search persistence.
 * No logging or recovery is attempted.
 *
 * @param parallelism How many write operations may run in parallel for this sink.
 * @param maxBulkSize How many writes to perform in one bulk.
 * @param writeInterval Delay between bulk operation requests. MongoDB backpressure is insufficient.
 * @return the sink.
 */
public Flow<Source<AbstractWriteModel, NotUsed>, WriteResultAndErrors, NotUsed> start(
        final int parallelism,
        final int maxBulkSize,
        final Duration writeInterval) {

    final Flow<List<AbstractWriteModel>, List<AbstractWriteModel>, NotUsed> throttleFlow;
    if (Duration.ZERO.minus(writeInterval).isNegative()) {
        throttleFlow = Flow.<List<AbstractWriteModel>>create()
                .delay(writeInterval, DelayOverflowStrategy.backpressure());
    } else {
        throttleFlow = Flow.create();
    }

    final Flow<Source<AbstractWriteModel, NotUsed>, List<AbstractWriteModel>, NotUsed> batchFlow =
            Flow.<Source<AbstractWriteModel, NotUsed>>create()
                    .flatMapConcat(source -> source.grouped(maxBulkSize))
                    .via(throttleFlow);

    final Flow<List<AbstractWriteModel>, WriteResultAndErrors, NotUsed> writeFlow =
            Flow.<List<AbstractWriteModel>>create()
                    .flatMapMerge(parallelism, this::executeBulkWrite)
                    // never initiate more than "parallelism" writes against the persistence
                    .withAttributes(Attributes.inputBuffer(parallelism, parallelism));

    return Flow.fromGraph(assembleFlows(batchFlow, writeFlow, createStartTimerFlow(), createStopTimerFlow()));
}
 

private <T> Flow<T, T, NotUsed> blockNamespaceFlow(final Function<T, String> namespaceExtractor) {
    return Flow.<T>create()
            .flatMapConcat(element -> {
                final String namespace = namespaceExtractor.apply(element);
                final CompletionStage<Boolean> shouldUpdate = blockedNamespaces.contains(namespace)
                        .handle((result, error) -> result == null || !result);
                return Source.fromCompletionStage(shouldUpdate)
                        .filter(Boolean::valueOf)
                        .map(bool -> element);
            });
}
 

private <T> Flow<Map<ThingId, T>, Map<ThingId, T>, NotUsed> filterMapKeysByBlockedNamespaces() {
    return Flow.<Map<ThingId, T>>create()
            .flatMapConcat(map ->
                    Source.fromIterator(map.entrySet()::iterator)
                            .via(blockNamespaceFlow(entry -> entry.getKey().getNamespace()))
                            .fold(new HashMap<>(), (accumulator, entry) -> {
                                accumulator.put(entry.getKey(), entry.getValue());
                                return accumulator;
                            })
            );
}
 

private <T> Source<T, NotUsed> processSearchPersistenceResult(Source<T, NotUsed> source,
        final DittoHeaders dittoHeaders) {

    final Flow<T, T, NotUsed> logAndFinishPersistenceSegmentFlow =
            Flow.fromFunction(result -> {
                // we know that the source provides exactly one ResultList
                LogUtil.enhanceLogWithCorrelationId(log, dittoHeaders.getCorrelationId());
                log.debug("Persistence returned: {}", result);
                return result;
            });

    return source.via(logAndFinishPersistenceSegmentFlow);
}
 

private <T> Flow<T, Object, NotUsed> stopTimerAndHandleError(final StartedTimer searchTimer,
        final WithDittoHeaders<?> command) {
    return Flow.<T, Object>fromFunction(
            element -> {
                stopTimer(searchTimer);
                return element;
            })
            .recoverWithRetries(1, new PFBuilder<Throwable, Graph<SourceShape<Object>, NotUsed>>()
                    .matchAny(error -> {
                        stopTimer(searchTimer);
                        return Source.single(asDittoRuntimeException(error, command));
                    })
                    .build()
            );
}
 

private <T extends Command> void executeCount(final T countCommand,
        final Function<T, Query> queryParseFunction,
        final boolean isSudo) {
    final DittoHeaders dittoHeaders = countCommand.getDittoHeaders();
    final Optional<String> correlationIdOpt = dittoHeaders.getCorrelationId();
    LogUtil.enhanceLogWithCorrelationId(log, correlationIdOpt);
    log.info("Processing CountThings command: {}", countCommand);
    final JsonSchemaVersion version = countCommand.getImplementedSchemaVersion();

    final String queryType = "count";

    final StartedTimer countTimer = startNewTimer(version, queryType);

    final StartedTimer queryParsingTimer = countTimer.startNewSegment(QUERY_PARSING_SEGMENT_NAME);

    final ActorRef sender = getSender();

    final Source<Object, ?> replySource = createQuerySource(queryParseFunction, countCommand)
            .flatMapConcat(query -> {
                LogUtil.enhanceLogWithCorrelationId(log, correlationIdOpt);
                stopTimer(queryParsingTimer);
                final StartedTimer databaseAccessTimer =
                        countTimer.startNewSegment(DATABASE_ACCESS_SEGMENT_NAME);

                final Source<Long, NotUsed> countResultSource = isSudo
                        ? searchPersistence.sudoCount(query)
                        : searchPersistence.count(query,
                        countCommand.getDittoHeaders().getAuthorizationContext().getAuthorizationSubjectIds());

                return processSearchPersistenceResult(countResultSource, dittoHeaders)
                        .via(Flow.fromFunction(result -> {
                            stopTimer(databaseAccessTimer);
                            return result;
                        }))
                        .map(count -> CountThingsResponse.of(count, dittoHeaders));
            })
            .via(stopTimerAndHandleError(countTimer, countCommand));

    Patterns.pipe(replySource.runWith(Sink.head(), materializer), getContext().dispatcher()).to(sender);
}
 

/**
 * Create an async flow for event sniffing.
 *
 * @param request the HTTP request that started the event stream.
 * @return flow to pass events through with a wiretap attached over an async barrier to the sink for sniffed events.
 */
default Flow<T, T, NotUsed> toAsyncFlow(final HttpRequest request) {
    return Flow.<T>create().wireTap(
            Flow.<T>create()
                    .async()
                    .to(Sink.lazyInitAsync(() -> CompletableFuture.completedFuture(
                            createSink(request)))));
}