<commit_before>package filodb.coordinator

import scala.collection.mutable.{HashMap => mHashMap, Queue => mQueue}
import scala.concurrent.Future
import scala.util.control.NonFatal

import akka.actor.{ActorRef, ActorSystem, Props}
import akka.dispatch.{Envelope, UnboundedStablePriorityMailbox}
import com.typesafe.config.Config
import kamon.Kamon
import kamon.tag.TagSet
import monix.execution.Scheduler
import net.ceedubs.ficus.Ficus._
import net.ceedubs.ficus.readers.ValueReader

import filodb.coordinator.queryplanner.SingleClusterPlanner
import filodb.core._
import filodb.core.memstore.{FiloSchedulers, MemStore, TermInfo}
import filodb.core.metadata.Schemas
import filodb.core.query.{QueryConfig, QueryContext, QuerySession}
import filodb.core.store.CorruptVectorException
import filodb.query._
import filodb.query.exec.ExecPlan


// A class to manage a queue and a set of currently executing queries, allowing no more than concurrentQueries
// at a time.  When the query and execplans finishes execution, then they are removed.  This has multiple benefits
// including limiting memory usage and tracking running queries.
// concurrentQueries defines the number of current queries where all execPlans belonging to the same queryID
// is counted as a single query.
// A queue keeps track of queries for which there is no space.  It is not expected for there to be more than
// one query in the queue belonging to the same queryID, this is because only high-level ExecPlans spawn more
// low-level ones, and spawned ones arrive in the Actor mailbox and will go striaght to executing if the parent
// plan is already executing.
// A single threaded scheduler is created for each query (all ExecPlans with the same query ID).  This reduces
// thread contention significantly and reduces task switching overhead, boosting query throughput especially
// when the tasks are small.
class QueryScheduler(concurrentQueries: Int, maxQueueLen: Int, tags: Map[String, String]) {
  //  Currently executing query IDs -> number of execPlans in that query ID.  Used to track when all plans finish
  //  so that cleanup can be done and more queries can be popped off the stack.
  val currentIDs = new mHashMap[String, Int]().withDefaultValue(0)
  //  Queue of queries waiting to be added to currently executing ones
  val queue = new mQueue[(ExecPlan, ActorRef)]
  //  Mapping of queryID to scheduler.  Each queryID gets its own scheduler.
  val idToSched = new mHashMap[String, Scheduler]()
  //  Queue of schedulers available for future queries
  val schedulers = new mQueue[Scheduler]()

  // Immediate ExecPlan execution, direct on current pool, no Queue, vs from queue
  private val countExecImmed = Kamon.counter("query-scheduler-execute-immediate").withTags(TagSet.from(tags))
  private val countExecQueue = Kamon.counter("query-scheduler-execute-queue").withTags(TagSet.from(tags))
  private val countAddQueue  = Kamon.counter("query-scheduler-add-queue").withTags(TagSet.from(tags))
  private val countQueueFull = Kamon.counter("query-scheduler-reject-queue-full").withTags(TagSet.from(tags))
  private val queueLen = Kamon.gauge("query-scheduler-queue-length").withTags(TagSet.from(tags))

  // Given a plan and replyTo, executes the plan if there is room or it is a currently executing query.
  // If there isn't room, add it to the queue.  If the queue has no room, return false; otherwise true.
  def execOrWait(plan: ExecPlan,
                 replyTo: ActorRef,
                 newPlanFunc: (ExecPlan, ActorRef, Scheduler) => Future[Unit]): Boolean = synchronized {
    val id = plan.queryContext.queryId
    if ((currentIDs contains id) || currentIDs.size < concurrentQueries) {
      countExecImmed.increment()
      execute(plan, replyTo, newPlanFunc)
      true
    } else if (queueSize < maxQueueLen) {
      countAddQueue.increment()
      queue.enqueue((plan, replyTo))
      true
    } else {
      countQueueFull.increment()
      false
    }
  }

  // Length of queue - not synchronized, so data might not be consistent.
  def queueSize: Int = queue.length

