akka.stream.javadsl.Sink Java Examples

private CompletionStage<Done> createNonExistingIndices(final String collectionName,
        final List<Index> indices) {
    if (indices.isEmpty()) {
        LOGGER.warn("No indices are defined, thus no indices are created.");
        return CompletableFuture.completedFuture(Done.getInstance());
    }
    return indexOperations.getIndicesExceptDefaultIndex(collectionName)
            .flatMapConcat(
                    existingIndices -> {
                        LOGGER.info("Create non-existing indices: Existing indices are: {}", existingIndices);
                        final List<Index> indicesToCreate = excludeIndices(indices, existingIndices);
                        LOGGER.info("Indices to create are: {}", indicesToCreate);
                        return createIndices(collectionName, indicesToCreate);
                    })
            .runWith(Sink.ignore(), materializer);
}
 

private Receive startBackup(long offset) {
    query
        .eventsByTag(tag, NoOffset.getInstance())
        // create backups of max [N] elements, or at least every [T] on activity
        // FIXME write a stage that, instead of buffering each chunk into memory, creates sub-streams instead.
        .groupedWithin(eventChunkSize, eventChunkDuration)
        .filter(list -> list.size() > 0)
        .mapAsync(4, list -> s3.store(tag, Vector.ofAll(list)).thenApply(done -> list.get(list.size() - 1).offset()))
        .runWith(Sink.actorRefWithAck(self(), "init", "ack", "done", Failure::new), materializer);
    
    return ReceiveBuilder.create()
        .matchEquals("init", msg -> sender().tell("ack", self()))
        .match(Long.class, l -> pipe(s3.saveOffset(l).thenApply(done -> "ack"), context().dispatcher()).to(sender()))
        .match(Failure.class, msg -> {
            log.error("Stream failed, rethrowing", msg.cause());
            throw new RuntimeException(msg.cause());
        })
        .matchEquals("done", msg -> { throw new IllegalStateException("eventsByTag completed, this should not happen. Killing actor, hoping for restart"); })
        .build();
}
 

private void insertTestThings() {
    final Thing baseThing = ThingsModelFactory.newThingBuilder()
            .setId(ThingId.of("thing", "00"))
            .setRevision(1234L)
            .setPermissions(AUTH_CONTEXT.getFirstAuthorizationSubject().orElseThrow(AssertionError::new),
                    Permission.READ)
            .setAttribute(JsonPointer.of("x"), JsonValue.of(5))
            .build();

    final Thing irrelevantThing = baseThing.toBuilder().removeAllAttributes().build();

    writePersistence.writeThingWithAcl(template(baseThing, 0, "a"))
            .concat(writePersistence.writeThingWithAcl(template(baseThing, 1, "b")))
            .concat(writePersistence.writeThingWithAcl(template(baseThing, 2, "a")))
            .concat(writePersistence.writeThingWithAcl(template(baseThing, 3, "b")))
            .concat(writePersistence.writeThingWithAcl(template(baseThing, 4, "c")))
            .concat(writePersistence.writeThingWithAcl(template(irrelevantThing, 5, "c")))
            .runWith(Sink.ignore(), materializer)
            .toCompletableFuture()
            .join();
}
 

@Test
public void extractJournalPidsFromEventsAndNotSnapshots() {
    insert("test_journal", new Document().append("pid", "pid3").append("to", 2L));
    insert("test_journal", new Document().append("pid", "pid4").append("to", 2L));
    insert("test_journal", new Document().append("pid", "pid1").append("to", 1L));
    insert("test_journal", new Document().append("pid", "pid2").append("to", 1L));
    insert("test_snaps", new Document().append("pid", "pid5").append("sn", 3L));
    insert("test_snaps", new Document().append("pid", "pid6").append("sn", 4L));

    final List<String> pids =
            readJournal.getJournalPids(2, Duration.ZERO, materializer)
                    .runWith(Sink.seq(), materializer)
                    .toCompletableFuture().join();

    assertThat(pids).containsExactly("pid1", "pid2", "pid3", "pid4");
}
 

