<commit_before>require 'thrift/server'

# thrift/server already imports fastthread/thread

module Thrift
  # this class expects to always use a FramedTransport for reading messages
  #--
  # this isn't very pretty, but we're working around the fact that FramedTransport
  # and the processors are all written in a synchronous manner.
  # So lets read data off the wire ourselves, check if we have a full frame, and
  # only then hand it to the transport to parse
  #
  # we inherit from ThreadPoolServer for the initialize/rescuable_serve methods
  class NonblockingServer < ThreadPoolServer
    def serve
      @server_thread = Thread.current
      @serverTransport.listen

      begin
        connections = {}
        running_connections = {}
        # the swapping_connections stuff is to ensure the thread doesn't
        # put the connection back into the regular list, then have the server
        # thread process it again, then have the first thread remove it from
        # the running_connections list
        swapping_connections = {}
        thread_group = ThreadGroup.new
        loop do
          break if @shutdown
          rd, = select([@serverTransport.handle, *connections.keys])
          next if rd.nil?
          rd.each do |socket|
            if socket == @serverTransport.handle
              client = @serverTransport.accept
              buffer = ''
              outtrans = @transportFactory.get_transport(client)
              outprot = @protocolFactory.get_protocol(outtrans)
              connections[client.handle] = [client, buffer, outtrans, outprot]
            else
              client, buffer, outtrans, outprot = connections[socket]
              if socket.eof?
                client.close
                connections.delete(socket)
              else
                buffer << client.read(4096, true)
                if has_full_frame?(buffer)
                  running_connections[socket] = connections.delete(socket)
                  @thread_q.push :token
                  t = Thread.new(Thread.current) do |master|
                    begin
                      membuf = MemoryBuffer.new(buffer)
                      intrans = @transportFactory.get_transport(membuf)
                      inprot = @protocolFactory.get_protocol(intrans)
                      @processor.process(inprot, outprot)
                      if @shutdown
                        client.close
                        running_connections.delete(socket)
                      else
                        swapping_connections[socket] = running_connections.delete(socket)
                        master.wakeup
                      end
                    rescue => e
                      outtrans.close
                      @exception_q.push e
                    ensure
                      running_connections.delete(socket)
                      connections[socket] = swapping_connections.delete(socket) if swapping_connections.include? socket
                      intrans.close
                      @thread_q.pop
                    end
                  end
                  thread_group.add t
                end
              end
            end
          end
        end
        if @shutdown
          @serverTransport.close
          connections.merge! running_connections
          connections.merge! swapping_connections
          connections.values.each do |client, buffer, outtrans, outprot|
            # can't close completely or we'll break active messages
            # but lets at least stop accepting input
            client.handle.close_read
          end
          start = Time.now.to_f
          until thread_group.list.empty?
            if @shutdown_timeout
              now = Time.now.to_f
              cur_timeout = @shutdown_timeout - (now - start)
              break if cur_timeout <= 0
              thread_group.list.first.join(cur_timeout)
            else
              thread_group.list.first.join
            end
          end
          thread_group.list.each { |t| t.kill } if @shutdown_kill
          # now kill connections completely if they still exists
          connections.values.each do |client, buffer, outtrans, outprot|
            client.close
          end
        end
      ensure
        @serverTransport.close
      end
    end

    # Stop accepting new messages and wait for active messages to finish
    # If the given timeout passes without the active messages finishing,
    # control will exit from #serve and leave the remaining threads active.
    # If you pass true for kill, the remaining threads will be reaped instead.
    # A false timeout means wait indefinitely
    def shutdown(timeout = nil, kill = false)<commit_msg>Add synchronization around shared resources in NonblockingServer<commit_after>require 'thrift/server'
require 'sync'
# thrift/server already imports fastthread/thread