  // Pop next item(s?) off top of queue so long as there is room to execute another query.
  def tryPop(newPlanFunc: (ExecPlan, ActorRef, Scheduler) => Future[Unit]): Unit = synchronized {
    if (currentIDs.size < concurrentQueries && queue.nonEmpty) {
      val (nextPlan, replyTo) = queue.dequeue
      countExecQueue.increment()
      execute(nextPlan, replyTo, newPlanFunc)
    }
  }

  // Executes newPlanFunc, updating currentIDs/plans as needed
  // The closure is called in the current thread, but the callback when Future/plan is completed will not be.
  // Note that an appropriate single threaded scheduler is pulled off the stack and used.
  private def execute(newPlan: ExecPlan,
                      replyTo: ActorRef,
                      newPlanFunc: (ExecPlan, ActorRef, Scheduler) => Future[Unit]): Unit = {
    val queryId = newPlan.queryContext.queryId

    // Add plan to data structures and obtain scheduler
    val sched = synchronized {
      currentIDs(queryId) += 1
      idToSched.getOrElseUpdate(queryId, {
        if (schedulers.nonEmpty) schedulers.dequeue
        else Scheduler.singleThread(s"Query-Queue-${currentIDs.size}")
      })
    }

    val planFut = newPlanFunc(newPlan, replyTo, sched)

    // Remove plan on Future complete.  If a slot opens up, recursively call myself to open up a new slot
    planFut.onComplete {
      case _ => synchronized {
                  currentIDs(queryId) -= 1
                  if (currentIDs(queryId) == 0) {
                    currentIDs.remove(queryId)
                    schedulers.enqueue(idToSched.remove(queryId).get)
                    queueLen.update(queueSize)
                    tryPop(newPlanFunc)
                  }
                }
    }(sched)
  }

  // Currently active query IDs.  Not synchronized, so data might not be consistent.
  def currentQueryIDs: Set[String] = currentIDs.keySet.toSet
}