private void startRestore() {
    s3
    .list(tag)
    // skip over entries until the one BEFORE entry where startTime >= offset (since the one before may have been only partially restored)
    .via(dropUntilNext(l -> S3.getStartInstant(l).toEpochMilli() >= offset, true))
    .flatMapConcat(entry -> s3.loadEvents(entry.key().substring(entry.key().lastIndexOf("/") + 1)))
    .mapAsync(maxInFlight, e -> {
        log.debug("Replaying {}:{}", e.getPersistenceId(), e.getSequenceNr());
        return ask(shardRegion, e, timeout);
    })
    .map(resp -> {
        log.debug("Responded {}", resp);
        return (Long) resp;
    })
    // only save one lastOffset update per minute, and only the lowest one
    .conflate((Long l1, Long l2) -> l1 < l2 ? l1 : l2)
    .runWith(Sink.actorRefWithAck(self(), "init", "ack", "done", Failure::new), materializer);
}
 

@Test
void verifyProcessingPipelineWorksAsExpectedForBogusMessages() {

	final TestKit catchAllProbe = new TestKit(system);

	final Source<Message, NotUsed> source = Source.single(new Strict("invalid message from ws"));
	final Sink<ProducerRecord<String, String>, NotUsed> sink = Sink.actorRef(catchAllProbe.getRef(), new ProducerRecord<String, String>("none", "completed"));

	WebsocketMessageToProducerRecordTranslator.eventProcessing()
			.on(system)
			.withHandler(() -> catchAllProbe.getRef().tell("Application replaced", catchAllProbe.getRef()))
			.withCallContextProvider(catchAllProbe.getRef())
			.withMetricsService(catchAllProbe.getRef())
			.from(source)
			.to(sink)
			.run();

	final ProducerRecord<String, String> completeMsg = catchAllProbe.expectMsgClass(ProducerRecord.class);
	assertThat(completeMsg.topic(), is("none"));
	assertThat(completeMsg.value(), is("completed"));
}
 

private Sink<Source<AbstractWriteModel, NotUsed>, NotUsed> createRestartSink() {
    final StreamConfig streamConfig = searchConfig.getStreamConfig();
    final PersistenceStreamConfig persistenceConfig = streamConfig.getPersistenceConfig();

    final int parallelism = persistenceConfig.getParallelism();
    final int maxBulkSize = persistenceConfig.getMaxBulkSize();
    final Duration writeInterval = streamConfig.getWriteInterval();
    final Sink<Source<AbstractWriteModel, NotUsed>, NotUsed> sink =
            mongoSearchUpdaterFlow.start(parallelism, maxBulkSize, writeInterval)
                    .via(bulkWriteResultAckFlow.start())
                    .log("SearchUpdaterStream/BulkWriteResult")
                    .withAttributes(Attributes.logLevels(
                            Attributes.logLevelInfo(),
                            Attributes.logLevelWarning(),
                            Attributes.logLevelError()))
                    .to(Sink.ignore());

    final ExponentialBackOffConfig backOffConfig = persistenceConfig.getExponentialBackOffConfig();

    return RestartSink.withBackoff(backOffConfig.getMax(), backOffConfig.getMax(), backOffConfig.getRandomFactor(),
            () -> sink);
}
 

