ruby-kafka-temp-fork 0.0.1

data/lib/kafka/ssl_socket_with_timeout.rb

@@ -0,0 +1,188 @@
+# frozen_string_literal: true
+
+require "socket"
+
+module Kafka
+
+  # Opens sockets in a non-blocking fashion, ensuring that we're not stalling
+  # for long periods of time.
+  #
+  # It's possible to set timeouts for connecting to the server, for reading data,
+  # and for writing data. Whenever a timeout is exceeded, Errno::ETIMEDOUT is
+  # raised.
+  #
+  class SSLSocketWithTimeout
+
+    # Opens a socket.
+    #
+    # @param host [String]
+    # @param port [Integer]
+    # @param connect_timeout [Integer] the connection timeout, in seconds.
+    # @param timeout [Integer] the read and write timeout, in seconds.
+    # @param ssl_context [OpenSSL::SSL::SSLContext] which SSLContext the ssl connection should use
+    # @raise [Errno::ETIMEDOUT] if the timeout is exceeded.
+    def initialize(host, port, connect_timeout: nil, timeout: nil, ssl_context:)
+      addr = Socket.getaddrinfo(host, nil)
+      sockaddr = Socket.pack_sockaddr_in(port, addr[0][3])
+
+      @connect_timeout = connect_timeout
+      @timeout = timeout
+
+      @tcp_socket = Socket.new(Socket.const_get(addr[0][0]), Socket::SOCK_STREAM, 0)
+      @tcp_socket.setsockopt(Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1)
+
+      # first initiate the TCP socket
+      begin
+        # Initiate the socket connection in the background. If it doesn't fail
+        # immediately it will raise an IO::WaitWritable (Errno::EINPROGRESS)
+        # indicating the connection is in progress.
+        @tcp_socket.connect_nonblock(sockaddr)
+      rescue IO::WaitWritable
+        # select will block until the socket is writable or the timeout
+        # is exceeded, whichever comes first.
+        unless select_with_timeout(@tcp_socket, :connect_write)
+          # select returns nil when the socket is not ready before timeout
+          # seconds have elapsed
+          @tcp_socket.close
+          raise Errno::ETIMEDOUT
+        end
+
+        begin
+          # Verify there is now a good connection.
+          @tcp_socket.connect_nonblock(sockaddr)
+        rescue Errno::EISCONN
+          # The socket is connected, we're good!
+        end
+      end
+
+      # once that's connected, we can start initiating the ssl socket
+      @ssl_socket = OpenSSL::SSL::SSLSocket.new(@tcp_socket, ssl_context)
+      @ssl_socket.hostname = host
+
+      begin
+        # Initiate the socket connection in the background. If it doesn't fail
+        # immediately it will raise an IO::WaitWritable (Errno::EINPROGRESS)
+        # indicating the connection is in progress.
+        # Unlike waiting for a tcp socket to connect, you can't time out ssl socket
+        # connections during the connect phase properly, because IO.select only partially works.
+        # Instead, you have to retry.
+        @ssl_socket.connect_nonblock
+      rescue Errno::EAGAIN, Errno::EWOULDBLOCK, IO::WaitReadable
+        if select_with_timeout(@ssl_socket, :connect_read)
+          retry
+        else
+          @ssl_socket.close
+          close
+          raise Errno::ETIMEDOUT
+        end
+      rescue IO::WaitWritable
+        if select_with_timeout(@ssl_socket, :connect_write)
+          retry
+        else
+          close
+          raise Errno::ETIMEDOUT
+        end
+      end
+    end
+
+    # Reads bytes from the socket, possible with a timeout.
+    #
+    # @param num_bytes [Integer] the number of bytes to read.
+    # @raise [Errno::ETIMEDOUT] if the timeout is exceeded.
+    # @return [String] the data that was read from the socket.
+    def read(num_bytes)
+      buffer = String.new
+
+      until buffer.length >= num_bytes
+        begin
+          # Unlike plain TCP sockets, SSL sockets don't support IO.select
+          # properly.
+          # Instead, timeouts happen on a per read basis, and we have to
+          # catch exceptions from read_nonblock and gradually build up
+          # our read buffer.
+          buffer << @ssl_socket.read_nonblock(num_bytes - buffer.length)
+        rescue IO::WaitReadable
+          if select_with_timeout(@ssl_socket, :read)
+            retry
+          else
+            raise Errno::ETIMEDOUT
+          end
+        rescue IO::WaitWritable
+          if select_with_timeout(@ssl_socket, :write)
+            retry
+          else
+            raise Errno::ETIMEDOUT
+          end
+        end
+      end
+
+      buffer
+    end
+
+    # Writes bytes to the socket, possible with a timeout.
+    #
+    # @param bytes [String] the data that should be written to the socket.
+    # @raise [Errno::ETIMEDOUT] if the timeout is exceeded.
+    # @return [Integer] the number of bytes written.
+    def write(bytes)
+      loop do
+        written = 0
+        begin
+          # unlike plain tcp sockets, ssl sockets don't support IO.select
+          # properly.
+          # Instead, timeouts happen on a per write basis, and we have to
+          # catch exceptions from write_nonblock, and gradually build up
+          # our write buffer.
+          written += @ssl_socket.write_nonblock(bytes)
+        rescue Errno::EFAULT => error
+          raise error
+        rescue OpenSSL::SSL::SSLError, Errno::EAGAIN, Errno::EWOULDBLOCK, IO::WaitWritable => error
+          if error.is_a?(OpenSSL::SSL::SSLError) && error.message == 'write would block'
+            if select_with_timeout(@ssl_socket, :write)
+              retry
+            else
+              raise Errno::ETIMEDOUT
+            end
+          else
+            raise error
+          end
+        end
+
+        # Fast, common case.
+        break if written == bytes.size
+
+        # This takes advantage of the fact that most ruby implementations
+        # have Copy-On-Write strings. Thusly why requesting a subrange
+        # of data, we actually don't copy data because the new string
+        # simply references a subrange of the original.
+        bytes = bytes[written, bytes.size]
+      end
+    end
+
+    def close
+      @tcp_socket.close
+      @ssl_socket.close
+    end
+
+    def closed?
+      @tcp_socket.closed? || @ssl_socket.closed?
+    end
+
+    def set_encoding(encoding)
+      @tcp_socket.set_encoding(encoding)
+    end
+
+    def select_with_timeout(socket, type)
+      case type
+      when :connect_read
+        IO.select([socket], nil, nil, @connect_timeout)
+      when :connect_write
+        IO.select(nil, [socket], nil, @connect_timeout)
+      when :read
+        IO.select([socket], nil, nil, @timeout)
+      when :write
+        IO.select(nil, [socket], nil, @timeout)
+      end
+    end
+  end
+end

