Synchronize instrumentation with the actor

engine/src/main/scala/cromwell/engine/workflow/tokens/JobExecutionTokenDispenserActor.scala

@@ -26,18 +26,19 @@ class JobExecutionTokenDispenserActor(override val serviceRegistryActor: ActorRe
   var currentTokenQueuePointer: Int = 0
   var tokenAssignments: Map[ActorRef, Lease[JobExecutionToken]] = Map.empty
 
-  scheduleInstrumentation {
+  val instrumentationAction = () => {
     sendGaugeJob(ExecutionStatus.Running.toString, tokenAssignments.size.toLong)
     sendGaugeJob(ExecutionStatus.QueuedInCromwell.toString, tokenQueues.values.map(_.size).sum.toLong)
   }
 
   override def preStart() = {
     ratePreStart()
     serviceRegistryActor ! ListenToLoadController
+    startInstrumentationTimer()
     super.preStart()
   }
 
-  override def receive: Actor.Receive = tokenDistributionReceive.orElse(rateReceive)
+  override def receive: Actor.Receive = tokenDistributionReceive.orElse(rateReceive).orElse(instrumentationReceive(instrumentationAction))
 
   private def tokenDistributionReceive: Receive = {
     case JobExecutionTokenRequest(hogGroup, tokenType) => enqueue(sender, hogGroup, tokenType)

engine/src/main/scala/cromwell/engine/workflow/workflowstore/WorkflowStoreEngineActor.scala

@@ -1,6 +1,6 @@
 package cromwell.engine.workflow.workflowstore
 
-import akka.actor.{ActorLogging, ActorRef, LoggingFSM, PoisonPill, Props}
+import akka.actor.{ActorLogging, ActorRef, LoggingFSM, PoisonPill, Props, Timers}
 import akka.pattern.ask
 import akka.util.Timeout
 import cats.data.NonEmptyList
@@ -26,14 +26,14 @@ final case class WorkflowStoreEngineActor private(store: WorkflowStore,
                                                   serviceRegistryActor: ActorRef,
                                                   abortAllJobsOnTerminate: Boolean,
                                                   workflowHeartbeatConfig: WorkflowHeartbeatConfig)
-  extends LoggingFSM[WorkflowStoreActorState, WorkflowStoreActorData] with ActorLogging with WorkflowInstrumentation with CromwellInstrumentationScheduler with WorkflowMetadataHelper {
+  extends LoggingFSM[WorkflowStoreActorState, WorkflowStoreActorData] with ActorLogging with WorkflowInstrumentation with CromwellInstrumentationScheduler with WorkflowMetadataHelper with Timers {
 
   implicit val ec: ExecutionContext = context.dispatcher
 
   startWith(Unstarted, WorkflowStoreActorData(None, List.empty))
   self ! InitializerCommand
 
-  scheduleInstrumentation {
+  val instrumentationAction = () => {
     store.stats map { stats: Map[WorkflowStoreState, Int] =>
       // Update the count for Submitted and Running workflows, defaulting to 0
       val statesMap = stats.withDefault(_ => 0)
@@ -45,6 +45,13 @@ final case class WorkflowStoreEngineActor private(store: WorkflowStore,
     ()
   }
 
+  override def preStart() = {
+    startInstrumentationTimer()
+    super.preStart()
+  }
+
+  override def receive = instrumentationReceive(instrumentationAction).orElse(super.receive)
+
   when(Unstarted) {
     case Event(InitializerCommand, _) =>
       val work = store.initialize map { _ =>

services/src/main/scala/cromwell/services/instrumentation/CromwellInstrumentation.scala

@@ -1,6 +1,7 @@
 package cromwell.services.instrumentation
 
-import akka.actor.{Actor, ActorRef, Cancellable}
+import akka.actor.{Actor, ActorRef, Timers}
+import akka.dispatch.ControlMessage
 import cats.data.NonEmptyList
 import com.typesafe.config.ConfigFactory
 import cromwell.services.instrumentation.CromwellInstrumentation._
@@ -110,17 +111,17 @@ trait CromwellInstrumentation {
 /**
   * Helper trait to provide a scheduler function that can be used for instrumentation purposes
   */
-trait CromwellInstrumentationScheduler { this: Actor =>
-  // The system dispatcher should be fine ? Do we need a special one for instrumentation ?
-  implicit private val instrumentationEc = context.system.dispatcher
+trait CromwellInstrumentationScheduler { this: Actor with Timers =>
+  private case object InstrumentationTimerKey
+  private case object InstrumentationTimerAction extends ControlMessage
 
-  private var scheduledInstrumentationTimers: Vector[Cancellable] = Vector.empty
-
-  def scheduleInstrumentation(f: => Unit) = {
-    scheduledInstrumentationTimers = scheduledInstrumentationTimers :+ context.system.scheduler.schedule(InstrumentationRate, InstrumentationRate)(f)
+  def startInstrumentationTimer() = {
+    timers.startSingleTimer(InstrumentationTimerKey, InstrumentationTimerAction, InstrumentationRate)
   }
 
-  override def postStop(): Unit = {
-    scheduledInstrumentationTimers foreach { _.cancel() }
+  protected def instrumentationReceive(instrumentationAction: () => Unit): Receive = {
+    case InstrumentationTimerAction => 
+      instrumentationAction()
+      timers.startSingleTimer(InstrumentationTimerKey, InstrumentationTimerAction, InstrumentationRate)
   }
 }

supportedBackends/google/pipelines/common/src/main/scala/cromwell/backend/google/pipelines/common/api/PipelinesApiRequestManager.scala

@@ -10,14 +10,14 @@ import com.google.api.client.googleapis.json.GoogleJsonError
 import com.google.api.client.http.{HttpHeaders, HttpRequest}
 import common.util.Backoff
 import cromwell.backend.BackendSingletonActorAbortWorkflow
+import cromwell.backend.google.pipelines.common.PapiInstrumentation
 import cromwell.backend.google.pipelines.common.api.PipelinesApiRequestManager._
 import cromwell.backend.google.pipelines.common.api.clients.PipelinesApiRunCreationClient.JobAbortedException
-import cromwell.backend.google.pipelines.common.PapiInstrumentation
 import cromwell.backend.standard.StandardAsyncJob
 import cromwell.core.Dispatcher.BackendDispatcher
 import cromwell.core.retry.SimpleExponentialBackoff
 import cromwell.core.{CromwellFatalExceptionMarker, LoadConfig, Mailbox, WorkflowId}
-import cromwell.services.instrumentation.{CromwellInstrumentation, CromwellInstrumentationScheduler}
+import cromwell.services.instrumentation.CromwellInstrumentationScheduler
 import cromwell.services.loadcontroller.LoadControllerService.{HighLoad, LoadMetric, NormalLoad}
 import cromwell.util.StopAndLogSupervisor
 import eu.timepit.refined.api.Refined
@@ -74,7 +74,6 @@ class PipelinesApiRequestManager(val qps: Int Refined Positive, requestWorkers:
   // 
   private lazy val workerBatchInterval = determineBatchInterval(qps) * nbWorkers.toLong
 
-  scheduleInstrumentation { updateQueueSize(workQueue.size) }
 
   // workQueue is protected for the unit tests, not intended to be generally overridden
   protected[api] var workQueue: Queue[PAPIApiRequest] = Queue.empty
@@ -90,18 +89,17 @@ class PipelinesApiRequestManager(val qps: Int Refined Positive, requestWorkers:
 
   override def preStart() = {
     log.info("{} Running with {} workers", self.path.name, requestWorkers.value)
-    timers.startSingleTimer(QueueMonitoringTimerKey, QueueMonitoringTimerAction, CromwellInstrumentation.InstrumentationRate)
+    startInstrumentationTimer()
     super.preStart()
   }
 
   def monitorQueueSize() = {
     val load = if (workQueue.size > LoadConfig.PAPIThreshold) HighLoad else NormalLoad
     serviceRegistryActor ! LoadMetric("PAPIQueryManager", load)
-    timers.startSingleTimer(QueueMonitoringTimerKey, QueueMonitoringTimerAction, CromwellInstrumentation.InstrumentationRate)
+    updateQueueSize(workQueue.size)
   }
 
-  override def receive = {
-    case QueueMonitoringTimerAction => monitorQueueSize()
+  val requestManagerReceive: Receive = {
     case BackendSingletonActorAbortWorkflow(id) => abort(id)
     case status: PAPIStatusPollRequest => workQueue :+= status
     case create: PAPIRunCreationRequest =>
@@ -117,6 +115,8 @@ class PipelinesApiRequestManager(val qps: Int Refined Positive, requestWorkers:
     case other => log.error(s"Unexpected message to JesPollingManager: $other")
   }
 
+  override def receive = instrumentationReceive(monitorQueueSize _).orElse(requestManagerReceive)
+
   private def abort(workflowId: WorkflowId) = {
     def aborted(query: PAPIRunCreationRequest) = query.requester ! PipelinesApiRunCreationQueryFailed(query, JobAbortedException)