<commit_before>require "active_support/core_ext/class/subclasses"
require "concurrent"
require "protobuf/rpc/server"
require "protobuf/rpc/service"

module Protobuf
  module Nats
    class Server
      include ::Protobuf::Rpc::Server
      include ::Protobuf::Logging

      attr_reader :nats, :thread_pool, :subscriptions

      def initialize(options)
        @options = options
        @running = true
        @stopped = false

        @nats = ::NATS::IO::Client.new
        @nats.connect(::Protobuf::Nats::Config.connection_options)

        # Don't let the queue grow beyond available threads.
        @thread_pool = ::Concurrent::FixedThreadPool.new(options[:threads], :max_queue => -1)

        @subscriptions = []
      end

      def service_klasses
        ::Protobuf::Rpc::Service.implemented_services.map(&:safe_constantize)
      end

      def execute_request_promise(request_data, reply_id)
        ::Concurrent::Promise.new(:executor => thread_pool).then do
          # Publish an ACK.
          nats.publish(reply_id, ::Protobuf::Nats::Messages::ACK)
          # Process request.
          response_data = handle_request(request_data)
          # Publish response.
          nats.publish(reply_id, response_data)
        end.on_error do |error|
          logger.error error
          if error.respond_to?(:backtrace) && error.backtrace.is_a?(::Array)
            logger.error error.backtrace.join("\n")
          end
        end.execute

        true
      rescue ::Concurrent::RejectedExecutionError
        false
      end

      def subscribe_to_services
        logger.info "Creating subscriptions:"

        service_klasses.each do |service_klass|
          service_klass.rpcs.each do |service_method, _|
            # Skip services that are not implemented.
            next unless service_klass.method_defined? service_method

            subscription_key_and_queue = ::Protobuf::Nats.subscription_key(service_klass, service_method)
            logger.info "  - #{subscription_key_and_queue}"

            subscriptions << nats.subscribe(subscription_key_and_queue, :queue => subscription_key_and_queue) do |request_data, reply_id, _subject|
              unless execute_request_promise(request_data, reply_id)
                logger.error { "Thread pool is full! Dropping message for: #{subscription_key_and_queue}" }
              end
            end
          end
        end
      end

      def run
        nats.on_reconnect do
          logger.warn "Reconnected to NATS server!"
          subscribe_to_services
        end

        nats.on_disconnect do
          logger.warn "Disconnected from NATS server!"
        end

        subscribe_to_services

        yield if block_given?

        loop do
          break unless @running
          sleep 1
        end

        subscriptions.each do |subscription_id|
          nats.unsubscribe(subscription_id)
        end

        thread_pool.shutdown
        thread_pool.wait_for_termination
      ensure
        @stopped = true
      end

      def running?
        @stopped
      end<commit_msg>NATS client will replay subscriptions<commit_after>require "active_support/core_ext/class/subclasses"
require "concurrent"
require "protobuf/rpc/server"
require "protobuf/rpc/service"

module Protobuf
  module Nats
    class Server
      include ::Protobuf::Rpc::Server
      include ::Protobuf::Logging

      attr_reader :nats, :thread_pool, :subscriptions

      def initialize(options)
        @options = options
        @running = true
        @stopped = false

        @nats = ::NATS::IO::Client.new
        @nats.connect(::Protobuf::Nats::Config.connection_options)

        # Don't let the queue grow beyond available threads.
        @thread_pool = ::Concurrent::FixedThreadPool.new(options[:threads], :max_queue => -1)

        @subscriptions = []
      end

      def service_klasses
        ::Protobuf::Rpc::Service.implemented_services.map(&:safe_constantize)
      end

      def execute_request_promise(request_data, reply_id)
        ::Concurrent::Promise.new(:executor => thread_pool).then do
          # Publish an ACK.
          nats.publish(reply_id, ::Protobuf::Nats::Messages::ACK)
          # Process request.
          response_data = handle_request(request_data)
          # Publish response.
          nats.publish(reply_id, response_data)
        end.on_error do |error|
          logger.error error
          if error.respond_to?(:backtrace) && error.backtrace.is_a?(::Array)
            logger.error error.backtrace.join("\n")
          end
        end.execute

        true
      rescue ::Concurrent::RejectedExecutionError
        false
      end

      def subscribe_to_services
        logger.info "Creating subscriptions:"

        service_klasses.each do |service_klass|
          service_klass.rpcs.each do |service_method, _|
            # Skip services that are not implemented.
            next unless service_klass.method_defined? service_method

            subscription_key_and_queue = ::Protobuf::Nats.subscription_key(service_klass, service_method)
            logger.info "  - #{subscription_key_and_queue}"

            subscriptions << nats.subscribe(subscription_key_and_queue, :queue => subscription_key_and_queue) do |request_data, reply_id, _subject|
              unless execute_request_promise(request_data, reply_id)
                logger.error { "Thread pool is full! Dropping message for: #{subscription_key_and_queue}" }
              end
            end
          end
        end
      end

      def run
        nats.on_reconnect do
          logger.warn "Reconnected to NATS server!"
        end

        nats.on_disconnect do
          logger.warn "Disconnected from NATS server!"
        end

        subscribe_to_services

        yield if block_given?

        loop do
          break unless @running
          sleep 1
        end

        subscriptions.each do |subscription_id|
          nats.unsubscribe(subscription_id)
        end

        thread_pool.shutdown
        thread_pool.wait_for_termination
      ensure
        @stopped = true
      end

      def running?
        @stopped
      end
{'fraction_non_alphanumeric': 0.08113391984359726, 'fraction_numerical': 0.0006516780710329097, 'mean_word_length': 2.2934549356223175, 'pattern_counts': {'":': 0, '<': 7, '<?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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2332000', 'n_tokens_mistral': 1723, 'n_tokens_neox': 1685, 'n_words': 370}