<commit_before>(ns onyx.peer.coordinator
  (:require [com.stuartsierra.component :as component]
            [onyx.schema :as os]
            [clojure.core.async :refer [>!! <!! poll! promise-chan sliding-buffer chan close! thread]]
            [onyx.static.planning :as planning]
            [taoensso.timbre :refer [debug info error warn trace fatal]]
            [schema.core :as s]
            [onyx.monitoring.measurements :refer [emit-latency emit-latency-value]]
            [com.stuartsierra.component :as component]
            [onyx.messaging.protocols.endpoint-status :as endpoint-status]
            [onyx.messaging.protocols.messenger :as m]
            [onyx.messaging.protocols.publisher :as pub]
            [onyx.messaging.messenger-state :as ms]
            [onyx.static.util :refer [ms->ns]]
            [onyx.peer.constants :refer [load-balance-slot-id]]
            [onyx.peer.status :refer [merge-statuses]]
            [onyx.checkpoint :as checkpoint :refer [read-checkpoint-coordinate
                                                    assume-checkpoint-coordinate
                                                    write-checkpoint-coordinate]]
            [onyx.extensions :as extensions]
            [onyx.log.replica]
            [onyx.static.default-vals :refer [arg-or-default]])
  (:import [org.apache.zookeeper KeeperException$BadVersionException]
           [java.util.concurrent.locks LockSupport]))

(defn input-publications [{:keys [peer-sites message-short-ids] :as replica} peer-id job-id]
  (let [allocations (get-in replica [:allocations job-id])
        input-tasks (get-in replica [:input-tasks job-id])
        coordinator-peer-id [:coordinator peer-id]]
    (->> input-tasks
         (mapcat (fn [task]
                   (->> (get allocations task)
                        (group-by (fn [input-peer]
                                    (get peer-sites input-peer)))
                        (map (fn [[site colocated-peers]]
                               {:src-peer-id coordinator-peer-id
                                :job-id job-id
                                :dst-task-id [job-id task]
                                :dst-peer-ids (set colocated-peers)
                                :batch-size 0
                                :short-id (get message-short-ids
                                               {:src-peer-type :coordinator
                                                :src-peer-id peer-id
                                                :job-id job-id
                                                :dst-task-id task
                                                :msg-slot-id load-balance-slot-id})
                                :slot-id load-balance-slot-id
                                :site site})))))
         (set))))

(defn offer-heartbeats
  [{:keys [messenger] :as state}]
  (run! pub/offer-heartbeat! (m/publishers messenger))
  (assoc state :last-heartbeat-time (System/nanoTime)))

(def coordinator-backoff-ms 1)

(defn offer-barriers [{:keys [messenger barrier] :as state}]
  (if (:offering? barrier)
    (let [offer-xf (comp (map (fn [pub]
                                [(m/offer-barrier messenger pub (:opts barrier))
                                 pub]))
                         (remove (comp pos? first))
                         (map second))
          new-remaining (sequence offer-xf (:remaining barrier))]
      (if (empty? new-remaining)
        (update state :barrier merge {:remaining nil
                                      :offering? false})
        (do ;; park and retry
            (LockSupport/parkNanos (ms->ns coordinator-backoff-ms))
            (update state :barrier merge {:remaining new-remaining}))))
    state))

