<commit_before>(ns onyx.plugin.core-async
  (:require [clojure.core.async :refer [chan >!! <!! alts!! timeout go <!]]
            [onyx.peer.function :as function]
            [onyx.peer.pipeline-extensions :as p-ext]
            [onyx.static.default-vals :refer [defaults]]
            [taoensso.timbre :refer [debug info] :as timbre]))

(defn inject-reader
  [event lifecycle]
  (assert (:core.async/chan event) 
          ":core.async/chan not found - add it using a :before-task-start lifecycle")
  (let [pipeline (:onyx.core/pipeline event)] 
    {:core.async/pending-messages (:pending-messages pipeline) 
     :core.async/drained (:drained pipeline)
     :core.async/retry-ch (:retry-ch pipeline)
     :core.async/retry-count (:retry-count pipeline)}))

(defn log-retry-count
  [event lifecycle]
  (info "core.async input plugin stopping. Retry count:" @(:core.async/retry-count event))
  {})

(defn inject-writer
  [event lifecycle]
  (assert (:core.async/chan event) ":core.async/chan not found - add it using a :before-task-start lifecycle")
  {})

(def reader-calls
  {:lifecycle/before-task-start inject-reader
   :lifecycle/after-task-stop log-retry-count})

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

(defrecord CoreAsyncInput [max-pending batch-size batch-timeout pending-messages 
                           drained retry-ch retry-count]
  p-ext/Pipeline
  (write-batch 
    [this event]
    (function/write-batch event))

  (read-batch [_ {:keys [core.async/chan] :as event}]
    (let [pending (count @pending-messages)
          max-segments (min (- max-pending pending) batch-size)
          timeout-ch (timeout batch-timeout)
          batch (if (pos? max-segments)
                  (loop [segments [] cnt 0]
                    (if (= cnt max-segments)
                      segments
                      (if-let [message (first (alts!! [retry-ch chan timeout-ch] :priority true))] 
                        (recur (conj segments 
                                     {:id (java.util.UUID/randomUUID)
                                      :input :core.async
                                      :message message})
                               (inc cnt))
                        segments)))
                  (<!! timeout-ch))]
      (doseq [m batch]
        (swap! pending-messages assoc (:id m) (:message m)))
      (when (and (= 1 (count @pending-messages))
                 (= (count batch) 1)
                 (= (:message (first batch)) :done))
        (reset! drained true))
      {:onyx.core/batch batch}))

  p-ext/PipelineInput

  (ack-message [_ _ message-id]
    (swap! pending-messages dissoc message-id))

  (retry-message 
    [_ _ message-id]
    (when-let [msg (get @pending-messages message-id)]
      (swap! pending-messages dissoc message-id)
      (when-not (= msg :done)
        (swap! retry-count inc))
      (>!! retry-ch msg)))

  (pending?
    [_ _ message-id]
    (get @pending-messages message-id))

  (drained? 
    [_ _]<commit_msg>Use input record type in the core.async plugin<commit_after>(ns onyx.plugin.core-async
  (:require [clojure.core.async :refer [chan >!! <!! alts!! timeout go <!]]
            [onyx.peer.function :as function]
            [onyx.peer.pipeline-extensions :as p-ext]
            [onyx.static.default-vals :refer [defaults]]
            [onyx.types :as t]
            [taoensso.timbre :refer [debug info] :as timbre]))

(defn inject-reader
  [event lifecycle]
  (assert (:core.async/chan event) 
          ":core.async/chan not found - add it using a :before-task-start lifecycle")
  (let [pipeline (:onyx.core/pipeline event)] 
    {:core.async/pending-messages (:pending-messages pipeline) 
     :core.async/drained (:drained pipeline)
     :core.async/retry-ch (:retry-ch pipeline)
     :core.async/retry-count (:retry-count pipeline)}))

(defn log-retry-count
  [event lifecycle]
  (info "core.async input plugin stopping. Retry count:" @(:core.async/retry-count event))
  {})

(defn inject-writer
  [event lifecycle]
  (assert (:core.async/chan event) ":core.async/chan not found - add it using a :before-task-start lifecycle")
  {})

(def reader-calls
  {:lifecycle/before-task-start inject-reader
   :lifecycle/after-task-stop log-retry-count})

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

(defrecord CoreAsyncInput [max-pending batch-size batch-timeout pending-messages 
                           drained retry-ch retry-count]
  p-ext/Pipeline
  (write-batch 
    [this event]
    (function/write-batch event))

  (read-batch [_ {:keys [core.async/chan] :as event}]
    (let [pending (count @pending-messages)
          max-segments (min (- max-pending pending) batch-size)
          timeout-ch (timeout batch-timeout)
          batch (if (pos? max-segments)
                  (loop [segments [] cnt 0]
                    (if (= cnt max-segments)
                      segments
                      (if-let [message (first (alts!! [retry-ch chan timeout-ch] :priority true))] 
                        (recur (conj segments 
                                     (t/input (java.util.UUID/randomUUID)
                                              message))
                               (inc cnt))
                        segments)))
                  (<!! timeout-ch))]
      (doseq [m batch]
        (swap! pending-messages assoc (:id m) (:message m)))
      (when (and (= 1 (count @pending-messages))
                 (= (count batch) 1)
                 (= (:message (first batch)) :done))
        (reset! drained true))
      {:onyx.core/batch batch}))

  p-ext/PipelineInput

  (ack-message [_ _ message-id]
    (swap! pending-messages dissoc message-id))

  (retry-message 
    [_ _ message-id]
    (when-let [msg (get @pending-messages message-id)]
      (swap! pending-messages dissoc message-id)
      (when-not (= msg :done)
        (swap! retry-count inc))
      (>!! retry-ch msg)))

  (pending?
    [_ _ message-id]
    (get @pending-messages message-id))

  (drained? 
    [_ _]
{'fraction_non_alphanumeric': 0.13681592039800994, 'fraction_numerical': 0.0009950248756218905, 'mean_word_length': 2.0676500508647, 'pattern_counts': {'":': 4, '<': 9, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6118443', 'n_tokens_mistral': 1884, 'n_tokens_neox': 1825, 'n_words': 447}