Repository: wvlet/scala-ai
Pull request #75: feature: Add wvlet.ai.core.control module
Author: login: xerial
Labels: name: feature
description: New feature
Created: 2025-05-19T04:06:45Z | Merged: 2025-05-19T04:30:32Z

**Description**

Introduced the `wvlet.ai.core.control` module to enhance functionality and modularity.

**Related Issue/Task**

N/A

**Checklist**

- [ ] This pull request focuses on a single task.
- [ ] The change does not contain security credentials


Diff:
diff --git a/ai-core/.js/src/main/scala/wvlet/ai/core/control/Compat.scala b/ai-core/.js/src/main/scala/wvlet/ai/core/control/Compat.scala
new file mode 100644
index 00000000..4ea1d3f9
--- /dev/null
+++ b/ai-core/.js/src/main/scala/wvlet/ai/core/control/Compat.scala
@@ -0,0 +1,23 @@
+/*
+ * 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.ai.core.control
+
+import scala.concurrent.{ExecutionContext, Future, Promise}
+
+/**
+  */
+object Compat:
+  def sleep(millis: Long): Unit = {
+    // do nothing in Scala.js as their is no way to implement sleep https://github.com/scala-js/scala-js/issues/1898
+  }
diff --git a/ai-core/.jvm/src/main/scala/wvlet/ai/core/control/Compat.scala b/ai-core/.jvm/src/main/scala/wvlet/ai/core/control/Compat.scala
new file mode 100644
index 00000000..01a9ac22
--- /dev/null
+++ b/ai-core/.jvm/src/main/scala/wvlet/ai/core/control/Compat.scala
@@ -0,0 +1,24 @@
+/*
+ * 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.ai.core.control
+
+import java.util.concurrent.{ScheduledThreadPoolExecutor, TimeUnit}
+import scala.concurrent.{ExecutionContext, Future, Promise}
+import scala.concurrent.{ExecutionContext, Future, Promise}
+
+/**
+  * Compat object provides compatibility functions for thread operations.
+  */
+object Compat:
+  def sleep(millis: Long): Unit = Thread.sleep(millis)
diff --git a/ai-core/.jvm/src/main/scala/wvlet/ai/core/log/AsyncHandler.scala b/ai-core/.jvm/src/main/scala/wvlet/ai/core/log/AsyncHandler.scala
index 216bd31c..7c31d777 100644
--- a/ai-core/.jvm/src/main/scala/wvlet/ai/core/log/AsyncHandler.scala
+++ b/ai-core/.jvm/src/main/scala/wvlet/ai/core/log/AsyncHandler.scala
@@ -1,7 +1,6 @@
 package wvlet.ai.core.log
 
 import wvlet.ai.core.control.{Guard, ThreadUtil}
-
 import java.io.Flushable
 import java.util
 import java.util.concurrent.atomic.AtomicBoolean
