<commit_before>;;
;; Copyright (c) Two Sigma Open Source, LLC
;;
;; Licensed under the Apache License, Version 2.0 (the "License");
;; you may not use this file except in compliance with the License.
;; You may obtain a copy of the License at
;;
;;  http://www.apache.org/licenses/LICENSE-2.0
;;
;; Unless required by applicable law or agreed to in writing, software
;; distributed under the License is distributed on an "AS IS" BASIS,
;; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
;; See the License for the specific language governing permissions and
;; limitations under the License.
;;
(ns cook.scheduler.scheduler
  (:require [chime :refer [chime-at chime-ch]]
            [clj-time.coerce :as tc]
            [clj-time.core :as time]
            [clojure.core.async :as async]
            [clojure.core.cache :as cache]
            [clojure.edn :as edn]
            [clojure.string :as str]
            [clojure.tools.logging :as log]
            [cook.compute-cluster :as cc]
            [cook.config :as config]
            [cook.datomic :as datomic]
            [cook.plugins.completion :as completion]
            [cook.plugins.definitions :as plugins]
            [cook.plugins.launch :as launch-plugin]
            [cook.scheduler.constraints :as constraints]
            [cook.scheduler.data-locality :as dl]
            [cook.scheduler.dru :as dru]
            [cook.scheduler.fenzo-utils :as fenzo]
            [cook.group :as group]
            [cook.pool :as pool]
            [cook.quota :as quota]
            [cook.mesos.reason :as reason]
            [cook.scheduler.share :as share]
            [cook.mesos.task :as task]
            [cook.tools :as util]
            [cook.rate-limit :as ratelimit]
            [cook.task]
            [datomic.api :as d :refer (q)]
            [mesomatic.scheduler :as mesos]
            [metatransaction.core :refer (db)]
            [metrics.counters :as counters]
            [metrics.gauges :as gauges]
            [metrics.histograms :as histograms]
            [metrics.meters :as meters]
            [metrics.timers :as timers]
            [plumbing.core :as pc])
  (import [com.netflix.fenzo TaskAssignmentResult TaskRequest TaskScheduler TaskScheduler$Builder
                             VirtualMachineLease VirtualMachineLease$Range
                             VirtualMachineCurrentState]
          [com.netflix.fenzo.functions Action1 Func1]))

(defn now
  []
  (tc/to-date (time/now)))

