<commit_before>require 'eventmachine'
require 'beefcake'

require 'riemann/query'
require 'riemann/attribute'
require 'riemann/state'
require 'riemann/event'
require 'riemann/message'

require_relative '../protocol'
require_relative '../plugin'
require_relative '../logging'

require_relative 'riemann/utils'
require_relative 'riemann/handler'

module EVD::Plugin::Riemann
  include EVD::Plugin
  include EVD::Logging

  register_plugin "riemann"

  class HandlerTCP
    include EVD::Logging
    include EVD::Plugin::Riemann::Utils
    include EVD::Plugin::Riemann::Handler

    def initialize
      setup_handler
    end

    def encode(m)
      m.encode_with_length
    end
  end

  class HandlerUDP
    include EVD::Logging
    include EVD::Plugin::Riemann::Utils
    include EVD::Plugin::Riemann::Handler

    def initialize
      setup_handler
    end

    def encode(m)
      m.encode
    end
  end

  class ConnectionBase < EM::Connection
    include EVD::Logging
    include EVD::Plugin::Riemann::Utils
    include EM::Protocols::ObjectProtocol

    module RiemannSerializer
      def self.dump(m)
        m.encode.to_s
      end

      def self.load(data)
        ::Riemann::Message.decode(data)
      end
    end

    def initialize(channel, log)
      @channel = channel
      @log = log
    end

    def serializer
      RiemannSerializer
    end

    def receive_object(m)
      unless m.events.nil? or m.events.empty?
        m.events.each do |e|
          @channel.event read_event(e)
        end
      end

      send_ok
    rescue => e
      @log.error "Failed to receive object", e
      send_error e
    end

    protected

    def send_ok; end
    def send_error(e); end
  end

  class ConnectionTCP < ConnectionBase
    def send_ok
      send_object(::Riemann::Message.new(
        :ok => true))
    end

    def send_error(e)
      send_object(::Riemann::Message.new(
        :ok => false, :error => e.to_s))
    end
  end

  class ConnectionUDP < ConnectionBase; end

  DEFAULT_HOST = "localhost"
  DEFAULT_PORT = 5555
  DEFAULT_PROTOCOL = 'tcp'

  HANDLERS = {:tcp => HandlerTCP, :udp => HandlerUDP}
  CONNECTIONS = {:tcp => ConnectionTCP, :udp => ConnectionUDP}

  def self.output_setup(opts={})
    opts[:host] ||= DEFAULT_HOST
    opts[:port] ||= DEFAULT_PORT

    protocol = EVD.parse_protocol(opts[:protocol] || DEFAULT_PROTOCOL)

    unless handler = HANDLERS[protocol.family]
      raise "No handler for protocol family: #{protocol.family}"
    end

    handler_instance = handler.new
    protocol.connect log, opts, handler_instance
  end

  def self.input_setup(opts={})
    opts[:host] ||= DEFAULT_HOST
    opts[:port] ||= DEFAULT_PORT
    protocol = EVD.parse_protocol(opts[:protocol] || DEFAULT_PROTOCOL)

    unless connection = CONNECTIONS[protocol.family]
      raise "No connection for protocol family: #{protocol.family}"
    end

    protocol.listen log, opts, connection, log
  end
end<commit_msg>riemann: Cleanup<commit_after>require 'eventmachine'
require 'beefcake'

require 'riemann/query'
require 'riemann/attribute'
require 'riemann/state'
require 'riemann/event'
require 'riemann/message'

require_relative '../protocol'
require_relative '../plugin'
require_relative '../logging'

require_relative 'riemann/utils'
require_relative 'riemann/handler'

module EVD::Plugin::Riemann
  include EVD::Plugin
  include EVD::Logging

  register_plugin "riemann"

  class HandlerTCP
    include EVD::Plugin::Riemann::Utils
    include EVD::Plugin::Riemann::Handler

    def initialize
      setup_handler
    end

    def encode(m)
      m.encode_with_length
    end
  end

  class HandlerUDP
    include EVD::Plugin::Riemann::Utils
    include EVD::Plugin::Riemann::Handler

    def initialize
      setup_handler
    end

    def encode(m)
      m.encode
    end
  end

  class ConnectionBase < EM::Connection
    include EVD::Plugin::Riemann::Utils
    include EM::Protocols::ObjectProtocol

    module RiemannSerializer
      def self.dump(m)
        m.encode.to_s
      end

      def self.load(data)
        ::Riemann::Message.decode(data)
      end
    end

    def initialize(channel, log)
      @channel = channel
      @log = log
    end

    def serializer
      RiemannSerializer
    end

    def receive_object(m)
      unless m.events.nil? or m.events.empty?
        m.events.each do |e|
          @channel.event read_event(e)
        end
      end

      send_ok
    rescue => e
      @log.error "Failed to receive object", e
      send_error e
    end

    protected

    def send_ok; end
    def send_error(e); end
  end

  class ConnectionTCP < ConnectionBase
    def send_ok
      send_object(::Riemann::Message.new(
        :ok => true))
    end

    def send_error(e)
      send_object(::Riemann::Message.new(
        :ok => false, :error => e.to_s))
    end
  end

  class ConnectionUDP < ConnectionBase; end

  DEFAULT_HOST = "localhost"
  DEFAULT_PORT = 5555
  DEFAULT_PROTOCOL = 'tcp'

  HANDLERS = {:tcp => HandlerTCP, :udp => HandlerUDP}
  CONNECTIONS = {:tcp => ConnectionTCP, :udp => ConnectionUDP}

  def self.output_setup(opts={})
    opts[:host] ||= DEFAULT_HOST
    opts[:port] ||= DEFAULT_PORT

    protocol = EVD.parse_protocol(opts[:protocol] || DEFAULT_PROTOCOL)

    unless handler = HANDLERS[protocol.family]
      raise "No handler for protocol family: #{protocol.family}"
    end

    instance = handler.new

    protocol.connect log, opts, instance
  end

  def self.input_setup(opts={})
    opts[:host] ||= DEFAULT_HOST
    opts[:port] ||= DEFAULT_PORT
    protocol = EVD.parse_protocol(opts[:protocol] || DEFAULT_PROTOCOL)

    unless connection = CONNECTIONS[protocol.family]
      raise "No connection for protocol family: #{protocol.family}"
    end

    protocol.listen log, opts, connection, log
  end
end
{'fraction_non_alphanumeric': 0.10167443466252374, 'fraction_numerical': 0.001380977041256689, 'mean_word_length': 2.9986197377501727, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 19, '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': '3796225', 'n_tokens_mistral': 1883, 'n_tokens_neox': 1760, 'n_words': 440}