feat: Add AutoCloseable shortcut on mapWithResource

Author: injae-kim

Diff for: stream-tests/src/test/java/org/apache/pekko/stream/javadsl/FlowTest.java

@@ -37,6 +37,7 @@
 
 import java.util.*;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
@@ -237,6 +238,21 @@ public void mustBeAbleToUseMapWithResource() {
     Assert.assertFalse(gate.get());
   }
 
+  @Test
+  public void mustBeAbleToUseMapWithAutoCloseableResource() {
+    final AtomicInteger closed = new AtomicInteger();
+    Source.from(Arrays.asList("1", "2", "3"))
+        .via(
+            Flow.of(String.class)
+                .mapWithResource(
+                    () -> (AutoCloseable) closed::incrementAndGet, (resource, elem) -> elem))
+        .runWith(TestSink.create(system), system)
+        .request(4)
+        .expectNext("1", "2", "3")
+        .expectComplete();
+    Assert.assertEquals(closed.get(), 3);
+  }
+
   @Test
   public void mustBeAbleToUseFoldWhile() throws Exception {
     final int result =

Diff for: stream-tests/src/test/java/org/apache/pekko/stream/javadsl/SourceTest.java

@@ -47,6 +47,7 @@
 import java.util.*;
 import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
@@ -815,6 +816,18 @@ public void mustBeAbleToUseMapWithResource() {
     Assert.assertFalse(gate.get());
   }
 
+  @Test
+  public void mustBeAbleToUseMapWithAutoCloseableResource() {
+    final AtomicInteger closed = new AtomicInteger();
+    Source.from(Arrays.asList("1", "2", "3"))
+        .mapWithResource(() -> (AutoCloseable) closed::incrementAndGet, (resource, elem) -> elem)
+        .runWith(TestSink.create(system), system)
+        .request(4)
+        .expectNext("1", "2", "3")
+        .expectComplete();
+    Assert.assertEquals(closed.get(), 3);
+  }
+
   @Test
   public void mustBeAbleToUseFoldWhile() throws Exception {
     final int result =

Diff for: stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/FlowMapWithResourceSpec.scala

@@ -410,6 +410,25 @@ class FlowMapWithResourceSpec extends StreamSpec(UnboundedMailboxConfig) {
       Await.result(promise.future, 3.seconds) shouldBe Done
     }
 
+    "will close the autocloseable resource" in {
+      val closedCounter = new AtomicInteger(0)
+      val create = () =>
+        new AutoCloseable {
+          override def close(): Unit = closedCounter.incrementAndGet()
+        }
+      val (pub, sub) = TestSource
+        .probe[Int]
+        .mapWithResource(create, (_, count) => count)
+        .toMat(TestSink.probe)(Keep.both)
+        .run()
+      sub.expectSubscription().request(2)
+      pub.sendNext(1)
+      sub.expectNext(1)
+      pub.sendComplete()
+      sub.expectComplete()
+      closedCounter.get shouldBe 1
+    }
+
   }
   override def afterTermination(): Unit = {
     fs.close()

Diff for: stream/src/main/scala/org/apache/pekko/stream/javadsl/Flow.scala

@@ -828,6 +828,45 @@ final class Flow[In, Out, Mat](delegate: scaladsl.Flow[In, Out, Mat]) extends Gr
         (resource, out) => f(resource, out),
         resource => close.apply(resource).toScala))
 
