Search code examples
scalaioakkaactor

How can I throttle messages to the IO(Tcp) actor in Akka


I have an actor like this

class TcpClientActor(target: Target) extends Actor with Logger {

  override def preStart(): Unit = {
    self ! TestConnection
  }

  override def receive: Receive = {
    case TestConnection =>
      IO(Tcp) ! Connect(remoteAddress = new InetSocketAddress(target.endpoint, target.port), localAddress = None, options = Nil, timeout = Some(timeout), pullMode = false)

    case failed@CommandFailed(_: Connect) =>
      info(s"Failure: $target.endpoint:$target.port")
      shutdown()

    case Connected(_, _) =>
      info(s"Success: $target.endpoint:$target.port")
      sender() ! Close
      shutdown()

  }

  def shutdown(): Unit = {
    context stop self
  }
}

I'm iterating over a file with endpoints to test against and creating one of these actors with each line as a constructor argument of type Target. I want to be able to throttle the number of parallel TCP connections to initiate to some set number, are there built-in mechanisms I can use in Akka to ensure I don't overload the system by just immediately creating a TcpClientActor for every line of input and kicking off a socket connection?


Solution

  • I would use an Akka Stream to throttle the messages

    import scala.concurrent.duration._
    import akka.NotUsed
    import akka.actor.ActorRef
    import akka.stream.{ ActorMaterializer, OverflowStrategy, ThrottleMode }
    import akka.stream.scaladsl.{ Sink, Source }
    
    object TcpThrottle {
      def throttler(ratePerSecond: Int, burstRate: Option[Int], bufferSize: Int = 1000)(implicit materializer: ActorMaterializer): ActorRef =
        Source.actorRef(bufferSize = bufferSize, OverflowStrategy.dropNew)
          .throttle(ratePerSecond, 1.second, burstRate.getOrElse(ratePerSecond), ThrottleMode.Shaping)
          .to(Sink.actorRef(IO(Tcp), NotUsed)
          .run()
     }
    
    class TcpClientActor(target: Target) extends Actor with Logger {
      val throttler = TcpThrottle.throttler(1, Some(5))
    
      // otherwise identical
    
      def receive: Receive = {
        case TestConnection => throttler ! Connect(remoteAddress = new InetSocketAddress(target.endpoint, target.port), localAddress = None, options = Nil, timeout = Some(timeout), pullMode = false)
    
        // other cases identical
      }
    }
    

    Adapted from The Akka 2.5 migration guide. It may be necessary