Publisher scanWith enhancements (#2640)

Motivation:
Publisher scanWith operators only allow a limited set of
mapped terminal transformations. Use cases like mapping
state and preserving the onError Throwable aren't possible.

Modifications:
- Introduce ScanMapper and ScanLifetimeMapper which
support MappedTerminal that allows all the terminal combinations
to be mapped.
- Deprecate ScanWithMapper and ScanWithLifetimeMapper in favor of
these new APIs.
- Deprecate Publisher.scanWith and Publisher.scanWithLiftime in
favor of the new extension operators.

servicetalk-concurrent-api/gradle/checkstyle/suppressions.xml

@@ -23,4 +23,5 @@
     files="docs[\\/]modules[\\/]ROOT[\\/]assets[\\/]images[\\/].+\.svg"/>
   <!-- mapOnError supports re-throwing a Throwable from onError(Throwable) -->
   <suppress checks="IllegalThrowsCheck" files="io[\\/]servicetalk[\\/]concurrent[\\/]api[\\/]ScanWithMapper.java"/>
+  <suppress checks="IllegalThrowsCheck" files="io[\\/]servicetalk[\\/]concurrent[\\/]api[\\/]ScanMapper.java"/>
 </suppressions>

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/Publisher.java

@@ -287,11 +287,54 @@ public final <R> Publisher<R> scanWith(Supplier<R> initial, BiFunction<R, ? supe
      * @param <R> Type of the items emitted by the returned {@link Publisher}.
      * @return A {@link Publisher} that transforms elements emitted by this {@link Publisher} into a different type.
      * @see <a href="https://reactivex.io/documentation/operators/scan.html">ReactiveX scan operator.</a>
+     * @deprecated Use {@link #scanWithMapper(Supplier)}.
      */
+    @Deprecated
     public final <R> Publisher<R> scanWith(Supplier<? extends ScanWithMapper<? super T, ? extends R>> mapperSupplier) {
         return new ScanWithPublisher<>(this, mapperSupplier);
     }
 
+    /**
+     * Apply a function to each {@link Subscriber#onNext(Object)} emitted by this {@link Publisher} as well as
+     * optionally concat one {@link Subscriber#onNext(Object)} signal before the terminal signal is emitted downstream.
+     * <p>
+     * This method provides a data transformation in sequential programming similar to:
+     * <pre>{@code
+     *     List<R> results = ...;
+     *     ScanWithLifetimeMapperExt<T, R> mapper = mapperSupplier.get();
+     *     MappedTerminal<R> mapped = null;
+     *     try {
+     *       for (T t : resultOfThisPublisher()) {
+     *         results.add(mapper.mapOnNext(t));
+     *       }
+     *     } catch (Throwable cause) {
+     *       mapped = mapper.mapOnError(cause);
+     *       if (mapped == null) {
+     *         throw cause;
+     *       }
+     *     }
+     *     if (mapped == null) {
+     *       mapped = mapper.mapOnComplete();
+     *     }
+     *     if (mapped.onNextValid()) {
+     *       results.add(mapped.onNext());
+     *     }
+     *     if (mapped.terminal() != null) {
+     *       throw mapped.terminal();
+     *     }
+     *     return results;
+     * }</pre>
+     * @param mapperSupplier Invoked on each {@link PublisherSource#subscribe(Subscriber)} and maintains any necessary
+     * state for the mapping/accumulation for each {@link Subscriber}.
+     * @param <R> Type of the items emitted by the returned {@link Publisher}.
+     * @return A {@link Publisher} that transforms elements emitted by this {@link Publisher} into a different type.
+     * @see <a href="https://reactivex.io/documentation/operators/scan.html">ReactiveX scan operator.</a>
+     */
+    public final <R> Publisher<R> scanWithMapper(
+            Supplier<? extends ScanMapper<? super T, ? extends R>> mapperSupplier) {
+        return new ScanWithPublisher<>(mapperSupplier, this);
+    }
+
     /**
      * Apply a function to each {@link Subscriber#onNext(Object)} emitted by this {@link Publisher} as well as
      * optionally concat one {@link Subscriber#onNext(Object)} signal before the terminal signal is emitted downstream.
@@ -329,12 +372,63 @@ public final <R> Publisher<R> scanWith(Supplier<? extends ScanWithMapper<? super
      * @param <R> Type of the items emitted by the returned {@link Publisher}.
      * @return A {@link Publisher} that transforms elements emitted by this {@link Publisher} into a different type.
      * @see <a href="https://reactivex.io/documentation/operators/scan.html">ReactiveX scan operator.</a>
+     * @deprecated Use {@link #scanWithLifetimeMapper(Supplier)}.
      */
+    @Deprecated
     public final <R> Publisher<R> scanWithLifetime(
             Supplier<? extends ScanWithLifetimeMapper<? super T, ? extends R>> mapperSupplier) {
         return new ScanWithLifetimePublisher<>(this, mapperSupplier);
     }
 
+    /**
+     * Apply a function to each {@link Subscriber#onNext(Object)} emitted by this {@link Publisher} as well as
+     * optionally concat one {@link Subscriber#onNext(Object)} signal before the terminal signal is emitted downstream.
+     * Additionally the {@link ScanLifetimeMapper#afterFinally()} method will be invoked on terminal or cancel
+     * signals which enables cleanup of state (if required). This provides a similar lifetime management as
+     * {@link TerminalSignalConsumer}.
+     *
+     * <p>
+     * This method provides a data transformation in sequential programming similar to:
+     * <pre>{@code
+     *     List<R> results = ...;
+     *     ScanWithLifetimeMapperExt<T, R> mapper = mapperSupplier.get();
+     *     try {
+     *       MappedTerminal<R> mapped = null;
+     *       try {
+     *         for (T t : resultOfThisPublisher()) {
+     *           results.add(mapper.mapOnNext(t));
+     *         }
+     *       } catch (Throwable cause) {
+     *         mapped = mapper.mapOnError(cause);
+     *         if (mapped == null) {
+     *           throw cause;
+     *         }
+     *       }
+     *       if (mapped == null) {
+     *         mapped = mapper.mapOnComplete();
+     *       }
+     *       if (mapped.onNextValid()) {
+     *         results.add(mapped.onNext());
+     *       }
+     *       if (mapped.terminal() != null) {
+     *          throw mapped.terminal();
+     *       }
+     *     } finally {
+     *       mapper.afterFinally();
+     *     }
+     *     return results;
+     * }</pre>
+     * @param mapperSupplier Invoked on each {@link PublisherSource#subscribe(Subscriber)} and maintains any necessary
+     * state for the mapping/accumulation for each {@link Subscriber}.
+     * @param <R> Type of the items emitted by the returned {@link Publisher}.
+     * @return A {@link Publisher} that transforms elements emitted by this {@link Publisher} into a different type.
+     * @see <a href="https://reactivex.io/documentation/operators/scan.html">ReactiveX scan operator.</a>
+     */
+    public final <R> Publisher<R> scanWithLifetimeMapper(
+            Supplier<? extends ScanLifetimeMapper<? super T, ? extends R>> mapperSupplier) {
+        return new ScanWithLifetimePublisher<>(mapperSupplier, this);
+    }
+
     /**
      * Transform errors emitted on this {@link Publisher} into a {@link Subscriber#onComplete()} signal
      * (e.g. swallows the error).

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/ScanLifetimeMapper.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright © 2023 Apple Inc. and the ServiceTalk project authors
+ *
+ * Licensed 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 io.servicetalk.concurrent.api;
+
+import io.servicetalk.concurrent.PublisherSource.Subscriber;
+import io.servicetalk.concurrent.PublisherSource.Subscription;
+
+import java.util.function.Supplier;
+
+/**
+ * Provides the ability to transform (aka map) signals emitted via
+ * the {@link Publisher#scanWithLifetimeMapper(Supplier)} operator, as well as the ability to cleanup state
+ * via {@link #afterFinally}.
+ * @param <T> Type of items emitted by the {@link Publisher} this operator is applied.
+ * @param <R> Type of items emitted by this operator.
+ */
+public interface ScanLifetimeMapper<T, R> extends ScanMapper<T, R> {
+    /**
+     * Invoked after a terminal signal {@link Subscriber#onError(Throwable)} or
+     * {@link Subscriber#onComplete()} or {@link Subscription#cancel()}.
+     * No further interaction will occur with the {@link ScanLifetimeMapper} to prevent use-after-free
+     * on internal state.
+     */
+    void afterFinally();
+}

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/ScanMapper.java

@@ -0,0 +1,106 @@
+/*
+ * Copyright © 2023 Apple Inc. and the ServiceTalk project authors
+ *
+ * Licensed 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 io.servicetalk.concurrent.api;
+
+import io.servicetalk.concurrent.PublisherSource.Subscriber;
+
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
+
+/**
+ * Provides the ability to transform (aka map) signals emitted via the {@link Publisher#scanWithMapper(Supplier)}
+ * operator.
+ * @param <T> Type of items emitted by the {@link Publisher} this operator is applied.
+ * @param <R> Type of items emitted by this operator.
+ */
+public interface ScanMapper<T, R> {
+    /**
+     * Invoked on each {@link Subscriber#onNext(Object)} signal and maps from type {@link T} to type {@link R}.
+     * @param next The next element emitted from {@link Subscriber#onNext(Object)}.
+     * @return The result of mapping {@code next}.
+     */
+    @Nullable
+    R mapOnNext(@Nullable T next);
+
+    /**
+     * Invoked when a {@link Subscriber#onError(Throwable)} signal is received and can map the current state into an
+     * object of type {@link R} which will be emitted downstream as {@link Subscriber#onNext(Object)}, followed by
+     * a terminal signal.
+     * <p>
+     * If this method throws the exception will be propagated downstream via {@link Subscriber#onError(Throwable)}.
+     * @param cause The cause from upstream {@link Subscriber#onError(Throwable)}.
+     * @return
+     * <ul>
+     *     <li>{@code null} if no mapping is required and {@code cause} is propagated to
+     *     {@link Subscriber#onError(Throwable)}</li>
+     *     <li>non-{@code null} will propagate {@link MappedTerminal#onNext()} to {@link Subscriber#onNext(Object)}
+     *     then will terminate with {@link MappedTerminal#terminal()}</li>
+     * </ul>
+     * @throws Throwable If an exception occurs, which will be propagated downstream via
+     * {@link Subscriber#onError(Throwable)}.
+     */
+    @Nullable
+    MappedTerminal<R> mapOnError(Throwable cause) throws Throwable;
+
+    /**
+     * Invoked when a {@link Subscriber#onComplete()} signal is received and can map the current state into an
+     * object of type {@link R} which will be emitted downstream as {@link Subscriber#onNext(Object)}, followed by
+     * a terminal signal.
+     * <p>
+     * If this method throws the exception will be propagated downstream via {@link Subscriber#onError(Throwable)}.
+     * @return
+     * <ul>
+     *     <li>{@code null} if no mapping is required and {@code cause} is propagated to
+     *     {@link Subscriber#onError(Throwable)}</li>
+     *     <li>non-{@code null} will propagate {@link MappedTerminal#onNext()} to {@link Subscriber#onNext(Object)}
+     *     then will terminate with {@link MappedTerminal#terminal()}</li>
+     * </ul>
+     * @throws Throwable If an exception occurs, which will be propagated downstream via
+     * {@link Subscriber#onError(Throwable)}.
+     */
+    @Nullable
+    MappedTerminal<R> mapOnComplete() throws Throwable;
+
+    /**
+     * Result of a mapping operation of a terminal signal.
+     * @param <R> The mapped result type.
+     */
+    interface MappedTerminal<R> {
+        /**
+         * Get the signal to be delivered to {@link Subscriber#onNext(Object)} if {@link #onNextValid()}.
+         * @return the signal to be delivered to {@link Subscriber#onNext(Object)} if {@link #onNextValid()}.
+         */
+        @Nullable
+        R onNext();
+
+        /**
+         * Determine if {@link #onNext()} is valid and should be propagated downstream.
+         * @return {@code true} to propagate {@link #onNext()}, {@code false} will only propagate {@link #terminal()}.
+         */
+        boolean onNextValid();
+
+        /**
+         * The terminal event to propagate.
+         * @return
+         * <ul>
+         *     <li>{@code null} means {@link Subscriber#onComplete()}</li>
+         *     <li>non-{@code null} will propagate as {@link Subscriber#onError(Throwable)}</li>
+         * </ul>
+         */
+        @Nullable
+        Throwable terminal();
+    }
+}

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/ScanWithLifetimeMapper.java

@@ -25,9 +25,10 @@
  * via {@link #afterFinally}.
  * @param <T> Type of items emitted by the {@link Publisher} this operator is applied.
  * @param <R> Type of items emitted by this operator.
+ * @deprecated Use {@link ScanLifetimeMapper}.
  */
+@Deprecated
 public interface ScanWithLifetimeMapper<T, R> extends ScanWithMapper<T, R> {
-
     /**
      * Invoked after a terminal signal {@link PublisherSource.Subscriber#onError(Throwable)} or
      * {@link PublisherSource.Subscriber#onComplete()} or {@link PublisherSource.Subscription#cancel()}.