private Route handleSudoCountThingsPerRequest(final RequestContext ctx, final SudoCountThings command) {
    final CompletableFuture<HttpResponse> httpResponseFuture = new CompletableFuture<>();

    Source.single(command)
            .to(Sink.actorRef(createHttpPerRequestActor(ctx, httpResponseFuture),
                    AbstractHttpRequestActor.COMPLETE_MESSAGE))
            .run(materializer);

    final CompletionStage<HttpResponse> allThingsCountHttpResponse = Source.fromCompletionStage(httpResponseFuture)
            .flatMapConcat(httpResponse -> httpResponse.entity().getDataBytes())
            .fold(ByteString.empty(), ByteString::concat)
            .map(ByteString::utf8String)
            .map(Integer::valueOf)
            .map(count -> JsonObject.newBuilder().set("allThingsCount", count).build())
            .map(jsonObject -> HttpResponse.create()
                    .withEntity(ContentTypes.APPLICATION_JSON, ByteString.fromString(jsonObject.toString()))
                    .withStatus(HttpStatusCode.OK.toInt()))
            .runWith(Sink.head(), materializer);

    return completeWithFuture(allThingsCountHttpResponse);
}
 

@Test()
void properlyHandleInvalidCommandMessage() {
	final TestKit kafkaProducer = new TestKit(system);
	final TestKit metricsService = new TestKit(system);
	final TestKit callContextProvider = new TestKit(system);

	final ConsumerRecord<String, String> consumerRecord = new ConsumerRecord<>("topic", 0, 0,
			"key", "NOT JSON");
	final Source<ConsumerRecord<String, String>, NotUsed> source = Source.single(consumerRecord);
	final Sink<ProducerRecord<String, String>, NotUsed> sink = Sink.<ProducerRecord<String, String>>ignore()
			.mapMaterializedValue(q -> NotUsed.getInstance());

	AriCommandResponseKafkaProcessor.commandResponseProcessing()
			.on(system)
			.withHandler(requestAndContext -> Http.get(system).singleRequest(requestAndContext._1))
			.withCallContextProvider(callContextProvider.getRef())
			.withMetricsService(metricsService.getRef())
			.from(source)
			.to(sink)
			.run();

	kafkaProducer.expectNoMsg(Duration.apply(250, TimeUnit.MILLISECONDS));
}
 

@Test
public void helloStream() throws Exception {
    // Important to concat our source with a maybe, this ensures the connection doesn't get closed once we've
    // finished feeding our elements in, and then also to take 3 from the response stream, this ensures our
    // connection does get closed once we've received the 3 elements.
    Source<String, ?> response = await(streamService.stream().invoke(
            Source.from(Arrays.asList("a", "b", "c"))
                    .concat(Source.maybe())));
    List<String> messages = await(response.take(3).runWith(Sink.seq(), mat));
    assertEquals(Arrays.asList("Hello, a!", "Hello, b!", "Hello, c!"), messages);
}
 

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);
}
 

private <T> List<T> parse(Protocol<CsvEvent, T> proto, String input)
    throws InterruptedException, ExecutionException, TimeoutException {
    List<T> output = Source.single(input)
        .via(new CsvParser(CsvSettings.RFC4180))
        .via(ProtocolReader.of(proto))
        .runWith(Sink.seq(), materializer)
        .toCompletableFuture()
        .get(1, TimeUnit.SECONDS);
    return output;
}
 

@Before
public void init() {
    system = ActorSystem.create();
    mat = ActorMaterializer.create(system);

    rematerializeSource();

    // materialize sink once - it never fails.
    final Sink<Integer, TestSubscriber.Probe<Integer>> sink = TestSink.probe(system);
    final Pair<TestSubscriber.Probe<Integer>, Sink<Integer, NotUsed>> sinkPair = sink.preMaterialize(mat);
    sinkProbe = sinkPair.first();
    testSink = sinkPair.second();
}
 

@Override
@SuppressWarnings("unchecked")
protected Publisher hide(Processor processor) {
    return (Publisher) Source
            .fromPublisher(processor)
            .runWith(Sink.asPublisher(AsPublisher.WITH_FANOUT), mat);
}
 