(defn write-coordinate [curr-version log tenancy-id job-id coordinate]
  (try (->> curr-version
            (write-checkpoint-coordinate log tenancy-id job-id coordinate)
            (:version))
       (catch KeeperException$BadVersionException bve
         (throw (Exception. "Coordinator failed to write coordinates.
                             This is likely due to the coordinator being quarantined, and another coordinator taking over.")))))

(defn next-replica
  [{:keys [log job-id peer-id messenger curr-replica tenancy-id] :as state}
   barrier-period-ns
   new-replica]
  (let [curr-version (get-in curr-replica [:allocation-version job-id])
        new-version (get-in new-replica [:allocation-version job-id])
        reallocated? (not= curr-version new-version)]
    (cond reallocated?
          (let [new-messenger (-> messenger
                                  (m/update-publishers (input-publications new-replica peer-id job-id))
                                  (m/set-replica-version! new-version)
                                  (m/set-epoch! 0))
                coordinates (read-checkpoint-coordinate log tenancy-id job-id)]
            (-> state
                (update :job merge {:completed? false
                                    :sealing? false})
                (update :barrier merge {:scheduled? false
                                        :offering? true
                                        :remaining (m/publishers new-messenger)
                                        :opts {:recover-coordinates coordinates}})
                (update :checkpoint merge {:epoch 0})
                (assoc :messenger new-messenger)
                (assoc :curr-replica new-replica)))

          :else
          (assoc state :curr-replica new-replica))))

(defn complete-job
  [{:keys [tenancy-id log job-id messenger checkpoint group-ch] :as state}]
  (info (format "Job %s completed, and final checkpoint has finished. Writing checkpoint coordinates." job-id))
  (let [replica-version (m/replica-version messenger)
        epoch (m/epoch messenger)]
    (let [coordinates {:tenancy-id tenancy-id 
                       :job-id job-id 
                       :replica-version replica-version 
                       :epoch epoch}
          {:keys [write-version]} checkpoint
          next-write-version (write-coordinate write-version log tenancy-id job-id coordinates)]
      (>!! group-ch [:send-to-outbox {:fn :complete-job :args {:job-id job-id}}])
      (-> state
          (update :job merge {:completed? true
                              :sealing? false})
          (update :checkpoint merge {:write-version next-write-version})))))

(defn merged-statuses [messenger]
  (->> (m/publishers messenger)
       (mapcat (comp endpoint-status/statuses pub/endpoint-status))
       (map val)
       (merge-statuses)))

(defn periodic-barrier
  [{:keys [tenancy-id workflow-depth log curr-replica job-id messenger barrier checkpoint] :as state}]
  (if (:offering? barrier)
    ;; No op because hasn't finished emitting last barrier, wait again
    state
    (let [;; always checkpoints at the moment
          start-checkpoint? (first (shuffle [true #_false]))
          messenger (m/set-epoch! messenger (inc (m/epoch messenger)))]
      (-> state
          (update :checkpoint merge {:initiated? true
                                     :epoch (if start-checkpoint? (m/epoch messenger) (:epoch checkpoint))})
          (update :barrier merge {:scheduled? false
                                  :offering? true
                                  :remaining (m/publishers messenger)
                                  :opts {:completed? (:sealing? (:job state))
                                         :checkpoint? start-checkpoint?}})
          (assoc :messenger messenger)))))

(defn shutdown [{:keys [messenger] :as state}]
  (assoc state :messenger (component/stop messenger)))

(defn initialise-state [{:keys [log job-id tenancy-id] :as state}]
  (let [write-version (assume-checkpoint-coordinate log tenancy-id job-id)]
    (-> state
        (assoc-in [:checkpoint :write-version] write-version)
        (assoc :last-heartbeat-time (System/nanoTime)))))

(defn checkpoint-complete?
  [{:keys [initiated? epoch] :as checkpoint} status]
  (or (not initiated?)
      (and (not (:checkpointing? status))
           (>= (:min-epoch status) epoch))))

(defn completed-checkpoint [{:keys [checkpoint messenger job-id tenancy-id log] :as state}]
  (let [{:keys [epoch write-version]} checkpoint
        write-coordinate? (> epoch 0)
        coordinates {:tenancy-id tenancy-id
                     :job-id job-id
                     :replica-version (m/replica-version messenger)
                     :epoch epoch}
        ;; get the next version of the zk node, so we can detect when there are other writers
        next-write-version (if write-coordinate?
                             (write-coordinate write-version log tenancy-id job-id coordinates)
                             write-version)]


    (-> state
        (update :barrier    merge {:scheduled? false})
        (update :checkpoint merge {:initiated? false
                                   :write-version next-write-version}))))

(defn schedule-next-barrier [state barrier-period-ns]
  (update state :barrier merge {:scheduled? true
                                :next-barrier-time (+ (System/nanoTime) barrier-period-ns)}))

(defn emit-final-completion-barrier [state]
  (-> state
      (assoc-in [:job :sealing?] true)
      (periodic-barrier)))

(defn coordinator-iteration
  [{:keys [messenger checkpoint last-heartbeat-time allocation-ch shutdown-ch barrier job job-id] :as state}
   coordinator-max-sleep-ns
   barrier-period-ns
   heartbeat-ns]
  (let [_ (run! pub/poll-heartbeats! (m/publishers messenger))
        status (merged-statuses messenger)
        {:keys [sealing? completed?]} job]
    (cond (> (System/nanoTime) (+ last-heartbeat-time heartbeat-ns))
          ;; Immediately offer heartbeats
          (offer-heartbeats state)
<commit_msg>Minor coordinator cleanup and defensive checks.<commit_after>(ns onyx.peer.coordinator
  (:require [com.stuartsierra.component :as component]
            [onyx.schema :as os]
            [clojure.core.async :refer [>!! <!! poll! promise-chan sliding-buffer chan close! thread]]
            [onyx.static.planning :as planning]
            [taoensso.timbre :refer [debug info error warn trace fatal]]
            [schema.core :as s]
            [onyx.monitoring.measurements :refer [emit-latency emit-latency-value]]
            [com.stuartsierra.component :as component]
            [onyx.messaging.protocols.endpoint-status :as endpoint-status]
            [onyx.messaging.protocols.messenger :as m]
            [onyx.messaging.protocols.publisher :as pub]
            [onyx.messaging.messenger-state :as ms]
            [onyx.static.util :refer [ms->ns]]
            [onyx.peer.constants :refer [load-balance-slot-id]]
            [onyx.peer.status :refer [merge-statuses]]
            [onyx.checkpoint :as checkpoint :refer [read-checkpoint-coordinate
                                                    assume-checkpoint-coordinate
                                                    write-checkpoint-coordinate]]
            [onyx.extensions :as extensions]
            [onyx.log.replica]
            [onyx.static.default-vals :refer [arg-or-default]])
  (:import [org.apache.zookeeper KeeperException$BadVersionException]
           [java.util.concurrent.locks LockSupport]))

(defn input-publications [{:keys [peer-sites message-short-ids] :as replica} peer-id job-id]
  (let [allocations (get-in replica [:allocations job-id])
        input-tasks (get-in replica [:input-tasks job-id])
        coordinator-peer-id [:coordinator peer-id]]
    (->> input-tasks
         (mapcat (fn [task]
                   (->> (get allocations task)
                        (group-by (fn [input-peer]
                                    (get peer-sites input-peer)))
                        (map (fn [[site colocated-peers]]
                               {:src-peer-id coordinator-peer-id
                                :job-id job-id
                                :dst-task-id [job-id task]
                                :dst-peer-ids (set colocated-peers)
                                :batch-size 0
                                :short-id (get message-short-ids
                                               {:src-peer-type :coordinator
                                                :src-peer-id peer-id
                                                :job-id job-id
                                                :dst-task-id task
                                                :msg-slot-id load-balance-slot-id})
                                :slot-id load-balance-slot-id
                                :site site})))))
         (set))))

(defn offer-heartbeats
  [{:keys [messenger] :as state}]
  (run! pub/offer-heartbeat! (m/publishers messenger))
  (assoc state :last-heartbeat-time (System/nanoTime)))

(def coordinator-backoff-ms 1)

(defn offer-barriers [{:keys [messenger barrier] :as state}]
  (if (:offering? barrier)
    (let [offer-xf (comp (map (fn [pub]
                                [(m/offer-barrier messenger pub (:opts barrier))
                                 pub]))
                         (remove (comp pos? first))
                         (map second))
          new-remaining (sequence offer-xf (:remaining barrier))]
      (if (empty? new-remaining)
        (update state :barrier merge {:remaining nil
                                      :offering? false})
        (do ;; park and retry
            (LockSupport/parkNanos (ms->ns coordinator-backoff-ms))
            (update state :barrier merge {:remaining new-remaining}))))
    state))

(defn write-coordinate [curr-version log tenancy-id job-id coordinate]
  (try (->> curr-version
            (write-checkpoint-coordinate log tenancy-id job-id coordinate)
            (:version))
       (catch KeeperException$BadVersionException bve
         (throw (Exception. "Coordinator failed to write coordinates.
                             This is likely due to the coordinator being quarantined, and another coordinator taking over.")))))

(defn next-replica
  [{:keys [log job-id peer-id messenger curr-replica tenancy-id] :as state}
   barrier-period-ns
   new-replica]
  (let [curr-version (get-in curr-replica [:allocation-version job-id])
        new-version (get-in new-replica [:allocation-version job-id])
        job-reallocated? (not= curr-version new-version)]
    (if job-reallocated?
      (let [new-messenger (-> messenger
                              (m/update-publishers (input-publications new-replica peer-id job-id))
                              (m/set-replica-version! new-version)
                              (m/set-epoch! 0))
            coordinates (read-checkpoint-coordinate log tenancy-id job-id)]
        (-> state
            (update :job merge {:completed? false
                                :sealing? false})
            (update :barrier merge {:scheduled? false
                                    :offering? true
                                    :remaining (m/publishers new-messenger)
                                    :opts {:recover-coordinates coordinates}})
            (update :checkpoint merge {:initiated? false
                                       :epoch 0})
            (assoc :messenger new-messenger)
            (assoc :curr-replica new-replica)))
      (assoc state :curr-replica new-replica))))

(defn complete-job
  [{:keys [tenancy-id log job-id messenger checkpoint group-ch] :as state}]
  (info (format "Job %s completed, and final checkpoint has finished. Writing checkpoint coordinates." job-id))
  (let [replica-version (m/replica-version messenger)
        epoch (m/epoch messenger)]
    (let [coordinates {:tenancy-id tenancy-id 
                       :job-id job-id 
                       :replica-version replica-version 
                       :epoch epoch}
          {:keys [write-version]} checkpoint
          next-write-version (write-coordinate write-version log tenancy-id job-id coordinates)]
      (>!! group-ch [:send-to-outbox {:fn :complete-job :args {:job-id job-id}}])
      (-> state
          (update :job merge {:completed? true
                              :sealing? false})
          (update :checkpoint merge {:write-version next-write-version})))))

(defn merged-statuses [messenger]
  (->> (m/publishers messenger)
       (mapcat (comp endpoint-status/statuses pub/endpoint-status))
       (map val)
       (merge-statuses)))

(defn periodic-barrier
  [{:keys [tenancy-id workflow-depth log curr-replica job-id messenger barrier checkpoint] :as state}]
  (if (:offering? barrier)
    ;; No op because hasn't finished emitting last barrier, wait again
    state
    (let [;; always checkpoints at the moment
          start-checkpoint? (first (shuffle [true #_false]))
          messenger (m/set-epoch! messenger (inc (m/epoch messenger)))]
      (-> state
          (update :checkpoint merge {:initiated? true
                                     :epoch (if start-checkpoint? (m/epoch messenger) (:epoch checkpoint))})
          (update :barrier merge {:scheduled? false
                                  :offering? true
                                  :remaining (m/publishers messenger)
                                  :opts {:completed? (:sealing? (:job state))
                                         :checkpoint? start-checkpoint?}})
          (assoc :messenger messenger)))))

(defn shutdown [{:keys [messenger] :as state}]
  (assoc state :messenger (component/stop messenger)))

(defn initialise-state [{:keys [log job-id tenancy-id] :as state}]
  (let [write-version (assume-checkpoint-coordinate log tenancy-id job-id)]
    (-> state
        (assoc-in [:checkpoint :write-version] write-version)
        (assoc :last-heartbeat-time (System/nanoTime)))))

(defn checkpoint-complete?
  [{:keys [initiated? epoch] :as checkpoint} status]
  (or (not initiated?)
      (and (not (:checkpointing? status))
           (>= (:min-epoch status) epoch))))

(defn completed-checkpoint [{:keys [checkpoint messenger job-id tenancy-id log] :as state}]
  (let [{:keys [epoch write-version]} checkpoint
        write-coordinate? (> epoch 0)
        coordinates {:tenancy-id tenancy-id
                     :job-id job-id
                     :replica-version (m/replica-version messenger)
                     :epoch epoch}
        ;; get the next version of the zk node, so we can detect when there are other writers
        next-write-version (if write-coordinate?
                             (write-coordinate write-version log tenancy-id job-id coordinates)
                             write-version)]
    (-> state
        (update :barrier    merge {:scheduled? false})
        (update :checkpoint merge {:initiated? false
                                   :write-version next-write-version}))))

(defn schedule-next-barrier [state barrier-period-ns]
  (update state :barrier merge {:scheduled? true
                                :next-barrier-time (+ (System/nanoTime) barrier-period-ns)}))

(defn emit-final-completion-barrier [state]
  (-> state
      (assoc-in [:job :sealing?] true)
      (periodic-barrier)))

(defn coordinator-iteration
  [{:keys [messenger checkpoint last-heartbeat-time allocation-ch shutdown-ch barrier job job-id] :as state}
   coordinator-max-sleep-ns
   barrier-period-ns
   heartbeat-ns]
  (let [_ (run! pub/poll-heartbeats! (m/publishers messenger))
        status (merged-statuses messenger)
        {:keys [sealing? completed?]} job]
    (cond (> (System/nanoTime) (+ last-heartbeat-time heartbeat-ns))
          ;; Immediately offer heartbeats
          (offer-heartbeats state)

{'fraction_non_alphanumeric': 0.11531402936378467, 'fraction_numerical': 0.0005097879282218597, 'mean_word_length': 1.761402027027027, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 47, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 61, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5912103', 'n_tokens_mistral': 5554, 'n_tokens_neox': 5294, 'n_words': 1376}