<commit_before>require 'rack'
require 'timeout'

module BackgroundStreamer
  class Worker
    CRLF = "\r\n".freeze
    RACK_HIJACK = 'rack.hijack'.freeze
    RACK_HIJACK_IO = 'rack.hijack_io'.freeze
    ACTION_DISPATCH_REQUEST_ID = 'action_dispatch.request_id'.freeze
    X_REQUEST_ID = 'X_REQUEST_ID'.freeze

    attr_reader :options, :timeout, :request_id, :status, :headers, :io, :body

    class << self
      def perform_async(env, body, options = {})
        threads.keep_if(&:alive?)

        if threads.size >= BackgroundStreamer.max_threads
          raise ThreadLimitExceeded, "Thread limit of #{BackgroundStreamer.max_threads} exceeded"
        end

        env[RACK_HIJACK].call

        threads << Thread.new do
          worker = new(env, body, options)
          worker.perform

          BackgroundStreamer.on_worker_exit.call if BackgroundStreamer.on_worker_exit
        end
      end

      private

      def threads
        @threads ||= []
      end
    end
      
    def initialize(env, body, options = {})
      @io         = env[RACK_HIJACK_IO]
      @request_id = env[ACTION_DISPATCH_REQUEST_ID] || env[X_REQUEST_ID]
      @options    = options
      @timeout    = @options.delete(:timeout) || 15.seconds

      app = -> _ { [200, create_headers, body] }
      @status, @headers, @body = Rack::Deflater.new(app).call(env)
    end

    def perform
      logger.push_tags("#{self.class.name}.#{request_id}") if logger.respond_to?(:push_tags)
      now = Time.now
      logger.info "Starting ..."
      
      stream
    ensure
      close
      logger.info "Finished in #{Time.now - now} seconds."
      logger.pop_tags if logger.respond_to?(:pop_tags)
    end

    private

    def stream
      write_headers

      Timeout::timeout(timeout) do
        logger.debug "Writing Body"

        body.each do |chunk|
          write_chunk chunk
        end
      end
    rescue => e
      logger.warn "Streaming has failed: #{e} => #{e.backtrace}"
    ensure
      write_chunk ""
    end

    def close
      logger.debug "Closing connection"

      unless io.closed?
        io.shutdown # in case of fork() in Rack app
        io.close # flush and uncork socket immediately, no keepalive
      end
    rescue => e
      logger.info "Closing connection failed: #{e} => #{e.backtrace}"
    end

    def write_chunk(chunk)
      io << [chunk.size.to_s(16), CRLF, chunk, CRLF].join
    end

    def create_headers
      {
        "Date"              => Time.now.strftime("%a %e %b %T %Y %Z"),
        "Status"            => "200 OK",
        "Connection"        => "close",
        "Content-Type"      => "application/json; charset=utf-8",
        "Cache-Control"     => "max-age=0, private, no-cache",
        "Transfer-Encoding" => "chunked",
        "X-Request-Id"      => request_id, 
      }
    end

    def write_headers
      logger.debug "Writing Headers"

      buffer = ""
      buffer << "HTTP/1.1 200 OK" << CRLF

      headers.each do |key, value|
        buffer << "#{key}: #{value}#{CRLF}"
      end

      buffer << CRLF

      io << buffer
    end

    def logger
      BackgroundStreamer.logger
    end
  end
end<commit_msg>always write headers on the main thread<commit_after>require 'rack'
require 'timeout'

module BackgroundStreamer
  class Worker
    CRLF = "\r\n".freeze
    RACK_HIJACK = 'rack.hijack'.freeze
    RACK_HIJACK_IO = 'rack.hijack_io'.freeze
    ACTION_DISPATCH_REQUEST_ID = 'action_dispatch.request_id'.freeze
    X_REQUEST_ID = 'X_REQUEST_ID'.freeze

    attr_reader :options, :timeout, :request_id, :status, :headers, :io, :body

    class << self
      def perform_async(env, body, options = {})
        threads.keep_if(&:alive?)

        if threads.size >= BackgroundStreamer.max_threads
          raise ThreadLimitExceeded, "Thread limit of #{BackgroundStreamer.max_threads} exceeded"
        end

        worker = new(env, body, options)

        threads << Thread.new do
          worker.perform
          BackgroundStreamer.on_worker_exit.call if BackgroundStreamer.on_worker_exit
        end
      end

      private

      def threads
        @threads ||= []
      end
    end
      
    def initialize(env, body, options = {})
      env[RACK_HIJACK].call

      @io         = env[RACK_HIJACK_IO]
      @request_id = env[ACTION_DISPATCH_REQUEST_ID] || env[X_REQUEST_ID]
      @options    = options
      @timeout    = @options.delete(:timeout) || 15.seconds

      app = -> _ { [200, create_headers, body] }
      @status, @headers, @body = Rack::Deflater.new(app).call(env)

      write_headers
    end

    def perform
      logger.push_tags("#{request_id}:#{self.class.name}") if logger.respond_to?(:push_tags)
      now = Time.now
      logger.info "Starting ..."
      
      write_body
    ensure
      close
      logger.info "Finished in #{Time.now - now} seconds."
      logger.pop_tags if logger.respond_to?(:pop_tags)
    end

    private

    def write_body
      Timeout::timeout(timeout) do
        logger.debug "Writing Body"

        body.each do |chunk|
          write_chunk chunk
        end
      end
    rescue => e
      logger.warn "Streaming has failed: #{e} => #{e.backtrace}"
    ensure
      write_chunk ""
    end

    def close
      logger.debug "Closing connection"

      unless io.closed?
        io.shutdown # in case of fork() in Rack app
        io.close # flush and uncork socket immediately, no keepalive
      end
    rescue => e
      logger.info "Closing connection failed: #{e} => #{e.backtrace}"
    end

    def write_chunk(chunk)
      io << [chunk.size.to_s(16), CRLF, chunk, CRLF].join
    end

    def create_headers
      {
        "Date"              => Time.now.strftime("%a %e %b %T %Y %Z"),
        "Status"            => "200 OK",
        "Connection"        => "close",
        "Content-Type"      => "application/json; charset=utf-8",
        "Cache-Control"     => "max-age=0, private, no-cache",
        "Transfer-Encoding" => "chunked",
        "X-Request-Id"      => request_id, 
      }
    end

    def write_headers
      logger.debug "Writing Headers"

      buffer = ""
      buffer << "HTTP/1.1 200 OK" << CRLF

      headers.each do |key, value|
        buffer << "#{key}: #{value}#{CRLF}"
      end

      buffer << CRLF

      io << buffer
    rescue => e
      logger.warn "Writing headers has failed: #{e} => #{e.backtrace}"
    end

    def logger
      BackgroundStreamer.logger
    end
  end
end
{'fraction_non_alphanumeric': 0.11612000621793875, 'fraction_numerical': 0.005285247940307788, 'mean_word_length': 2.137006338371526, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 31, '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': '3847360', 'n_tokens_mistral': 2096, 'n_tokens_neox': 2001, 'n_words': 480}