feat: support the flattening syntax for supervising

actor-typed-tests/src/test/java/jdocs/org/apache/pekko/typed/RouterTest.java

@@ -133,7 +133,7 @@ static Behavior<Void> showPoolRouting() {
               Routers.pool(
                   poolSize,
                   // make sure the workers are restarted if they fail
-                  Behaviors.supervise(Worker.create()).onFailure(SupervisorStrategy.restart()));
+                  Behaviors.supervise(Worker.create()).onNotFatalFailure(SupervisorStrategy.restart()));
           ActorRef<Worker.Command> router = context.spawn(pool, "worker-pool");
 
           for (int i = 0; i < 10; i++) {

actor-typed-tests/src/test/java/jdocs/org/apache/pekko/typed/supervision/SupervisionCompileOnlyTest.java

@@ -43,7 +43,7 @@ public Got(int n) {
 
     // #top-level
     public static Behavior<Command> create() {
-      return Behaviors.supervise(counter(1)).onFailure(SupervisorStrategy.restart());
+      return Behaviors.supervise(counter(1)).onNotFatalFailure(SupervisorStrategy.restart());
     }
     // #top-level
 
@@ -86,9 +86,8 @@ public void supervision() {
     // #restart-limit
 
     // #multiple
-    Behaviors.supervise(
-            Behaviors.supervise(behavior)
-                .onFailure(IllegalStateException.class, SupervisorStrategy.restart()))
+    Behaviors.supervise(behavior)
+        .onFailure(IllegalStateException.class, SupervisorStrategy.restart())
         .onFailure(IllegalArgumentException.class, SupervisorStrategy.stop());
     // #multiple
 
@@ -115,7 +114,7 @@ static Behavior<String> parent() {
                         return Behaviors.same();
                       });
                 }))
-        .onFailure(SupervisorStrategy.restart());
+        .onNotFatalFailure(SupervisorStrategy.restart());
   }
   // #restart-stop-children
 
@@ -136,7 +135,7 @@ static Behavior<String> parent2() {
                         child2.tell(parts[1]);
                         return Behaviors.same();
                       }))
-              .onFailure(SupervisorStrategy.restart().withStopChildren(false));
+              .onNotFatalFailure(SupervisorStrategy.restart().withStopChildren(false));
         });
   }
   // #restart-keep-children

actor-typed-tests/src/test/java/org/apache/pekko/actor/typed/javadsl/ActorCompile.java

@@ -183,9 +183,8 @@ public Behavior<MyMsg> receive(TypedActorContext<MyMsg> context, MyMsg message)
     SupervisorStrategy strategy7 = strategy6.withResetBackoffAfter(Duration.ofSeconds(2));
 
     Behavior<MyMsg> behv =
-        Behaviors.supervise(
-                Behaviors.supervise(Behaviors.<MyMsg>ignore())
-                    .onFailure(IllegalStateException.class, strategy6))
+        Behaviors.supervise(Behaviors.<MyMsg>ignore())
+            .onFailure(IllegalStateException.class, strategy6)
             .onFailure(RuntimeException.class, strategy1);
   }
 

actor-typed-tests/src/test/scala/docs/org/apache/pekko/typed/supervision/SupervisionCompileOnly.scala

