Serialized TCP server using Akka streams that accepts at most one client

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 Server extends App {
  implicit val system = ActorSystem("Sys")
  implicit val materializer = ActorMaterializer()

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

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

  case class Done(token: Any)

  connections runForeach { connection =>
    val src = Source.actorRef[ByteString](10, OverflowStrategy.fail)
    val sink = Sink.actorRef(manager, Done(connection))

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

    manager ! Connection(connection, m)
  }

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

  case class Connection(token: Any, client: ActorRef)

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

    startWith(Idle, None)

    when(Idle) {
      case Event(c: Connection, None) =>
        goto(Running) using Some(c)
    }

    when(Running) {
      case Event(newConnection: Connection, Some(Connection(_, oldClient))) =>
        oldClient ! PoisonPill
        stay using Some(newConnection)
      case Event(bs: ByteString, Some(Connection(_, client))) =>
        resource tell (bs, client)
        stay
      case Event(Done(token), Some(connection)) =>
        if (token == connection.token) {
          connection.client ! PoisonPill
          goto(Idle) using None
        } else
          stay
    }
  }

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