Reduce receive-timeout state allocations

actor-tests/src/test/scala/org/apache/pekko/actor/ReceiveTimeoutSpec.scala

@@ -96,18 +96,22 @@ class ReceiveTimeoutSpec extends PekkoSpec() {
 
     "reschedule timeout after regular receive" taggedAs TimingTest in {
       val timeoutLatch = TestLatch()
+      val messageCount = 100
+      val processedLatch = TestLatch(messageCount)
 
       val timeoutActor = system.actorOf(Props(new Actor {
-        context.setReceiveTimeout(500.milliseconds)
+        context.setReceiveTimeout(1.second)
 
         def receive = {
-          case Tick           => ()
+          case Tick           => processedLatch.countDown()
           case ReceiveTimeout => timeoutLatch.open()
         }
       }))
 
-      timeoutActor ! Tick
+      (1 to messageCount).foreach(_ => timeoutActor ! Tick)
 
+      Await.ready(processedLatch, TestLatch.DefaultTimeout)
+      intercept[TimeoutException] { Await.ready(timeoutLatch, 500.milliseconds) }
       Await.ready(timeoutLatch, TestLatch.DefaultTimeout)
       system.stop(timeoutActor)
     }

actor/src/main/scala/org/apache/pekko/actor/dungeon/ReceiveTimeout.scala

@@ -22,63 +22,83 @@ import pekko.actor.Cancellable
 import pekko.actor.NotInfluenceReceiveTimeout
 
 private[pekko] object ReceiveTimeout {
-  final val emptyReceiveTimeoutData: (Duration, Cancellable) = (Duration.Undefined, ActorCell.emptyCancellable)
+  // Compared only across one message invocation, so wrapping over an actor's lifetime is harmless.
+  private final class State(var timeout: Duration, var task: Cancellable, var version: Int)
 }
 
 private[pekko] trait ReceiveTimeout { this: ActorCell =>
 
   import ActorCell._
   import ReceiveTimeout._
 
-  private var receiveTimeoutData: (Duration, Cancellable) = emptyReceiveTimeoutData
+  private var receiveTimeoutData: State = _
 
-  final def receiveTimeout: Duration = receiveTimeoutData._1
+  final def receiveTimeout: Duration =
+    if (receiveTimeoutData eq null) Duration.Undefined else receiveTimeoutData.timeout
 
-  final def setReceiveTimeout(timeout: Duration): Unit = receiveTimeoutData = receiveTimeoutData.copy(_1 = timeout)
+  final def setReceiveTimeout(timeout: Duration): Unit = {
+    val data = receiveTimeoutData
+    if (data eq null)
+      receiveTimeoutData = new State(timeout, emptyCancellable, version = 1)
+    else {
+      data.timeout = timeout
+      data.version += 1
+    }
+  }
 
   /** Called after `ActorCell.receiveMessage` or `ActorCell.autoReceiveMessage`. */
-  protected def checkReceiveTimeoutIfNeeded(message: Any, beforeReceive: (Duration, Cancellable)): Unit =
-    if (hasTimeoutData || receiveTimeoutChanged(beforeReceive))
-      checkReceiveTimeout(!message.isInstanceOf[NotInfluenceReceiveTimeout] || receiveTimeoutChanged(beforeReceive))
+  protected def checkReceiveTimeoutIfNeeded(message: Any, beforeReceiveVersion: Int): Unit = {
+    val timeoutChanged = receiveTimeoutChanged(beforeReceiveVersion)
+    if (hasTimeoutData || timeoutChanged)
+      checkReceiveTimeout(!message.isInstanceOf[NotInfluenceReceiveTimeout] || timeoutChanged)
+  }
 
   final def checkReceiveTimeout(reschedule: Boolean): Unit = {
-    val (recvTimeout, task) = receiveTimeoutData
-    recvTimeout match {
-      case f: FiniteDuration =>
-        // The fact that recvTimeout is FiniteDuration and task is emptyCancellable
-        // means that a user called `context.setReceiveTimeout(...)`
-        // while sending the ReceiveTimeout message is not scheduled yet.
-        // We have to handle the case and schedule sending the ReceiveTimeout message
-        // ignoring the reschedule parameter.
-        if (reschedule || (task eq emptyCancellable))
-          rescheduleReceiveTimeout(f)
-
-      case _ => cancelReceiveTimeoutTask()
+    val data = receiveTimeoutData
+    if (data ne null) {
+      data.timeout match {
+        case f: FiniteDuration =>
+          // The fact that timeout is FiniteDuration and task is emptyCancellable
+          // means that a user called `context.setReceiveTimeout(...)`
+          // while sending the ReceiveTimeout message is not scheduled yet.
+          // We have to handle the case and schedule sending the ReceiveTimeout message
+          // ignoring the reschedule parameter.
+          if (reschedule || (data.task eq emptyCancellable))
+            rescheduleReceiveTimeout(data, f)
+
+        case _ => cancelReceiveTimeoutTask()
+      }
     }
   }
 
-  private def rescheduleReceiveTimeout(f: FiniteDuration): Unit = {
-    receiveTimeoutData._2.cancel() // Cancel any ongoing future
-    val task = system.scheduler.scheduleOnce(f, self, pekko.actor.ReceiveTimeout)(this.dispatcher)
-    receiveTimeoutData = (f, task)
+  private def rescheduleReceiveTimeout(data: State, timeout: FiniteDuration): Unit = {
+    data.task.cancel() // Cancel any ongoing future
+    data.task = system.scheduler.scheduleOnce(timeout, self, pekko.actor.ReceiveTimeout)(this.dispatcher)
+    data.version += 1
   }
 
-  private def hasTimeoutData: Boolean = receiveTimeoutData ne emptyReceiveTimeoutData
+  private def hasTimeoutData: Boolean = receiveTimeoutData ne null
 
-  private def receiveTimeoutChanged(beforeReceive: (Duration, Cancellable)): Boolean =
-    receiveTimeoutData ne beforeReceive
+  private def receiveTimeoutVersion: Int =
+    if (receiveTimeoutData eq null) 0 else receiveTimeoutData.version
 
-  protected def cancelReceiveTimeoutIfNeeded(message: Any): (Duration, Cancellable) = {
+  private def receiveTimeoutChanged(beforeReceiveVersion: Int): Boolean =
+    receiveTimeoutVersion != beforeReceiveVersion
+
+  protected def cancelReceiveTimeoutIfNeeded(message: Any): Int = {
     if (hasTimeoutData && !message.isInstanceOf[NotInfluenceReceiveTimeout])
       cancelReceiveTimeoutTask()
 
-    receiveTimeoutData
+    receiveTimeoutVersion
   }
 
-  private[pekko] def cancelReceiveTimeoutTask(): Unit =
-    if (receiveTimeoutData._2 ne emptyCancellable) {
-      receiveTimeoutData._2.cancel()
-      receiveTimeoutData = (receiveTimeoutData._1, emptyCancellable)
+  private[pekko] def cancelReceiveTimeoutTask(): Unit = {
+    val data = receiveTimeoutData
+    if ((data ne null) && (data.task ne emptyCancellable)) {
+      data.task.cancel()
+      data.task = emptyCancellable
+      data.version += 1
     }
+  }
 
 }

bench-jmh/src/main/scala/org/apache/pekko/actor/ReceiveTimeoutBenchmark.scala

@@ -0,0 +1,96 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.pekko.actor
+
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
+
+import scala.concurrent.duration._
+
+import com.typesafe.config.ConfigFactory
+import org.openjdk.jmh.annotations._
+
+@State(Scope.Benchmark)
+@BenchmarkMode(Array(Mode.Throughput))
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+@Fork(3)
+@Threads(1)
+@Warmup(iterations = 5, time = 1)
+@Measurement(iterations = 5, time = 1)
+class ReceiveTimeoutBenchmark {
+  import ReceiveTimeoutBenchmark._
+
+  private var system: ActorSystem = _
+  private var target: ActorRef = _
+
+  @Setup(Level.Trial)
+  def setup(): Unit = {
+    system = ActorSystem(
+      "ReceiveTimeoutBenchmark",
+      ConfigFactory.parseString("""
+        pekko {
+          log-dead-letters = off
+          scheduler.tick-duration = 1 ms
+          actor.default-dispatcher {
+            fork-join-executor {
+              parallelism-min = 1
+              parallelism-factor = 1.0
+              parallelism-max = 1
+            }
+            throughput = 1000
+          }
+        }
+      """))
+    target = system.actorOf(Props[TimeoutActor]())
+  }
+
+  @TearDown(Level.Trial)
+  def shutdown(): Unit =
+    system.close()
+
+  @Benchmark
+  @OperationsPerInvocation(MessageCount)
+  def influencingMessages(): Unit = {
+    val latch = new CountDownLatch(1)
+    var i = 0
+    while (i < MessageCount) {
+      target ! Message
+      i += 1
+    }
+    target ! new Finished(latch)
+    latch.await()
+  }
+}
+
+object ReceiveTimeoutBenchmark {
+  final val MessageCount = 10000
+
+  case object Message
+
+  final class Finished(val latch: CountDownLatch) extends NotInfluenceReceiveTimeout
+
+  final class TimeoutActor extends Actor {
+    context.setReceiveTimeout(100.millis)
+
+    override def receive: Receive = {
+      case Message            =>
+      case finished: Finished => finished.latch.countDown()
+      case ReceiveTimeout     =>
+    }
+  }
+}