Basic Akka Streams Source for receiving UDP datagrams

Udp.scala

package test

import java.net.InetSocketAddress

import akka.actor.{ActorLogging, ActorRef, Props}
import akka.io
import akka.stream.ActorMaterializer
import akka.stream.actor.ActorPublisher
import akka.stream.scaladsl.Source

import scala.annotation.tailrec
import scala.concurrent.{Future, Promise}


/**
  * Basic demo of receiving UDP datagrams via an Akka Streams Source.
  *
  * The main method uses the UDP source to just print UTF-8 text received in datagrams
  * to stdout.
  *
  * Test via Netcat or ncat:
  *
  * `$ echo 'Hello, world' | ncat -u 127.0.0.1 6000`
  *
  * This is very unlikely to be production quality!
  */
object Udp {

  final class ServerBinding(val localAddress: InetSocketAddress, unbindAction: () => Future[Unit]) {
    def unbind(): Future[Unit] = unbindAction()
  }

  object HandlerActor {

    def props(listenOn: InetSocketAddress, maxBuffer: Int): Props = Props(new HandlerActor(listenOn, maxBuffer))

  }

  private final class HandlerActor(listenOn: InetSocketAddress, maxBuffer: Int) extends ActorPublisher[io.Udp.Received] with ActorLogging {

    require(maxBuffer > 0, "Datagram buffer must hold at least one datagram")

    import akka.stream.actor.ActorPublisherMessage._
    import context.system

    io.IO(io.Udp) ! io.Udp.Bind(self, listenOn)

    private val bindingPromise = Promise[ServerBinding]
    private val unbindPromise = Promise[Unit]
    private var buf = Vector.empty[io.Udp.Received]
    private var listener: ActorRef = _
    private var suspended = false

    def receive = {
      case io.Udp.Bound(local) =>
        log.info("UDP listener bound to {}", local)
        listener = sender()
        val target = self
        bindingPromise.success(new ServerBinding(local, () => {
          target ! io.Udp.Unbind
          unbindPromise.future
        }))
        context.become(ready())
      case Request(_) =>
        // can't send anything to a stream until we're bound
      case Cancel =>
        context.stop(self)
      case io.Udp.CommandFailed(c) =>
        log.error("UDP command {} failed", c)
        onCompleteThenStop()
    }

    private def ready(): Receive = {
      case d@io.Udp.Received(_, _) =>
        // if we're at max buffer size, ignore the request
        // for UDP, clients should be able to gracefully handle undelivered datagrams anyway
        if (buf.size < maxBuffer) {
          buf :+= d
          deliverToStream()
        } else if (!suspended) {
          log.warning("UDP datagram buffer maximum size {} reached, suspending reading", maxBuffer)
          suspended = true
          listener ! io.Udp.SuspendReading
        }
      case Request(_) =>
        deliverToStream()
      case Cancel =>
        listener ! io.Udp.Unbind
      case io.Udp.CommandFailed(c) =>
        log.error("UDP command {} failed", c)
        onCompleteThenStop()
      case io.Udp.Unbound =>
        log.debug("Received unbind request")
        unbindPromise.success(())
        onCompleteThenStop()
        sender ! io.Udp.Unbound
    }

    @tailrec
    private def deliverToStream(): Unit = {
      if (totalDemand > 0) {
        if (suspended) {
          log.info("UDP datagram buffer space now available, resuming reading")
          suspended = false
          listener ! io.Udp.ResumeReading
        }
        log.debug("Delivering for total demand of {}, with {} received datagrams", totalDemand, buf.size)
        if (totalDemand <= Int.MaxValue) {
          val (use, keep) = buf.splitAt(totalDemand.toInt)
          buf = keep
          use foreach onNext
        } else {
          val (use, keep) = buf.splitAt(Int.MaxValue)
          buf = keep
          use foreach onNext
          deliverToStream()
        }
      }
    }

  }

  def source(listenOn: InetSocketAddress, maxBuffer: Int = 10): Source[io.Udp.Received, ActorRef] = {
    Source.actorPublisher[io.Udp.Received](HandlerActor.props(listenOn, maxBuffer))
  }

  def main(args: Array[String]): Unit = {
    implicit val system = akka.actor.ActorSystem("udptest")
    implicit val materializer = ActorMaterializer()

    val udpSource = source(
      listenOn = new InetSocketAddress(java.net.InetAddress.getLoopbackAddress, 6000),
      maxBuffer = 10)
    udpSource.runForeach {
      case io.Udp.Received(data, remote) =>
        println(s"Received from $remote: [${data.utf8String}]")
    }
    println("Press Ctrl-C to quit")
  }
}