<commit_before>module Mouth
  
  class SuckerConnection < EM::Connection
    attr_accessor :sucker
  
    def receive_data(data)
      Mouth.logger.debug "UDP packet: '#{data}'"
  
      sucker.store!(data)
    end
  end
  
  class Sucker
    
    # Host/Port to suck UDP packets on
    attr_accessor :host
    attr_accessor :port

    # Actual EM::Mongo connection
    attr_accessor :mongo_db
    
    # Info to connect to mongo
    attr_accessor :mongo_db_name
    attr_accessor :mongo_hosts
    
    # Accumulators of our data
    attr_accessor :counters
    attr_accessor :timers
    
    def initialize(options = {})
      self.host = options[:host] || "localhost"
      self.port = options[:port] || 8889
      self.mongo_db_name = options[:mongo_db_name] || "mouth"
      self.mongo_hosts = options[:mongo_hosts] || ["localhost"]
      
      self.counters = {}
      self.timers = {}
    end
    
    def suck!
      EM.run do
        # Connect to mongo now
        self.mongo_db

        EM.open_datagram_socket host, port, SuckerConnection do |conn|
          conn.sucker = self
        end
        
        EM.add_periodic_timer(10) do
          Mouth.logger.info "Counters: #{self.counters.inspect}"
          Mouth.logger.info "Timers: #{self.timers.inspect}"
          self.flush!
        end

        EM.next_tick do
          Mouth.logger.info "Mouth reactor started..."
        end
      end
    end
    
    def store!(data)
      command, key, value = data.split(":")
      key = sanitize_key(key)
      value = value.to_f
      
      ts = minute_timestamps
      
      if command == "ms"
        self.timers[ts] ||= {}
        self.timers[ts][key] ||= []
        self.timers[ts][key] << value
      elsif command == "c"
        self.counters[ts] ||= {}
        self.counters[ts][key] ||= 0.0
        self.counters[ts][key] += value
      end
    end
    
    def flush!
      ts = minute_timestamps
      limit_ts = ts - 1
      mongo_docs = {}
      
      # We're going to construct mongo_docs which look like this:
      # "mycollections:234234": {  # NOTE: this timpstamp will be popped into .t = 234234
      #   c: {
      #     poopings: 37,
      #     shittings: 3
      #   },
      #   ms: {
      #     peeings: {...}
      #   }
      # }
      
      self.counters.each do |cur_ts, counters_to_save|
        if cur_ts <= limit_ts
          counters_to_save.each do |counter_key, value|
            ns, sub_key = parse_key(counter_key)
            mongo_key = "#{ns}:#{ts}"
            mongo_docs[mongo_key] ||= {}
            
            cur_mongo_doc = mongo_docs[mongo_key]
            cur_mongo_doc["c"] ||= {}
            cur_mongo_doc["c"][sub_key] = value
          end
          
          self.counters.delete(cur_ts)
        end
      end
      
      self.timers.each do |cur_ts, timers_to_save|
        if cur_ts <= limit_ts
          timers_to_save.each do |timer_key, values|
            ns, sub_key = parse_key(timer_key)
            mongo_key = "#{ns}:#{ts}"
            mongo_docs[mongo_key] ||= {}
            
            cur_mongo_doc = mongo_docs[mongo_key]
            cur_mongo_doc["ms"] ||= {}
            cur_mongo_doc["ms"][sub_key] = analyze_timer(values)
          end
          
          self.timers.delete(cur_ts)
        end
      end
      
      save_documents!(mongo_docs)
    end
    
    def save_documents!(mongo_docs)
      Mouth.logger.info "Saving Docs: #{mongo_docs.inspect}"
      
      mongo_docs.each do |key, doc|
        ns, ts = key.split(":")
        collection_name = "mouth_#{ns}"
        doc["t"] = ts
        
        self.mongo_db.collection(collection_name).insert(doc)
      end
    end
    
    def mongo_db
      @mongo_db ||= begin
        if self.mongo_hosts.length == 1
          EM::Mongo::Connection.new(self.mongo_hosts.first).db(self.mongo_db_name)
        else
          raise "TODO: ability to connect to a replica set."
        end
      end
    end
    
    private
    
    def minute_timestamps
      Time.now.to_i / 60
    end
    
    def sanitize_key(key)
      key.gsub(/\s+/, '_').gsub(/\//, '-').gsub(/[^a-zA-Z_\-0-9\.]/, '')<commit_msg>Store timestamps in integer format<commit_after>module Mouth
  
  class SuckerConnection < EM::Connection
    attr_accessor :sucker
  
    def receive_data(data)
      Mouth.logger.debug "UDP packet: '#{data}'"
  
      sucker.store!(data)
    end
  end
  
  class Sucker
    
    # Host/Port to suck UDP packets on
    attr_accessor :host
    attr_accessor :port

    # Actual EM::Mongo connection
    attr_accessor :mongo_db
    
    # Info to connect to mongo
    attr_accessor :mongo_db_name
    attr_accessor :mongo_hosts
    
    # Accumulators of our data
    attr_accessor :counters
    attr_accessor :timers
    
    def initialize(options = {})
      self.host = options[:host] || "localhost"
      self.port = options[:port] || 8889
      self.mongo_db_name = options[:mongo_db_name] || "mouth"
      self.mongo_hosts = options[:mongo_hosts] || ["localhost"]
      
      self.counters = {}
      self.timers = {}
    end
    
    def suck!
      EM.run do
        # Connect to mongo now
        self.mongo_db

        EM.open_datagram_socket host, port, SuckerConnection do |conn|
          conn.sucker = self
        end
        
        EM.add_periodic_timer(10) do
          Mouth.logger.info "Counters: #{self.counters.inspect}"
          Mouth.logger.info "Timers: #{self.timers.inspect}"
          self.flush!
        end

        EM.next_tick do
          Mouth.logger.info "Mouth reactor started..."
        end
      end
    end
    
    def store!(data)
      command, key, value = data.split(":")
      key = sanitize_key(key)
      value = value.to_f
      
      ts = minute_timestamps
      
      if command == "ms"
        self.timers[ts] ||= {}
        self.timers[ts][key] ||= []
        self.timers[ts][key] << value
      elsif command == "c"
        self.counters[ts] ||= {}
        self.counters[ts][key] ||= 0.0
        self.counters[ts][key] += value
      end
    end
    
    def flush!
      ts = minute_timestamps
      limit_ts = ts - 1
      mongo_docs = {}
      
      # We're going to construct mongo_docs which look like this:
      # "mycollections:234234": {  # NOTE: this timpstamp will be popped into .t = 234234
      #   c: {
      #     poopings: 37,
      #     shittings: 3
      #   },
      #   ms: {
      #     peeings: {...}
      #   }
      # }
      
      self.counters.each do |cur_ts, counters_to_save|
        if cur_ts <= limit_ts
          counters_to_save.each do |counter_key, value|
            ns, sub_key = parse_key(counter_key)
            mongo_key = "#{ns}:#{ts}"
            mongo_docs[mongo_key] ||= {}
            
            cur_mongo_doc = mongo_docs[mongo_key]
            cur_mongo_doc["c"] ||= {}
            cur_mongo_doc["c"][sub_key] = value
          end
          
          self.counters.delete(cur_ts)
        end
      end
      
      self.timers.each do |cur_ts, timers_to_save|
        if cur_ts <= limit_ts
          timers_to_save.each do |timer_key, values|
            ns, sub_key = parse_key(timer_key)
            mongo_key = "#{ns}:#{ts}"
            mongo_docs[mongo_key] ||= {}
            
            cur_mongo_doc = mongo_docs[mongo_key]
            cur_mongo_doc["ms"] ||= {}
            cur_mongo_doc["ms"][sub_key] = analyze_timer(values)
          end
          
          self.timers.delete(cur_ts)
        end
      end
      
      save_documents!(mongo_docs)
    end
    
    def save_documents!(mongo_docs)
      Mouth.logger.info "Saving Docs: #{mongo_docs.inspect}"
      
      mongo_docs.each do |key, doc|
        ns, ts = key.split(":")
        collection_name = "mouth_#{ns}"
        doc["t"] = ts.to_i
        
        self.mongo_db.collection(collection_name).insert(doc)
      end
    end
    
    def mongo_db
      @mongo_db ||= begin
        if self.mongo_hosts.length == 1
          EM::Mongo::Connection.new(self.mongo_hosts.first).db(self.mongo_db_name)
        else
          raise "TODO: ability to connect to a replica set."
        end
      end
    end
    
    private
    
    def minute_timestamps
      Time.now.to_i / 60
    end
    
    def sanitize_key(key)
      key.gsub(/\s+/, '_').gsub(/\//, '-').gsub(/[^a-zA-Z_\-0-9\.]/, '')
{'fraction_non_alphanumeric': 0.11443323727185399, 'fraction_numerical': 0.006964457252641691, 'mean_word_length': 1.7344057780695994, 'pattern_counts': {'":': 6, '<': 13, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1832627', 'n_tokens_mistral': 2903, 'n_tokens_neox': 2826, 'n_words': 550}