Combine stream completion

bitrockteam · Jun 1, 2020 · 3f4a9f0 · 3f4a9f0
1 parent 51731d9
commit 3f4a9f0
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 7 deletions.
diff --git a/src/main/scala/it/bitrock/dvs/producer/aviationedge/Main.scala b/src/main/scala/it/bitrock/dvs/producer/aviationedge/Main.scala
@@ -20,11 +20,11 @@ object Main extends App with LazyLogging {
 
   bindingFuture.map(serverBinding => logger.info(s"Exposing to ${serverBinding.localAddress}"))
 
-  val (cancellableFlight, flightCompletion, _, _)     = runAviationEdgeStream[FlightStream.type]()
-  val (cancellableAirplane, airplaneCompletion, _, _) = runAviationEdgeStream[AirplaneStream.type]()
-  val (cancellableAirport, airportCompletion, _, _)   = runAviationEdgeStream[AirportStream.type]()
-  val (cancellableAirline, airlineCompletion, _, _)   = runAviationEdgeStream[AirlineStream.type]()
-  val (cancellableCity, cityCompletion, _, _)         = runAviationEdgeStream[CityStream.type]()
+  val (cancellableFlight, flightCompletion)           = runAviationEdgeStream[FlightStream.type]()
+  val (cancellableAirplane, airplaneCompletion)       = runAviationEdgeStream[AirplaneStream.type]()
+  val (cancellableAirport, airportCompletion)         = runAviationEdgeStream[AirportStream.type]()
+  val (cancellableAirline, airlineCompletion)         = runAviationEdgeStream[AirlineStream.type]()
+  val (cancellableCity, cityCompletion)               = runAviationEdgeStream[CityStream.type]()
   val (cancellableFlightState, flightStateCompletion) = runOpenSkyStream[FlightStateStream.type]()
 
   val streamsCompletion =

diff --git a/src/main/scala/it/bitrock/dvs/producer/aviationedge/services/MainFunctions.scala b/src/main/scala/it/bitrock/dvs/producer/aviationedge/services/MainFunctions.scala
@@ -39,7 +39,7 @@ object MainFunctions {
   def runAviationEdgeStream[A: ApiProviderStreamContext]()(
       implicit system: ActorSystem,
       ec: ExecutionContext
-  ): (Cancellable, Future[Done], Future[Done], Future[Done]) = {
+  ): (Cancellable, Future[Done]) = {
     val config = AviationStreamContext[A].config(apiProviderConfig)
 
     val tickSource = new TickSource(config.pollingStart, config.pollingInterval, aviationConfig.tickSource).source
@@ -54,7 +54,7 @@ object MainFunctions {
 
     val jsonSource = tickSource.via(aviationFlow).via(monitoringGraph(monitoringSink)).mapConcat(identity)
 
-    mainGraph(jsonSource, rawSink, errorSink, invalidFlightSink).run()
+    mainGraph(jsonSource, rawSink, errorSink, invalidFlightSink).mapMaterializedValue(combineCompletionValue).run()
   }
 
   def runOpenSkyStream[A: ApiProviderStreamContext]()(
@@ -73,4 +73,9 @@ object MainFunctions {
 
     collectRightMessagesGraph(jsonSource, rawSink).run()
   }
+
+  private def combineCompletionValue[T](value: (Cancellable, Future[T], Future[T], Future[T]))(implicit ec: ExecutionContext) =
+    value match {
+      case (cancel, one, two, three) => (cancel, Future.firstCompletedOf(List(one, two, three)))
+    }
 }