diff --git a/ai-core/.native/src/main/scala/wvlet/ai/core/control/Compat.scala b/ai-core/.native/src/main/scala/wvlet/ai/core/control/Compat.scala
new file mode 100644
index 00000000..f785f0b7
--- /dev/null
+++ b/ai-core/.native/src/main/scala/wvlet/ai/core/control/Compat.scala
@@ -0,0 +1,23 @@
+/*
+ * 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.ai.core.control
+
+import java.util.concurrent.{ScheduledThreadPoolExecutor, TimeUnit}
+import scala.concurrent.{ExecutionContext, Future, Promise}
+
+/**
+  * Compatibility layer for Scala Native
+  */
+object Compat:
+  def sleep(millis: Long): Unit = Thread.sleep(millis)
diff --git a/ai-core/src/main/scala/wvlet/ai/core/control/CircuitBreaker.scala b/ai-core/src/main/scala/wvlet/ai/core/control/CircuitBreaker.scala
new file mode 100644
index 00000000..1afc5f49
--- /dev/null
+++ b/ai-core/src/main/scala/wvlet/ai/core/control/CircuitBreaker.scala
@@ -0,0 +1,361 @@
+/*
+ * 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.ai.core.control
+
+import java.util.concurrent.atomic.{AtomicInteger, AtomicLong, AtomicReference}
+
+import scala.util.{Failure, Success, Try}
+import wvlet.ai.core.log.LogSupport
+import wvlet.ai.core.control.Retry.{Jitter, RetryPolicy, RetryPolicyConfig, RetryableFailure}
+
+import scala.reflect.ClassTag
+
+/**
+  * An exception thrown when the circuit breaker is open.
+  */
+case class CircuitBreakerOpenException(context: CircuitBreakerContext) extends Exception
+
+sealed trait CircuitBreakerState
+
+/**
+  */
+object CircuitBreaker extends LogSupport:
+
+  case object OPEN      extends CircuitBreakerState
+  case object HALF_OPEN extends CircuitBreakerState
+  case object CLOSED    extends CircuitBreakerState
+
+  def default: CircuitBreaker                         = new CircuitBreaker()
+  def newCircuitBreaker(name: String): CircuitBreaker = new CircuitBreaker().withName(name)
+
+  /**
+    * Create a CircuitBreaker that will be open after observing numFailures out of numExecutions.
+    */
+  def withFailureThreshold(numFailures: Int, numExecutions: Int = 10): CircuitBreaker = default
+    .withHealthCheckPolicy(HealthCheckPolicy.markDeadOnFailureThreshold(numFailures, numExecutions))
+
+  /**
+    * Create a CircuitBreaker that will be open if the failure rate in a time window exceeds the
+    * given threshold. The failure rate will be decayed exponentially as time goes.
+    */
+  def withFailureRate(failureRate: Double, timeWindowMillis: Int = 60000): CircuitBreaker = default
+    .withHealthCheckPolicy(
+      HealthCheckPolicy.markDeadOnRecentFailureRate(failureRate, timeWindowMillis)
+    )
+
+  /**
+    * Create a CircuitBreaker that will be open if the number of consecutive failures exceeds the
+    * given threshold.
+    */
+  def withConsecutiveFailures(numFailures: Int): CircuitBreaker = default.withHealthCheckPolicy(
+    HealthCheckPolicy.markDeadOnConsecutiveFailures(numFailures)
+  )
+
+  def alwaysClosed: CircuitBreaker = default.withHealthCheckPolicy(HealthCheckPolicy.alwaysHealthy)
+
+  private[control] def throwOpenException: CircuitBreakerContext => Unit =
+    (ctx: CircuitBreakerContext) => throw CircuitBreakerOpenException(ctx)
+
+  private[control] def reportStateChange =
+    (ctx: CircuitBreakerContext) => info(s"CircuitBreaker[${ctx.name}] is changed to ${ctx.state}")
+
+end CircuitBreaker
+
+import CircuitBreaker.*
+
+/**
+  * A safe interface for accessing CircuitBreaker states when handling events.
+  */
+trait CircuitBreakerContext:
+  def name: String
+  def state: CircuitBreakerState
+  def lastFailure: Option[Throwable]
+
+trait CircuitBreakerRecoveryPolicy:
+
+  /**
+    * Called when a request succeeds
+    */
+  def recordSuccess(): Unit
+
+  /**
+    * Called when request is failed.
+    */
+  def recordFailure(): Unit
+
+  /**
+    * Called when this policy is needed to reset.
+    */
+  def reset(): Unit
+
+  /**
+    * Check if CircuitBreaker can recover from HALF_OPEN to CLOSE.
+    */
+  def canRecover: Boolean
+
+object CircuitBreakerRecoveryPolicy:
+  def recoverImmediately =
+    new CircuitBreakerRecoveryPolicy():
+      override def recordSuccess(): Unit = {}
+      override def recordFailure(): Unit = {}
+      override def reset(): Unit         = {}
+      override def canRecover: Boolean   = true
+
+  def recoverAfterConsecutiveSuccesses(numberOfSuccess: Int) =
+    new CircuitBreakerRecoveryPolicy():
+      private val counter                = new AtomicInteger(0)
+      override def recordSuccess(): Unit = counter.incrementAndGet()
+      override def recordFailure(): Unit = counter.set(0)
+      override def reset(): Unit         = counter.set(0)
+      override def canRecover: Boolean   = counter.get() >= numberOfSuccess
+
+  def recoverAfterWait(elapsedTimeMillis: Int) =
+    new CircuitBreakerRecoveryPolicy():
+      private val timestamp = new AtomicLong(Long.MaxValue)
+      override def recordSuccess(): Unit = timestamp.compareAndSet(
+        Long.MaxValue,
+        System.currentTimeMillis()
+      )
+      override def recordFailure(): Unit = timestamp.set(Long.MaxValue)
+      override def reset(): Unit         = timestamp.set(Long.MaxValue)
+      override def canRecover: Boolean =
+        timestamp.get() <= System.currentTimeMillis() - elapsedTimeMillis
+
+case class CircuitBreaker(
+    name: String = "default",
+    healthCheckPolicy: HealthCheckPolicy = HealthCheckPolicy.markDeadOnConsecutiveFailures(3),
+    resultClassifier: Any => ResultClass = ResultClass.ALWAYS_SUCCEED,
+    errorClassifier: Throwable => ResultClass.Failed = ResultClass.ALWAYS_RETRY,
+    onOpenFailureHandler: CircuitBreakerContext => Unit = CircuitBreaker.throwOpenException,
+    onStateChangeListener: CircuitBreakerContext => Unit = CircuitBreaker.reportStateChange,
+    fallbackHandler: Throwable => Any = t => throw t,
+    delayAfterMarkedDead: RetryPolicy =
+      new Jitter(new RetryPolicyConfig(initialIntervalMillis = 30000)), // 30 seconds
+    recoveryPolicy: CircuitBreakerRecoveryPolicy = CircuitBreakerRecoveryPolicy.recoverImmediately,
+    private var nextProvingTimeMillis: Long = Long.MaxValue,
+    private var provingWaitTimeMillis: Long = 0L,
+    var lastFailure: Option[Throwable] = None,
+    private val currentState: AtomicReference[CircuitBreakerState] =
+      new AtomicReference(CircuitBreaker.CLOSED)
+) extends CircuitBreakerContext
+    with LogSupport:
+  def state: CircuitBreakerState = currentState.get()
+
+  /**
+    * Set the name of this CircuitBreaker
+    */
+  def withName(newName: String): CircuitBreaker = this.copy(name = newName)
+
+  /**
+    * Set a health check policy, which will be used to determine the state of the target service.
+    */
+  def withHealthCheckPolicy(newHealthCheckPolicy: HealthCheckPolicy): CircuitBreaker = this.copy(
+    healthCheckPolicy = newHealthCheckPolicy
+  )
+
+  /**
+    * Set a classifier to determine whether the execution result of the code block is successful or
+    * not.
+    */
+  def withResultClassifier(newResultClassifier: Any => ResultClass): CircuitBreaker = this.copy(
+    resultClassifier = newResultClassifier
+  )
+
+  /**
+    * Set a classifier to determine whether the exception happened in the code block can be
+    * ignoreable or not for the accessing the target service.
+    */
+  def withErrorClassifier(newErrorClassifier: Throwable => ResultClass.Failed): CircuitBreaker =
+    this.copy(errorClassifier = newErrorClassifier)
+
+  /**
+    * Set a delay policy until moving the state from OPEN to HALF_OPEN (probing) state. The default
+    * is Jittered-exponential backoff delay with the initial interval of 30 seconds.
+    */
+  def withDelayAfterMarkedDead(retryPolicy: RetryPolicy): CircuitBreaker = this.copy(
+    delayAfterMarkedDead = retryPolicy
+  )
+
+  /**
+    * Set a fallback handler which process the exception happened in the code block. The default is
+    * just throwing the exception as it is.
+    */
+  def withFallbackHandler(handler: Throwable => Any): CircuitBreaker = this.copy(fallbackHandler =
+    handler
+  )
+
+  /**
+    * Set an event listener that monitors CircuitBreaker state changes
+    */
+  def onStateChange(listener: CircuitBreakerContext => Unit): CircuitBreaker = this.copy(
+    onStateChangeListener = listener
+  )
+
+  /**
+    * Set a recovery policiy which determine if ths circuit breaker can recover from HALF_OPEN to
+    * CLOSED. The default policy recovers immediately if health check is once successful.
+    */
+  def withRecoveryPolicy(recoveryPolicy: CircuitBreakerRecoveryPolicy): CircuitBreaker = this.copy(
+    recoveryPolicy = recoveryPolicy
+  )
+
+  /**
+    * Defines the action when trying to use the open circuit. The default behavior is to throw
+    * CircuitBreakerOpenException
+    */
+  def onOpenFailure(handler: CircuitBreakerContext => Unit): CircuitBreaker = this.copy(
+    onOpenFailureHandler = handler
+  )
+
+  /**
+    * Reset the lastFailure and close the circuit
+    */
+  def reset: Unit =
+    lastFailure = None
+    currentState.set(CLOSED)
+    nextProvingTimeMillis = Long.MaxValue
+    provingWaitTimeMillis = 0L
+    healthCheckPolicy.recovered()
+    recoveryPolicy.reset()
+
+  /**
+    * Force setting the current state.
+    */
+  def setState(newState: CircuitBreakerState): this.type =
+    if currentState.get() != newState then
+      currentState.set(newState)
+      onStateChangeListener(this)
+    this
+
+  def open(): this.type     = setState(OPEN)
+  def halfOpen(): this.type = setState(HALF_OPEN)
+  def close(): this.type    = setState(CLOSED)
+
+  /**
+    * Returns true when the circuit can execute the code ( OPEN or HALF_OPEN state)
+    */
+  def isConnected: Boolean =
+    val s = currentState.get()
+    s == CLOSED || s == HALF_OPEN
+
+  /**
+    * Note: Use this method only for the standalone mode. Generally, using CircuiteBreaker.run is
+    * sufficient.
+    *
+    * If the connection is open, perform the specified action. The default behavior is fail-fast,
+    * i.e., throwing CircuitBreakerOpenException
+    */
+  def verifyConnection(): Unit =
+    if !isConnected then
+      val currentTime = System.currentTimeMillis()
+      if currentTime > nextProvingTimeMillis then
+        halfOpen()
+      else
+        onOpenFailureHandler(this)
+
+  /**
+    * Note: Use this method only for the standalone mode. Generally, using CircuitBreaker.run is
+    * sufficient.
+    *
+    * This method reports a successful state to the CircuitBreaker.
+    */
+  def recordSuccess(): Unit =
+    healthCheckPolicy.recordSuccess()
+    recoveryPolicy.recordSuccess()
+    val isDead = healthCheckPolicy.isMarkedDead
+    currentState.get() match
+      case HALF_OPEN if (recoveryPolicy.canRecover) =>
+        // Probe request succeeds, so move to CLOSED state
+        healthCheckPolicy.recovered()
+        close()
+      case CLOSED if isDead =>
+        open()
+      case OPEN if !isDead =>
+        // Service is not marked dead, so try proving at HALF_OPEN state
+        halfOpen()
+      case _ =>
+
+  /**
+    * Note: Use this method only for the standalone mode. Generally, using CircuitBreaker.run is
+    * sufficient.
+    *
+    * This method reports a failure state to the CircuitBreaker.
+    */
+  def recordFailure(e: Throwable): Unit =
+    lastFailure = Some(e)
+    healthCheckPolicy.recordFailure()
+    recoveryPolicy.recordFailure()
+    if healthCheckPolicy.isMarkedDead then
+      val baseWaitMillis =
+        provingWaitTimeMillis
+          .max(delayAfterMarkedDead.retryPolicyConfig.initialIntervalMillis)
+          .toInt
+      val nextWaitMillis = delayAfterMarkedDead.nextWait(baseWaitMillis)
+      provingWaitTimeMillis = delayAfterMarkedDead.updateBaseWait(baseWaitMillis)
+      nextProvingTimeMillis = System.currentTimeMillis() + nextWaitMillis
+      open()
+
+  /**
+    * Execute the body block through the CircuitBreaker.
+    *
+    * If the state is OPEN, this will throw CircuitBreakerOpenException (fail-fast). The state will
+    * move to HALF_OPEN state after a certain amount of delay, determined by the
+    * delayAfterMarkedDead policy.
+    *
+    * If the state is HALF_OPEN, this method allows running the code block once, and if the result
+    * is successful, the state will move to CLOSED. If not, the state will be OPEN again.
+    *
+    * If the state is CLOSED, the code block will be executed normally. If the result is marked
+    * failure or nonRetryable exception is thrown, it will report to the failure to the
+    * HealthCheckPolicy. If this policy determines the target service is dead, the circuit will
+    * shift to OPEN state to block the future execution.
+    */
+  def run[A: ClassTag](body: => A): A =
+    verifyConnection()
+
+    val result = Try(body)
+
+    val resultClass =
+      result match
+        case Success(x) =>
+          resultClassifier(x)
+        case Failure(RetryableFailure(e)) =>
+          ResultClass.retryableFailure(e)
+        case Failure(e) =>
+          errorClassifier(e)
+
+    resultClass match
+      case ResultClass.Succeeded =>
+        recordSuccess()
+        result.get
+      case ResultClass.Failed(retryable, cause, _) =>
+        recordFailure(cause)
+        if retryable then
+          // If the error is retryable, rethrow as it is then the caller (maybe Retryer) should handle it.
+          throw cause
+        else
+          // If the error is not retryable, apply fallbackHandler
+          val x     = fallbackHandler(cause)
+          val clazz = implicitly[ClassTag[A]].runtimeClass
+          if clazz.isAssignableFrom(x.getClass) then
+            x.asInstanceOf[A]
+          else
+            throw new ClassCastException(
+              s"The fallback handler is returning ${x}, which is not an instance of ${clazz
+                  .getName}"
+            )
+
+  end run
+
+end CircuitBreaker
diff --git a/ai-core/src/main/scala/wvlet/ai/core/control/Control.scala b/ai-core/src/main/scala/wvlet/ai/core/control/Control.scala
new file mode 100644
index 00000000..57b64af8
--- /dev/null
+++ b/ai-core/src/main/scala/wvlet/ai/core/control/Control.scala
@@ -0,0 +1,80 @@
+/*
+ * 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.ai.core.control
+
+import scala.util.control.NonFatal
+import scala.concurrent.{ExecutionContext, Future}
+import scala.util.{Failure, Success, Try}
+
+/**
+  */
+object Control:
+  def withResource[R <: AutoCloseable, U](resource: R)(body: R => U): U =
+    try body(resource)
+    finally
+      if resource != null then
+        resource.close()
+
+  def withResources[R1 <: AutoCloseable, R2 <: AutoCloseable, U](resource1: R1, resource2: R2)(
+      body: (R1, R2) => U
+  ): U =
+    try body(resource1, resource2)
+    finally closeResources(resource1, resource2)
+
+  /**
+    * A loan pattern for Future[U].
+    *
+    * TODO: Test this after async test is available in 22.5.0
+    */
+  def withResourceAsync[R <: AutoCloseable, U](
+      resource: R
+  )(body: R => Future[U])(implicit sc: ExecutionContext): Future[U] = Future
+    .apply(resource)
+    .flatMap(r => body(r))
+    .transform { case any =>
+      closeSafely(any, () => resource.close())
+    }
+
+  /**
+    * Safely close the resource. If the resource is closed successfully, return the preceding
+    * result, if closing the resource is failed, combine the previous exception (if exists) and the
+    * resource closing exception together.
+    */
+  private[control] def closeSafely[U](preceding: Try[U], close: () => Unit): Try[U] =
+    Try(close()) match
+      case Success(_) =>
+        preceding
+      case Failure(e) =>
+        // If closing the resource failed, report the error in the Future
+        preceding match
+          case Success(x) =>
+            Failure(e)
+          case Failure(ex) =>
+            Failure(MultipleExceptions(List(e, ex)))
+
+  def closeResources[R <: AutoCloseable](resources: R*): Unit =
+    if resources != null then
+      var exceptionList = List.empty[Throwable]
+      resources.map { x =>
+        try
+          if x != null then
+            x.close()
+        catch
+          case NonFatal(e) =>
+            exceptionList = e :: exceptionList
+      }
+      if exceptionList.nonEmpty then
+        throw MultipleExceptions(exceptionList)
+
+end Control
diff --git a/ai-core/src/main/scala/wvlet/ai/core/control/ExponentialMovingAverage.scala b/ai-core/src/main/scala/wvlet/ai/core/control/ExponentialMovingAverage.scala
new file mode 100644
index 00000000..50ca0745
--- /dev/null
+++ b/ai-core/src/main/scala/wvlet/ai/core/control/ExponentialMovingAverage.scala
@@ -0,0 +1,71 @@
+/*
+ * 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.ai.core.control.util
+
+/**
+  * This code is based on com.twitter.finagle.util.Ema.
+  *
+  * This class is non-thread safe, so callers of the update method need to manage synchronization.
+  */
+private[control] class ExponentialMovingAverage(windowSize: Long):
+  private var time = Long.MinValue
+
+  // this is volatile to allow read-only calls to `last`
+  // without needing synchronization.
+  @volatile
+  private var ema = 0.0
+
+  /**
+    * Update the average with observed value `x`, and return the new average.
+    *
+    * Since `update` requires monotonic timestamps, it is up to the caller to ensure that calls to
+    * update do not race.
+    */
+  def update(timeStamp: Long, x: Double): Double =
+    if time == Long.MinValue then
+      time = timeStamp
+      ema = x
+      x
+    else
+      val td = timeStamp - time
+      if td < 0 then
+        throw new IllegalArgumentException(
+          s"Time must be monotonic. Previous time: ${time}, new time: ${timeStamp}"
+        )
+      time = timeStamp
+      val w =
+        if windowSize == 0.0 then
+          0.0
+        else
+          math.exp(-td.toDouble / windowSize)
+      val newEma = x * (1 - w) + ema * w
+      ema = newEma
+      newEma
+
+  /**
+    * Return the last observation.
+    *
+    * @note
+    *   This is safe to call without synchronization.
+    */
+  def last: Double = ema
+
+  /**
+    * Reset the average to 0 and erase all observations.
+    */
+  def reset(): Unit =
+    time = Long.MinValue
+    ema = 0.0
+
+end ExponentialMovingAverage
diff --git a/ai-core/src/main/scala/wvlet/ai/core/control/HealthCheckPolicy.scala b/ai-core/src/main/scala/wvlet/ai/core/control/HealthCheckPolicy.scala
new file mode 100644
index 00000000..77736ddf
--- /dev/null
+++ b/ai-core/src/main/scala/wvlet/ai/core/control/HealthCheckPolicy.scala
@@ -0,0 +1,145 @@
+/*
+ * 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.ai.core.control
+
+import java.util.concurrent.atomic.AtomicLong
+
+import wvlet.ai.core.control.util.ExponentialMovingAverage
+import wvlet.ai.core.log.LogSupport
+
+/**
+  */
+trait HealthCheckPolicy:
+  def isAlive: Boolean = !isMarkedDead
+  def isMarkedDead: Boolean
+
+  /**
+    * Called when a request succeeds
+    */
+  def recordSuccess(): Unit
+
+  /**
+    * Called when request is failed.
+    */
+  def recordFailure(): Unit
+
+  /**
+    * Called when the target service is recovered()
+    */
+  def recovered(): Unit
+
+object HealthCheckPolicy extends LogSupport:
+
+  /**
+    * A special health check policy used for disabling circuit breaker
+    */
+  def alwaysHealthy: HealthCheckPolicy =
+    new HealthCheckPolicy:
+      override def isMarkedDead: Boolean = false
+      override def recordSuccess(): Unit = {}
+      override def recordFailure(): Unit = {}
+      override def recovered(): Unit     = {}
+
+  /**
+    * A policy for marking the service dead upon consecutive failures
+    */
+  def markDeadOnConsecutiveFailures(numFailures: Int): HealthCheckPolicy =
+    new HealthCheckPolicy:
+      private val consecutiveFailures = new AtomicLong(0L)
+
+      override def isMarkedDead: Boolean = consecutiveFailures.get() >= numFailures
+
+      override def recordSuccess(): Unit = consecutiveFailures.set(0)
+
+      /**
+        * Called when request is failed. Returns delay
+        */
+      override def recordFailure(): Unit = consecutiveFailures.incrementAndGet()
+
+      /**
+        * Called when the target service is recovered()
+        */
+      override def recovered(): Unit = consecutiveFailures.set(0)
+
+  def markDeadOnRecentFailureRate(failureRate: Double, timeWindowMillis: Long): HealthCheckPolicy =
+    new HealthCheckPolicy:
+      private val failureRateEMA = new ExponentialMovingAverage(timeWindowMillis)
+
+      override def isMarkedDead: Boolean =
+        // logger.warn(s"failure rate: ${failureRateEMA.last}, ${failureRate}")
+        failureRateEMA.last > failureRate
+
+      /**
+        * Called when a request succeeds
+        */
+      override def recordSuccess(): Unit = failureRateEMA.update(System.currentTimeMillis(), 0)
+
+      /**
+        * Called when request is failed.
+        */
+      override def recordFailure(): Unit = failureRateEMA.update(System.currentTimeMillis(), 1)
+
+      /**
+        * Called when the target service is recovered()
+        */
+      override def recovered(): Unit = failureRateEMA.reset()
+
+  def markDeadOnFailureThreshold(numFailures: Int, numExecutions: Int) =
+    require(numExecutions > 0, s"numExecusions ${numExecutions} should be larger than 0")
+    require(
+      numFailures <= numExecutions,
+      s"numFailures ${numFailures} should be less than numExections(${numExecutions})"
+    )
+
+    new HealthCheckPolicy:
+      private val arraySize = (numExecutions + 64 - 1) / 64
+      // Circular bit vector of execution hisory 0 (success) or 1 (failure)
+      private val executionHistory: Array[Long] = Array.fill[Long](arraySize)(0L)
+      private var executionCount: Long          = 0
+
+      private def failureCount =
+        executionHistory
+          .map {
+            java.lang.Long.bitCount(_)
+          }
+          .sum
+
+      override def isMarkedDead: Boolean =
+        if executionCount < numExecutions then
+          false
+        else
+          failureCount >= numFailures
+
+      private def setAndMove(v: Boolean): Unit =
+        val i    = (executionCount % numExecutions).toInt
+        val mask = 1L << (63 - i % 64)
+        if v == true then
+          executionHistory(i / 64) |= mask
+        else
+          executionHistory(i / 64) &= ~mask
+        executionCount += 1
+        if executionCount < 0 then
+          // Reset upon overflow
+          executionCount = numExecutions
+
+      override def recordSuccess(): Unit = setAndMove(false)
+      override def recordFailure(): Unit = setAndMove(true)
+
+      override def recovered(): Unit = executionCount = 0
+
+    end new
+
+  end markDeadOnFailureThreshold
+
+end HealthCheckPolicy
diff --git a/ai-core/src/main/scala/wvlet/ai/core/control/IO.scala b/ai-core/src/main/scala/wvlet/ai/core/control/IO.scala
new file mode 100644
index 00000000..e14ba638
--- /dev/null
+++ b/ai-core/src/main/scala/wvlet/ai/core/control/IO.scala
@@ -0,0 +1,52 @@
+/*
+ * 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.ai.core.control
+
+import java.io.{ByteArrayOutputStream, File, InputStream}
+import java.net.URL
+import java.nio.charset.StandardCharsets
+
+import wvlet.ai.core.control.Control.withResource
+
+/**
+  */
+object IO:
+
+  def readAsString(f: File): String = readAsString(f.toURI.toURL)
+
+  def readAsString(url: URL): String =
+    withResource(url.openStream()) { in =>
+      readAsString(in)
+    }
+
+  def readAsString(in: InputStream): String = new String(readFully(in), StandardCharsets.UTF_8)
+
+  def readFully(in: InputStream): Array[Byte] =
+    val byteArray =
+      if in == null then
+        Array.emptyByteArray
+      else
+        withResource(new ByteArrayOutputStream) { b =>
+          val buf = new Array[Byte](8192)
+          withResource(in) { src =>
+            var readBytes = 0
+            while {
+              readBytes = src.read(buf);
+              readBytes != -1
+            } do
+              b.write(buf, 0, readBytes)
+          }
+          b.toByteArray
+        }
+    byteArray
diff --git a/ai-core/src/main/scala/wvlet/ai/core/control/MultipleExceptions.scala b/ai-core/src/main/scala/wvlet/ai/core/control/MultipleExceptions.scala
new file mode 100644
index 00000000..7a1e4a5f
--- /dev/null
+++ b/ai-core/src/main/scala/wvlet/ai/core/control/MultipleExceptions.scala
@@ -0,0 +1,24 @@
+/*
+ * 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.ai.core.control
+
+/**
+  */
+case class MultipleExceptions(causes: Seq[Throwable]) extends Exception:
+  causes.foreach { ex =>
+    this.addSuppressed(ex)
+  }
+
+  override def getMessage: String =
+    s"Multiple failures occurred:\n${causes.map(x => s"  - ${x.getMessage}").mkString("\n")}"
diff --git a/ai-core/src/main/scala/wvlet/ai/core/control/OSType.scala b/ai-core/src/main/scala/wvlet/ai/core/control/OSType.scala
new file mode 100644
index 00000000..eb9a0777
--- /dev/null
+++ b/ai-core/src/main/scala/wvlet/ai/core/control/OSType.scala
@@ -0,0 +1,28 @@
+/*
+ * Copyright 2012 Taro L. Saito
+ *
+ * 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.ai.core.control
+
+/**
+  * operating system information
+  * @author
+  *   leo
+  */
+enum OSType:
+  case Windows
+  case Mac
+  case Linux
+  case Other
diff --git a/ai-core/src/main/scala/wvlet/ai/core/control/Resource.scala b/ai-core/src/main/scala/wvlet/ai/core/control/Resource.scala
new file mode 100644
index 00000000..2d93a297
--- /dev/null
+++ b/ai-core/src/main/scala/wvlet/ai/core/control/Resource.scala
@@ -0,0 +1,88 @@
+/*
+ * 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.ai.core.control
+
+import java.io.File
+import scala.concurrent.{ExecutionContext, Future}
+
+/**
+  * Resource that can be closed.
+  * @tparam A
+  */
+trait Resource[A] extends AutoCloseable:
+  def get: A
+  def close(): Unit
+
+  /**
+    * Use the resource within the limited scope. After existing the scope, the resource will be
+    * closed
+    * @param body
+    * @tparam U
+    * @return
+    */
+  def use[U](body: A => U): U =
+    try body(get)
+    finally close()
+
+  /**
+    * Wrap a Future with this resource. After the future completes, the resource will be closed
+    */
+  def wrapFuture[U](body: A => Future[U])(implicit ec: ExecutionContext): Future[U] = Future
+    .apply(get)
+    .flatMap { (a: A) =>
+      body(a)
+    }
+    .transform { case any =>
+      Control.closeSafely(any, () => close())
+    }
+
+object Resource:
+
+  /**
+    * Create a resource for a temporary file, which will be deleted after closing the resource
+    * @param name
+    * @param suffix
+    * @param dir
+    * @return
+    */
+  def newTempFile(name: String, suffix: String = ".tmp", dir: String = "target"): Resource[File] =
+    newResource[File](
+      resource =
+        new File(dir)
+        // Create the target directory if not exists
+        val d = new File(dir)
+        d.mkdirs()
+        File.createTempFile(name, suffix, d)
+      ,
+      onClose = (file: File) => file.delete
+    )
+
+  /**
+    * Create a new Resource from an AutoClosable object
+    * @param resource
+    * @tparam R
+    * @return
+    */
+  def newResource[R](
+      resource: R,
+      onInit: R => Unit = { (x: R) =>
+      },
+      onClose: R => Unit
+  ): Resource[R] =
+    new Resource[R]:
+      onInit(resource)
+      override def get: R        = resource
+      override def close(): Unit = onClose(resource)
+
+end Resource
diff --git a/ai-core/src/main/scala/wvlet/ai/core/control/ResultClass.scala b/ai-core/src/main/scala/wvlet/ai/core/control/ResultClass.scala
new file mode 100644
index 00000000..051a372d
--- /dev/null
+++ b/ai-core/src/main/scala/wvlet/ai/core/control/ResultClass.scala
@@ -0,0 +1,51 @@
+/*
+ * 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.ai.core.control
+
+import wvlet.ai.core.control.Retry.ExtraWait
+
+/**
+  * A classification of the code execution result
+  */
+sealed trait ResultClass
+
+object ResultClass:
+
+  /**
+    * A label for the successful code execution result
+    */
+  case object Succeeded extends ResultClass
+
+  /**
+    * A label for the failed code execution result
+    */
+  case class Failed(
+      isRetryable: Boolean,
+      cause: Throwable,
+      extraWait: ExtraWait = Retry.noExtraWait
+  ) extends ResultClass:
+    def withExtraWaitMillis(extraWaitMillis: Int): Failed = this.copy(extraWait =
+      ExtraWait(maxExtraWaitMillis = extraWaitMillis)
+    )
+
+    def withExtraWaitFactor(factor: Double, maxExtraWaitMillis: Int = 5000): Failed = this.copy(
+      extraWait = ExtraWait(maxExtraWaitMillis = maxExtraWaitMillis, factor = factor)
+    )
+
+  def retryableFailure(e: Throwable): Failed    = Retry.retryableFailure(e)
+  def nonRetryableFailure(e: Throwable): Failed = Retry.nonRetryableFailure(e)
+
+  def ALWAYS_SUCCEED: Any => ResultClass = (x: Any) => Succeeded
+
+  def ALWAYS_RETRY: Throwable => ResultClass.Failed = (e: Throwable) => retryableFailure(e)
diff --git a/ai-core/src/main/scala/wvlet/ai/core/control/Retry.scala b/ai-core/src/main/scala/wvlet/ai/core/control/Retry.scala
new file mode 100644
index 00000000..d3f97cbd
--- /dev/null
+++ b/ai-core/src/main/scala/wvlet/ai/core/control/Retry.scala
@@ -0,0 +1,358 @@
+/*
+ * 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.ai.core.control
+
+import wvlet.ai.core.control.ResultClass.Failed
+import wvlet.ai.core.log.LogSupport
+import wvlet.ai.core.rx.Rx
+
+import java.util.concurrent.TimeUnit
+import scala.concurrent.{ExecutionContext, Future}
+import scala.util.{Failure, Random, Success, Try}
+
+/**
+  * Retry logic implementation helper
+  */
+object Retry extends LogSupport:
+  def retryableFailure(e: Throwable)    = Failed(isRetryable = true, e)
+  def nonRetryableFailure(e: Throwable) = Failed(isRetryable = false, e)
+
+  def withBackOff(
+      maxRetry: Int = 3,
+      initialIntervalMillis: Int = 100,
+      maxIntervalMillis: Int = 15000,
+      multiplier: Double = 1.5
+  ): RetryContext = defaultRetryContext
+    .withMaxRetry(maxRetry)
+    .withBackOff(initialIntervalMillis, maxIntervalMillis, multiplier)
+
+  def withBoundedBackoff(
+      initialIntervalMillis: Int = 100,
+      maxTotalWaitMillis: Int = 180000,
+      multiplier: Double = 1.5
+  ): RetryContext =
+    require(initialIntervalMillis > 0, s"initialWaitMillis must be > 0: ${initialIntervalMillis}")
+
+    // S = totalWaitMillis = w * r^0 + w * r^1 + w * r^2 + ...  + w * r^n
+    // S = w * (1-r^n) / (1-r)
+    // r^n = 1 - S * (1-r)/w
+    // n * log(r) = log(1 - S * (1-r) / w)
+    val N =
+      math.log(1 - (maxTotalWaitMillis * (1 - multiplier) / initialIntervalMillis)) /
+        math.log(multiplier)
+
+    def total(n: Int) = initialIntervalMillis * (1 - math.pow(multiplier, n)) / (1 - multiplier)
+
+    var maxRetry = N.ceil.toInt
+    while maxRetry > 0 && total(maxRetry) > maxTotalWaitMillis do
+      maxRetry -= 1
+    var maxIntervalMillis = initialIntervalMillis * math.pow(multiplier, N).toInt
+    withBackOff(
+      maxRetry = maxRetry.max(0),
+      initialIntervalMillis = initialIntervalMillis,
+      maxIntervalMillis = maxIntervalMillis,
+      multiplier = multiplier
+    )
+
+  def withJitter(
+      maxRetry: Int = 3,
+      initialIntervalMillis: Int = 100,
+      maxIntervalMillis: Int = 15000,
+      multiplier: Double = 1.5
+  ): RetryContext = defaultRetryContext
+    .withMaxRetry(maxRetry)
+    .withJitter(initialIntervalMillis, maxIntervalMillis, multiplier)
+
+  private val defaultRetryContext: RetryContext =
+    val retryConfig = RetryPolicyConfig()
+    RetryContext(
+      context = None,
+      lastError = NOT_STARTED,
+      retryCount = 0,
+      maxRetry = 3,
+      retryWaitStrategy = new Jitter(retryConfig),
+      nextWaitMillis = retryConfig.initialIntervalMillis,
+      baseWaitMillis = retryConfig.initialIntervalMillis,
+      extraWaitMillis = 0
+    )
+
+  case class MaxRetryException(retryContext: RetryContext)
+      extends Exception(
+        s"Reached the max retry count ${retryContext.retryCount}/${retryContext
+            .maxRetry}: ${retryContext.lastError.getMessage}",
+        retryContext.lastError
+      )
+
+  // Throw this to force retry the execution
+  case class RetryableFailure(e: Throwable) extends Exception(e)
+
+  case object NOT_STARTED extends Exception("Code is not executed")
+
+  private def REPORT_RETRY_COUNT: RetryContext => Unit =
+    (ctx: RetryContext) =>
+      warn(
+        f"[${ctx.retryCount}/${ctx.maxRetry}] Execution failed: ${ctx
+            .lastError
+            .getMessage}. Retrying in ${ctx.nextWaitMillis / 1000.0}%.2f sec."
+      )
+
+  private def RETHROW_ALL: Throwable => ResultClass.Failed = (e: Throwable) => throw e
+
+  private[control] val noExtraWait = ExtraWait()
+
+  case class ExtraWait(maxExtraWaitMillis: Int = 0, factor: Double = 0.0):
+    require(maxExtraWaitMillis >= 0)
+    require(factor >= 0)
+
+    def hasNoWait: Boolean = maxExtraWaitMillis == 0 && factor == 0.0
+
+    // Compute the extra wait millis based on the next wait millis
+    def extraWaitMillis(nextWaitMillis: Int): Int =
+      if maxExtraWaitMillis == 0 then
+        if factor == 0.0 then
+          0
+        else
+          (nextWaitMillis * factor).toInt
+      else if factor == 0.0 then
+        maxExtraWaitMillis
+      else
+        (nextWaitMillis * factor).toInt.min(maxExtraWaitMillis)
+
+  case class RetryContext(
+      context: Option[Any],
+      lastError: Throwable,
+      retryCount: Int,
+      maxRetry: Int,
+      retryWaitStrategy: RetryPolicy,
+      nextWaitMillis: Int,
+      baseWaitMillis: Int,
+      extraWaitMillis: Int,
+      resultClassifier: Any => ResultClass = ResultClass.ALWAYS_SUCCEED,
+      errorClassifier: Throwable => ResultClass.Failed = ResultClass.ALWAYS_RETRY,
+      beforeRetryAction: RetryContext => Any = REPORT_RETRY_COUNT
+  ):
+    def init(context: Option[Any] = None): RetryContext = this.copy(
+      context = context,
+      lastError = NOT_STARTED,
+      retryCount = 0,
+      nextWaitMillis = retryWaitStrategy.retryPolicyConfig.initialIntervalMillis,
+      baseWaitMillis = retryWaitStrategy.retryPolicyConfig.initialIntervalMillis,
+      extraWaitMillis = 0
+    )
+
+    def canContinue: Boolean = retryCount < maxRetry
+
+    /**
+      * Update the retry context, including retry count, last error, next wait time, etc.
+      *
+      * @param retryReason
+      * @return
+      *   the next retry context
+      */
+    def nextRetry(retryReason: Throwable): RetryContext =
+      val nextRetryCtx = this.copy(
+        lastError = retryReason,
+        retryCount = retryCount + 1,
+        nextWaitMillis = retryWaitStrategy.nextWait(baseWaitMillis) + extraWaitMillis,
+        baseWaitMillis = retryWaitStrategy.updateBaseWait(baseWaitMillis),
+        extraWaitMillis = 0
+      )
+      beforeRetryAction(nextRetryCtx)
+      nextRetryCtx
+
+    def withExtraWait(extraWait: ExtraWait): RetryContext =
+      if extraWait.hasNoWait && this.extraWaitMillis == 0 then
+        this
+      else
+        this.copy(extraWaitMillis = extraWait.extraWaitMillis(nextWaitMillis))
+
+    def withRetryWaitStrategy(newRetryWaitStrategy: RetryPolicy): RetryContext = this.copy(
+      retryWaitStrategy = newRetryWaitStrategy
+    )
+
+    def withMaxRetry(newMaxRetry: Int): RetryContext = this.copy(maxRetry = newMaxRetry)
+
+    def noRetry: RetryContext = this.copy(maxRetry = 0)
+
+    def withBackOff(
+        initialIntervalMillis: Int = 100,
+        maxIntervalMillis: Int = 15000,
+        multiplier: Double = 1.5
+    ): RetryContext =
+      val config = RetryPolicyConfig(initialIntervalMillis, maxIntervalMillis, multiplier)
+      this.copy(retryWaitStrategy = new ExponentialBackOff(config))
+
+    def withJitter(
+        initialIntervalMillis: Int = 100,
+        maxIntervalMillis: Int = 15000,
+        multiplier: Double = 1.5
+    ): RetryContext =
+      val config = RetryPolicyConfig(initialIntervalMillis, maxIntervalMillis, multiplier)
+      this.copy(retryWaitStrategy = new Jitter(config))
+
+    def withResultClassifier[U](newResultClassifier: U => ResultClass): RetryContext = this.copy(
+      resultClassifier = newResultClassifier.asInstanceOf[Any => ResultClass]
+    )
+
+    /**
+      * Set a detailed error handler upon Exception. If the given exception is not retryable, just
+      * rethrow the exception. Otherwise, consume the exception.
+      */
+    def withErrorClassifier(errorClassifier: Throwable => ResultClass.Failed): RetryContext = this
+      .copy(errorClassifier = errorClassifier)
+
+    def beforeRetry[U](handler: RetryContext => U): RetryContext = this.copy(beforeRetryAction =
+      handler
+    )
+
+    /**
+      * Clear the default beforeRetry action
+      */
+    def noRetryLogging: RetryContext = this.copy(beforeRetryAction = (x: RetryContext) => ())
+
+    /**
+      * Add a partial function that accepts exceptions that need to be retried.
+      *
+      * @param errorClassifier
+      * @return
+      */
+    def retryOn(errorClassifier: PartialFunction[Throwable, ResultClass.Failed]): RetryContext =
+      this.copy(errorClassifier = (e: Throwable) => errorClassifier.applyOrElse(e, RETHROW_ALL))
+
+    def run[A](body: => A): A = runInternal(None)(body)
+
+    def runWithContext[A](
+        context: Any,
+        circuitBreaker: CircuitBreaker = CircuitBreaker.alwaysClosed
+    )(body: => A): A = runInternal(Option(context), circuitBreaker)(body)
+
+    private def classifyResult[A](result: A): ResultClass =
+      val resultClass =
+        result match
+          case Success(x) =>
+            // Test whether the code block execution is succeeded or failed
+            resultClassifier(x)
+          case Failure(RetryableFailure(e)) =>
+            ResultClass.retryableFailure(e)
+          case Failure(e) =>
+            errorClassifier(e)
+      resultClass
+
+    protected def runInternal[A](
+        context: Option[Any],
+        circuitBreaker: CircuitBreaker = CircuitBreaker.alwaysClosed
+    )(body: => A): A =
+      var result: Option[A]          = None
+      var retryContext: RetryContext = init(context)
+
+      var isFirst: Boolean = true
+
+      while isFirst || (result.isEmpty && retryContext.canContinue) do
+        isFirst = false
+
+        val ret = Try {
+          circuitBreaker.verifyConnection()
+          body
+        }
+        val resultClass = classifyResult(ret)
+        resultClass match
+          case ResultClass.Succeeded =>
+            circuitBreaker.recordSuccess()
+            // OK. Exit the loop
+            result = Some(ret.get)
+          case ResultClass.Failed(isRetryable, cause, extraWait) if isRetryable =>
+            circuitBreaker.recordFailure(cause)
+            // Retryable error
+            retryContext = retryContext.withExtraWait(extraWait).nextRetry(cause)
+            // Wait until the next retry
+            Compat.sleep(retryContext.nextWaitMillis)
+          case ResultClass.Failed(_, cause, _) =>
+            // For regular non-retryable failures, we need to treat them as successful responses
+            circuitBreaker.recordSuccess()
+            // Non-retryable error. Exit the loop by throwing the exception
+            throw cause
+
+      result match
+        case Some(a) =>
+          a
+        case None =>
+          throw MaxRetryException(retryContext)
+
+    end runInternal
+
+    def runAsyncWithContext[A](
+        context: Any,
+        circuitBreaker: CircuitBreaker = CircuitBreaker.alwaysClosed
+    )(body: => Rx[A]): Rx[A] =
+      def loop(retryContext: RetryContext, isFirst: Boolean): Rx[A] =
+        if !isFirst && !retryContext.canContinue then
+          Rx.exception(MaxRetryException(retryContext))
+        else
+          Rx.fromTry(Try(circuitBreaker.verifyConnection()))
+            .flatMap(_ => body)
+            .transformRx { (ret: Try[A]) =>
+              val resultClass = classifyResult(ret)
+              resultClass match
+                case ResultClass.Succeeded =>
+                  circuitBreaker.recordSuccess()
+                  // Exit the loop
+                  Rx.fromTry(ret)
+                case ResultClass.Failed(isRetryable, cause, extraWait) if isRetryable =>
+                  // Retryable error
+                  circuitBreaker.recordFailure(cause)
+                  // Add retry wait
+                  val nextRetry = retryContext.withExtraWait(extraWait).nextRetry(cause)
+                  Rx.delay(nextRetry.nextWaitMillis, TimeUnit.MILLISECONDS)
+                    .flatMap(_ => loop(nextRetry, isFirst = false))
+                case ResultClass.Failed(_, cause, _) =>
+                  // For regular non-retryable failures, we need to treat them as successful responses
+                  circuitBreaker.recordSuccess()
+                  // Non-retryable error. Exit the loop with the exception
+                  Rx.exception(cause)
+            }
+
+      loop(retryContext = init(Option(context)), isFirst = true)
+
+    end runAsyncWithContext
+
+  end RetryContext
+
+  case class RetryPolicyConfig(
+      initialIntervalMillis: Int = 100,
+      maxIntervalMillis: Int = 15000,
+      multiplier: Double = 1.5
+  ):
+    require(initialIntervalMillis >= 0)
+    require(maxIntervalMillis >= 0)
+    require(multiplier >= 0)
+
+  trait RetryPolicy:
+    def retryPolicyConfig: RetryPolicyConfig
+    def updateBaseWait(waitMillis: Int): Int = math
+      .round(waitMillis * retryPolicyConfig.multiplier)
+      .toInt
+      .min(retryPolicyConfig.maxIntervalMillis)
+
+    def nextWait(baseWaitMillis: Int): Int
+
+  class ExponentialBackOff(val retryPolicyConfig: RetryPolicyConfig) extends RetryPolicy:
+    override def nextWait(baseWaitMillis: Int): Int = baseWaitMillis
+
+  class Jitter(val retryPolicyConfig: RetryPolicyConfig, rand: Random = new Random())
+      extends RetryPolicy:
+    override def nextWait(baseWaitMillis: Int): Int = (baseWaitMillis.toDouble * rand.nextDouble())
+      .round
+      .toInt
+
+end Retry
diff --git a/ai-core/src/main/scala/wvlet/ai/core/control/Ticker.scala b/ai-core/src/main/scala/wvlet/ai/core/control/Ticker.scala
new file mode 100644
index 00000000..b9215e02
--- /dev/null
+++ b/ai-core/src/main/scala/wvlet/ai/core/control/Ticker.scala
@@ -0,0 +1,36 @@
+/*
+ * 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.ai.core.control
+
+/**
+  * Ticker is for measuring the elapsed time.
+  */
+trait Ticker:
+  // Return the number of nanoseconds elapsed
+  def read: Long
+
+/**
+  * A Ticker implementation that can be incremented for test
+  */
+class ManualTicker(private var counter: Long = 0) extends Ticker:
+  def tick(n: Long): Unit = counter += n
+  def read: Long          = counter
+
+object Ticker:
+  // A ticker that reads the current time using System.nanoTime()
+  def systemTicker: Ticker =
+    new Ticker:
+      override def read: Long = System.nanoTime()
+
+  def manualTicker: ManualTicker = new ManualTicker
diff --git a/ai-core/src/test/scala/wvlet/ai/core/control/CircuitBreakerTest.scala b/ai-core/src/test/scala/wvlet/ai/core/control/CircuitBreakerTest.scala
new file mode 100644
index 00000000..89b93093
--- /dev/null
+++ b/ai-core/src/test/scala/wvlet/ai/core/control/CircuitBreakerTest.scala
@@ -0,0 +1,134 @@
+package wvlet.ai.core.control
+
+import wvlet.airspec.*
+import java.util.concurrent.TimeoutException
+
+class CircuitBreakerTest extends AirSpec:
+
+  test("support changing states") {
+    val cb = CircuitBreaker.default
+    cb.state shouldBe CircuitBreaker.CLOSED
+    cb.isConnected shouldBe true
+
+    cb.open()
+    cb.state shouldBe CircuitBreaker.OPEN
+    cb.isConnected shouldBe false
+
+    cb.halfOpen()
+    cb.state shouldBe CircuitBreaker.HALF_OPEN
+    cb.isConnected shouldBe true
+
+    cb.close()
+    cb.state shouldBe CircuitBreaker.CLOSED
+    cb.isConnected shouldBe true
+  }
+
+  test("support standalone usage") {
+    val cb = CircuitBreaker.default
+
+    cb.verifyConnection()
+    try
+      cb.recordSuccess()
+    catch
+      case e: Throwable =>
+        cb.recordFailure(e)
+
+    // Capture the open exception
+    intercept[CircuitBreakerOpenException] {
+      cb.open()
+      cb.verifyConnection()
+    }
+  }
+
+  test("support failure threshold") {
+    val cb = CircuitBreaker.withFailureThreshold(2, 5)
+    val e  = new TimeoutException()
+    cb.isConnected shouldBe true
+
+    // 0/0
+    cb.recordSuccess()
+    cb.isConnected shouldBe true
+
+    // 0/1
+    cb.recordSuccess()
+    cb.isConnected shouldBe true
+
+    // 1/2
+    cb.recordFailure(e)
+    cb.isConnected shouldBe true
+
+    // 2/3
+    cb.recordFailure(e)
+    cb.isConnected shouldBe true
+
+    // 2/4
+    cb.recordSuccess()
+    cb.isConnected shouldBe false
+
+    // 2/5 -> open the circuit breaker
+    cb.recordSuccess()
+    cb.isConnected shouldBe false
+
+    // 2/5
+    cb.recordSuccess()
+    cb.isConnected shouldBe false
+
+    // 2/5
+    cb.recordSuccess()
+    cb.isConnected shouldBe true
+
+    // 1/5
+    cb.recordSuccess()
+    cb.isConnected shouldBe true
+  }
+
+  test("support consecutive failure health checker") {
+    val cb = CircuitBreaker.withConsecutiveFailures(2)
+    val e  = new TimeoutException()
+    cb.isConnected shouldBe true
+
+    // 1/1
+    cb.recordSuccess()
+    cb.isConnected shouldBe true
+
+    // 1/2
+    cb.recordFailure(e)
+    cb.isConnected shouldBe true
+
+    // 1/3
+    cb.recordFailure(e)
+    cb.isConnected shouldBe false
+
+    // Force probing
+    cb.halfOpen()
+
+    // 1/4
+    cb.recordSuccess()
+    cb.isConnected shouldBe true
+  }
+
+  test("run code with circuit") {
+    val cb = CircuitBreaker.withFailureThreshold(1, 2)
+
+    cb.run {}
+
+    intercept[TimeoutException] {
+      cb.run[Any] {
+        throw new TimeoutException()
+      }
+    }
+
+    intercept[CircuitBreakerOpenException] {
+      cb.run {}
+    }
+  }
+
+  test("disable circuit breaker") {
+    val cb = CircuitBreaker.alwaysClosed
+    val e  = new TimeoutException()
+    for i <- 0 to 10 do
+      cb.recordFailure(e)
+      cb.isConnected shouldBe true
+  }
+
+end CircuitBreakerTest
diff --git a/ai-core/src/test/scala/wvlet/ai/core/control/ControlTest.scala b/ai-core/src/test/scala/wvlet/ai/core/control/ControlTest.scala
new file mode 100644
index 00000000..5a0ad7c9
--- /dev/null
+++ b/ai-core/src/test/scala/wvlet/ai/core/control/ControlTest.scala
@@ -0,0 +1,84 @@
+/*
+ * 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.ai.core.control
+
+import wvlet.airspec.AirSpec
+
+object ControlTest:
+  class A extends AutoCloseable:
+    var closed: Boolean        = false
+    override def close(): Unit = closed = true
+
+/**
+  */
+class ControlTest extends AirSpec:
+
+  test("have loan pattern") {
+    val out = new ControlTest.A
+    out.closed shouldBe false
+    Control.withResource(out) { o =>
+      // do nothing
+    }
+    out.closed shouldBe true
+  }
+
+  test("have loan pattern for two resources") {
+    val in  = new ControlTest.A
+    val out = new ControlTest.A
+    out.closed shouldBe false
+    Control.withResources(in, out) { (i, o) =>
+      // do nothing
+    }
+    in.closed shouldBe true
+    out.closed shouldBe true
+  }
+
+  test("not cause error for null resource") {
+    Control.withResource(null) { o =>
+      // do nothing
+    }
+    Control.withResources(null, null) { (i, o) =>
+      // do nothing
+    }
+  }
+
+  test("report resource closing errors") {
+    class FirstException  extends RuntimeException
+    class SecondException extends RuntimeException
+
+    intercept[FirstException] {
+      Control.withResource(
+        new AutoCloseable:
+          override def close(): Unit = throw new FirstException()
+      ) { o =>
+        // do nothing
+      }
+    }
+
+    val m = intercept[MultipleExceptions] {
+      Control.withResources(
+        new AutoCloseable:
+          override def close(): Unit = throw new FirstException()
+        ,
+        new AutoCloseable:
+          override def close(): Unit = throw new SecondException()
+      ) { (i, o) =>
+        // do nothing
+      }
+    }
+    m.causes.exists(_.getClass == classOf[FirstException]) shouldBe true
+    m.causes.exists(_.getClass == classOf[SecondException]) shouldBe true
+  }
+
+end ControlTest
diff --git a/ai-core/src/test/scala/wvlet/ai/core/control/RetryTest.scala b/ai-core/src/test/scala/wvlet/ai/core/control/RetryTest.scala
new file mode 100644
index 00000000..75b27027
--- /dev/null
+++ b/ai-core/src/test/scala/wvlet/ai/core/control/RetryTest.scala
@@ -0,0 +1,162 @@
+/*
+ * 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.ai.core.control
+
+import wvlet.ai.core.control.Retry.{MaxRetryException, RetryContext}
+import wvlet.airspec.AirSpec
+
+import scala.concurrent.TimeoutException
+
+/**
+  */
+class RetryTest extends AirSpec:
+
+  test("support backoff retry") {
+    var count = 0
+
+    val r = Retry.withBackOff(maxRetry = 3)
+
+    val e = r.run {
+      count += 1
+      logger.info("hello retry")
+      if count < 3 then
+        throw new IllegalStateException("retry test")
+      else
+        "success"
+    }
+
+    e shouldBe "success"
+    count shouldBe 3
+  }
+
+  test("support bounded backoff retry") {
+    val maxWait = 50000
+    var r =
+      Retry
+        .withBoundedBackoff(initialIntervalMillis = 1000, maxTotalWaitMillis = maxWait)
+        .noRetryLogging
+
+    r = r.init()
+    var waitTotal = 0
+    while r.canContinue do
+      waitTotal += r.nextWaitMillis
+      r = r.nextRetry(new IllegalStateException())
+    waitTotal <= maxWait shouldBe true
+  }
+
+  test("support jitter retry") {
+    var count = 0
+
+    val r = Retry
+      .withJitter(maxRetry = 3)
+      .run {
+        count += 1
+        logger.info("hello retry")
+        if count < 2 then
+          throw new IllegalStateException("retry test")
+        else
+          "success"
+      }
+
+    r shouldBe "success"
+    count shouldBe 2
+  }
+
+  test("throw max retry exception") {
+    val e = intercept[MaxRetryException] {
+      Retry
+        .withBackOff(maxRetry = 3)
+        .retryOn { case e: IllegalStateException =>
+          warn(e.getMessage)
+          ResultClass.retryableFailure(e)
+        }
+        .run {
+          logger.info("hello retry")
+          throw new IllegalStateException("retry test")
+        }
+    }
+
+    e.retryContext.maxRetry shouldBe 3
+    e.retryContext.retryCount shouldBe 3
+    e.retryContext.lastError.getClass shouldBe classOf[IllegalStateException]
+  }
+
+  test("change retry wait strategy") {
+    val r = Retry.withJitter().withBackOff(initialIntervalMillis = 3)
+
+    r.retryWaitStrategy.retryPolicyConfig.initialIntervalMillis shouldBe 3
+
+    val j = r.withJitter(initialIntervalMillis = 20)
+    j.retryWaitStrategy.retryPolicyConfig.initialIntervalMillis shouldBe 20
+
+    val m = j.withMaxRetry(100)
+    m.maxRetry shouldBe 100
+  }
+
+  test("pass the execution context") {
+    val r = Retry.withBackOff(initialIntervalMillis = 0)
+
+    var count   = 0
+    var checked = false
+    r.beforeRetry { (ctx: RetryContext) =>
+        ctx.context shouldBe Some("hello world")
+        checked = true
+      }
+      .runWithContext("hello world") {
+        if count == 0 then
+          count += 1
+          throw new Exception()
+      }
+
+    count shouldBe 1
+    checked shouldBe true
+  }
+
+  test("add extra wait") {
+    intercept[MaxRetryException] {
+      Retry
+        .withBackOff(initialIntervalMillis = 10)
+        .withMaxRetry(1)
+        .withErrorClassifier { case e: TimeoutException =>
+          ResultClass.retryableFailure(e).withExtraWaitMillis(100)
+        }
+        .beforeRetry { (ctx: RetryContext) =>
+          debug(s"${ctx.retryCount} ${ctx.nextWaitMillis}")
+          ctx.nextWaitMillis shouldBe 10 + 100
+        }
+        .run {
+          throw new TimeoutException()
+        }
+    }
+  }
+
+  test("add extra wait factor") {
+    intercept[MaxRetryException] {
+      Retry
+        .withBackOff(initialIntervalMillis = 10)
+        .withMaxRetry(1)
+        .withErrorClassifier { case e: TimeoutException =>
+          ResultClass.retryableFailure(e).withExtraWaitFactor(0.2)
+        }
+        .beforeRetry { (ctx: RetryContext) =>
+          debug(s"${ctx.retryCount} ${ctx.nextWaitMillis}")
+          ctx.nextWaitMillis shouldBe 10 + 2
+        }
+        .run {
+          throw new TimeoutException()
+        }
+    }
+  }
+
+end RetryTest
diff --git a/ai-core/src/test/scala/wvlet/ai/core/rx/RxTest.scala b/ai-core/src/test/scala/wvlet/ai/core/rx/RxTest.scala
index b52aceab..2432fd8a 100644
--- a/ai-core/src/test/scala/wvlet/ai/core/rx/RxTest.scala
+++ b/ai-core/src/test/scala/wvlet/ai/core/rx/RxTest.scala
@@ -13,7 +13,7 @@
  */
 package wvlet.ai.core.rx
 
