<commit_before>(ns onyx.plugin.bench-plugin
  (:require [clojure.core.async :refer [chan >!! <!! close! alts!! timeout]]
            [clojure.data.fressian :as fressian]
            [onyx.static.default-vals :refer [defaults]]
            [onyx.peer.pipeline-extensions :as p-ext]))

(def hundred-bytes 
  (into-array Byte/TYPE (range 100)))

(defn inject-reader
  [event lifecycle]
  (assert (:core.async/chan event) ":core.async/chan not found - add it via inject-lifecycle-resources.")
  {:generator/pending-messages (atom {})
   :generator/retry (atom [])})

(def reader-calls
  {:lifecycle/before-task-start inject-reader})

(defn flush-swap! [a f-read f-swap]
  (loop []
    (let [v (deref a)]
      (if (compare-and-set! a v (f-swap v))
        (f-read v)
        (recur)))))

(defmethod p-ext/read-batch :generator
  [{:keys [onyx.core/task-map
           generator/pending-messages
           generator/retry] :as event}]
  (let [pending (count @pending-messages)
        max-pending (or (:onyx/max-pending task-map) (:onyx/max-pending defaults))
        batch-size (:onyx/batch-size task-map)
        max-segments (min (- max-pending pending) batch-size)
        segments (->> (flush-swap! retry 
                                   #(take max-segments %)
                                   #(subvec % (min max-segments (count %))))
                      (map (fn [m] {:id (java.util.UUID/randomUUID)
                                    :input :generator
                                    :message m})))
        batch (loop [n (count segments) <commit_msg>Channel not required for benchmark plugin<commit_after>(ns onyx.plugin.bench-plugin
  (:require [clojure.core.async :refer [chan >!! <!! close! alts!! timeout]]
            [clojure.data.fressian :as fressian]
            [onyx.static.default-vals :refer [defaults]]
            [onyx.peer.pipeline-extensions :as p-ext]))

(def hundred-bytes 
  (into-array Byte/TYPE (range 100)))

(defn inject-reader
  [event lifecycle]
  {:generator/pending-messages (atom {})
   :generator/retry (atom [])})

(def reader-calls
  {:lifecycle/before-task-start inject-reader})

(defn flush-swap! [a f-read f-swap]
  (loop []
    (let [v (deref a)]
      (if (compare-and-set! a v (f-swap v))
        (f-read v)
        (recur)))))

(defmethod p-ext/read-batch :generator
  [{:keys [onyx.core/task-map
           generator/pending-messages
           generator/retry] :as event}]
  (let [pending (count @pending-messages)
        max-pending (or (:onyx/max-pending task-map) (:onyx/max-pending defaults))
        batch-size (:onyx/batch-size task-map)
        max-segments (min (- max-pending pending) batch-size)
        segments (->> (flush-swap! retry 
                                   #(take max-segments %)
                                   #(subvec % (min max-segments (count %))))
                      (map (fn [m] {:id (java.util.UUID/randomUUID)
                                    :input :generator
                                    :message m})))
        batch (loop [n (count segments) 
{'fraction_non_alphanumeric': 0.14790575916230367, 'fraction_numerical': 0.001963350785340314, 'mean_word_length': 2.0177690029615003, 'pattern_counts': {'":': 1, '<': 5, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4130325', 'n_tokens_mistral': 978, 'n_tokens_neox': 959, 'n_words': 225}