object QueryCommandPriority extends java.util.Comparator[Envelope] {
  override def compare(o1: Envelope, o2: Envelope): Int = {
    (o1.message, o2.message) match {<commit_msg> fix(coordinator): Query threads should start with "query-sched" for compatibility. (#841)<commit_after>package filodb.coordinator

import scala.collection.mutable.{HashMap => mHashMap, Queue => mQueue}
import scala.concurrent.Future
import scala.util.control.NonFatal

import akka.actor.{ActorRef, ActorSystem, Props}
import akka.dispatch.{Envelope, UnboundedStablePriorityMailbox}
import com.typesafe.config.Config
import kamon.Kamon
import kamon.tag.TagSet
import monix.execution.Scheduler
import net.ceedubs.ficus.Ficus._
import net.ceedubs.ficus.readers.ValueReader

import filodb.coordinator.queryplanner.SingleClusterPlanner
import filodb.core._
import filodb.core.memstore.{FiloSchedulers, MemStore, TermInfo}
import filodb.core.metadata.Schemas
import filodb.core.query.{QueryConfig, QueryContext, QuerySession}
import filodb.core.store.CorruptVectorException
import filodb.query._
import filodb.query.exec.ExecPlan


// A class to manage a queue and a set of currently executing queries, allowing no more than concurrentQueries
// at a time.  When the query and execplans finishes execution, then they are removed.  This has multiple benefits
// including limiting memory usage and tracking running queries.
// concurrentQueries defines the number of current queries where all execPlans belonging to the same queryID
// is counted as a single query.
// A queue keeps track of queries for which there is no space.  It is not expected for there to be more than
// one query in the queue belonging to the same queryID, this is because only high-level ExecPlans spawn more
// low-level ones, and spawned ones arrive in the Actor mailbox and will go striaght to executing if the parent
// plan is already executing.
// A single threaded scheduler is created for each query (all ExecPlans with the same query ID).  This reduces
// thread contention significantly and reduces task switching overhead, boosting query throughput especially
// when the tasks are small.
class QueryScheduler(concurrentQueries: Int, maxQueueLen: Int, tags: Map[String, String]) {
  //  Currently executing query IDs -> number of execPlans in that query ID.  Used to track when all plans finish
  //  so that cleanup can be done and more queries can be popped off the stack.
  val currentIDs = new mHashMap[String, Int]().withDefaultValue(0)
  //  Queue of queries waiting to be added to currently executing ones
  val queue = new mQueue[(ExecPlan, ActorRef)]
  //  Mapping of queryID to scheduler.  Each queryID gets its own scheduler.
  val idToSched = new mHashMap[String, Scheduler]()
  //  Queue of schedulers available for future queries
  val schedulers = new mQueue[Scheduler]()

  // Immediate ExecPlan execution, direct on current pool, no Queue, vs from queue
  private val countExecImmed = Kamon.counter("query-scheduler-execute-immediate").withTags(TagSet.from(tags))
  private val countExecQueue = Kamon.counter("query-scheduler-execute-queue").withTags(TagSet.from(tags))
  private val countAddQueue  = Kamon.counter("query-scheduler-add-queue").withTags(TagSet.from(tags))
  private val countQueueFull = Kamon.counter("query-scheduler-reject-queue-full").withTags(TagSet.from(tags))
  private val queueLen = Kamon.gauge("query-scheduler-queue-length").withTags(TagSet.from(tags))

  // Given a plan and replyTo, executes the plan if there is room or it is a currently executing query.
  // If there isn't room, add it to the queue.  If the queue has no room, return false; otherwise true.
  def execOrWait(plan: ExecPlan,
                 replyTo: ActorRef,
                 newPlanFunc: (ExecPlan, ActorRef, Scheduler) => Future[Unit]): Boolean = synchronized {
    val id = plan.queryContext.queryId
    if ((currentIDs contains id) || currentIDs.size < concurrentQueries) {
      countExecImmed.increment()
      execute(plan, replyTo, newPlanFunc)
      true
    } else if (queueSize < maxQueueLen) {
      countAddQueue.increment()
      queue.enqueue((plan, replyTo))
      true
    } else {
      countQueueFull.increment()
      false
    }
  }

  // Length of queue - not synchronized, so data might not be consistent.
  def queueSize: Int = queue.length

  // Pop next item(s?) off top of queue so long as there is room to execute another query.
  def tryPop(newPlanFunc: (ExecPlan, ActorRef, Scheduler) => Future[Unit]): Unit = synchronized {
    if (currentIDs.size < concurrentQueries && queue.nonEmpty) {
      val (nextPlan, replyTo) = queue.dequeue
      countExecQueue.increment()
      execute(nextPlan, replyTo, newPlanFunc)
    }
  }

  // Executes newPlanFunc, updating currentIDs/plans as needed
  // The closure is called in the current thread, but the callback when Future/plan is completed will not be.
  // Note that an appropriate single threaded scheduler is pulled off the stack and used.
  private def execute(newPlan: ExecPlan,
                      replyTo: ActorRef,
                      newPlanFunc: (ExecPlan, ActorRef, Scheduler) => Future[Unit]): Unit = {
    val queryId = newPlan.queryContext.queryId

    // Add plan to data structures and obtain scheduler
    val sched = synchronized {
      currentIDs(queryId) += 1
      idToSched.getOrElseUpdate(queryId, {
        if (schedulers.nonEmpty) schedulers.dequeue
        else Scheduler.singleThread(s"${FiloSchedulers.QuerySchedName}-${currentIDs.size}")
      })
    }

    val planFut = newPlanFunc(newPlan, replyTo, sched)

    // Remove plan on Future complete.  If a slot opens up, recursively call myself to open up a new slot
    planFut.onComplete {
      case _ => synchronized {
                  currentIDs(queryId) -= 1
                  if (currentIDs(queryId) == 0) {
                    currentIDs.remove(queryId)
                    schedulers.enqueue(idToSched.remove(queryId).get)
                    queueLen.update(queueSize)
                    tryPop(newPlanFunc)
                  }
                }
    }(sched)
  }

  // Currently active query IDs.  Not synchronized, so data might not be consistent.
  def currentQueryIDs: Set[String] = currentIDs.keySet.toSet
}

object QueryCommandPriority extends java.util.Comparator[Envelope] {
  override def compare(o1: Envelope, o2: Envelope): Int = {
    (o1.message, o2.message) match {
{'fraction_non_alphanumeric': 0.08352162685603615, 'fraction_numerical': 0.0015332472562943835, 'mean_word_length': 3.453108156665469, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3954922', 'n_tokens_mistral': 3550, 'n_tokens_neox': 3400, 'n_words': 1210}