Skip to content

Commit

Permalink
=str Fix maybe throw for MinimalStage.
Browse files Browse the repository at this point in the history
  • Loading branch information
He-Pin committed Dec 2, 2023
1 parent 8fefca1 commit 485c413
Show file tree
Hide file tree
Showing 2 changed files with 69 additions and 14 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
/*
* 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.pekko.stream.javadsl;

import org.apache.pekko.japi.Pair;
import org.apache.pekko.stream.StreamTest;
import org.apache.pekko.testkit.PekkoJUnitActorSystemResource;
import org.apache.pekko.testkit.PekkoSpec;
import org.junit.Assert;
import org.junit.ClassRule;
import org.junit.Test;

import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;

public class FlowUnfoldAsyncTest extends StreamTest {
@ClassRule
public static PekkoJUnitActorSystemResource actorSystemResource =
new PekkoJUnitActorSystemResource("SourceTest", PekkoSpec.testConf());

public FlowUnfoldAsyncTest() {
super(actorSystemResource);
}

@Test
public void testFoldAsync() throws Exception {
final Integer result = Source.unfoldAsync(
0,
idx -> {
if (idx >= 10) {
return CompletableFuture.completedStage(Optional.empty());
} else {
return CompletableFuture.completedStage(Optional.of(Pair.create(idx + 1, idx)));
}
})
.runFold(0, Integer::sum, system)
.toCompletableFuture()
.get(3, TimeUnit.SECONDS);
Assert.assertEquals(45, result.intValue());
}
}
26 changes: 12 additions & 14 deletions stream/src/main/scala/org/apache/pekko/stream/impl/Unfold.scala
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ import pekko.stream.impl.Stages.DefaultAttributes
import pekko.stream.stage.{ GraphStage, GraphStageLogic, OutHandler }

import java.util.Optional
import java.util.concurrent.CompletableFuture
import java.util.concurrent.CompletionStage
import scala.concurrent.Future
import scala.util.{ Failure, Success, Try }
Expand Down Expand Up @@ -124,21 +123,20 @@ import scala.util.{ Failure, Success, Try }
}

def onPull(): Unit = {
f.apply(state) match {
case cf: CompletableFuture[Optional[Pair[S, E]] @unchecked] if cf.isDone && !cf.isCompletedExceptionally =>
handle(cf.join())
case future =>
future.handle((r, ex) => {
if (ex != null) {
asyncHandler(Failure(ex))
} else {
asyncHandler(Success(r))
}
null
})
val future = f.apply(state).toCompletableFuture
if (future.isDone && !future.isCompletedExceptionally) {
handle(future.getNow(null))
} else {
future.handle((r, ex) => {
if (ex != null) {
asyncHandler(Failure(ex))
} else {
asyncHandler(Success(r))
}
null
})
}
}

setHandler(out, this)
}
}

0 comments on commit 485c413

Please sign in to comment.