module Thrift
  # this class expects to always use a FramedTransport for reading messages
  #--
  # this isn't very pretty, but we're working around the fact that FramedTransport
  # and the processors are all written in a synchronous manner.
  # So lets read data off the wire ourselves, check if we have a full frame, and
  # only then hand it to the transport to parse
  #
  # we inherit from ThreadPoolServer for the initialize/rescuable_serve methods
  class NonblockingServer < ThreadPoolServer
    def initialize(processor, serverTransport, transportFactory=nil, protocolFactory=nil, num=20)
      super
      @sync = Sync.new
    end

    def serve
      @server_thread = Thread.current
      @serverTransport.listen

      begin
        connections = {}
        running_connections = {}
        # the swapping_connections stuff is to ensure the thread doesn't
        # put the connection back into the regular list, then have the server
        # thread process it again, then have the first thread remove it from
        # the running_connections list
        swapping_connections = {}
        thread_group = ThreadGroup.new
        loop do
          break if @shutdown
          handles = [@serverTransport.handle]
          @sync.synchronize(Sync::SH) do
            handles.concat connections.keys
          end
          rd, = select(handles)
          next if rd.nil?
          rd.each do |socket|
            if socket == @serverTransport.handle
              client = @serverTransport.accept
              buffer = ''
              outtrans = @transportFactory.get_transport(client)
              outprot = @protocolFactory.get_protocol(outtrans)
              @sync.synchronize(Sync::EX) do
                connections[client.handle] = [client, buffer, outtrans, outprot]
              end
            else
              client, buffer, outtrans, outprot = nil # for scope
              @sync.synchronize(Sync::SH) do
                client, buffer, outtrans, outprot = connections[socket]
              end
              if socket.eof?
                client.close
                @sync.synchronize(Sync::EX) do
                  connections.delete(socket)
                end
              else
                buffer << client.read(4096, true)
                if has_full_frame?(buffer)
                  @sync.synchronize(Sync::EX) do
                    running_connections[socket] = connections.delete(socket)
                  end
                  @thread_q.push :token
                  t = Thread.new(Thread.current) do |master|
                    begin
                      membuf = MemoryBuffer.new(buffer)
                      intrans = @transportFactory.get_transport(membuf)
                      inprot = @protocolFactory.get_protocol(intrans)
                      @processor.process(inprot, outprot)
                      if @shutdown
                        client.close
                        @sync.synchronize(Sync::EX) do
                          running_connections.delete(socket)
                        end
                      else
                        @sync.synchronize(Sync::EX) do
                          swapping_connections[socket] = running_connections.delete(socket)
                        end
                      end
                    rescue => e
                      outtrans.close
                      @exception_q.push e
                    ensure
                      should_wakeup = false
                      @sync.synchronize(Sync::EX) do
                        running_connections.delete(socket)
                        if swapping_connections.include? socket
                          connections[socket] = swapping_connections.delete(socket)
                          should_wakeup = true
                        end
                      end
                      master.wakeup if should_wakeup
                      intrans.close
                      @thread_q.pop
                    end
                  end
                  thread_group.add t
                end
              end
            end
          end
        end
        if @shutdown
          @serverTransport.close
          handles = []
          @sync.synchronize(Sync::SH) do
            handles = connections
            handles.merge! running_connections
            handles.merge! swapping_connections
          end
          handles.values.each do |client, buffer, outtrans, outprot|
            # can't close completely or we'll break active messages
            # but lets at least stop accepting input
            client.handle.close_read
          end
          start = Time.now.to_f
          until thread_group.list.empty?
            if @shutdown_timeout
              now = Time.now.to_f
              cur_timeout = @shutdown_timeout - (now - start)
              break if cur_timeout <= 0
              thread_group.list.first.join(cur_timeout)
            else
              thread_group.list.first.join
            end
          end
          thread_group.list.each { |t| t.kill } if @shutdown_kill
          # now kill connections completely if they still exists
          handles.values.each do |client, buffer, outtrans, outprot|
            client.close
          end
        end
      ensure
        @serverTransport.close
      end
    end

    # Stop accepting new messages and wait for active messages to finish
    # If the given timeout passes without the active messages finishing,
    # control will exit from #serve and leave the remaining threads active.
    # If you pass true for kill, the remaining threads will be reaped instead.
    # A false timeout means wait indefinitely
    def shutdown(timeout = nil, kill = false)
{'fraction_non_alphanumeric': 0.05730740181268882, 'fraction_numerical': 0.0011329305135951663, 'mean_word_length': 1.5476190476190477, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 5, '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': '2392831', 'n_tokens_mistral': 2527, 'n_tokens_neox': 2400, 'n_words': 766}