rx (feature): Add Rx.runOn for side-effect

Author: xerial

airframe-http/src/test/scala/wvlet/airframe/http/client/RPCErrorHandlingTest.scala

@@ -0,0 +1,78 @@
+/*
+ * 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 wvlet.airframe.http.client
+
+import wvlet.airframe.http.HttpMessage.Response
+import wvlet.airframe.http.{Http, HttpMessage, HttpStatus, RPCMethod, RPCStatus, RxHttpEndpoint, RxHttpFilter}
+import wvlet.airframe.rx.Rx
+import wvlet.airframe.surface.Surface
+import wvlet.airspec.AirSpec
+
+import java.util.concurrent.atomic.AtomicInteger
+import scala.util.{Failure, Success}
+
+object RPCErrorHandlingTest extends AirSpec {
+
+  private class DummyHttpChannel extends HttpChannel {
+    private val requestCount = new AtomicInteger(0)
+
+    override def send(req: HttpMessage.Request, channelConfig: HttpChannelConfig): Response = ???
+    override def sendAsync(req: HttpMessage.Request, channelConfig: HttpChannelConfig): Rx[Response] = {
+      requestCount.getAndIncrement() match {
+        case 0 =>
+          // For retry test
+          Rx.single(Http.response(HttpStatus.InternalServerError_500))
+        case _ =>
+          Rx.single(RPCStatus.INVALID_REQUEST_U1.newException("RPC error").toResponse)
+      }
+    }
+    override def close(): Unit = {}
+  }
+
+  initDesign {
+    _.bind[AsyncClient].toInstance {
+      val config = Http.client // .withClientFilter(null)
+      new AsyncClientImpl(new DummyHttpChannel, config)
+    }
+  }
+
+  test("test error response") { (client: AsyncClient) =>
+    val dummyRPCMerthod =
+      RPCMethod("/rpc_method", "demo.RPCClass", "hello", Surface.of[Map[String, Any]], Surface.of[String])
+
+    var observedRetry = false
+    val myClient = client.withClientFilter(new RxHttpFilter {
+      override def apply(request: HttpMessage.Request, next: RxHttpEndpoint): Rx[Response] = {
+        next(request).runOn {
+          case Success(x) =>
+            x.status match {
+              case HttpStatus.InternalServerError_500 =>
+                observedRetry = true
+              case _ =>
+            }
+          case Failure(e) =>
+            logger.info(e)
+        }
+      }
+    })
+
+    myClient.rpc[Map[String, Any], String](dummyRPCMerthod, Map("message" -> "world")).transform {
+      case scala.util.Failure(e) =>
+        observedRetry shouldBe true
+        e.getMessage shouldContain "RPC error"
+      case _ =>
+        fail("should fail")
+    }
+  }
+}

airframe-rx/src/main/scala/wvlet/airframe/rx/Rx.scala

@@ -14,9 +14,7 @@
 package wvlet.airframe.rx
 
 import java.util.concurrent.TimeUnit
-import wvlet.airframe.rx.Rx.{RecoverOp, RecoverWithOp}
 import wvlet.log.LogSupport
-
 import scala.concurrent.{ExecutionContext, Future}
 import scala.language.higherKinds
 import scala.util.{Failure, Success, Try}
@@ -26,19 +24,39 @@ import scala.util.{Failure, Success, Try}
   * @tparam A
   */
 trait RxOps[+A] { self =>
+  import Rx._
+
   def parents: Seq[RxOps[_]]
 
   def toRx: Rx[A]
 
   /**
     * Recover from a known error and emit a replacement value
     */
-  def recover[U](f: PartialFunction[Throwable, U]): Rx[U] = RecoverOp(this.toRx, f)
+  def recover[U](f: PartialFunction[Throwable, U]): Rx[U] = RecoverOp(this, f)
 
   /**
     * Recover from a known error and emit replacement values from a given Rx
     */
-  def recoverWith[A](f: PartialFunction[Throwable, RxOps[A]]): Rx[A] = RecoverWithOp(this.toRx, f)
+  def recoverWith[A](f: PartialFunction[Throwable, RxOps[A]]): Rx[A] = RecoverWithOp(this, f)
+
+  /**
+    * A utility method for running the given effect function when the Rx value is available and bypass the original
+    * input value as is.
+    *
+    * This method is useful for debugging Rx chains. For example:
+    * {{{
+    *   rx.runOn {
+    *     case Success(v) => debug(s"received ${v}")
+    *     case Failure(e) => error(s"request failed", e)
+    *   }
+    * }}}
+    *
+    * @param f
+    * @return
+    *   the original input value as is
+    */
+  def runOn(f: PartialFunction[Try[A], Unit]): Rx[A] = RunOnOp(this, f)
 
   /**
     * Evaluate this Rx[A] and apply the given effect function. Once OnError(e) or OnCompletion is observed, it will stop
@@ -416,6 +434,7 @@ object Rx extends LogSupport {
   }
   case class RecoverOp[A, U](input: RxOps[A], f: PartialFunction[Throwable, U])            extends UnaryRx[A, U]
   case class RecoverWithOp[A, U](input: RxOps[A], f: PartialFunction[Throwable, RxOps[U]]) extends UnaryRx[A, U]
+  case class RunOnOp[A](input: RxOps[A], f: PartialFunction[Try[A], Unit])                 extends UnaryRx[A, A]
 
   case class IntervalOp(interval: Long, unit: TimeUnit) extends Rx[Long] {
     override def parents: Seq[RxOps[_]] = Seq.empty

airframe-rx/src/main/scala/wvlet/airframe/rx/RxRunner.scala

@@ -2,13 +2,11 @@ package wvlet.airframe.rx
 
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicBoolean
-
 import wvlet.log.LogSupport
 
 import scala.annotation.tailrec
 import scala.collection.immutable.Queue
 import scala.util.{Failure, Success, Try}
-import Rx._
 
 /**
   * States for propagating the result of the downstream operators.
@@ -75,6 +73,8 @@ class RxRunner(
     continuous: Boolean
 ) extends LogSupport { runner =>
 
+  import Rx._
+
   /**
     * Build an executable chain of Rx operators. The resulting chain will be registered as a subscriber to the root node
     * (see RxVar.foreach). If the root value changes, the effect code block will be executed.
@@ -459,6 +459,19 @@ class RxRunner(
         Cancelable { () =>
           c1.cancel; c2.cancel
         }
+      case RunOnOp(in, f) =>
+        val f0 = f.asInstanceOf[PartialFunction[Try[_], Unit]]
+        run(in) { ev =>
+          ev match {
+            case OnNext(v) =>
+              // Skip the error handling so as not to discard the input event
+              Try(f0.applyOrElse(Success(v), (_: Try[_]) => ()))
+            case OnError(e) =>
+              Try(f0.applyOrElse(Failure(e), (_: Try[_]) => ()))
+            case _ =>
+          }
+          effect(ev)
+        }
       case SingleOp(v) =>
         Try(effect(OnNext(v.eval))) match {
           case Success(c) =>