+  /**
+   * Transform each stream element with the help of a [[AutoCloseable]] resource and close it.
+   *
+   * The resource creation function is invoked once when the stream is materialized and the returned resource is passed to
+   * the mapping function for mapping the first element. The mapping function returns a mapped element to emit
+   * downstream. The returned `T` MUST NOT be `null` as it is illegal as stream element - according to the Reactive Streams specification.
+   *
+   * The [[AutoCloseable]] resource is closed only once when the upstream or downstream finishes or fails.
+   *
+   * Early completion can be done with combination of the [[takeWhile]] operator.
+   *
+   * Adheres to the [[ActorAttributes.SupervisionStrategy]] attribute.
+   *
+   * You can configure the default dispatcher for this Source by changing the `akka.stream.materializer.blocking-io-dispatcher` or
+   * set it for a given Source by using [[ActorAttributes]].
+   *
+   * '''Emits when''' the mapping function returns an element and downstream is ready to consume it
+   *
+   * '''Backpressures when''' downstream backpressures
+   *
+   * '''Completes when''' upstream completes
+   *
+   * '''Cancels when''' downstream cancels
+   *
+   * @tparam R the type of the resource
+   * @tparam T the type of the output elements
+   * @param create function that creates the resource
+   * @param f function that transforms the upstream element and the resource to output element
+   * @since 1.1.0
+   */
+  def mapWithResource[R <: AutoCloseable, T](
+      create: function.Creator[R],
+      f: function.Function2[R, Out, T]): javadsl.Flow[In, T, Mat] =
+    mapWithResource(create, f,
+      (resource: AutoCloseable) => {
+        resource.close()
+        Optional.empty()
+      })
+
   /**
    * Transform each input element into an `Iterable` of output elements that is
    * then flattened into the output stream. The transformation is meant to be stateful,

Diff for: stream/src/main/scala/org/apache/pekko/stream/javadsl/Source.scala

@@ -2541,6 +2541,45 @@ final class Source[Out, Mat](delegate: scaladsl.Source[Out, Mat]) extends Graph[
         (resource, out) => f(resource, out),
         resource => close.apply(resource).toScala))
 
+  /**
+   * Transform each stream element with the help of a [[AutoCloseable]] resource and close it.
+   *
+   * The resource creation function is invoked once when the stream is materialized and the returned resource is passed to
+   * the mapping function for mapping the first element. The mapping function returns a mapped element to emit
+   * downstream. The returned `T` MUST NOT be `null` as it is illegal as stream element - according to the Reactive Streams specification.
+   *
+   * The [[AutoCloseable]] resource is closed only once when the upstream or downstream finishes or fails.
+   *
+   * Early completion can be done with combination of the [[takeWhile]] operator.
+   *
+   * Adheres to the [[ActorAttributes.SupervisionStrategy]] attribute.
+   *
+   * You can configure the default dispatcher for this Source by changing the `akka.stream.materializer.blocking-io-dispatcher` or
+   * set it for a given Source by using [[ActorAttributes]].
+   *
+   * '''Emits when''' the mapping function returns an element and downstream is ready to consume it
+   *
+   * '''Backpressures when''' downstream backpressures
+   *
+   * '''Completes when''' upstream completes
+   *
+   * '''Cancels when''' downstream cancels
+   *
+   * @tparam R the type of the resource
+   * @tparam T the type of the output elements
+   * @param create function that creates the resource
+   * @param f function that transforms the upstream element and the resource to output element
+   * @since 1.1.0
+   */
+  def mapWithResource[R <: AutoCloseable, T](
+      create: function.Creator[R],
+      f: function.Function2[R, Out, T]): javadsl.Source[T, Mat] =
+    mapWithResource(create, f,
+      (resource: AutoCloseable) => {
+        resource.close()
+        Optional.empty()
+      })
+
   /**
    * Transform each input element into an `Iterable` of output elements that is
    * then flattened into the output stream. The transformation is meant to be stateful,

Diff for: stream/src/main/scala/org/apache/pekko/stream/javadsl/SubFlow.scala

@@ -285,6 +285,45 @@ class SubFlow[In, Out, Mat](
         (resource, out) => f(resource, out),
         resource => close.apply(resource).toScala))
 
+  /**
+   * Transform each stream element with the help of a [[AutoCloseable]] resource and close it.
+   *
+   * The resource creation function is invoked once when the stream is materialized and the returned resource is passed to
+   * the mapping function for mapping the first element. The mapping function returns a mapped element to emit
+   * downstream. The returned `T` MUST NOT be `null` as it is illegal as stream element - according to the Reactive Streams specification.
+   *
+   * The [[AutoCloseable]] resource is closed only once when the upstream or downstream finishes or fails.
+   *
+   * Early completion can be done with combination of the [[takeWhile]] operator.
+   *
+   * Adheres to the [[ActorAttributes.SupervisionStrategy]] attribute.
+   *
+   * You can configure the default dispatcher for this Source by changing the `akka.stream.materializer.blocking-io-dispatcher` or
+   * set it for a given Source by using [[ActorAttributes]].
+   *
+   * '''Emits when''' the mapping function returns an element and downstream is ready to consume it
+   *
+   * '''Backpressures when''' downstream backpressures
+   *
+   * '''Completes when''' upstream completes
+   *
+   * '''Cancels when''' downstream cancels
+   *
+   * @tparam R the type of the resource
+   * @tparam T the type of the output elements
+   * @param create function that creates the resource
+   * @param f function that transforms the upstream element and the resource to output element
+   * @since 1.1.0
+   */
+  def mapWithResource[R <: AutoCloseable, T](
+      create: function.Creator[R],
+      f: function.Function2[R, Out, T]): javadsl.SubFlow[In, T, Mat] =
+    mapWithResource(create, f,
+      (resource: AutoCloseable) => {
+        resource.close()
+        Optional.empty()
+      })
+
   /**
    * Transform each input element into an `Iterable` of output elements that is
    * then flattened into the output stream. The transformation is meant to be stateful,

Diff for: stream/src/main/scala/org/apache/pekko/stream/javadsl/SubSource.scala

@@ -276,6 +276,45 @@ class SubSource[Out, Mat](
         (resource, out) => f(resource, out),
         resource => close.apply(resource).toScala))
 
+  /**
+   * Transform each stream element with the help of a [[AutoCloseable]] resource and close it.
+   *
+   * The resource creation function is invoked once when the stream is materialized and the returned resource is passed to
+   * the mapping function for mapping the first element. The mapping function returns a mapped element to emit
+   * downstream. The returned `T` MUST NOT be `null` as it is illegal as stream element - according to the Reactive Streams specification.
+   *
+   * The [[AutoCloseable]] resource is closed only once when the upstream or downstream finishes or fails.
+   *
+   * Early completion can be done with combination of the [[takeWhile]] operator.
+   *
+   * Adheres to the [[ActorAttributes.SupervisionStrategy]] attribute.
+   *
+   * You can configure the default dispatcher for this Source by changing the `akka.stream.materializer.blocking-io-dispatcher` or
+   * set it for a given Source by using [[ActorAttributes]].
+   *
+   * '''Emits when''' the mapping function returns an element and downstream is ready to consume it
+   *
+   * '''Backpressures when''' downstream backpressures
+   *
+   * '''Completes when''' upstream completes
+   *
+   * '''Cancels when''' downstream cancels
+   *
+   * @tparam R the type of the resource
+   * @tparam T the type of the output elements
+   * @param create function that creates the resource
+   * @param f function that transforms the upstream element and the resource to output element
+   * @since 1.1.0
+   */
+  def mapWithResource[R <: AutoCloseable, T](
+      create: function.Creator[R],
+      f: function.Function2[R, Out, T]): javadsl.SubSource[T, Mat] =
+    mapWithResource(create, f,
+      (resource: AutoCloseable) => {
+        resource.close()
+        Optional.empty()
+      })
+
   /**
    * Transform each input element into an `Iterable` of output elements that is
    * then flattened into the output stream. The transformation is meant to be stateful,

Diff for: stream/src/main/scala/org/apache/pekko/stream/scaladsl/Flow.scala

@@ -1142,6 +1142,43 @@ trait FlowOps[+Out, +Mat] {
         resource => close(resource))
         .withAttributes(DefaultAttributes.mapWithResource))
 
+  /**
+   * Transform each stream element with the help of a [[AutoCloseable]] resource and close it.
+   *
+   * The resource creation function is invoked once when the stream is materialized and the returned resource is passed to
+   * the mapping function for mapping the first element. The mapping function returns a mapped element to emit
+   * downstream. The returned `T` MUST NOT be `null` as it is illegal as stream element - according to the Reactive Streams specification.
+   *
+   * The [[AutoCloseable]] resource is closed only once when the upstream or downstream finishes or fails.
+   *
+   * Early completion can be done with combination of the [[takeWhile]] operator.
+   *
+   * Adheres to the [[ActorAttributes.SupervisionStrategy]] attribute.
+   *
+   * You can configure the default dispatcher for this Source by changing the `akka.stream.materializer.blocking-io-dispatcher` or
+   * set it for a given Source by using [[ActorAttributes]].
+   *
+   * '''Emits when''' the mapping function returns an element and downstream is ready to consume it
+   *
+   * '''Backpressures when''' downstream backpressures
+   *
+   * '''Completes when''' upstream completes
+   *
+   * '''Cancels when''' downstream cancels
+   *
+   * @tparam R the type of the resource
+   * @tparam T the type of the output elements
+   * @param create function that creates the resource
+   * @param f function that transforms the upstream element and the resource to output element
+   * @since 1.1.0
+   */
+  def mapWithResource[R <: AutoCloseable, T](create: () => R, f: (R, Out) => T): Repr[T] =
+    mapWithResource(create)(f,
+      (resource: AutoCloseable) => {
+        resource.close()
+        None
+      })
+
   /**
    * Transform each input element into an `Iterable` of output elements that is
    * then flattened into the output stream. The transformation is meant to be stateful,