-import wvlet.airframe.Design
+import wvlet.ai.core.design.Design
 import wvlet.airspec.AirSpec
 
 import java.util.concurrent.atomic.{AtomicInteger, AtomicReference}


Reviews:
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:11Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:12Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:13Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:14Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:15Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:16Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:17Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:18Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:19Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:20Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:44Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:44Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:45Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:46Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:47Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:48Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:49Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:50Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:51Z
user: login: amazon-q-developer[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T04:08:52Z

Review comments:
diff_hunk: @@ -0,0 +1,88 @@
+/*
+ * 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.ai.core.control
+
+import java.io.File
+import scala.concurrent.{ExecutionContext, Future}
+
+/**
+  * Resource that can be closed.
+  * @tparam A
+  */
+trait Resource[A] extends AutoCloseable:
+  def get: A
+  def close(): Unit
+
+  /**
+    * Use the resource within the limited scope. After existing the scope, the resource will be
+    * closed
+    * @param body
+    * @tparam U
+    * @return
+    */
+  def use[U](body: A => U): U =
+    try body(get)
+    finally close()
+
+  /**
+    * Wrap a Future with this resource. After the future completes, the resource will be closed
+    */
+  def wrapFuture[U](body: A => Future[U])(implicit ec: ExecutionContext): Future[U] = Future
+    .apply(get)
+    .flatMap { (a: A) =>
+      body(a)
+    }
+    .transform { case any =>
+      Control.closeSafely(any, () => close())
+    }
+
+object Resource:
+
+  /**
+    * Create a resource for a temporary file, which will be deleted after closing the resource
+    * @param name
+    * @param suffix
+    * @param dir
+    * @return
+    */
+  def newTempFile(name: String, suffix: String = ".tmp", dir: String = "target"): Resource[File] =
+    newResource[File](
+      resource =
+        new File(dir)
+        // Create the target directory if not exists
+        val d = new File(dir)
+        d.mkdirs()
+        File.createTempFile(name, suffix, d)
path: ai-core/src/main/scala/wvlet/ai/core/control/Resource.scala
user: login: amazon-q-developer[bot]
body: > [!WARNING]
> ****Description:**** The `newTempFile` method doesn't handle potential exceptions when creating temporary files or directories. Wrap the file creation logic in a try-catch block to handle potential IOExceptions.

****Severity:**** High


created_at: 2025-05-19T04:08:11Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,68 @@
+/*
+ * 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.ai.core.control.util
+
+/**
+  * This code is based on com.twitter.finagle.util.Ema.
+  *
+  * This class is non-thread safe, so callers of the update method need to manage synchronization.
+  */
+private[control] class ExponentialMovingAverage(windowSize: Long):
+  private var time = Long.MinValue
+
+  // this is volatile to allow read-only calls to `last`
+  // without needing synchronization.
+  @volatile
+  private var ema = 0.0
+
+  /**
+    * Update the average with observed value `x`, and return the new average.
+    *
+    * Since `update` requires monotonic timestamps, it is up to the caller to ensure that calls to
+    * update do not race.
+    */
+  def update(timeStamp: Long, x: Double): Double =
+    if time == Long.MinValue then
+      time = timeStamp
+      ema = x
+      x
+    else
+      val td = timeStamp - time
+      assert(td >= 0, "Nonmonotonic timestamp")
path: ai-core/src/main/scala/wvlet/ai/core/control/ExponentialMovingAverage.scala
user: login: amazon-q-developer[bot]
body: > [!WARNING]
> ****Description:**** The assert statement on line 42 may throw an AssertionError in production, which is not ideal for error handling. Consider replacing the assert with a more robust error handling mechanism, such as throwing a custom exception or returning an error value.

****Severity:**** High


created_at: 2025-05-19T04:08:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,52 @@
+/*
+ * 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.ai.core.control
+
+import java.io.{ByteArrayOutputStream, File, InputStream}
+import java.net.URL
+import java.nio.charset.StandardCharsets
+
+import wvlet.ai.core.control.Control.withResource
+
+/**
+  */
+object IO:
+
+  def readAsString(f: File): String = readAsString(f.toURI.toURL)
+
+  def readAsString(url: URL): String =
+    withResource(url.openStream()) { in =>
+      readAsString(in)
+    }
+
+  def readAsString(in: InputStream): String = new String(readFully(in), StandardCharsets.UTF_8)
+
+  def readFully(in: InputStream): Array[Byte] =
path: ai-core/src/main/scala/wvlet/ai/core/control/IO.scala
user: login: amazon-q-developer[bot]
body: > [!WARNING]
> ****Description:**** The readFully method doesn't handle potential IOException when reading from the input stream. Wrap the reading operation in a try-catch block to handle IOException. Consider logging the error or rethrowing as a custom exception.

****Severity:**** High


created_at: 2025-05-19T04:08:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,22 @@
+/*
+ * 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.ai.core.control
+
+import java.util.concurrent.{ScheduledThreadPoolExecutor, TimeUnit}
+import scala.concurrent.{ExecutionContext, Future, Promise}
+
+/**
+  */
+object Compat:
path: ai-core/.jvm/src/main/scala/wvlet/ai/core/control/Compat.scala
user: login: amazon-q-developer[bot]
body: ****Description:**** The object Compat lacks a descriptive comment explaining its purpose and functionality. Add a brief comment describing the purpose and functionality of the Compat object.

****Severity:**** Medium


created_at: 2025-05-19T04:08:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,358 @@
+/*
+ * 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.ai.core.control
+
+import wvlet.ai.core.control.ResultClass.Failed
+import wvlet.ai.core.log.LogSupport
+import wvlet.ai.core.rx.Rx
+
+import java.util.concurrent.TimeUnit
+import scala.concurrent.{ExecutionContext, Future}
+import scala.util.{Failure, Random, Success, Try}
+
+/**
+  * Retry logic implementation helper
+  */
+object Retry extends LogSupport:
+  def retryableFailure(e: Throwable)    = Failed(isRetryable = true, e)
+  def nonRetryableFailure(e: Throwable) = Failed(isRetryable = false, e)
+
+  def withBackOff(
+      maxRetry: Int = 3,
+      initialIntervalMillis: Int = 100,
+      maxIntervalMillis: Int = 15000,
+      multiplier: Double = 1.5
+  ): RetryContext = defaultRetryContext
+    .withMaxRetry(maxRetry)
+    .withBackOff(initialIntervalMillis, maxIntervalMillis, multiplier)
+
+  def withBoundedBackoff(
+      initialIntervalMillis: Int = 100,
+      maxTotalWaitMillis: Int = 180000,
+      multiplier: Double = 1.5
+  ): RetryContext =
+    require(initialIntervalMillis > 0, s"initialWaitMillis must be > 0: ${initialIntervalMillis}")
+
+    // S = totalWaitMillis = w * r^0 + w * r^1 + w * r^2 + ...  + w * r^n
+    // S = w * (1-r^n) / (1-r)
+    // r^n = 1 - S * (1-r)/w
+    // n * log(r) = log(1 - S * (1-r) / w)
+    val N =
+      math.log(1 - (maxTotalWaitMillis * (1 - multiplier) / initialIntervalMillis)) /
+        math.log(multiplier)
+
+    def total(n: Int) = initialIntervalMillis * (1 - math.pow(multiplier, n)) / (1 - multiplier)
+
+    var maxRetry = N.ceil.toInt
+    while maxRetry > 0 && total(maxRetry) > maxTotalWaitMillis do
path: ai-core/src/main/scala/wvlet/ai/core/control/Retry.scala
user: login: amazon-q-developer[bot]
body: ****Description:**** The `withBoundedBackoff` method uses a while loop to calculate `maxRetry`, which may be inefficient for large values. Consider using a more efficient algorithm to calculate `maxRetry` without a loop.

****Severity:**** Medium


created_at: 2025-05-19T04:08:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,80 @@
+/*
+ * 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.ai.core.control
+
+import scala.util.control.NonFatal
+import scala.concurrent.{ExecutionContext, Future}
+import scala.util.{Failure, Success, Try}
+
+/**
+  */
+object Control:
+  def withResource[R <: AutoCloseable, U](resource: R)(body: R => U): U =
path: ai-core/src/main/scala/wvlet/ai/core/control/Control.scala
user: login: amazon-q-developer[bot]
body: ****Description:**** The `withResource` method doesn't handle the case where `resource` is null before trying to use it. Add a null check for `resource` at the beginning of the `withResource` method and throw an appropriate exception if it's null.

****Severity:**** Medium


created_at: 2025-05-19T04:08:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,68 @@
+/*
+ * 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.ai.core.control.util
+
+/**
+  * This code is based on com.twitter.finagle.util.Ema.
+  *
+  * This class is non-thread safe, so callers of the update method need to manage synchronization.
+  */
+private[control] class ExponentialMovingAverage(windowSize: Long):
+  private var time = Long.MinValue
+
+  // this is volatile to allow read-only calls to `last`
+  // without needing synchronization.
+  @volatile
+  private var ema = 0.0
+
+  /**
+    * Update the average with observed value `x`, and return the new average.
+    *
+    * Since `update` requires monotonic timestamps, it is up to the caller to ensure that calls to
+    * update do not race.
+    */
+  def update(timeStamp: Long, x: Double): Double =
+    if time == Long.MinValue then
+      time = timeStamp
+      ema = x
+      x
+    else
+      val td = timeStamp - time
+      assert(td >= 0, "Nonmonotonic timestamp")
+      time = timeStamp
+      val w =
+        if windowSize == 0.0 then
path: ai-core/src/main/scala/wvlet/ai/core/control/ExponentialMovingAverage.scala
user: login: amazon-q-developer[bot]
body: ****Description:**** The comparison windowSize == 0.0 on line 45 uses a Double literal, which may lead to precision issues. Consider using windowSize == 0L for comparing with a Long value, or explicitly cast windowSize to Double if that's the intended comparison.

****Severity:**** Medium


created_at: 2025-05-19T04:08:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,84 @@
+/*
+ * 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.ai.core.control
+
+import wvlet.airspec.AirSpec
+
+object ControlTest:
+  class A extends AutoCloseable:
+    var closed: Boolean        = false
+    override def close(): Unit = closed = true
+
+/**
+  */
+class ControlTest extends AirSpec:
+
+  test("have loan pattern") {
+    val out = new ControlTest.A
+    out.closed shouldBe false
+    Control.withResource(out) { o =>
+      // do nothing
path: ai-core/src/test/scala/wvlet/ai/core/control/ControlTest.scala
user: login: amazon-q-developer[bot]
body: ****Description:**** The test methods contain empty code blocks with comments "// do nothing", which may be unclear to readers. Consider adding a brief comment explaining why these blocks are intentionally left empty or if they represent placeholder code for future implementation.

****Severity:**** Low


created_at: 2025-05-19T04:08:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,23 @@
+/*
+ * 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.ai.core.control
+
+import scala.concurrent.{ExecutionContext, Future, Promise}
+
+/**
+  */
+object Compat:
+  def sleep(millis: Long): Unit = {
+    // do nothing in Scala.js as their is no way to implement sleep https://github.com/scala-js/scala-js/issues/1898
path: ai-core/.js/src/main/scala/wvlet/ai/core/control/Compat.scala
user: login: amazon-q-developer[bot]
body: ****Description:**** The sleep method contains a comment explaining why it does nothing, but this could be clearer. Consider rephrasing the comment to clarify that this is a no-op implementation for Scala.js compatibility.

****Severity:**** Low


created_at: 2025-05-19T04:08:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,80 @@
+/*
+ * 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.ai.core.control
+
+import scala.util.control.NonFatal
+import scala.concurrent.{ExecutionContext, Future}
+import scala.util.{Failure, Success, Try}
+
+/**
+  */
+object Control:
+  def withResource[R <: AutoCloseable, U](resource: R)(body: R => U): U =
+    try body(resource)
+    finally
+      if resource != null then
+        resource.close()
+
+  def withResources[R1 <: AutoCloseable, R2 <: AutoCloseable, U](resource1: R1, resource2: R2)(
+      body: (R1, R2) => U
+  ): U =
+    try body(resource1, resource2)
+    finally closeResources(resource1, resource2)
+
+  /**
+    * A loan pattern for Future[U].
+    *
+    * TODO: Test this after async test is available in 22.5.0
+    */
+  def withResourceAsync[R <: AutoCloseable, U](
+      resource: R
+  )(body: R => Future[U])(implicit sc: ExecutionContext): Future[U] = Future
+    .apply(resource)
+    .flatMap(r => body(r))
+    .transform { case any =>
+      closeSafely(any, () => resource.close())
+    }
+
+  /**
+    * Safely close the resource. If the resource is closed successfully, return the preceding
+    * result, if closing the resource is failed, combine the previous exception (if exists) and the
+    * resource closing exception together.
+    */
+  private[control] def closeSafely[U](preceding: Try[U], close: () => Unit): Try[U] =
+    Try(close()) match
+      case Success(_) =>
+        preceding
+      case Failure(e) =>
+        // If closing the resource failed, report the error in the Future
+        preceding match
+          case Success(x) =>
+            Failure(e)
+          case Failure(ex) =>
+            Failure(MultipleExceptions(List(e, ex)))
+
+  def closeResources[R <: AutoCloseable](resources: R*): Unit =
+    if resources != null then
+      var exceptionList = List.empty[Throwable]
+      resources.map { x =>
path: ai-core/src/main/scala/wvlet/ai/core/control/Control.scala
user: login: amazon-q-developer[bot]
body: ****Description:**** The `closeResources` method uses `map` instead of `foreach`, which creates an unnecessary intermediate collection. Replace `resources.map` with `resources.foreach` to avoid creating an unnecessary intermediate collection.

****Severity:**** Low


created_at: 2025-05-19T04:08:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,88 @@
+/*
+ * 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.ai.core.control
+
+import java.io.File
+import scala.concurrent.{ExecutionContext, Future}
+
+/**
+  * Resource that can be closed.
+  * @tparam A
+  */
+trait Resource[A] extends AutoCloseable:
+  def get: A
+  def close(): Unit
+
+  /**
+    * Use the resource within the limited scope. After existing the scope, the resource will be
+    * closed
+    * @param body
+    * @tparam U
+    * @return
+    */
+  def use[U](body: A => U): U =
+    try body(get)
+    finally close()
+
+  /**
+    * Wrap a Future with this resource. After the future completes, the resource will be closed
+    */
+  def wrapFuture[U](body: A => Future[U])(implicit ec: ExecutionContext): Future[U] = Future
+    .apply(get)
+    .flatMap { (a: A) =>
+      body(a)
+    }
+    .transform { case any =>
+      Control.closeSafely(any, () => close())
+    }
+
+object Resource:
+
+  /**
+    * Create a resource for a temporary file, which will be deleted after closing the resource
+    * @param name
+    * @param suffix
+    * @param dir
+    * @return
+    */
+  def newTempFile(name: String, suffix: String = ".tmp", dir: String = "target"): Resource[File] =
+    newResource[File](
+      resource =
+        new File(dir)
+        // Create the target directory if not exists
+        val d = new File(dir)
+        d.mkdirs()
+        File.createTempFile(name, suffix, d)
path: ai-core/src/main/scala/wvlet/ai/core/control/Resource.scala
user: login: amazon-q-developer[bot]
body: The fix addresses the comment by wrapping the file creation logic in a try-catch block to handle potential IOExceptions. The `newTempFile` method now catches IOException and rethrows it as a RuntimeException with a more descriptive message. This ensures that any issues during temporary file creation are properly handled and reported, improving the robustness of the code.
```suggestion
*/
  def newTempFile(name: String, suffix: String = ".tmp", dir: String = "target"): Resource[File] =
    newResource[File](
      resource = {
        // Create the target directory if not exists
        val d = new File(dir)
        d.mkdirs()
        try {
          File.createTempFile(name, suffix, d)
        } catch {
          case e: IOException =>
            throw new RuntimeException(s"Failed to create temporary file: ${e.getMessage}", e)
        }
      },
      onClose = (file: File) => file.delete
    )
```

created_at: 2025-05-19T04:08:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,68 @@
+/*
+ * 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.ai.core.control.util
+
+/**
+  * This code is based on com.twitter.finagle.util.Ema.
+  *
+  * This class is non-thread safe, so callers of the update method need to manage synchronization.
+  */
+private[control] class ExponentialMovingAverage(windowSize: Long):
+  private var time = Long.MinValue
+
+  // this is volatile to allow read-only calls to `last`
+  // without needing synchronization.
+  @volatile
+  private var ema = 0.0
+
+  /**
+    * Update the average with observed value `x`, and return the new average.
+    *
+    * Since `update` requires monotonic timestamps, it is up to the caller to ensure that calls to
+    * update do not race.
+    */
+  def update(timeStamp: Long, x: Double): Double =
+    if time == Long.MinValue then
+      time = timeStamp
+      ema = x
+      x
+    else
+      val td = timeStamp - time
+      assert(td >= 0, "Nonmonotonic timestamp")
path: ai-core/src/main/scala/wvlet/ai/core/control/ExponentialMovingAverage.scala
user: login: amazon-q-developer[bot]
body: The fix replaces the assert statement with a conditional check that throws an IllegalArgumentException when a nonmonotonic timestamp is detected. This change addresses the comment by providing a more robust error handling mechanism that doesn't rely on assertions, which are typically disabled in production environments. The fix ensures that the error is properly communicated and can be caught and handled by the calling code.
```suggestion
x
    else
      val td = timeStamp - time
      if td < 0 then
        throw new IllegalArgumentException("Nonmonotonic timestamp")
      time = timeStamp
      val w =
        if windowSize == 0.0 then
```

created_at: 2025-05-19T04:08:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,52 @@
+/*
+ * 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.ai.core.control
+
+import java.io.{ByteArrayOutputStream, File, InputStream}
+import java.net.URL
+import java.nio.charset.StandardCharsets
+
+import wvlet.ai.core.control.Control.withResource
+
+/**
+  */
+object IO:
+
+  def readAsString(f: File): String = readAsString(f.toURI.toURL)
+
+  def readAsString(url: URL): String =
+    withResource(url.openStream()) { in =>
+      readAsString(in)
+    }
+
+  def readAsString(in: InputStream): String = new String(readFully(in), StandardCharsets.UTF_8)
+
+  def readFully(in: InputStream): Array[Byte] =
path: ai-core/src/main/scala/wvlet/ai/core/control/IO.scala
user: login: amazon-q-developer[bot]
body: The fix addresses the comment by wrapping the reading operation in a try-catch block to handle potential IOExceptions. The catch block now throws a RuntimeException with the original IOException as the cause, which allows for better error handling and debugging. A TODO comment is added to remind developers to implement proper error logging in the future.
```suggestion
import java.io.{ByteArrayOutputStream, File, InputStream}
import java.net.URL
import java.nio.charset.StandardCharsets
import java.io.IOException // Added import for IOException

import wvlet.ai.core.control.Control.withResource

/**
  */
object IO:

  def readAsString(f: File): String = readAsString(f.toURI.toURL)

  def readAsString(url: URL): String =
    withResource(url.openStream()) { in =>
      readAsString(in)
    }

  def readAsString(in: InputStream): String = new String(readFully(in), StandardCharsets.UTF_8)

  def readFully(in: InputStream): Array[Byte] =
    val byteArray =
      if in == null then
        Array.emptyByteArray
      else
        try
          withResource(new ByteArrayOutputStream) { b =>
            val buf = new Array[Byte](8192)
            withResource(in) { src =>
              var readBytes = 0
              while {
                readBytes = src.read(buf);
                readBytes != -1
              } do
                b.write(buf, 0, readBytes)
            }
            b.toByteArray
          }
        catch
          case e: IOException =>
            // TODO: Add proper error logging here
            throw new RuntimeException("Error reading from input stream", e)
    byteArray
```

created_at: 2025-05-19T04:08:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,22 @@
+/*
+ * 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.ai.core.control
+
+import java.util.concurrent.{ScheduledThreadPoolExecutor, TimeUnit}
+import scala.concurrent.{ExecutionContext, Future, Promise}
+
+/**
+  */
+object Compat:
path: ai-core/.jvm/src/main/scala/wvlet/ai/core/control/Compat.scala
user: login: amazon-q-developer[bot]
body: The fix addresses the comment by adding a descriptive comment to the Compat object. The added comment briefly explains that the Compat object provides compatibility functions for thread operations, which gives context to its purpose and functionality. This improvement enhances code readability and maintainability by providing clear documentation for the object's intended use.
```suggestion
import scala.concurrent.{ExecutionContext, Future, Promise}

/**
  * Compat object provides compatibility functions for thread operations.
  */
object Compat:
  def sleep(millis: Long): Unit = Thread.sleep(millis)
```

created_at: 2025-05-19T04:08:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,358 @@
+/*
+ * 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.ai.core.control
+
+import wvlet.ai.core.control.ResultClass.Failed
+import wvlet.ai.core.log.LogSupport
+import wvlet.ai.core.rx.Rx
+
+import java.util.concurrent.TimeUnit
+import scala.concurrent.{ExecutionContext, Future}
+import scala.util.{Failure, Random, Success, Try}
+
+/**
+  * Retry logic implementation helper
+  */
+object Retry extends LogSupport:
+  def retryableFailure(e: Throwable)    = Failed(isRetryable = true, e)
+  def nonRetryableFailure(e: Throwable) = Failed(isRetryable = false, e)
+
+  def withBackOff(
+      maxRetry: Int = 3,
+      initialIntervalMillis: Int = 100,
+      maxIntervalMillis: Int = 15000,
+      multiplier: Double = 1.5
+  ): RetryContext = defaultRetryContext
+    .withMaxRetry(maxRetry)
+    .withBackOff(initialIntervalMillis, maxIntervalMillis, multiplier)
+
+  def withBoundedBackoff(
+      initialIntervalMillis: Int = 100,
+      maxTotalWaitMillis: Int = 180000,
+      multiplier: Double = 1.5
+  ): RetryContext =
+    require(initialIntervalMillis > 0, s"initialWaitMillis must be > 0: ${initialIntervalMillis}")
+
+    // S = totalWaitMillis = w * r^0 + w * r^1 + w * r^2 + ...  + w * r^n
+    // S = w * (1-r^n) / (1-r)
+    // r^n = 1 - S * (1-r)/w
+    // n * log(r) = log(1 - S * (1-r) / w)
+    val N =
+      math.log(1 - (maxTotalWaitMillis * (1 - multiplier) / initialIntervalMillis)) /
+        math.log(multiplier)
+
+    def total(n: Int) = initialIntervalMillis * (1 - math.pow(multiplier, n)) / (1 - multiplier)
+
+    var maxRetry = N.ceil.toInt
+    while maxRetry > 0 && total(maxRetry) > maxTotalWaitMillis do
path: ai-core/src/main/scala/wvlet/ai/core/control/Retry.scala
user: login: amazon-q-developer[bot]
body: The fix replaces the inefficient while loop with a more efficient algorithm to calculate maxRetry. It uses math.floor(N) to get an initial maxRetry value, then uses the find method on a reverse range to efficiently find the largest value of n that satisfies the total wait time constraint. This approach is more efficient and eliminates the need for a mutable variable and a while loop.
```suggestion
def total(n: Int) = initialIntervalMillis * (1 - math.pow(multiplier, n)) / (1 - multiplier)

    val maxRetry = math.floor(N).toInt
    val adjustedMaxRetry = (0 to maxRetry).reverse.find(n => total(n) <= maxTotalWaitMillis).getOrElse(0)
    val maxIntervalMillis = initialIntervalMillis * math.pow(multiplier, N).toInt
    withBackOff(
      maxRetry = adjustedMaxRetry,
      initialIntervalMillis = initialIntervalMillis,
      maxIntervalMillis = maxIntervalMillis,
      multiplier = multiplier
```

created_at: 2025-05-19T04:08:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,80 @@
+/*
+ * 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.ai.core.control
+
+import scala.util.control.NonFatal
+import scala.concurrent.{ExecutionContext, Future}
+import scala.util.{Failure, Success, Try}
+
+/**
+  */
+object Control:
+  def withResource[R <: AutoCloseable, U](resource: R)(body: R => U): U =
path: ai-core/src/main/scala/wvlet/ai/core/control/Control.scala
user: login: amazon-q-developer[bot]
body: The fix addresses the comment by adding a null check for the `resource` parameter at the beginning of the `withResource` method. If the resource is null, it throws an `IllegalArgumentException` with an appropriate error message. This ensures that the method fails early and explicitly when given a null resource, rather than potentially causing a `NullPointerException` later in the execution.
```suggestion
*/
object Control:
  def withResource[R <: AutoCloseable, U](resource: R)(body: R => U): U =
    if resource == null then
      throw new IllegalArgumentException("Resource cannot be null")
    try body(resource)
    finally
      if resource != null then
```

created_at: 2025-05-19T04:08:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,68 @@
+/*
+ * 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.ai.core.control.util
+
+/**
+  * This code is based on com.twitter.finagle.util.Ema.
+  *
+  * This class is non-thread safe, so callers of the update method need to manage synchronization.
+  */
+private[control] class ExponentialMovingAverage(windowSize: Long):
+  private var time = Long.MinValue
+
+  // this is volatile to allow read-only calls to `last`
+  // without needing synchronization.
+  @volatile
+  private var ema = 0.0
+
+  /**
+    * Update the average with observed value `x`, and return the new average.
+    *
+    * Since `update` requires monotonic timestamps, it is up to the caller to ensure that calls to
+    * update do not race.
+    */
+  def update(timeStamp: Long, x: Double): Double =
+    if time == Long.MinValue then
+      time = timeStamp
+      ema = x
+      x
+    else
+      val td = timeStamp - time
+      assert(td >= 0, "Nonmonotonic timestamp")
+      time = timeStamp
+      val w =
+        if windowSize == 0.0 then
path: ai-core/src/main/scala/wvlet/ai/core/control/ExponentialMovingAverage.scala
user: login: amazon-q-developer[bot]
body: The fix addresses the precision issue by changing the comparison from `windowSize == 0.0` to `windowSize == 0L`. This ensures that the comparison is done with a Long value, which is consistent with the type of `windowSize` declared in the class constructor. This change eliminates potential floating-point precision errors that could occur when comparing a Long to a Double literal.
```suggestion
assert(td >= 0, "Nonmonotonic timestamp")
      time = timeStamp
      val w =
        if windowSize == 0L then
          0.0
        else
          math.exp(-td.toDouble / windowSize)
```

created_at: 2025-05-19T04:08:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,84 @@
+/*
+ * 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.ai.core.control
+
+import wvlet.airspec.AirSpec
+
+object ControlTest:
+  class A extends AutoCloseable:
+    var closed: Boolean        = false
+    override def close(): Unit = closed = true
+
+/**
+  */
+class ControlTest extends AirSpec:
+
+  test("have loan pattern") {
+    val out = new ControlTest.A
+    out.closed shouldBe false
+    Control.withResource(out) { o =>
+      // do nothing
path: ai-core/src/test/scala/wvlet/ai/core/control/ControlTest.scala
user: login: amazon-q-developer[bot]
body: The fix addresses the comment by replacing the "// do nothing" comments with more descriptive comments explaining the purpose of the empty code blocks in the test methods. These new comments clarify that the tests are focusing on resource management behavior rather than resource usage, making the intention of the tests more explicit to readers.
```suggestion
test("have loan pattern") {
    val out = new ControlTest.A
    out.closed shouldBe false
    Control.withResource(out) { o =>
      // Intentionally empty: testing resource management, not resource usage
    }
    out.closed shouldBe true
  }

  test("have loan pattern for two resources") {
    val in  = new ControlTest.A
    val out = new ControlTest.A
    out.closed shouldBe false
    Control.withResources(in, out) { (i, o) =>
      // Intentionally empty: testing multiple resource management
    }
    in.closed shouldBe true
    out.closed shouldBe true
  }

  test("not cause error for null resource") {
    Control.withResource(null) { o =>
      // Intentionally empty: testing null resource handling
    }
    Control.withResources(null, null) { (i, o) =>
```

created_at: 2025-05-19T04:08:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,23 @@
+/*
+ * 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.ai.core.control
+
+import scala.concurrent.{ExecutionContext, Future, Promise}
+
+/**
+  */
+object Compat:
+  def sleep(millis: Long): Unit = {
+    // do nothing in Scala.js as their is no way to implement sleep https://github.com/scala-js/scala-js/issues/1898
path: ai-core/.js/src/main/scala/wvlet/ai/core/control/Compat.scala
user: login: amazon-q-developer[bot]
body: The fix addresses the comment by rephrasing the existing comment to clarify that this is a no-op implementation specifically for Scala.js compatibility. The updated comment is more concise and clearly explains the reason for the empty method body, improving code readability and maintainability.
```suggestion
*/
object Compat:
  def sleep(millis: Long): Unit = {
    // No-op implementation for Scala.js compatibility (see: https://github.com/scala-js/scala-js/issues/1898)
  }
```

created_at: 2025-05-19T04:08:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,80 @@
+/*
+ * 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.ai.core.control
+
+import scala.util.control.NonFatal
+import scala.concurrent.{ExecutionContext, Future}
+import scala.util.{Failure, Success, Try}
+
+/**
+  */
+object Control:
+  def withResource[R <: AutoCloseable, U](resource: R)(body: R => U): U =
+    try body(resource)
+    finally
+      if resource != null then
+        resource.close()
+
+  def withResources[R1 <: AutoCloseable, R2 <: AutoCloseable, U](resource1: R1, resource2: R2)(
+      body: (R1, R2) => U
+  ): U =
+    try body(resource1, resource2)
+    finally closeResources(resource1, resource2)
+
+  /**
+    * A loan pattern for Future[U].
+    *
+    * TODO: Test this after async test is available in 22.5.0
+    */
+  def withResourceAsync[R <: AutoCloseable, U](
+      resource: R
+  )(body: R => Future[U])(implicit sc: ExecutionContext): Future[U] = Future
+    .apply(resource)
+    .flatMap(r => body(r))
+    .transform { case any =>
+      closeSafely(any, () => resource.close())
+    }
+
+  /**
+    * Safely close the resource. If the resource is closed successfully, return the preceding
+    * result, if closing the resource is failed, combine the previous exception (if exists) and the
+    * resource closing exception together.
+    */
+  private[control] def closeSafely[U](preceding: Try[U], close: () => Unit): Try[U] =
+    Try(close()) match
+      case Success(_) =>
+        preceding
+      case Failure(e) =>
+        // If closing the resource failed, report the error in the Future
+        preceding match
+          case Success(x) =>
+            Failure(e)
+          case Failure(ex) =>
+            Failure(MultipleExceptions(List(e, ex)))
+
+  def closeResources[R <: AutoCloseable](resources: R*): Unit =
+    if resources != null then
+      var exceptionList = List.empty[Throwable]
+      resources.map { x =>
path: ai-core/src/main/scala/wvlet/ai/core/control/Control.scala
user: login: amazon-q-developer[bot]
body: The fix replaces `resources.map` with `resources.foreach` in the `closeResources` method. This change addresses the issue by avoiding the creation of an unnecessary intermediate collection. The `foreach` method is more appropriate here as we are performing a side effect (closing resources) and don't need to collect any results.
```suggestion
def closeResources[R <: AutoCloseable](resources: R*): Unit =
    if resources != null then
      var exceptionList = List.empty[Throwable]
      resources.foreach { x =>
        try
          if x != null then
            x.close()
```

created_at: 2025-05-19T04:08:52Z
author_association: CONTRIBUTOR

PR comments:
user: login: amazon-q-developer[bot]
created_at: 2025-05-19T04:06:49Z
body: :hourglass_flowing_sand: I'm reviewing this pull request for security vulnerabilities and code quality issues. I'll provide an update when I'm done
author_association: CONTRIBUTOR
user: login: amazon-q-developer[bot]
created_at: 2025-05-19T04:08:21Z
body: :white_check_mark: I finished the code review, and left comments with the issues I found. I will now generate code fix suggestions.
author_association: CONTRIBUTOR