data/lib/kafka/statsd.rb

@@ -0,0 +1,296 @@
+# frozen_string_literal: true
+
+begin
+  require "statsd"
+rescue LoadError
+  $stderr.puts "In order to report Kafka client metrics to Statsd you need to install the `statsd-ruby` gem."
+  raise
+end
+
+require "active_support/subscriber"
+
+module Kafka
+  # Reports operational metrics to a Statsd agent.
+  #
+  #     require "kafka/statsd"
+  #
+  #     # Default is "ruby_kafka".
+  #     Kafka::Statsd.namespace = "custom-namespace"
+  #
+  #     # Default is "127.0.0.1".
+  #     Kafka::Statsd.host = "statsd.something.com"
+  #
+  #     # Default is 8125.
+  #     Kafka::Statsd.port = 1234
+  #
+  # Once the file has been required, no further configuration is needed – all operational
+  # metrics are automatically emitted.
+  module Statsd
+    DEFAULT_NAMESPACE = "ruby_kafka"
+    DEFAULT_HOST = '127.0.0.1'
+    DEFAULT_PORT = 8125
+
+    def self.statsd
+      @statsd ||= ::Statsd.new(DEFAULT_HOST, DEFAULT_PORT).tap { |sd| sd.namespace = DEFAULT_NAMESPACE }
+    end
+
+    def self.host=(host)
+      statsd.host = host
+      statsd.connect if statsd.respond_to?(:connect)
+    end
+
+    def self.port=(port)
+      statsd.port = port
+      statsd.connect if statsd.respond_to?(:connect)
+    end
+
+    def self.namespace=(namespace)
+      statsd.namespace = namespace
+    end
+
+    class StatsdSubscriber < ActiveSupport::Subscriber
+      private
+
+      %w[increment count timing gauge].each do |type|
+        define_method(type) do |*args|
+          Kafka::Statsd.statsd.send(type, *args)
+        end
+      end
+    end
+
+    class ConnectionSubscriber < StatsdSubscriber
+      def request(event)
+        client = event.payload.fetch(:client_id)
+        api = event.payload.fetch(:api, "unknown")
+        request_size = event.payload.fetch(:request_size, 0)
+        response_size = event.payload.fetch(:response_size, 0)
+        broker = event.payload.fetch(:broker_host)
+
+        timing("api.#{client}.#{api}.#{broker}.latency", event.duration)
+        increment("api.#{client}.#{api}.#{broker}.calls")
+
+        timing("api.#{client}.#{api}.#{broker}.request_size", request_size)
+        timing("api.#{client}.#{api}.#{broker}.response_size", response_size)
+
+        if event.payload.key?(:exception)
+          increment("api.#{client}.#{api}.#{broker}.errors")
+        end
+      end
+
+      attach_to "connection.kafka"
+    end
+
+    class ConsumerSubscriber < StatsdSubscriber
+      def process_message(event)
+        offset_lag = event.payload.fetch(:offset_lag)
+        create_time = event.payload.fetch(:create_time)
+        client = event.payload.fetch(:client_id)
+        group_id = event.payload.fetch(:group_id)
+        topic = event.payload.fetch(:topic)
+        partition = event.payload.fetch(:partition)
+
+        time_lag = create_time && ((Time.now - create_time) * 1000).to_i
+
+        if event.payload.key?(:exception)
+          increment("consumer.#{client}.#{group_id}.#{topic}.#{partition}.process_message.errors")
+        else
+          timing("consumer.#{client}.#{group_id}.#{topic}.#{partition}.process_message.latency", event.duration)
+          increment("consumer.#{client}.#{group_id}.#{topic}.#{partition}.messages")
+        end
+
+        gauge("consumer.#{client}.#{group_id}.#{topic}.#{partition}.lag", offset_lag)
+
+        # Not all messages have timestamps.
+        if time_lag
+          gauge("consumer.#{client}.#{group_id}.#{topic}.#{partition}.time_lag", time_lag)
+        end
+      end
+
+      def process_batch(event)
+        messages = event.payload.fetch(:message_count)
+        client = event.payload.fetch(:client_id)
+        group_id = event.payload.fetch(:group_id)
+        topic = event.payload.fetch(:topic)
+        partition = event.payload.fetch(:partition)
+
+        if event.payload.key?(:exception)
+          increment("consumer.#{client}.#{group_id}.#{topic}.#{partition}.process_batch.errors")
+        else
+          timing("consumer.#{client}.#{group_id}.#{topic}.#{partition}.process_batch.latency", event.duration)
+          count("consumer.#{client}.#{group_id}.#{topic}.#{partition}.messages", messages)
+        end
+      end
+
+      def fetch_batch(event)
+        lag = event.payload.fetch(:offset_lag)
+        batch_size = event.payload.fetch(:message_count)
+        client = event.payload.fetch(:client_id)
+        group_id = event.payload.fetch(:group_id)
+        topic = event.payload.fetch(:topic)
+        partition = event.payload.fetch(:partition)
+
+        count("consumer.#{client}.#{group_id}.#{topic}.#{partition}.batch_size", batch_size)
+        gauge("consumer.#{client}.#{group_id}.#{topic}.#{partition}.lag", lag)
+      end
+
+      def join_group(event)
+        client = event.payload.fetch(:client_id)
+        group_id = event.payload.fetch(:group_id)
+
+        timing("consumer.#{client}.#{group_id}.join_group", event.duration)
+
+        if event.payload.key?(:exception)
+          increment("consumer.#{client}.#{group_id}.join_group.errors")
+        end
+      end
+
+      def sync_group(event)
+        client = event.payload.fetch(:client_id)
+        group_id = event.payload.fetch(:group_id)
+
+        timing("consumer.#{client}.#{group_id}.sync_group", event.duration)
+
+        if event.payload.key?(:exception)
+          increment("consumer.#{client}.#{group_id}.sync_group.errors")
+        end
+      end
+
+      def leave_group(event)
+        client = event.payload.fetch(:client_id)
+        group_id = event.payload.fetch(:group_id)
+
+        timing("consumer.#{client}.#{group_id}.leave_group", event.duration)
+
+        if event.payload.key?(:exception)
+          increment("consumer.#{client}.#{group_id}.leave_group.errors")
+        end
+      end
+
+      def pause_status(event)
+        client = event.payload.fetch(:client_id)
+        group_id = event.payload.fetch(:group_id)
+        topic = event.payload.fetch(:topic)
+        partition = event.payload.fetch(:partition)
+
+        duration = event.payload.fetch(:duration)
+
+        gauge("consumer.#{client}.#{group_id}.#{topic}.#{partition}.pause.duration", duration)
+      end
+
+      attach_to "consumer.kafka"
+    end
+
+    class ProducerSubscriber < StatsdSubscriber
+      def produce_message(event)
+        client = event.payload.fetch(:client_id)
+        topic = event.payload.fetch(:topic)
+        message_size = event.payload.fetch(:message_size)
+        buffer_size = event.payload.fetch(:buffer_size)
+        max_buffer_size = event.payload.fetch(:max_buffer_size)
+        buffer_fill_ratio = buffer_size.to_f / max_buffer_size.to_f
+        buffer_fill_percentage = buffer_fill_ratio * 100.0
+
+        # This gets us the write rate.
+        increment("producer.#{client}.#{topic}.produce.messages")
+
+        timing("producer.#{client}.#{topic}.produce.message_size", message_size)
+
+        # This gets us the avg/max buffer size per producer.
+        timing("producer.#{client}.buffer.size", buffer_size)
+
+        # This gets us the avg/max buffer fill ratio per producer.
+        timing("producer.#{client}.buffer.fill_ratio", buffer_fill_ratio)
+        timing("producer.#{client}.buffer.fill_percentage", buffer_fill_percentage)
+      end
+
+      def buffer_overflow(event)
+        client = event.payload.fetch(:client_id)
+        topic = event.payload.fetch(:topic)
+
+        increment("producer.#{client}.#{topic}.produce.errors")
+      end
+
+      def deliver_messages(event)
+        client = event.payload.fetch(:client_id)
+        message_count = event.payload.fetch(:delivered_message_count)
+        attempts = event.payload.fetch(:attempts)
+
+        if event.payload.key?(:exception)
+          increment("producer.#{client}.deliver.errors")
+        end
+
+        timing("producer.#{client}.deliver.latency", event.duration)
+
+        # Messages delivered to Kafka:
+        count("producer.#{client}.deliver.messages", message_count)
+
+        # Number of attempts to deliver messages:
+        timing("producer.#{client}.deliver.attempts", attempts)
+      end
+
+      def ack_message(event)
+        client = event.payload.fetch(:client_id)
+        topic = event.payload.fetch(:topic)
+
+        # Number of messages ACK'd for the topic.
+        increment("producer.#{client}.#{topic}.ack.messages")
+
+        # Histogram of delay between a message being produced and it being ACK'd.
+        timing("producer.#{client}.#{topic}.ack.delay", event.payload.fetch(:delay))
+      end
+
+      def topic_error(event)
+        client = event.payload.fetch(:client_id)
+        topic = event.payload.fetch(:topic)
+
+        increment("producer.#{client}.#{topic}.ack.errors")
+      end
+
+      attach_to "producer.kafka"
+    end
+
+    class AsyncProducerSubscriber < StatsdSubscriber
+      def enqueue_message(event)
+        client = event.payload.fetch(:client_id)
+        topic = event.payload.fetch(:topic)
+        queue_size = event.payload.fetch(:queue_size)
+        max_queue_size = event.payload.fetch(:max_queue_size)
+        queue_fill_ratio = queue_size.to_f / max_queue_size.to_f
+
+        # This gets us the avg/max queue size per producer.
+        timing("async_producer.#{client}.#{topic}.queue.size", queue_size)
+
+        # This gets us the avg/max queue fill ratio per producer.
+        timing("async_producer.#{client}.#{topic}.queue.fill_ratio", queue_fill_ratio)
+      end
+
+      def buffer_overflow(event)
+        client = event.payload.fetch(:client_id)
+        topic = event.payload.fetch(:topic)
+
+        increment("async_producer.#{client}.#{topic}.produce.errors")
+      end
+
+      def drop_messages(event)
+        client = event.payload.fetch(:client_id)
+        message_count = event.payload.fetch(:message_count)
+
+        count("async_producer.#{client}.dropped_messages", message_count)
+      end
+
+      attach_to "async_producer.kafka"
+    end
+
+    class FetcherSubscriber < StatsdSubscriber
+      def loop(event)
+        queue_size = event.payload.fetch(:queue_size)
+        client = event.payload.fetch(:client_id)
+        group_id = event.payload.fetch(:group_id)
+
+        gauge("fetcher.#{client}.#{group_id}.queue_size", queue_size)
+      end
+
+      attach_to "fetcher.kafka"
+    end
+  end
+end