/**
 * 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)))));
}
 

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"))));
}
 

private <T> List<T> runBlockingWithReturn(final Source<T, NotUsed> publisher) {
    final CompletionStage<List<T>> done = publisher.runWith(Sink.seq(), mat);
    try {
        return done.toCompletableFuture().get(WAIT_TIMEOUT, TimeUnit.SECONDS);
    } catch (final InterruptedException | ExecutionException | TimeoutException e) {
        throw new IllegalStateException(e);
    }
}
 

private static Sink<ProducerMessage.Results<String, String, PassThrough>, CompletionStage<Done>> publishSuccessSink() {

        // basically, we don't know if the 'publish' will succeed or fail. We would need to write our own
        // GraphStage actor for Kafka and MQTT, since alpakka doesn't provide this useful information for us.
        return Sink.foreach(results -> {
            final ConnectionMonitor connectionMonitor = results.passThrough().connectionMonitor;
            connectionMonitor.success(results.passThrough().externalMessage);
        });
    }
 

private void stream(final StreamThings streamThings) {
    log.withCorrelationId(streamThings)
            .info("Processing StreamThings command: {}", streamThings);
    final JsonSchemaVersion version = streamThings.getImplementedSchemaVersion();
    final String queryType = "query"; // same as queryThings
    final StartedTimer searchTimer = startNewTimer(version, queryType);
    final StartedTimer queryParsingTimer = searchTimer.startNewSegment(QUERY_PARSING_SEGMENT_NAME);
    final ActorRef sender = getSender();
    final Set<String> namespaces = streamThings.getNamespaces().orElse(null);
    final Source<Optional<ThingsSearchCursor>, NotUsed> cursorSource =
            ThingsSearchCursor.extractCursor(streamThings);
    final Source<SourceRef<String>, NotUsed> sourceRefSource = cursorSource.flatMapConcat(cursor -> {
        cursor.ifPresent(c -> c.logCursorCorrelationId(log, streamThings));
        return createQuerySource(queryParser::parse, streamThings).flatMapConcat(parsedQuery -> {
            final Query query = ThingsSearchCursor.adjust(cursor, parsedQuery, queryParser.getCriteriaFactory());
            stopTimer(queryParsingTimer);
            searchTimer.startNewSegment(DATABASE_ACCESS_SEGMENT_NAME); // segment stopped by stopTimerAndHandleError
            final List<String> subjectIds = streamThings.getDittoHeaders().getAuthorizationSubjects();
            final CompletionStage<SourceRef<String>> sourceRefFuture =
                    searchPersistence.findAllUnlimited(query, subjectIds, namespaces)
                            .map(ThingId::toString) // for serialization???
                            .runWith(StreamRefs.sourceRef(), materializer);
            return Source.fromCompletionStage(sourceRefFuture);
        });
    });
    final Source<Object, NotUsed> replySourceWithErrorHandling =
            sourceRefSource.via(stopTimerAndHandleError(searchTimer, streamThings));

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

private String extractBody(HttpResponse httpResponse) {
    return httpResponse.entity()
            .getDataBytes()
            .fold(ByteString.empty(), ByteString::concat)
            .map(ByteString::utf8String)
            .runWith(Sink.head(), actorSystem)
            .toCompletableFuture()
            .join();
}
 

public AkkaHubProxy(ActorMaterializer mat) {
    this.mat = mat;
    //  Obtain a Sink and Source which will publish and receive from the "bus" respectively.
    Pair<Sink<Object, NotUsed>, Source<Object, NotUsed>> sinkAndSource =
            MergeHub.of(Object.class, 16)
                    .toMat(BroadcastHub.of(Object.class, 256), Keep.both())
                    .run(mat);

    Sink<Object, NotUsed> sink = sinkAndSource.first();
    source = sinkAndSource.second().takeWhile((Predicate<Object>) o -> o != Done.getInstance());
    //source.runWith(Sink.ignore(), mat);
    busFlow = Flow.fromSinkAndSource(sink, source)
            .joinMat(KillSwitches.singleBidi(), Keep.right());
}
 

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 Graph<SinkShape<Either<StreamControlMessage, Signal>>, ?> getStreamControlOrSignalSink() {

        return Sink.foreach(either -> {
            final Object streamControlMessageOrSignal = either.isLeft() ? either.left().get() : either.right().get();
            streamingActor.tell(streamControlMessageOrSignal, ActorRef.noSender());
        });
    }
 

@Test
public void ensureCollectionIsCapped() throws Exception {
    final MongoCollection<Document> collection =
            syncPersistence.getCollection().runWith(Sink.head(), materializer).toCompletableFuture().get();

    runBlocking(syncPersistence.setTimestamp(Instant.now()));
    runBlocking(syncPersistence.setTimestamp(Instant.now()));

    assertThat(runBlocking(Source.fromPublisher(collection.count()))).containsExactly(1L);
}
 

@Override
public Removable addUpstream(Source<Object, NotUsed> publisher) {
    UniqueKillSwitch killSwitch =
            publisher.viaMat(busFlow, Keep.right())
                    .to(Sink.ignore())
                    .run(mat);
    subscriptions.put(publisher, killSwitch);
    return () -> AkkaHubProxy.this.removeUpstream(publisher);
}
 

private Route handleDevOpsPerRequest(final RequestContext ctx,
        final Source<ByteString, ?> payloadSource,
        final Function<String, DevOpsCommand<?>> requestJsonToCommandFunction) {
    final CompletableFuture<HttpResponse> httpResponseFuture = new CompletableFuture<>();

    payloadSource
            .fold(ByteString.empty(), ByteString::concat)
            .map(ByteString::utf8String)
            .map(requestJsonToCommandFunction)
            .to(Sink.actorRef(createHttpPerRequestActor(ctx, httpResponseFuture),
                    AbstractHttpRequestActor.COMPLETE_MESSAGE))
            .run(materializer);

    return completeWithFuture(httpResponseFuture);
}
 

public StreamedHttpResponse createStreamedResponse(HttpVersion version, List<String> body, long contentLength) {
    List<HttpContent> content = new ArrayList<>();
    for (String chunk: body) {
        content.add(new DefaultHttpContent(Unpooled.copiedBuffer(chunk, Charset.forName("utf-8"))));
    }
    Publisher<HttpContent> publisher = Source.from(content).runWith(Sink.<HttpContent>asPublisher(AsPublisher.WITH_FANOUT), materializer);
    StreamedHttpResponse response = new DefaultStreamedHttpResponse(version, HttpResponseStatus.OK, publisher);
    HttpUtil.setContentLength(response, contentLength);
    return response;
}
 

@Test
public void encodeAndDecodeAreInverse() {
    final ThingsSearchCursor input = randomCursor();
    final ThingsSearchCursor decoded =
            ThingsSearchCursor.decode(input.encode(), materializer)
                    .runWith(Sink.head(), materializer)
                    .toCompletableFuture().join();

    assertThat(decoded).isEqualTo(input);
}
 

/**
  * Loads the last known written offset from S3, or returns 0 if not found
  */
 public CompletionStage<Long> loadOffset() {
     return download("_lastOffset")
 		.reduce((bs1, bs2) -> bs1.concat(bs2))
 		.map(bs -> Long.parseLong(bs.utf8String()))
 		.recoverWith(new PFBuilder<Throwable, Source<Long,NotUsed>>()
 			.matchAny(x -> Source.single(0L)) // not found -> start at 0
	.build()
)
 		.runWith(Sink.head(), materializer);
 }
 

@Override
@SuppressWarnings("unchecked")
protected <T> Publisher<T> filter(Processor processor, final Class<T> filterClass) {
    Source src = Source.fromPublisher(processor)
            .filter(o -> filterClass.isAssignableFrom(o.getClass()));
    return (Publisher<T>) src.runWith(Sink.asPublisher(AsPublisher.WITH_FANOUT), mat);
}