(defn offer-resource-values
  [offer resource-name value-type]
  (->> offer :resources (filter #(= (:name %) resource-name)) (map value-type)))

(defn offer-resource-scalar
  [offer resource-name]
  (reduce + 0.0 (offer-resource-values offer resource-name :scalar)))

(defn offer-resource-ranges
  [offer resource-name]
  (reduce into [] (offer-resource-values offer resource-name :ranges)))

(defn get-offer-attr-map
  "Gets all the attributes from an offer and puts them in a simple, less structured map of the form
   name->value"
  [offer]
  (let [mesos-attributes (->> offer
                              :attributes
                              (map #(vector (:name %) (case (:type %)
                                                        :value-scalar (:scalar %)
                                                        :value-ranges (:ranges %)
                                                        :value-set (:set %)
                                                        :value-text (:text %)
                                                        ; Default
                                                        (:value %))))
                              (into {}))
        cook-attributes {"HOSTNAME" (:hostname offer)
                         "COOK_GPU?" (-> offer
                                         (offer-resource-scalar "gpus")
                                         (or 0.0)
                                         pos?)}]
    (merge mesos-attributes cook-attributes)))

(timers/deftimer [cook-mesos scheduler handle-status-update-duration])
(timers/deftimer [cook-mesos scheduler handle-framework-message-duration])
(meters/defmeter [cook-mesos scheduler handle-framework-message-rate])

(timers/deftimer [cook-mesos scheduler generate-user-usage-map-duration])

(defn metric-title
  [metric-name pool]
  ["cook-mesos" "scheduler" metric-name (str "pool-" pool)])

(defn completion-rate-meter
  [status pool]
  (let [metric-name (case status
                      :succeeded "tasks-succeeded"
                      :failed "tasks-failed"
                      :completed "tasks-completed")]
    (meters/meter (metric-title metric-name pool))))

(defn completion-mem-meter
  [status pool]
  (let [metric-name (case status
                      :succeeded "tasks-succeeded-mem"
                      :failed "tasks-failed-mem"
                      :completed "tasks-completed-mem")]
    (meters/meter (metric-title metric-name pool))))

(defn completion-cpus-meter
  [status pool]
  (let [metric-name (case status
                      :succeeded "tasks-succeeded-cpus"
                      :failed "tasks-failed-cpus"
                      :completed "tasks-completed-cpus")]
    (meters/meter (metric-title metric-name pool))))

(defn completion-run-times-histogram
  [status pool]
  (let [metric-name (case status
                      :succeeded "hist-task-succeed-times"
                      :failed "hist-task-fail-times"
                      :completed "hist-task-complete-times")]
    (histograms/histogram (metric-title metric-name pool))))

(defn completion-run-times-meter
  [status pool]
  (let [metric-name (case status
                      :succeeded "task-succeed-times"
                      :failed "task-fail-times"
                      :completed "task-complete-times")]
    (meters/meter (metric-title metric-name pool))))

(defn handle-throughput-metrics [job-resources run-time status pool]
  (let [completion-rate (completion-rate-meter status pool)
        completion-mem (completion-mem-meter status pool)
        completion-cpus (completion-cpus-meter status pool)
        completion-hist-run-times (completion-run-times-histogram status pool)
        completion-MA-run-times (completion-run-times-meter status pool)]
    (meters/mark! completion-rate)
    (meters/mark!
      completion-mem
      (:mem job-resources))
    (meters/mark!
      completion-cpus
      (:cpus job-resources))
    (histograms/update!
      completion-hist-run-times
      run-time)
    (meters/mark!
      completion-MA-run-times
      run-time)))

(defn interpret-task-status
  "Converts the status packet from Mesomatic into a more friendly data structure"
  [s]
  {:task-id (-> s :task-id :value)
   :reason (:reason s)
   :task-state (:state s)
   :progress (try
               (when (:data s)
                 (:percent (edn/read-string (String. (.toByteArray (:data s))))))
               (catch Exception e
                 (try (log/debug e (str "Error parsing mesos status data to edn."
                                        "Is it in the format we expect?"
                                        "String representation: "
                                        (String. (.toByteArray (:data s)))))
                      (catch Exception e
                        (log/debug e "Error reading a string from mesos status data. Is it in the format we expect?")))))})

(defn update-reason-metrics!
  "Updates histograms and counters for run time, cpu time, and memory time,
  where the histograms have the failure reason in the title"
  [db mesos-reason instance-runtime {:keys [cpus mem]}]
  (let [reason (->> mesos-reason
                    (reason/mesos-reason->cook-reason-entity-id db)
                    (d/entity db)
                    :reason/name
                    name)
        update-metrics! (fn update-metrics! [s v]
                          (histograms/update!
                            (histograms/histogram
                              ["cook-mesos" "scheduler" "hist-task-fail" reason s])
                            v)
                          (counters/inc!
                            (counters/counter
                              ["cook-mesos" "scheduler" "hist-task-fail" reason s "total"])
                            v))
        instance-runtime-seconds (/ instance-runtime 1000)
        mem-gb (/ mem 1024)]
    (update-metrics! "times" instance-runtime-seconds)
    (update-metrics! "cpu-times" (* instance-runtime-seconds cpus))
    (update-metrics! "mem-times" (* instance-runtime-seconds mem-gb))))

(defn write-status-to-datomic
  "Takes a status update from mesos."
  [conn pool->fenzo status]
  (log/info "Mesos status is:" status)
  (timers/time!
    handle-status-update-duration
    (try (let [db (db conn)
               {:keys [task-id reason task-state progress]} (interpret-task-status status)
               _ (when-not task-id
                   (throw (ex-info "task-id is nil. Something unexpected has happened."
                                   {:status status
                                    :task-id task-id
                                    :reason reason
                                    :task-state task-state
                                    :progress progress})))
               [job instance prior-instance-status] (first (q '[:find ?j ?i ?status
                                                                :in $ ?task-id
                                                                :where
                                                                [?i :instance/task-id ?task-id]
                                                                [?i :instance/status ?s]
                                                                [?s :db/ident ?status]
                                                                [?j :job/instance ?i]]
                                                              db task-id))
               job-ent (d/entity db job)
               instance-ent (d/entity db instance)
               previous-reason (reason/instance-entity->reason-entity db instance-ent)
               instance-status (condp contains? task-state
                                 #{:task-staging} :instance.status/unknown
                                 #{:task-starting
                                   :task-running} :instance.status/running
                                 #{:task-finished} :instance.status/success
                                 #{:task-failed
                                   :task-killed
                                   :task-lost
                                   :task-error} :instance.status/failed)
               prior-job-state (:job/state (d/entity db job))
               current-time (now)
               instance-runtime (- (.getTime current-time) ; Used for reporting
                                   (.getTime (or (:instance/start-time instance-ent) current-time)))
               job-resources (util/job-ent->resources job-ent)
               pool-name (util/job->pool-name job-ent)
               ^TaskScheduler fenzo (get pool->fenzo pool-name)]
           (when (#{:instance.status/success :instance.status/failed} instance-status)
             (log/debug "Unassigning task" task-id "from" (:instance/hostname instance-ent))
             (try
               (locking fenzo
                 (.. fenzo
                     (getTaskUnAssigner)
                     (call task-id (:instance/hostname instance-ent))))
               (catch Exception e
                 (log/error e "Failed to unassign task" task-id "from" (:instance/hostname instance-ent)))))
           (when (= instance-status :instance.status/success)
             (handle-throughput-metrics job-resources instance-runtime :succeeded pool-name)
             (handle-throughput-metrics job-resources instance-runtime :completed pool-name))
           (when (= instance-status :instance.status/failed)
             (handle-throughput-metrics job-resources instance-runtime :failed pool-name)
             (handle-throughput-metrics job-resources instance-runtime :completed pool-name)
             (when-not previous-reason
               (update-reason-metrics! db reason instance-runtime job-resources)))
           (when-not (nil? instance)
             ;; (println "update:" task-id task-state job instance instance-status prior-job-state)
             (log/debug "Transacting updated state for instance" instance "to status" instance-status)
             ;; The database can become inconsistent if we make multiple calls to :instance/update-state in a single
             ;; transaction; see the comment in the definition of :instance/update-state for more details
             (let [transaction-chan (datomic/transact-with-retries
                                     conn
                                     (reduce
                                      into
                                      [[:instance/update-state instance instance-status (or (:db/id previous-reason)
                                                                                            (reason/mesos-reason->cook-reason-entity-id db reason)
                                                                                            [:reason.name :unknown])]] ; Warning: Default is not mea-culpa
                                      [(when (and (#{:instance.status/failed} instance-status) (not previous-reason) reason)
                                         [[:db/add instance :instance/reason (reason/mesos-reason->cook-reason-entity-id db reason)]])
                                       (when (and (#{:instance.status/success
                                                     :instance.status/failed} instance-status)
                                                  (nil? (:instance/end-time instance-ent)))
                                         [[:db/add instance :instance/end-time (now)]])
                                       (when (and (#{:task-starting :task-running} task-state)
                                                  (nil? (:instance/mesos-start-time instance-ent)))
                                         [[:db/add instance :instance/mesos-start-time (now)]])
                                       (when progress
                                         [[:db/add instance :instance/progress progress]])]))]
               (async/go
                 ; Wait for the transcation to complete before running the plugin
                 (let [chan-result (async/<! transaction-chan)]
                   (when (#{:instance.status/success :instance.status/failed} instance-status)
                     (let [db (d/db conn)
                           updated-job (d/entity db job)
                           updated-instance (d/entity db instance)]
                       (try
                         (plugins/on-instance-completion completion/plugin updated-job updated-instance)
                         (catch Exception e
                           (log/error e "Error while running instance completion plugin.")))))
                   chan-result)))))
         (catch Exception e
           (log/error e "Mesos scheduler status update error")))))

(defn- task-id->instance-id
  "Retrieves the instance-id given a task-id"
  [db task-id]
  (-> (d/entity db [:instance/task-id task-id])
      :db/id))

(defn handle-framework-message
  "Processes a framework message from Mesos."
  [conn {:keys [handle-exit-code handle-progress-message]}
   {:strs [exit-code progress-message progress-percent progress-sequence task-id] :as message}]
  (log/info "Received framework message:" {:task-id task-id, :message message})
  (timers/time!
    handle-framework-message-duration
    (try
      (when (str/blank? task-id)
        (throw (ex-info "task-id is empty in framework message" {:message message})))
      (let [instance-id (task-id->instance-id (db conn) task-id)]
        (if (nil? instance-id)
          (throw (ex-info "No instance found!" {:task-id task-id}))
          (do
            (when (or progress-message progress-percent)
              (log/debug "Updating instance" instance-id "progress to" progress-percent progress-message)
              (handle-progress-message {:instance-id instance-id
                                        :progress-message progress-message
                                        :progress-percent progress-percent
                                        :progress-sequence progress-sequence}))
            (when exit-code
              (log/info "Updating instance" instance-id "exit-code to" exit-code)
              (handle-exit-code task-id exit-code)))))
      (catch Exception e
        (log/error e "Mesos scheduler framework message error")))))

(timers/deftimer [cook-mesos scheduler tx-report-queue-processing-duration])
(meters/defmeter [cook-mesos scheduler tx-report-queue-datoms])
(meters/defmeter [cook-mesos scheduler tx-report-queue-update-job-state])
(meters/defmeter [cook-mesos scheduler tx-report-queue-job-complete])
(meters/defmeter [cook-mesos scheduler tx-report-queue-tasks-killed])

(defn monitor-tx-report-queue
  "Takes an async channel that will have tx report queue elements on it"
  [tx-report-chan conn]
  (log/info "Starting tx-report-queue")
  (let [kill-chan (async/chan)
        query-db (d/db conn)
        query-basis (d/basis-t query-db)
        tasks-to-kill (q '[:find ?i
                           :in $ [?status ...]
                           :where
                           [?i :instance/status ?status]
                           [?job :job/instance ?i]
                           [?job :job/state :job.state/completed]]
                         query-db [:instance.status/unknown :instance.status/running])]
    (doseq [[task-entity-id] tasks-to-kill]
      (let [task-id (cook.task/task-entity-id->task-id query-db task-entity-id)
            compute-cluster-name (cook.task/task-entity-id->compute-cluster-name query-db task-entity-id)]
        (if-let [compute-cluster (cc/compute-cluster-name->ComputeCluster compute-cluster-name)]
          (try
            (log/info "Attempting to kill task" task-id "in" compute-cluster-name "from already completed job")
            (meters/mark! tx-report-queue-tasks-killed)
            (cc/safe-kill-task compute-cluster task-id)
            (catch Exception e
              (log/error e (str "Failed to kill task" task-id))))
          (log/warn "Unable to kill task" task-id "with unknown cluster" compute-cluster-name))))
    (async/go
      (loop []
        (async/alt!
          tx-report-chan ([tx-report]
                           (async/go
                             (timers/start-stop-time! ; Use this in go blocks, time! doesn't play nice
                               tx-report-queue-processing-duration
                               (let [{:keys [tx-data db-after]} tx-report]
                                 (when (< query-basis (d/basis-t db-after))
                                   (let [db (db conn)]
                                     (meters/mark! tx-report-queue-datoms (count tx-data))
                                     ;; Monitoring whether a job is completed.
                                     (doseq [{:keys [e a v]} tx-data]
                                       (try
                                         (when (and (= a (d/entid db :job/state))
                                                    (= v (d/entid db :job.state/completed)))
                                           (meters/mark! tx-report-queue-job-complete)
                                           (doseq [[task-entity-id]
                                                   (q '[:find ?i
                                                        :in $ ?job [?status ...]
                                                        :where
                                                        [?job :job/instance ?i]
                                                        [?i :instance/status ?status]]
                                                      db e [:instance.status/unknown
                                                            :instance.status/running])]
                                             (let [task-id (cook.task/task-entity-id->task-id db task-entity-id)
                                                   compute-cluster-name (cook.task/task-entity-id->compute-cluster-name db task-entity-id)]
                                               (if-let [compute-cluster (cc/compute-cluster-name->ComputeCluster compute-cluster-name)]
                                                 (do (log/info "Attempting to kill task" task-id "in" compute-cluster-name "due to job completion")
                                                     (meters/mark! tx-report-queue-tasks-killed)
                                                     (cc/safe-kill-task compute-cluster task-id))
                                                 (log/error "Couldn't kill task" task-id "due to no Mesos driver for compute cluster" compute-cluster-name "!")))))
                                         (catch Exception e
                                           (log/error e "Unexpected exception on tx report queue processor")))))))))
                           (recur))
          kill-chan ([_] nil))))
    #(async/close! kill-chan)))

;;; ===========================================================================
;;; API for matcher
;;; ===========================================================================

(defrecord VirtualMachineLeaseAdapter [offer time]
  VirtualMachineLease
  (cpuCores [_] (or (offer-resource-scalar offer "cpus") 0.0))
  (diskMB [_] (or (offer-resource-scalar offer "disk") 0.0))
  (getScalarValue [_ name] (or (double (offer-resource-scalar offer name)) 0.0))
  (getScalarValues [_]
    (reduce (fn [result resource]
              (if-let [value (:scalar resource)]
                ;; Do not remove the following fnil--either arg to + can be nil!
                (update-in result [(:name resource)] (fnil + 0.0 0.0) value)
                result))
            {}
            (:resources offer)))
  ; Some Fenzo plugins (which are included with fenzo, such as host attribute constraints) expect the "HOSTNAME"
  ; attribute to contain the hostname of this virtual machine.
  (getAttributeMap [_] (get-offer-attr-map offer))
  (getId [_] (-> offer :id :value))
  (getOffer [_] (throw (UnsupportedOperationException.)))
  (getOfferedTime [_] time)
  (getVMID [_] (-> offer :slave-id :value))
  (hostname [_] (:hostname offer))
  (memoryMB [_] (or (offer-resource-scalar offer "mem") 0.0))
  (networkMbps [_] 0.0)
  (portRanges [_] (mapv (fn [{:keys [begin end]}]
                          (VirtualMachineLease$Range. begin end))
                        (offer-resource-ranges offer "ports"))))


(defrecord TaskRequestAdapter [job resources task-id assigned-resources guuid->considerable-cotask-ids constraints needs-gpus? scalar-requests]
  TaskRequest
  (getCPUs [_] (:cpus resources))
  (getDisk [_] 0.0)
  (getHardConstraints [_] constraints)
  (getId [_] task-id)
  (getScalarRequests [_] scalar-requests)
  (getAssignedResources [_] @assigned-resources)
  (setAssignedResources [_ v] (reset! assigned-resources v))
  (getCustomNamedResources [_] {})
  (getMemory [_] (:mem resources))
  (getNetworkMbps [_] 0.0)
  (getPorts [_] (:ports resources))
  (getSoftConstraints [_] [])
  (taskGroupName [_] (str (:job/uuid job))))

(defn make-task-request
  "Helper to create a TaskRequest using TaskRequestAdapter. TaskRequestAdapter implements Fenzo's TaskRequest interface
   given a job, its resources, its task-id and a function assigned-cotask-getter. assigned-cotask-getter should be a
   function that takes a group uuid and returns a set of task-ids, which correspond to the tasks that will be assigned
   during the same Fenzo scheduling cycle as the newly created TaskRequest."
  [db job & {:keys [resources task-id assigned-resources guuid->considerable-cotask-ids reserved-hosts running-cotask-cache]
             :or {resources (util/job-ent->resources job)
                  task-id (str (d/squuid))
                  assigned-resources (atom nil)
                  guuid->considerable-cotask-ids (constantly #{})
                  running-cotask-cache (atom (cache/fifo-cache-factory {} :threshold 1))
                  reserved-hosts #{}}}]
  (let [constraints (-> (constraints/make-fenzo-job-constraints job)
                        (conj (constraints/build-rebalancer-reservation-constraint reserved-hosts))
                        (into
                          (remove nil?
                                  (mapv (fn make-group-constraints [group]
                                          (constraints/make-fenzo-group-constraint
                                            db group #(guuid->considerable-cotask-ids (:group/uuid group)) running-cotask-cache))
                                        (:group/_job job)))))
        needs-gpus? (constraints/job-needs-gpus? job)
        scalar-requests (reduce (fn [result resource]
                                  (if-let [value (:resource/amount resource)]
                                    (assoc result (name (:resource/type resource)) value)
                                    result))
                                {}
                                (:job/resource job))]
    (->TaskRequestAdapter job resources task-id assigned-resources guuid->considerable-cotask-ids constraints needs-gpus? scalar-requests)))

(defn match-offer-to-schedule
  "Given an offer and a schedule, computes all the tasks should be launched as a result.

   A schedule is just a sorted list of tasks, and we're going to greedily assign them to
   the offer.

   Returns {:matches (list of tasks that got matched to the offer)
            :failures (list of unmatched tasks, and why they weren't matched)}"
  [db ^TaskScheduler fenzo considerable offers rebalancer-reservation-atom pool-name]
  (log/info "In" pool-name "pool, matching" (count offers) "offers to" (count considerable) "jobs with fenzo")
  (log/debug "In" pool-name "pool, tasks to scheduleOnce" considerable)
  (dl/update-cost-staleness-metric considerable)
  (let [t (System/currentTimeMillis)
        leases (mapv #(->VirtualMachineLeaseAdapter % t) offers)
        considerable->task-id (plumbing.core/map-from-keys (fn [_] (str (d/squuid))) considerable)
        guuid->considerable-cotask-ids (util/make-guuid->considerable-cotask-ids considerable->task-id)
        running-cotask-cache (atom (cache/fifo-cache-factory {} :threshold (max 1 (count considerable))))
        job-uuid->reserved-host (or (:job-uuid->reserved-host @rebalancer-reservation-atom) {})
        reserved-hosts (into (hash-set) (vals job-uuid->reserved-host))
        ; Important that requests maintains the same order as considerable
        requests (mapv (fn [job]
                         (make-task-request db job
                                            :guuid->considerable-cotask-ids guuid->considerable-cotask-ids
                                            :reserved-hosts (disj reserved-hosts (job-uuid->reserved-host (:job/uuid job)))
                                            :running-cotask-cache running-cotask-cache
                                            :task-id (considerable->task-id job)))
                       considerable)
        ;; Need to lock on fenzo when accessing scheduleOnce because scheduleOnce and
        ;; task assigner can not be called at the same time.
        ;; task assigner may be called when reconciling
        result (locking fenzo
                 (.scheduleOnce fenzo requests leases))
        failure-results (.. result getFailures values)
        assignments (.. result getResultMap values)]
    (doall (map (fn [^VirtualMachineLease lease]
                  (when (-> lease :offer :reject-after-match-attempt)
                    (log/info "In" pool-name "pool, retracting lease" (-> lease :offer :id))
                    (locking fenzo
                      (.expireLease fenzo (.getId lease)))))
                leases))

    (log/debug "In" pool-name "pool, found this assignment:" result)

    {:matches (mapv (fn [assignment]
                      {:leases (.getLeasesUsed assignment)
                       :tasks (.getTasksAssigned assignment)
                       :hostname (.getHostname assignment)})
                    assignments)
     :failures failure-results}))

(meters/defmeter [cook-mesos scheduler scheduler-offer-declined])

(defn decline-offers
  "declines a collection of offer ids"
  [compute-cluster offer-ids]
  (log/debug "Declining offers:" offer-ids)
  (meters/mark! scheduler-offer-declined (count offer-ids))
  (try
    (cc/decline-offers compute-cluster offer-ids)
    (catch Throwable t
      (log/error t "Error declining offers for" compute-cluster))))

(histograms/defhistogram [cook-mesos scheduler number-tasks-matched])
(histograms/defhistogram [cook-mesos-scheduler number-offers-matched])
(meters/defmeter [cook-mesos scheduler scheduler-offer-matched])
(meters/defmeter [cook-mesos scheduler handle-resource-offer!-errors])
(def front-of-job-queue-mem-atom (atom 0))
(def front-of-job-queue-cpus-atom (atom 0))
(gauges/defgauge [cook-mesos scheduler front-of-job-queue-mem] (fn [] @front-of-job-queue-mem-atom))
(gauges/defgauge [cook-mesos scheduler front-of-job-queue-cpus] (fn [] @front-of-job-queue-cpus-atom))

(defn generate-user-usage-map
  "Returns a mapping from user to usage stats"
  [unfiltered-db pool-name]
  (timers/time!
    generate-user-usage-map-duration
    (->> (util/get-running-task-ents unfiltered-db)
         (map :job/_instance)
         (remove #(not= pool-name (util/job->pool-name %)))
         (group-by :job/user)
         (pc/map-vals (fn [jobs]
                        (->> jobs
                             (map util/job->usage)
                             (reduce (partial merge-with +))))))))

;Shared as we use this for unscheduled too.
(defonce pool->user->number-jobs (atom {}))

(defn pending-jobs->considerable-jobs
  "Limit the pending jobs to considerable jobs based on usage and quota.
   Further limit the considerable jobs to a maximum of num-considerable jobs."
  [db pending-jobs user->quota user->usage num-considerable pool-name]
  (log/debug "In" pool-name "pool, there are" (count pending-jobs) "pending jobs:" pending-jobs)
  (let [enforcing-job-launch-rate-limit? (ratelimit/enforce? ratelimit/job-launch-rate-limiter)
        user->number-jobs (atom {})
        user->rate-limit-count (atom {})
        user-within-launch-rate-limit?-fn
        (fn
          [{:keys [job/user]}]
          ; Account for each time we see a job for a user.
          (swap! user->number-jobs update user #(inc (or % 0)))
          (let [tokens (ratelimit/get-token-count! ratelimit/job-launch-rate-limiter user)
                number-jobs-for-user-so-far (@user->number-jobs user)
                is-rate-limited? (> number-jobs-for-user-so-far tokens)]
            (when is-rate-limited?
              (swap! user->rate-limit-count update user #(inc (or % 0))))
            (not (and is-rate-limited? enforcing-job-launch-rate-limit?))))
        quota-filter-jobs (util/filter-based-on-quota user->quota user->usage pending-jobs)
        allowed-to-start-jobs (filter (fn [job] (util/job-allowed-to-start? db job)) quota-filter-jobs)
        within-launch-rate-jobs (filter user-within-launch-rate-limit?-fn allowed-to-start-jobs)
        launch-plugin-jobs (filter launch-plugin/filter-job-launches within-launch-rate-jobs)
        considerable-jobs (take num-considerable launch-plugin-jobs)
        ; Force this to be taken eagerly so that the log line is accurate.
        considerable-jobs* (doall considerable-jobs)]
    (swap! pool->user->number-jobs update pool-name (constantly @user->number-jobs))
    (log/info "Users whose job launches are rate-limited " @user->rate-limit-count "( enforcing =" enforcing-job-launch-rate-limit? ")")
    (log/info "In" pool-name "pool, there are"
              (count pending-jobs) "pending ->"
              (count quota-filter-jobs) "quota-filtered ->"
              (count allowed-to-start-jobs) "allowed-to-start ->"
              (count within-launch-rate-jobs) "within-launch-rate ->"
              (count launch-plugin-jobs) "allowed by launch-plugin ->"
              "and"
              (count considerable-jobs*) "considerable jobs")
    considerable-jobs*))


(defn matches->job-uuids
  "Returns the matched job uuids."
  [matches pool-name]
  (let [jobs (->> matches
                  (mapcat #(-> % :tasks))
                  (map #(-> % .getRequest :job)))
        job-uuids (set (map :job/uuid jobs))]
    (log/debug "In" pool-name "pool, matched jobs:" (count job-uuids))
    (when (seq matches)
      (let [matched-normal-jobs-resource-requirements (util/sum-resources-of-jobs jobs)]
        (meters/mark! (meters/meter (metric-title "matched-tasks-cpus" pool-name))
                      (:cpus matched-normal-jobs-resource-requirements))
        (meters/mark! (meters/meter (metric-title "matched-tasks-mem" pool-name))
                      (:mem matched-normal-jobs-resource-requirements))))
    job-uuids))

(defn remove-matched-jobs-from-pending-jobs
  "Removes matched jobs from pool->pending-jobs."
  [pool-name->pending-jobs matched-job-uuids pool-name]<commit_msg>Remove extra logging that makes computing the result more expen… (#1326)<commit_after>;;
;; Copyright (c) Two Sigma Open Source, LLC
;;
;; Licensed under the Apache License, Version 2.0 (the "License");
;; you may not use this file except in compliance with the License.
;; You may obtain a copy of the License at
;;
;;  http://www.apache.org/licenses/LICENSE-2.0
;;
;; Unless required by applicable law or agreed to in writing, software
;; distributed under the License is distributed on an "AS IS" BASIS,
;; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
;; See the License for the specific language governing permissions and
;; limitations under the License.
;;
(ns cook.scheduler.scheduler
  (:require [chime :refer [chime-at chime-ch]]
            [clj-time.coerce :as tc]
            [clj-time.core :as time]
            [clojure.core.async :as async]
            [clojure.core.cache :as cache]
            [clojure.edn :as edn]
            [clojure.string :as str]
            [clojure.tools.logging :as log]
            [cook.compute-cluster :as cc]
            [cook.config :as config]
            [cook.datomic :as datomic]
            [cook.plugins.completion :as completion]
            [cook.plugins.definitions :as plugins]
            [cook.plugins.launch :as launch-plugin]
            [cook.scheduler.constraints :as constraints]
            [cook.scheduler.data-locality :as dl]
            [cook.scheduler.dru :as dru]
            [cook.scheduler.fenzo-utils :as fenzo]
            [cook.group :as group]
            [cook.pool :as pool]
            [cook.quota :as quota]
            [cook.mesos.reason :as reason]
            [cook.scheduler.share :as share]
            [cook.mesos.task :as task]
            [cook.tools :as util]
            [cook.rate-limit :as ratelimit]
            [cook.task]
            [datomic.api :as d :refer (q)]
            [mesomatic.scheduler :as mesos]
            [metatransaction.core :refer (db)]
            [metrics.counters :as counters]
            [metrics.gauges :as gauges]
            [metrics.histograms :as histograms]
            [metrics.meters :as meters]
            [metrics.timers :as timers]
            [plumbing.core :as pc])
  (import [com.netflix.fenzo TaskAssignmentResult TaskRequest TaskScheduler TaskScheduler$Builder
                             VirtualMachineLease VirtualMachineLease$Range
                             VirtualMachineCurrentState]
          [com.netflix.fenzo.functions Action1 Func1]))

(defn now
  []
  (tc/to-date (time/now)))

(defn offer-resource-values
  [offer resource-name value-type]
  (->> offer :resources (filter #(= (:name %) resource-name)) (map value-type)))

(defn offer-resource-scalar
  [offer resource-name]
  (reduce + 0.0 (offer-resource-values offer resource-name :scalar)))

(defn offer-resource-ranges
  [offer resource-name]
  (reduce into [] (offer-resource-values offer resource-name :ranges)))

(defn get-offer-attr-map
  "Gets all the attributes from an offer and puts them in a simple, less structured map of the form
   name->value"
  [offer]
  (let [mesos-attributes (->> offer
                              :attributes
                              (map #(vector (:name %) (case (:type %)
                                                        :value-scalar (:scalar %)
                                                        :value-ranges (:ranges %)
                                                        :value-set (:set %)
                                                        :value-text (:text %)
                                                        ; Default
                                                        (:value %))))
                              (into {}))
        cook-attributes {"HOSTNAME" (:hostname offer)
                         "COOK_GPU?" (-> offer
                                         (offer-resource-scalar "gpus")
                                         (or 0.0)
                                         pos?)}]
    (merge mesos-attributes cook-attributes)))

(timers/deftimer [cook-mesos scheduler handle-status-update-duration])
(timers/deftimer [cook-mesos scheduler handle-framework-message-duration])
(meters/defmeter [cook-mesos scheduler handle-framework-message-rate])

(timers/deftimer [cook-mesos scheduler generate-user-usage-map-duration])

(defn metric-title
  [metric-name pool]
  ["cook-mesos" "scheduler" metric-name (str "pool-" pool)])

(defn completion-rate-meter
  [status pool]
  (let [metric-name (case status
                      :succeeded "tasks-succeeded"
                      :failed "tasks-failed"
                      :completed "tasks-completed")]
    (meters/meter (metric-title metric-name pool))))

(defn completion-mem-meter
  [status pool]
  (let [metric-name (case status
                      :succeeded "tasks-succeeded-mem"
                      :failed "tasks-failed-mem"
                      :completed "tasks-completed-mem")]
    (meters/meter (metric-title metric-name pool))))

(defn completion-cpus-meter
  [status pool]
  (let [metric-name (case status
                      :succeeded "tasks-succeeded-cpus"
                      :failed "tasks-failed-cpus"
                      :completed "tasks-completed-cpus")]
    (meters/meter (metric-title metric-name pool))))

(defn completion-run-times-histogram
  [status pool]
  (let [metric-name (case status
                      :succeeded "hist-task-succeed-times"
                      :failed "hist-task-fail-times"
                      :completed "hist-task-complete-times")]
    (histograms/histogram (metric-title metric-name pool))))

(defn completion-run-times-meter
  [status pool]
  (let [metric-name (case status
                      :succeeded "task-succeed-times"
                      :failed "task-fail-times"
                      :completed "task-complete-times")]
    (meters/meter (metric-title metric-name pool))))

(defn handle-throughput-metrics [job-resources run-time status pool]
  (let [completion-rate (completion-rate-meter status pool)
        completion-mem (completion-mem-meter status pool)
        completion-cpus (completion-cpus-meter status pool)
        completion-hist-run-times (completion-run-times-histogram status pool)
        completion-MA-run-times (completion-run-times-meter status pool)]
    (meters/mark! completion-rate)
    (meters/mark!
      completion-mem
      (:mem job-resources))
    (meters/mark!
      completion-cpus
      (:cpus job-resources))
    (histograms/update!
      completion-hist-run-times
      run-time)
    (meters/mark!
      completion-MA-run-times
      run-time)))

(defn interpret-task-status
  "Converts the status packet from Mesomatic into a more friendly data structure"
  [s]
  {:task-id (-> s :task-id :value)
   :reason (:reason s)
   :task-state (:state s)
   :progress (try
               (when (:data s)
                 (:percent (edn/read-string (String. (.toByteArray (:data s))))))
               (catch Exception e
                 (try (log/debug e (str "Error parsing mesos status data to edn."
                                        "Is it in the format we expect?"
                                        "String representation: "
                                        (String. (.toByteArray (:data s)))))
                      (catch Exception e
                        (log/debug e "Error reading a string from mesos status data. Is it in the format we expect?")))))})

(defn update-reason-metrics!
  "Updates histograms and counters for run time, cpu time, and memory time,
  where the histograms have the failure reason in the title"
  [db mesos-reason instance-runtime {:keys [cpus mem]}]
  (let [reason (->> mesos-reason
                    (reason/mesos-reason->cook-reason-entity-id db)
                    (d/entity db)
                    :reason/name
                    name)
        update-metrics! (fn update-metrics! [s v]
                          (histograms/update!
                            (histograms/histogram
                              ["cook-mesos" "scheduler" "hist-task-fail" reason s])
                            v)
                          (counters/inc!
                            (counters/counter
                              ["cook-mesos" "scheduler" "hist-task-fail" reason s "total"])
                            v))
        instance-runtime-seconds (/ instance-runtime 1000)
        mem-gb (/ mem 1024)]
    (update-metrics! "times" instance-runtime-seconds)
    (update-metrics! "cpu-times" (* instance-runtime-seconds cpus))
    (update-metrics! "mem-times" (* instance-runtime-seconds mem-gb))))

(defn write-status-to-datomic
  "Takes a status update from mesos."
  [conn pool->fenzo status]
  (log/info "Mesos status is:" status)
  (timers/time!
    handle-status-update-duration
    (try (let [db (db conn)
               {:keys [task-id reason task-state progress]} (interpret-task-status status)
               _ (when-not task-id
                   (throw (ex-info "task-id is nil. Something unexpected has happened."
                                   {:status status
                                    :task-id task-id
                                    :reason reason
                                    :task-state task-state
                                    :progress progress})))
               [job instance prior-instance-status] (first (q '[:find ?j ?i ?status
                                                                :in $ ?task-id
                                                                :where
                                                                [?i :instance/task-id ?task-id]
                                                                [?i :instance/status ?s]
                                                                [?s :db/ident ?status]
                                                                [?j :job/instance ?i]]
                                                              db task-id))
               job-ent (d/entity db job)
               instance-ent (d/entity db instance)
               previous-reason (reason/instance-entity->reason-entity db instance-ent)
               instance-status (condp contains? task-state
                                 #{:task-staging} :instance.status/unknown
                                 #{:task-starting
                                   :task-running} :instance.status/running
                                 #{:task-finished} :instance.status/success
                                 #{:task-failed
                                   :task-killed
                                   :task-lost
                                   :task-error} :instance.status/failed)
               prior-job-state (:job/state (d/entity db job))
               current-time (now)
               instance-runtime (- (.getTime current-time) ; Used for reporting
                                   (.getTime (or (:instance/start-time instance-ent) current-time)))
               job-resources (util/job-ent->resources job-ent)
               pool-name (util/job->pool-name job-ent)
               ^TaskScheduler fenzo (get pool->fenzo pool-name)]
           (when (#{:instance.status/success :instance.status/failed} instance-status)
             (log/debug "Unassigning task" task-id "from" (:instance/hostname instance-ent))
             (try
               (locking fenzo
                 (.. fenzo
                     (getTaskUnAssigner)
                     (call task-id (:instance/hostname instance-ent))))
               (catch Exception e
                 (log/error e "Failed to unassign task" task-id "from" (:instance/hostname instance-ent)))))
           (when (= instance-status :instance.status/success)
             (handle-throughput-metrics job-resources instance-runtime :succeeded pool-name)
             (handle-throughput-metrics job-resources instance-runtime :completed pool-name))
           (when (= instance-status :instance.status/failed)
             (handle-throughput-metrics job-resources instance-runtime :failed pool-name)
             (handle-throughput-metrics job-resources instance-runtime :completed pool-name)
             (when-not previous-reason
               (update-reason-metrics! db reason instance-runtime job-resources)))
           (when-not (nil? instance)
             ;; (println "update:" task-id task-state job instance instance-status prior-job-state)
             (log/debug "Transacting updated state for instance" instance "to status" instance-status)
             ;; The database can become inconsistent if we make multiple calls to :instance/update-state in a single
             ;; transaction; see the comment in the definition of :instance/update-state for more details
             (let [transaction-chan (datomic/transact-with-retries
                                     conn
                                     (reduce
                                      into
                                      [[:instance/update-state instance instance-status (or (:db/id previous-reason)
                                                                                            (reason/mesos-reason->cook-reason-entity-id db reason)
                                                                                            [:reason.name :unknown])]] ; Warning: Default is not mea-culpa
                                      [(when (and (#{:instance.status/failed} instance-status) (not previous-reason) reason)
                                         [[:db/add instance :instance/reason (reason/mesos-reason->cook-reason-entity-id db reason)]])
                                       (when (and (#{:instance.status/success
                                                     :instance.status/failed} instance-status)
                                                  (nil? (:instance/end-time instance-ent)))
                                         [[:db/add instance :instance/end-time (now)]])
                                       (when (and (#{:task-starting :task-running} task-state)
                                                  (nil? (:instance/mesos-start-time instance-ent)))
                                         [[:db/add instance :instance/mesos-start-time (now)]])
                                       (when progress
                                         [[:db/add instance :instance/progress progress]])]))]
               (async/go
                 ; Wait for the transcation to complete before running the plugin
                 (let [chan-result (async/<! transaction-chan)]
                   (when (#{:instance.status/success :instance.status/failed} instance-status)
                     (let [db (d/db conn)
                           updated-job (d/entity db job)
                           updated-instance (d/entity db instance)]
                       (try
                         (plugins/on-instance-completion completion/plugin updated-job updated-instance)
                         (catch Exception e
                           (log/error e "Error while running instance completion plugin.")))))
                   chan-result)))))
         (catch Exception e
           (log/error e "Mesos scheduler status update error")))))

(defn- task-id->instance-id
  "Retrieves the instance-id given a task-id"
  [db task-id]
  (-> (d/entity db [:instance/task-id task-id])
      :db/id))

(defn handle-framework-message
  "Processes a framework message from Mesos."
  [conn {:keys [handle-exit-code handle-progress-message]}
   {:strs [exit-code progress-message progress-percent progress-sequence task-id] :as message}]
  (log/info "Received framework message:" {:task-id task-id, :message message})
  (timers/time!
    handle-framework-message-duration
    (try
      (when (str/blank? task-id)
        (throw (ex-info "task-id is empty in framework message" {:message message})))
      (let [instance-id (task-id->instance-id (db conn) task-id)]
        (if (nil? instance-id)
          (throw (ex-info "No instance found!" {:task-id task-id}))
          (do
            (when (or progress-message progress-percent)
              (log/debug "Updating instance" instance-id "progress to" progress-percent progress-message)
              (handle-progress-message {:instance-id instance-id
                                        :progress-message progress-message
                                        :progress-percent progress-percent
                                        :progress-sequence progress-sequence}))
            (when exit-code
              (log/info "Updating instance" instance-id "exit-code to" exit-code)
              (handle-exit-code task-id exit-code)))))
      (catch Exception e
        (log/error e "Mesos scheduler framework message error")))))

(timers/deftimer [cook-mesos scheduler tx-report-queue-processing-duration])
(meters/defmeter [cook-mesos scheduler tx-report-queue-datoms])
(meters/defmeter [cook-mesos scheduler tx-report-queue-update-job-state])
(meters/defmeter [cook-mesos scheduler tx-report-queue-job-complete])
(meters/defmeter [cook-mesos scheduler tx-report-queue-tasks-killed])

(defn monitor-tx-report-queue
  "Takes an async channel that will have tx report queue elements on it"
  [tx-report-chan conn]
  (log/info "Starting tx-report-queue")
  (let [kill-chan (async/chan)
        query-db (d/db conn)
        query-basis (d/basis-t query-db)
        tasks-to-kill (q '[:find ?i
                           :in $ [?status ...]
                           :where
                           [?i :instance/status ?status]
                           [?job :job/instance ?i]
                           [?job :job/state :job.state/completed]]
                         query-db [:instance.status/unknown :instance.status/running])]
    (doseq [[task-entity-id] tasks-to-kill]
      (let [task-id (cook.task/task-entity-id->task-id query-db task-entity-id)
            compute-cluster-name (cook.task/task-entity-id->compute-cluster-name query-db task-entity-id)]
        (if-let [compute-cluster (cc/compute-cluster-name->ComputeCluster compute-cluster-name)]
          (try
            (log/info "Attempting to kill task" task-id "in" compute-cluster-name "from already completed job")
            (meters/mark! tx-report-queue-tasks-killed)
            (cc/safe-kill-task compute-cluster task-id)
            (catch Exception e
              (log/error e (str "Failed to kill task" task-id))))
          (log/warn "Unable to kill task" task-id "with unknown cluster" compute-cluster-name))))
    (async/go
      (loop []
        (async/alt!
          tx-report-chan ([tx-report]
                           (async/go
                             (timers/start-stop-time! ; Use this in go blocks, time! doesn't play nice
                               tx-report-queue-processing-duration
                               (let [{:keys [tx-data db-after]} tx-report]
                                 (when (< query-basis (d/basis-t db-after))
                                   (let [db (db conn)]
                                     (meters/mark! tx-report-queue-datoms (count tx-data))
                                     ;; Monitoring whether a job is completed.
                                     (doseq [{:keys [e a v]} tx-data]
                                       (try
                                         (when (and (= a (d/entid db :job/state))
                                                    (= v (d/entid db :job.state/completed)))
                                           (meters/mark! tx-report-queue-job-complete)
                                           (doseq [[task-entity-id]
                                                   (q '[:find ?i
                                                        :in $ ?job [?status ...]
                                                        :where
                                                        [?job :job/instance ?i]
                                                        [?i :instance/status ?status]]
                                                      db e [:instance.status/unknown
                                                            :instance.status/running])]
                                             (let [task-id (cook.task/task-entity-id->task-id db task-entity-id)
                                                   compute-cluster-name (cook.task/task-entity-id->compute-cluster-name db task-entity-id)]
                                               (if-let [compute-cluster (cc/compute-cluster-name->ComputeCluster compute-cluster-name)]
                                                 (do (log/info "Attempting to kill task" task-id "in" compute-cluster-name "due to job completion")
                                                     (meters/mark! tx-report-queue-tasks-killed)
                                                     (cc/safe-kill-task compute-cluster task-id))
                                                 (log/error "Couldn't kill task" task-id "due to no Mesos driver for compute cluster" compute-cluster-name "!")))))
                                         (catch Exception e
                                           (log/error e "Unexpected exception on tx report queue processor")))))))))
                           (recur))
          kill-chan ([_] nil))))
    #(async/close! kill-chan)))

;;; ===========================================================================
;;; API for matcher
;;; ===========================================================================

(defrecord VirtualMachineLeaseAdapter [offer time]
  VirtualMachineLease
  (cpuCores [_] (or (offer-resource-scalar offer "cpus") 0.0))
  (diskMB [_] (or (offer-resource-scalar offer "disk") 0.0))
  (getScalarValue [_ name] (or (double (offer-resource-scalar offer name)) 0.0))
  (getScalarValues [_]
    (reduce (fn [result resource]
              (if-let [value (:scalar resource)]
                ;; Do not remove the following fnil--either arg to + can be nil!
                (update-in result [(:name resource)] (fnil + 0.0 0.0) value)
                result))
            {}
            (:resources offer)))
  ; Some Fenzo plugins (which are included with fenzo, such as host attribute constraints) expect the "HOSTNAME"
  ; attribute to contain the hostname of this virtual machine.
  (getAttributeMap [_] (get-offer-attr-map offer))
  (getId [_] (-> offer :id :value))
  (getOffer [_] (throw (UnsupportedOperationException.)))
  (getOfferedTime [_] time)
  (getVMID [_] (-> offer :slave-id :value))
  (hostname [_] (:hostname offer))
  (memoryMB [_] (or (offer-resource-scalar offer "mem") 0.0))
  (networkMbps [_] 0.0)
  (portRanges [_] (mapv (fn [{:keys [begin end]}]
                          (VirtualMachineLease$Range. begin end))
                        (offer-resource-ranges offer "ports"))))


(defrecord TaskRequestAdapter [job resources task-id assigned-resources guuid->considerable-cotask-ids constraints needs-gpus? scalar-requests]
  TaskRequest
  (getCPUs [_] (:cpus resources))
  (getDisk [_] 0.0)
  (getHardConstraints [_] constraints)
  (getId [_] task-id)
  (getScalarRequests [_] scalar-requests)
  (getAssignedResources [_] @assigned-resources)
  (setAssignedResources [_ v] (reset! assigned-resources v))
  (getCustomNamedResources [_] {})
  (getMemory [_] (:mem resources))
  (getNetworkMbps [_] 0.0)
  (getPorts [_] (:ports resources))
  (getSoftConstraints [_] [])
  (taskGroupName [_] (str (:job/uuid job))))

(defn make-task-request
  "Helper to create a TaskRequest using TaskRequestAdapter. TaskRequestAdapter implements Fenzo's TaskRequest interface
   given a job, its resources, its task-id and a function assigned-cotask-getter. assigned-cotask-getter should be a
   function that takes a group uuid and returns a set of task-ids, which correspond to the tasks that will be assigned
   during the same Fenzo scheduling cycle as the newly created TaskRequest."
  [db job & {:keys [resources task-id assigned-resources guuid->considerable-cotask-ids reserved-hosts running-cotask-cache]
             :or {resources (util/job-ent->resources job)
                  task-id (str (d/squuid))
                  assigned-resources (atom nil)
                  guuid->considerable-cotask-ids (constantly #{})
                  running-cotask-cache (atom (cache/fifo-cache-factory {} :threshold 1))
                  reserved-hosts #{}}}]
  (let [constraints (-> (constraints/make-fenzo-job-constraints job)
                        (conj (constraints/build-rebalancer-reservation-constraint reserved-hosts))
                        (into
                          (remove nil?
                                  (mapv (fn make-group-constraints [group]
                                          (constraints/make-fenzo-group-constraint
                                            db group #(guuid->considerable-cotask-ids (:group/uuid group)) running-cotask-cache))
                                        (:group/_job job)))))
        needs-gpus? (constraints/job-needs-gpus? job)
        scalar-requests (reduce (fn [result resource]
                                  (if-let [value (:resource/amount resource)]
                                    (assoc result (name (:resource/type resource)) value)
                                    result))
                                {}
                                (:job/resource job))]
    (->TaskRequestAdapter job resources task-id assigned-resources guuid->considerable-cotask-ids constraints needs-gpus? scalar-requests)))

(defn match-offer-to-schedule
  "Given an offer and a schedule, computes all the tasks should be launched as a result.

   A schedule is just a sorted list of tasks, and we're going to greedily assign them to
   the offer.

   Returns {:matches (list of tasks that got matched to the offer)
            :failures (list of unmatched tasks, and why they weren't matched)}"
  [db ^TaskScheduler fenzo considerable offers rebalancer-reservation-atom pool-name]
  (log/info "In" pool-name "pool, matching" (count offers) "offers to" (count considerable) "jobs with fenzo")
  (log/debug "In" pool-name "pool, tasks to scheduleOnce" considerable)
  (dl/update-cost-staleness-metric considerable)
  (let [t (System/currentTimeMillis)
        leases (mapv #(->VirtualMachineLeaseAdapter % t) offers)
        considerable->task-id (plumbing.core/map-from-keys (fn [_] (str (d/squuid))) considerable)
        guuid->considerable-cotask-ids (util/make-guuid->considerable-cotask-ids considerable->task-id)
        running-cotask-cache (atom (cache/fifo-cache-factory {} :threshold (max 1 (count considerable))))
        job-uuid->reserved-host (or (:job-uuid->reserved-host @rebalancer-reservation-atom) {})
        reserved-hosts (into (hash-set) (vals job-uuid->reserved-host))
        ; Important that requests maintains the same order as considerable
        requests (mapv (fn [job]
                         (make-task-request db job
                                            :guuid->considerable-cotask-ids guuid->considerable-cotask-ids
                                            :reserved-hosts (disj reserved-hosts (job-uuid->reserved-host (:job/uuid job)))
                                            :running-cotask-cache running-cotask-cache
                                            :task-id (considerable->task-id job)))
                       considerable)
        ;; Need to lock on fenzo when accessing scheduleOnce because scheduleOnce and
        ;; task assigner can not be called at the same time.
        ;; task assigner may be called when reconciling
        result (locking fenzo
                 (.scheduleOnce fenzo requests leases))
        failure-results (.. result getFailures values)
        assignments (.. result getResultMap values)]
    (doall (map (fn [^VirtualMachineLease lease]
                  (when (-> lease :offer :reject-after-match-attempt)
                    (log/info "In" pool-name "pool, retracting lease" (-> lease :offer :id))
                    (locking fenzo
                      (.expireLease fenzo (.getId lease)))))
                leases))

    (log/debug "In" pool-name "pool, found this assignment:" result)

    {:matches (mapv (fn [assignment]
                      {:leases (.getLeasesUsed assignment)
                       :tasks (.getTasksAssigned assignment)
                       :hostname (.getHostname assignment)})
                    assignments)
     :failures failure-results}))

(meters/defmeter [cook-mesos scheduler scheduler-offer-declined])

(defn decline-offers
  "declines a collection of offer ids"
  [compute-cluster offer-ids]
  (log/debug "Declining offers:" offer-ids)
  (meters/mark! scheduler-offer-declined (count offer-ids))
  (try
    (cc/decline-offers compute-cluster offer-ids)
    (catch Throwable t
      (log/error t "Error declining offers for" compute-cluster))))

(histograms/defhistogram [cook-mesos scheduler number-tasks-matched])
(histograms/defhistogram [cook-mesos-scheduler number-offers-matched])
(meters/defmeter [cook-mesos scheduler scheduler-offer-matched])
(meters/defmeter [cook-mesos scheduler handle-resource-offer!-errors])
(def front-of-job-queue-mem-atom (atom 0))
(def front-of-job-queue-cpus-atom (atom 0))
(gauges/defgauge [cook-mesos scheduler front-of-job-queue-mem] (fn [] @front-of-job-queue-mem-atom))
(gauges/defgauge [cook-mesos scheduler front-of-job-queue-cpus] (fn [] @front-of-job-queue-cpus-atom))

(defn generate-user-usage-map
  "Returns a mapping from user to usage stats"
  [unfiltered-db pool-name]
  (timers/time!
    generate-user-usage-map-duration
    (->> (util/get-running-task-ents unfiltered-db)
         (map :job/_instance)
         (remove #(not= pool-name (util/job->pool-name %)))
         (group-by :job/user)
         (pc/map-vals (fn [jobs]
                        (->> jobs
                             (map util/job->usage)
                             (reduce (partial merge-with +))))))))

;Shared as we use this for unscheduled too.
(defonce pool->user->number-jobs (atom {}))

(defn pending-jobs->considerable-jobs
  "Limit the pending jobs to considerable jobs based on usage and quota.
   Further limit the considerable jobs to a maximum of num-considerable jobs."
  [db pending-jobs user->quota user->usage num-considerable pool-name]
  (log/debug "In" pool-name "pool, there are" (count pending-jobs) "pending jobs:" pending-jobs)
  (let [enforcing-job-launch-rate-limit? (ratelimit/enforce? ratelimit/job-launch-rate-limiter)
        user->number-jobs (atom {})
        user->rate-limit-count (atom {})
        user-within-launch-rate-limit?-fn
        (fn
          [{:keys [job/user]}]
          ; Account for each time we see a job for a user.
          (swap! user->number-jobs update user #(inc (or % 0)))
          (let [tokens (ratelimit/get-token-count! ratelimit/job-launch-rate-limiter user)
                number-jobs-for-user-so-far (@user->number-jobs user)
                is-rate-limited? (> number-jobs-for-user-so-far tokens)]
            (when is-rate-limited?
              (swap! user->rate-limit-count update user #(inc (or % 0))))
            (not (and is-rate-limited? enforcing-job-launch-rate-limit?))))
        considerable-jobs
        (->> pending-jobs
             (util/filter-based-on-quota user->quota user->usage)
             (filter (fn [job] (util/job-allowed-to-start? db job)))
             (filter user-within-launch-rate-limit?-fn)
             (filter launch-plugin/filter-job-launches)
             (take num-considerable)
             ; Force this to be taken eagerly so that the log line is accurate.
             (doall))]
    (swap! pool->user->number-jobs update pool-name (constantly @user->number-jobs))
    (log/info "Users whose job launches are rate-limited " @user->rate-limit-count "( enforcing =" enforcing-job-launch-rate-limit? ")")
    considerable-jobs))


(defn matches->job-uuids
  "Returns the matched job uuids."
  [matches pool-name]
  (let [jobs (->> matches
                  (mapcat #(-> % :tasks))
                  (map #(-> % .getRequest :job)))
        job-uuids (set (map :job/uuid jobs))]
    (log/debug "In" pool-name "pool, matched jobs:" (count job-uuids))
    (when (seq matches)
      (let [matched-normal-jobs-resource-requirements (util/sum-resources-of-jobs jobs)]
        (meters/mark! (meters/meter (metric-title "matched-tasks-cpus" pool-name))
                      (:cpus matched-normal-jobs-resource-requirements))
        (meters/mark! (meters/meter (metric-title "matched-tasks-mem" pool-name))
                      (:mem matched-normal-jobs-resource-requirements))))
    job-uuids))

(defn remove-matched-jobs-from-pending-jobs
  "Removes matched jobs from pool->pending-jobs."
  [pool-name->pending-jobs matched-job-uuids pool-name]
{'fraction_non_alphanumeric': 0.11702608276609842, 'fraction_numerical': 0.0013252213722064936, 'mean_word_length': 1.6683677569717914, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 170, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 292, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3619596', 'n_tokens_mistral': 17834, 'n_tokens_neox': 16901, 'n_words': 4545}