Use InMemoryPersistenceStorage for persisting pod records

.travis.yml

@@ -7,7 +7,7 @@ language: scala
 install:
   - true
 scala:
-   - 2.12.7
+  - 2.12.7
 script:
   - ./gradlew checkScalaFmtAll && sudo ./gradlew provision && ./gradlew ci --info
 before_cache:

core-models/src/main/resources/reference.conf

@@ -4,4 +4,11 @@ akka {
   loggers = ["akka.event.slf4j.Slf4jLogger"]
   loglevel = "INFO"
   logging-filter = "akka.event.slf4j.Slf4jLoggingFilter"
-}
+}
+
+scheduler {
+  persistence {
+    pipeline-limit = 128
+    load-timeout = 60 # seconds
+  }
+}

core-models/src/main/scala/com/mesosphere/usi/core/models/StateEvent.scala

@@ -12,14 +12,9 @@ sealed trait PodStateEvent extends StateEvent {
   def id: PodId
 }
 
-case class StateSnapshot(
-    podStatuses: Seq[PodStatus],
-    podRecords: Seq[PodRecord],
-    agentRecords: Seq[AgentRecord],
-    reservationStatuses: Seq[ReservationStatus])
-    extends StateEvent
+case class StateSnapshot(podRecords: Seq[PodRecord] = Nil, agentRecords: Seq[AgentRecord] = Nil) extends StateEvent
 object StateSnapshot {
-  def empty = StateSnapshot(Nil, Nil, Nil, Nil)
+  val empty = StateSnapshot()
 }
 
 /**

core/build.gradle

@@ -1,5 +1,6 @@
 dependencies {
     compile project(':core-models')
+    compile project(':persistence')
     compile project(':mesos-client')
     testCompile project(':test-utils')
 }

core/src/main/scala/com/mesosphere/usi/core/CallerThreadExecutionContext.scala

@@ -0,0 +1,12 @@
+package com.mesosphere.usi.core
+
+import java.util.concurrent.Executor
+import scala.concurrent.ExecutionContext
+import scala.concurrent.ExecutionContextExecutor
+
+// TODO : This is a duplicated in com.mesosphere.usi.async. Ideally, this needs to be in a shared module.
+private[usi] object CallerThreadExecutionContext {
+  val executor: Executor = (command: Runnable) => command.run()
+
+  implicit val context: ExecutionContextExecutor = ExecutionContext.fromExecutor(executor)
+}

core/src/main/scala/com/mesosphere/usi/core/FrameResultBuilder.scala

@@ -27,12 +27,14 @@ case class FrameResultBuilder(
     if (schedulerEvents == SchedulerEvents.empty)
       this
     else {
-      val newDirty = dirtyPodIds ++ schedulerEvents.stateEvents.iterator.collect {
-        case podEvent: PodStateEvent => podEvent.id
-        case _: StateSnapshot =>
-          // We need to handle status snapshots and create a mechanism to signal that all cache should be recomputed
-          ???
-      }
+      val newDirty: Set[PodId] = dirtyPodIds ++ schedulerEvents.stateEvents.iterator
+        .collect[Set[PodId]] {
+          case podEvent: PodStateEvent => Set(podEvent.id)
+          // We need to handle status snapshots with a mechanism to signal that all cache should be recomputed
+          case snapshot: StateSnapshot => snapshot.podRecords.map(_.podId).toSet
+        }
+        .flatten
+
       copy(
         state = state.applyStateIntents(schedulerEvents.stateEvents),
         dirtyPodIds = newDirty,

core/src/main/scala/com/mesosphere/usi/core/Scheduler.scala

@@ -4,11 +4,24 @@ import akka.{Done, NotUsed}
 import akka.stream.scaladsl.{BidiFlow, Broadcast, Flow, GraphDSL, Sink, SinkQueueWithCancel, Source}
 import akka.stream.{BidiShape, FlowShape, KillSwitches, Materializer, OverflowStrategy, QueueOfferResult}
 import com.mesosphere.mesos.client.{MesosCalls, MesosClient}
-import com.mesosphere.usi.core.models.{SpecEvent, SpecUpdated, SpecsSnapshot, StateEvent, StateSnapshot}
+import com.mesosphere.usi.core.conf.SchedulerSettings
+import com.mesosphere.usi.core.models.PodId
+import com.mesosphere.usi.core.models.PodRecord
+import com.mesosphere.usi.core.models.{
+  PodRecordUpdated,
+  SpecEvent,
+  SpecUpdated,
+  SpecsSnapshot,
+  StateEvent,
+  StateSnapshot
+}
+import com.mesosphere.usi.repository.PodRecordRepository
+import com.typesafe.config.ConfigFactory
 import org.apache.mesos.v1.Protos.FrameworkInfo
 import org.apache.mesos.v1.scheduler.Protos.{Call => MesosCall, Event => MesosEvent}
-
 import scala.collection.JavaConverters._
+import scala.concurrent.Await
+import scala.concurrent.duration._
 import scala.concurrent.{Future, Promise}
 import scala.util.{Failure, Success}
 
@@ -56,20 +69,18 @@ object Scheduler {
 
   type StateOutput = (StateSnapshot, Source[StateEvent, Any])
 
-  def fromSnapshot(specsSnapshot: SpecsSnapshot, client: MesosClient): Flow[SpecUpdated, StateOutput, NotUsed] =
-    Flow[SpecUpdated].prefixAndTail(0).map { case (_, rest) => specsSnapshot -> rest }.via(fromClient(client))
+  private val schedulerSettings = SchedulerSettings.fromConfig(ConfigFactory.load().getConfig("scheduler"))
 
-  def asFlow(specsSnapshot: SpecsSnapshot, client: MesosClient)(
+  def asFlow(specsSnapshot: SpecsSnapshot, client: MesosClient, podRecordRepository: PodRecordRepository)(
       implicit materializer: Materializer): Future[(StateSnapshot, Flow[SpecUpdated, StateEvent, NotUsed])] = {
 
-    implicit val ec = scala.concurrent.ExecutionContext.Implicits.global //only for ultra-fast non-blocking onComplete
+    implicit val ec = scala.concurrent.ExecutionContext.Implicits.global //only for ultra-fast non-blocking map
 
-    val (snap, source, sink) = asSourceAndSink(specsSnapshot, client)
+    val (snap, source, sink) = asSourceAndSink(specsSnapshot, client, podRecordRepository)
 
     snap.map { snapshot =>
       (snapshot, Flow.fromSinkAndSourceCoupled(sink, source))
     }
-
   }
 
   /**
@@ -80,9 +91,9 @@ object Scheduler {
     * @param specsSnapshot Snapshot of the current specs
     * @return Snapshot of the current state, as well as Source which produces StateEvents and Sink which accepts SpecEvents
     */
