Serialized TP server using Akka streams and actor to manage state

Server.scala

package io

import scala.concurrent.Future
import akka.actor.ActorRef
import akka.actor.ActorSystem
import akka.actor.FSM
import akka.actor.PoisonPill
import akka.actor.Props
import akka.stream.ActorMaterializer
import akka.stream.OverflowStrategy
import akka.stream.scaladsl.Flow
import akka.stream.scaladsl.Sink
import akka.stream.scaladsl.Source
import akka.stream.scaladsl.Tcp
import akka.stream.scaladsl.Tcp.IncomingConnection
import akka.stream.scaladsl.Tcp.ServerBinding
import akka.util.ByteString
import akka.actor.Actor

/**
 * TCP server over a serialized resource that permits interaction with at most one client at a time
 */
object ServerNoToken extends App {
  implicit val system = ActorSystem("Sys")
  implicit val materializer = ActorMaterializer()

  val manager = system.actorOf(Props(new ResourceManager(system.actorOf(Props[Resource]))),
    name = "manager")

  val connections: Source[IncomingConnection, Future[ServerBinding]] =
    Tcp().bind("127.0.0.1", 9999)

  case class Done(token: Any)

  connections runForeach { connection =>
    val connectionHandler = system.actorOf(Props(new ConnectionHandler(manager)),
      name = "connectionHandler" + connection.remoteAddress.getPort)

    val src = Source.actorRef[ByteString](10, OverflowStrategy.fail)
    val sink = Sink.actorRef(connectionHandler, Done)

    val handler = Flow.wrap(sink, src)((m1, m2) => m2)
    val m = connection.handleWith(handler)

    connectionHandler ! Attach(m)
  }

  class ConnectionHandler(manager: ActorRef) extends FSM[State, Option[ActorRef]] {
    manager ! Arrived(self)

    startWith(Idle, None)

    when(Idle) {
      case Event(Attach(connection), None) =>
        goto(Running) using Some(connection)
    }

    when(Running) {
      case Event(Done, None) =>
        stop
      case Event(Done, Some(connection)) =>
        manager ! Gone(self)
        connection ! PoisonPill
        stop()
      case Event(Die, Some(connection)) =>
        connection ! PoisonPill
        stay using None
      case Event(data: ByteString, Some(connection)) =>
        manager ! Receive(data)
        stay
      case Event(Write(data), Some(connection)) =>
        connection ! data
        stay
    }
  }

  case object Done

  case object Die

  case class Attach(connection: ActorRef)

  case class Receive(data: ByteString)
  case class Write(data: ByteString)

  case class Arrived(client: ActorRef)
  case class Gone(client: ActorRef)

  sealed trait State
  case object Idle extends State
  case object Running extends State

  class ResourceManager(resource: ActorRef) extends FSM[State, Option[ActorRef]] {

    startWith(Idle, None)

    when(Idle) {
      case Event(Arrived(newClient), None) =>
        goto(Running) using Some(newClient)
    }

    when(Running) {
      case Event(Arrived(newClient), Some(oldClient)) =>
        oldClient ! Die
        stay using Some(newClient)
      case Event(Receive(data), Some(client)) =>
        resource tell (data, client)
        stay
      case Event(Gone(oldClient), Some(client)) if (oldClient == client) =>
        goto(Idle) using None
      case Event(Gone(oldClient), Some(client)) =>
        stay
    }
  }

  class Resource extends Actor {
    def receive = {
      case bs: ByteString => sender() ! Write(bs)
    }
  }
}