@@ -44,7 +44,8 @@ object SupervisionCompileOnly {
 
   // #multiple
   Behaviors
-    .supervise(Behaviors.supervise(behavior).onFailure[IllegalStateException](SupervisorStrategy.restart))
+    .supervise(behavior)
+    .onFailure[IllegalStateException](SupervisorStrategy.restart)
     .onFailure[IllegalArgumentException](SupervisorStrategy.stop)
   // #multiple
 

actor-typed-tests/src/test/scala/org/apache/pekko/actor/typed/SupervisionSpec.scala

@@ -844,7 +844,8 @@ class SupervisionSpec extends ScalaTestWithActorTestKit("""
     "support nesting to handle different exceptions" in {
       val probe = TestProbe[Event]("evt")
       val behv = Behaviors
-        .supervise(Behaviors.supervise(targetBehavior(probe.ref)).onFailure[Exc2](SupervisorStrategy.resume))
+        .supervise(targetBehavior(probe.ref))
+        .onFailure[Exc2](SupervisorStrategy.resume)
         .onFailure[Exc3](SupervisorStrategy.restart)
       val ref = spawn(behv)
       ref ! IncrementState

actor-typed/src/main/scala/org/apache/pekko/actor/typed/Behavior.scala

@@ -19,11 +19,9 @@ import scala.reflect.ClassTag
 
 import org.apache.pekko
 import pekko.actor.InvalidMessageException
-import pekko.actor.typed.internal.BehaviorImpl
+import pekko.actor.typed.internal.{ BehaviorImpl, BehaviorTags, InterceptorImpl, Supervisor }
 import pekko.actor.typed.internal.BehaviorImpl.DeferredBehavior
 import pekko.actor.typed.internal.BehaviorImpl.StoppedBehavior
-import pekko.actor.typed.internal.BehaviorTags
-import pekko.actor.typed.internal.InterceptorImpl
 import pekko.annotation.DoNotInherit
 import pekko.annotation.InternalApi
 
@@ -71,6 +69,43 @@ abstract class Behavior[T](private[pekko] val _tag: Int) { behavior =>
 
 }
 
+/**
+ * TODO scala doc
+ */
+final class SuperviseBehavior[T] private[pekko] (
+    val wrapped: Behavior[T]) extends Behavior[T](BehaviorTags.SuperviseBehavior) {
+  private final val ThrowableClassTag = ClassTag(classOf[Throwable])
+
+  /** Specify the [[SupervisorStrategy]] to be invoked when the wrapped behavior throws. */
+  def onFailure[Thr <: Throwable](strategy: SupervisorStrategy)(
+      implicit tag: ClassTag[Thr] = ThrowableClassTag): SuperviseBehavior[T] = {
+    val effectiveTag = if (tag == ClassTag.Nothing) ThrowableClassTag else tag
+    new SuperviseBehavior[T](Supervisor(Behavior.validateAsInitial(wrapped), strategy)(effectiveTag))
+  }
+
+  /**
+   * Java API:
+   * Specify the [[SupervisorStrategy]] to be invoked when the wrapped behavior throws.
+   *
+   * Only exceptions of the given type (and their subclasses) will be handled by this supervision behavior.
+   */
+  def onFailure[Thr <: Throwable](clazz: Class[Thr], strategy: SupervisorStrategy): SuperviseBehavior[T] = {
+    onFailure(strategy)(ClassTag(clazz))
+  }
+
+  /**
+   * Java API:
+   * Specify the [[SupervisorStrategy]] to be invoked when the wrapped behavior throws.
+   *
+   * Only exceptions of the given type (and their subclasses) will be handled by this supervision behavior.
+   */
+  def onNotFatalFailure[Thr <: Throwable](strategy: SupervisorStrategy): SuperviseBehavior[T] = {
+    onFailure(classOf[Exception], strategy)
+  }
+
+  private[pekko] def unwrap: Behavior[T] = wrapped
+}
+
 /**
  * Extension point for implementing custom behaviors in addition to the existing
  * set of behaviors available through the DSLs in [[pekko.actor.typed.scaladsl.Behaviors]] and [[pekko.actor.typed.javadsl.Behaviors]]
@@ -180,7 +215,8 @@ object Behavior {
         val startedInner = start(wrapped.nestedBehavior, ctx.asInstanceOf[TypedActorContext[Any]])
         if (startedInner eq wrapped.nestedBehavior) wrapped
         else wrapped.replaceNested(startedInner)
-      case _ => behavior
+      case supervise: SuperviseBehavior[T] => start(supervise.unwrap, ctx)
+      case _                               => behavior
     }
   }
 
@@ -266,6 +302,8 @@ object Behavior {
         throw new IllegalArgumentException(s"cannot execute with [$behavior] as behavior")
       case BehaviorTags.DeferredBehavior =>
         throw new IllegalArgumentException(s"deferred [$behavior] should not be passed to interpreter")
+      case BehaviorTags.SuperviseBehavior =>
+        throw new IllegalArgumentException(s"supervise [$behavior] should not be passed to interpreter")
       case BehaviorTags.IgnoreBehavior =>
         BehaviorImpl.same[T]
       case BehaviorTags.StoppedBehavior =>

actor-typed/src/main/scala/org/apache/pekko/actor/typed/internal/BehaviorImpl.scala

@@ -39,6 +39,7 @@ private[pekko] object BehaviorTags {
   final val SameBehavior = 6
   final val FailedBehavior = 7
   final val StoppedBehavior = 8
+  final val SuperviseBehavior = 9
 
 }
 

actor-typed/src/main/scala/org/apache/pekko/actor/typed/javadsl/Behaviors.scala

@@ -261,27 +261,8 @@ object Behaviors {
    *      .onFailure[IndexOutOfBoundsException](SupervisorStrategy.resume) // resume for IndexOutOfBoundsException exceptions
    * }}}
    */
-  def supervise[T](wrapped: Behavior[T]): Supervise[T] =
-    new Supervise[T](wrapped)
-
-  final class Supervise[T] private[pekko] (wrapped: Behavior[T]) {
-
-    /**
-     * Specify the [[SupervisorStrategy]] to be invoked when the wrapped behavior throws.
-     *
-     * Only exceptions of the given type (and their subclasses) will be handled by this supervision behavior.
-     */
-    def onFailure[Thr <: Throwable](clazz: Class[Thr], strategy: SupervisorStrategy): Behavior[T] =
-      Supervisor(Behavior.validateAsInitial(wrapped), strategy)(ClassTag(clazz))
-
-    /**
-     * Specify the [[SupervisorStrategy]] to be invoked when the wrapped behavior throws.
-     *
-     * All non-fatal (see [[scala.util.control.NonFatal]]) exceptions types will be handled using the given strategy.
-     */
-    def onFailure(strategy: SupervisorStrategy): Behavior[T] =
-      onFailure(classOf[Exception], strategy)
-  }
+  def supervise[T](wrapped: Behavior[T]): SuperviseBehavior[T] =
+    scaladsl.Behaviors.supervise(wrapped)
 
   /**
    * Transform the incoming messages by placing a funnel in front of the wrapped `Behavior`: the supplied

actor-typed/src/main/scala/org/apache/pekko/actor/typed/scaladsl/Behaviors.scala

@@ -225,19 +225,8 @@ object Behaviors {
    *      .onFailure[IndexOutOfBoundsException](SupervisorStrategy.resume) // resume for IndexOutOfBoundsException exceptions
    * }}}
    */
-  def supervise[T](wrapped: Behavior[T]): Supervise[T] =
-    new Supervise[T](wrapped)
-
-  private final val ThrowableClassTag = ClassTag(classOf[Throwable])
-  final class Supervise[T] private[pekko] (val wrapped: Behavior[T]) extends AnyVal {
-
-    /** Specify the [[SupervisorStrategy]] to be invoked when the wrapped behavior throws. */
-    def onFailure[Thr <: Throwable](strategy: SupervisorStrategy)(
-        implicit tag: ClassTag[Thr] = ThrowableClassTag): Behavior[T] = {
-      val effectiveTag = if (tag == ClassTag.Nothing) ThrowableClassTag else tag
-      Supervisor(Behavior.validateAsInitial(wrapped), strategy)(effectiveTag)
-    }
-  }
+  def supervise[T](wrapped: Behavior[T]): SuperviseBehavior[T] =
+    new SuperviseBehavior[T](wrapped)
 
   /**
    * Support for scheduled `self` messages in an actor.

cluster-typed/src/test/java/jdocs/org/apache/pekko/cluster/typed/SingletonCompileOnlyTest.java

@@ -120,7 +120,7 @@ public static void backoff() {
         singleton.init(
             SingletonActor.of(
                 Behaviors.supervise(Counter.create())
-                    .onFailure(
+                    .onNotFatalFailure(
                         SupervisorStrategy.restartWithBackoff(
                             Duration.ofSeconds(1), Duration.ofSeconds(10), 0.2)),
                 "GlobalCounter"));

docs/src/test/java/jdocs/typed/tutorial_1/ActorHierarchyExperiments.java

@@ -120,7 +120,7 @@ private SupervisingActor(ActorContext<String> context) {
     super(context);
     child =
         context.spawn(
-            Behaviors.supervise(SupervisedActor.create()).onFailure(SupervisorStrategy.restart()),
+            Behaviors.supervise(SupervisedActor.create()).onNotFatalFailure(SupervisorStrategy.restart()),
             "supervised-actor");
   }
 

persistence-typed-tests/src/test/java/org/apache/pekko/persistence/typed/javadsl/EventSourcedBehaviorJavaDslTest.java

@@ -398,7 +398,7 @@ protected void log() {
   public void workWhenWrappedInOtherBehavior() {
     Behavior<Command> behavior =
         Behaviors.supervise(counter(PersistenceId.ofUniqueId("c6")))
-            .onFailure(
+            .onNotFatalFailure(
                 SupervisorStrategy.restartWithBackoff(
                     Duration.ofSeconds(1), Duration.ofSeconds(10), 0.1));
     ActorRef<Command> c = testKit.spawn(behavior);