Better way to setup graceful shutdown with Akka Streams - scala

Question

I have pretty simple application that has an Akka HTTP endpoint, does some processing and writes results to either of the output files. The code to ensure graceful shutdown looks a bit complicated, is there a way to make it more succinct?
val bindingFuture = Http().newServerAt("localhost", config.port).bind(route)
val validQueue: BoundedSourceQueue[ByteString] = ???
val invalidQueue: BoundedSourceQueue[ByteString] = ???
val validDone: Future[Done] = ???
val invalidDone: Future[Done] = ???
val allDone = Future.sequence(validDone, invalidDone)
bindingFuture.onComplete {
case Success(binding) =>
logger.info("Server started on port {}", config.port)
binding.addToCoordinatedShutdown(5.seconds)
case Failure(ex) =>
logger.error("Can't start server", ex)
system.terminate()
}
allDone.onComplete { result =>
result match {
case Failure(ex) =>
logger.error("Streams completed with error", ex)
case Success(_) =>
logger.info("Streams completed successfully")
}
system.terminate()
}
sys.addShutdownHook {
logger.info("Shutting down...")
validQueue.complete()
invalidQueue.complete()
}

Answer 1

Akka installs JVM shutdown hooks by default, there is no need to add a shutdown hook yourself, you can remove sys.addShutdownHook { ... }
Calling ActorSystem.terminate will also terminate all streams. (Streams can be terminated abruptly, in most applications this is not a problem though, and it sounds like this would also not be a problem in your case.)
Minor cleanup, you could consider using map, recoverWith and andThen.
allDone.map { _ =>
logger.info("Streams completed successfully")
}.recoverWith {
case ex =>
logger.error("Streams completed with error", ex)
}.andThen {
case _ => system.terminate()
}
You could use CoordinatedShutdown:
CoordinatedShutdown(context.system).addTask(CoordinatedShutdown.PhaseServiceRequestsDone, "complete hdfs sinks") { () =>
validQueue.complete()
invalidQueue.complete()
}
You could also use a shared kill switch (https://doc.akka.io/docs/akka/current/stream/stream-dynamic.html#sharedkillswitch), which you can put in your flows with .via(sharedKillSwitch.flow), you could shutdown the switch from the CoordinatedShutdown:
// create it somewhere, use in your flows
val sharedKillSwitch = KillSwitches.shared("hdfs-switch")
// use switch in CoordinatedShutdown
CoordinatedShutdown(context.system).addTask(CoordinatedShutdown.PhaseServiceRequestsDone, "complete hdfs sinks") { () =>
sharedKillSwitch.shutdown()
}