-  def asSourceAndSink(specsSnapshot: SpecsSnapshot, client: MesosClient)(
+  def asSourceAndSink(specsSnapshot: SpecsSnapshot, client: MesosClient, podRecordRepository: PodRecordRepository)(
       implicit mat: Materializer): (Future[StateSnapshot], Source[StateEvent, NotUsed], Sink[SpecUpdated, NotUsed]) = {
-    val flow = fromClient(client)
+    val flow = fromClient(client, podRecordRepository)
     asSourceAndSink(specsSnapshot, flow)(mat)
   }
 
@@ -190,30 +201,32 @@ object Scheduler {
 
   }
 
-  def fromClient(client: MesosClient): Flow[SpecInput, StateOutput, NotUsed] = {
+  def fromClient(
+      client: MesosClient,
+      podRecordRepository: PodRecordRepository): Flow[SpecInput, StateOutput, NotUsed] = {
     if (!isMultiRoleFramework(client.frameworkInfo)) {
       throw new IllegalArgumentException(
-        "USI scheduler provides support for MULTI_ROLE frameworks only. Please provide create MesosClient with FrameworkInfo that has capability MULTI_ROLE")
+        "USI scheduler provides support for MULTI_ROLE frameworks only. " +
+          "Please provide a MesosClient with FrameworkInfo that has capability MULTI_ROLE")
     }
-    fromFlow(client.calls, Flow.fromSinkAndSource(client.mesosSink, client.mesosSource))
+    fromFlow(client.calls, podRecordRepository, Flow.fromSinkAndSource(client.mesosSink, client.mesosSource))
   }
 
-  private def isMultiRoleFramework(frameworkInfo: FrameworkInfo): Boolean =
-    frameworkInfo.getCapabilitiesList.asScala.exists(_.getType == FrameworkInfo.Capability.Type.MULTI_ROLE)
-
   def fromFlow(
       mesosCallFactory: MesosCalls,
+      podRecordRepository: PodRecordRepository,
       mesosFlow: Flow[MesosCall, MesosEvent, Any]): Flow[SpecInput, StateOutput, NotUsed] = {
     Flow.fromGraph {
-      GraphDSL.create(unconnectedGraph(mesosCallFactory), mesosFlow)((_, _) => NotUsed) { implicit builder =>
-        { (graph, mesos) =>
-          import GraphDSL.Implicits._
+      GraphDSL.create(unconnectedGraph(mesosCallFactory, podRecordRepository), mesosFlow)((_, _) => NotUsed) {
+        implicit builder =>
+          { (graph, mesos) =>
+            import GraphDSL.Implicits._
 
-          mesos ~> graph.in2
-          graph.out2 ~> mesos
+            mesos ~> graph.in2
+            graph.out2 ~> mesos
 
-          FlowShape(graph.in1, graph.out1)
-        }
+            FlowShape(graph.in1, graph.out1)
+          }
       }
     }
   }
@@ -230,30 +243,36 @@ object Scheduler {
       rest.prepend(Source.single(snapshot))
   }
 
-  // TODO (DCOS-47476) use actual prefixAndTail and expect first event to be a Snapshot; change the prefixAndTail param from 0 value to 1, let fail, etc.
-  private val stateOutputBreakoutFlow: Flow[StateEvent, StateOutput, NotUsed] = Flow[StateEvent].prefixAndTail(0).map {
-    case (_, stateEvents) =>
+  private val stateOutputBreakoutFlow: Flow[StateEvent, StateOutput, NotUsed] = Flow[StateEvent].prefixAndTail(1).map {
+    case (Seq(snapshot), stateEvents) =>
+      val stateSnapshot = snapshot match {
+        case x: StateSnapshot => x
+        case _ => throw new IllegalStateException("First event is allowed to be only a state snapshot")
+      }
       val stateUpdates = stateEvents.map {
         case _: StateSnapshot =>
           throw new IllegalStateException("Only the first event is allowed to be a state snapshot")
         case event => event
       }
-      (StateSnapshot.empty, stateUpdates)
+      (stateSnapshot, stateUpdates)
   }
 
-  def unconnectedGraph(
-      mesosCallFactory: MesosCalls): BidiFlow[SpecInput, StateOutput, MesosEvent, MesosCall, NotUsed] = {
+  private[core] def unconnectedGraph(
+      mesosCallFactory: MesosCalls,
+      podRecordRepository: PodRecordRepository): BidiFlow[SpecInput, StateOutput, MesosEvent, MesosCall, NotUsed] = {
+    val schedulerLogicGraph = new SchedulerLogicGraph(mesosCallFactory, loadPodRecords(podRecordRepository))
     BidiFlow.fromGraph {
-      GraphDSL.create(new SchedulerLogicGraph(mesosCallFactory)) { implicit builder => (schedulerLogic) =>
+      GraphDSL.create(schedulerLogicGraph) { implicit builder => (schedulerLogic) =>
         {
           import GraphDSL.Implicits._
 
           val broadcast = builder.add(Broadcast[SchedulerEvents](2, eagerCancel = true))
           val specInputFlattening = builder.add(specInputFlatteningFlow)
           val stateOutputBreakout = builder.add(stateOutputBreakoutFlow)
 
+          val persistenceStorageFlow = builder.add(persistenceFlow(podRecordRepository))
           specInputFlattening ~> schedulerLogic.in0
-          schedulerLogic.out ~> broadcast.in
+          schedulerLogic.out ~> persistenceStorageFlow ~> broadcast.in
 
           val mesosCalls = broadcast.out(0).mapConcat { frameResult =>
             frameResult.mesosCalls
@@ -269,4 +288,40 @@ object Scheduler {
       }
     }
   }
+
+  private[core] def persistenceFlow(
+      podRecordRepository: PodRecordRepository): Flow[SchedulerEvents, SchedulerEvents, NotUsed] = {
+    Flow[SchedulerEvents]
+      .mapConcat(persistEvents(_, podRecordRepository))
+      .mapAsync(schedulerSettings.persistencePipelineLimit)(call => call())
+      .collect { case Some(events) => events }
+  }
+
+  private def persistEvents(
+      events: SchedulerEvents,
+      podRecordRepository: PodRecordRepository): List[() => Future[Option[SchedulerEvents]]] = {
+    val ops: List[() => Future[Option[SchedulerEvents]]] = events.stateEvents.collect {
+      case PodRecordUpdated(_, Some(podRecord)) =>
+        () =>
+          podRecordRepository.store(podRecord).map(_ => None)(CallerThreadExecutionContext.context)
+      case PodRecordUpdated(podId, None) =>
+        () =>
+          podRecordRepository.delete(podId).map(_ => None)(CallerThreadExecutionContext.context)
+    }
+    ops :+ (() => Future.successful(Some(events)))
+  }
+
+  /*
+   * We don't start processing any commands until we've finished loading the entire set of podRecords
+   * This code delays building a scheduler stage until this podRecord snapshot is available.
+   *
+   * Block for IO - If the IO call fails or a timeout occurs, we should not make any progress.
+   */
+  private def loadPodRecords(podRecordRepository: PodRecordRepository): Map[PodId, PodRecord] = {
+    // Add error handling (and maybe a retry mechanism).
+    Await.result(podRecordRepository.readAll(), schedulerSettings.persistenceLoadTimeout.seconds)
+  }
+
+  private def isMultiRoleFramework(frameworkInfo: FrameworkInfo): Boolean =
+    frameworkInfo.getCapabilitiesList.asScala.exists(_.getType == FrameworkInfo.Capability.Type.MULTI_ROLE)
 }

core/src/main/scala/com/mesosphere/usi/core/SchedulerLogicGraph.scala

@@ -3,9 +3,8 @@ package com.mesosphere.usi.core
 import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
 import akka.stream.{Attributes, FanInShape2, Inlet, Outlet}
 import com.mesosphere.mesos.client.MesosCalls
-import com.mesosphere.usi.core.models.SpecEvent
+import com.mesosphere.usi.core.models.{PodId, PodRecord, SpecEvent, StateSnapshot}
 import org.apache.mesos.v1.scheduler.Protos.{Event => MesosEvent}
-
 import scala.collection.mutable
 
 object SchedulerLogicGraph {
@@ -41,22 +40,24 @@ object SchedulerLogicGraph {
   * It's existence is only warranted by forecasted future needs. It's kept as a graph with an internal buffer as we will
   * likely need timers, other callbacks, and additional output ports (such as an offer event stream?).
   */
-class SchedulerLogicGraph(mesosCallFactory: MesosCalls)
+private[core] class SchedulerLogicGraph(mesosCallFactory: MesosCalls, initialPodRecords: Map[PodId, PodRecord])
     extends GraphStage[FanInShape2[SpecEvent, MesosEvent, SchedulerEvents]] {
   import SchedulerLogicGraph.BUFFER_SIZE
 
-  val mesosEventsInlet = Inlet[MesosEvent]("mesos-events")
-  val specEventsInlet = Inlet[SpecEvent]("specs")
-  val frameResultOutlet = Outlet[SchedulerEvents]("effects")
+  private val mesosEventsInlet = Inlet[MesosEvent]("mesos-events")
+  private val specEventsInlet = Inlet[SpecEvent]("specs")
+  private val frameResultOutlet = Outlet[SchedulerEvents]("effects")
+
   // Define the shape of this stage, which is SourceShape with the port we defined above
   override val shape: FanInShape2[SpecEvent, MesosEvent, SchedulerEvents] =
     new FanInShape2(specEventsInlet, mesosEventsInlet, frameResultOutlet)
 
   // This is where the actual (possibly stateful) logic will live
   override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = {
-    val handler = new SchedulerLogicHandler(mesosCallFactory)
 
     new GraphStageLogic(shape) {
+      private[this] var handler: SchedulerLogicHandler = _
+
       val pendingEffects: mutable.Queue[SchedulerEvents] = mutable.Queue.empty
 
       setHandler(mesosEventsInlet, new InHandler {
@@ -83,9 +84,9 @@ class SchedulerLogicGraph(mesosCallFactory: MesosCalls)
       })
 
       override def preStart(): Unit = {
-        // Start the stream
-        pull(specEventsInlet)
-        pull(mesosEventsInlet)
+        handler = new SchedulerLogicHandler(mesosCallFactory, initialPodRecords)
+        // Publish the initial state snapshot event; podStatuses will not be in the snapshot in the future
+        pushOrQueueIntents(SchedulerEvents(List(StateSnapshot(podRecords = initialPodRecords.values.toSeq))))
       }
 
       def pushOrQueueIntents(effects: SchedulerEvents): Unit = {

core/src/main/scala/com/mesosphere/usi/core/SchedulerLogicHandler.scala

@@ -2,6 +2,7 @@ package com.mesosphere.usi.core
 
 import com.mesosphere.mesos.client.MesosCalls
 import com.mesosphere.usi.core.logic.{MesosEventsLogic, SpecLogic}
+import com.mesosphere.usi.core.models.PodRecord
 import com.mesosphere.usi.core.models.{PodId, PodInvalid, PodSpec, PodSpecUpdated, SpecEvent, SpecsSnapshot}
 import org.apache.mesos.v1.scheduler.Protos.{Call, Event => MesosEvent}
 
@@ -62,7 +63,7 @@ import org.apache.mesos.v1.scheduler.Protos.{Call, Event => MesosEvent}
   * through the application of several functions. The result of that event will lead to the emission of several
   * stateEvents and Mesos intents, which will be accumulated and emitted via a data structure we call the FrameResult.
   */
-private[core] class SchedulerLogicHandler(mesosCallFactory: MesosCalls) {
+private[core] class SchedulerLogicHandler(mesosCallFactory: MesosCalls, initialPodRecords: Map[PodId, PodRecord]) {
 
   private val schedulerLogic = new SpecLogic(mesosCallFactory)
   private val mesosEventsLogic = new MesosEventsLogic(mesosCallFactory)
@@ -77,7 +78,7 @@ private[core] class SchedulerLogicHandler(mesosCallFactory: MesosCalls) {
     * contain persistent, non-recoverable facts from Mesos, such as pod-launched time, agentId on which a pod was
     * launched, agent information or the time at which a pod was first seen as unhealthy or unreachable.
     */
-  private var state: SchedulerState = SchedulerState.empty
+  private var state: SchedulerState = SchedulerState(podStatuses = Map.empty, podRecords = initialPodRecords)
 
   /**
     * Cached view of SchedulerState and SpecificationState. We incrementally update this computation at the end of each
@@ -110,6 +111,20 @@ private[core] class SchedulerLogicHandler(mesosCallFactory: MesosCalls) {
     }
   }
 
+  /**
+    * Process a Mesos event and update internal state.
+    *
+    * @param event
+    * @return The events describing state changes as Mesos call intents
+    */
+  def handleMesosEvent(event: MesosEvent): SchedulerEvents = {
+    handleFrame { builder =>
+      builder.process { (specs, state, _) =>
+        mesosEventsLogic.processEvent(specs, state, cachedPendingLaunch.pendingLaunch)(event)
+      }
+    }
+  }
+
   /**
     * Instantiate a frameResultBuilder instance, call the handler, then follow up with housekeeping:
     *
@@ -171,20 +186,6 @@ private[core] class SchedulerLogicHandler(mesosCallFactory: MesosCalls) {
     SchedulerEvents(mesosCalls = reviveCalls ++ suppressCalls)
   }
 
-  /**
-    * Process a Mesos event and update internal state.
-    *
-    * @param event
-    * @return The events describing state changes as Mesos call intents
-    */
-  def handleMesosEvent(event: MesosEvent): SchedulerEvents = {
-    handleFrame { builder =>
-      builder.process { (specs, state, _) =>
-        mesosEventsLogic.processEvent(specs, state, cachedPendingLaunch.pendingLaunch)(event)
-      }
-    }
-  }
-
   /**
     * We remove a task if it is not reachable and running, and it has no podSpec defined
     *