| # |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # |
| |
| require 'logger' |
| require 'thread' |
| |
| module Thrift |
| # this class expects to always use a FramedTransport for reading messages |
| class NonblockingServer < BaseServer |
| def initialize(processor, server_transport, transport_factory=nil, protocol_factory=nil, num=20, logger=nil) |
| super(processor, server_transport, transport_factory, protocol_factory) |
| @num_threads = num |
| if logger.nil? |
| @logger = Logger.new(STDERR) |
| @logger.level = Logger::WARN |
| else |
| @logger = logger |
| end |
| @shutdown_semaphore = Mutex.new |
| @transport_semaphore = Mutex.new |
| end |
| |
| def serve |
| @logger.info "Starting #{self}" |
| @server_transport.listen |
| @io_manager = start_io_manager |
| |
| begin |
| loop do |
| break if @server_transport.closed? |
| rd, = select([@server_transport], nil, nil, 0.1) |
| next if rd.nil? |
| socket = @server_transport.accept |
| @logger.debug "Accepted socket: #{socket.inspect}" |
| @io_manager.add_connection socket |
| end |
| rescue IOError => e |
| end |
| # we must be shutting down |
| @logger.info "#{self} is shutting down, goodbye" |
| ensure |
| @transport_semaphore.synchronize do |
| @server_transport.close |
| end |
| @io_manager.ensure_closed unless @io_manager.nil? |
| end |
| |
| def shutdown(timeout = 0, block = true) |
| @shutdown_semaphore.synchronize do |
| return if @is_shutdown |
| @is_shutdown = true |
| end |
| # nonblocking is intended for calling from within a Handler |
| # but we can't change the order of operations here, so lets thread |
| shutdown_proc = lambda do |
| @io_manager.shutdown(timeout) |
| @transport_semaphore.synchronize do |
| @server_transport.close # this will break the accept loop |
| end |
| end |
| if block |
| shutdown_proc.call |
| else |
| Thread.new &shutdown_proc |
| end |
| end |
| |
| private |
| |
| def start_io_manager |
| iom = IOManager.new(@processor, @server_transport, @transport_factory, @protocol_factory, @num_threads, @logger) |
| iom.spawn |
| iom |
| end |
| |
| class IOManager # :nodoc: |
| DEFAULT_BUFFER = 2**20 |
| |
| def initialize(processor, server_transport, transport_factory, protocol_factory, num, logger) |
| @processor = processor |
| @server_transport = server_transport |
| @transport_factory = transport_factory |
| @protocol_factory = protocol_factory |
| @num_threads = num |
| @logger = logger |
| @connections = [] |
| @buffers = Hash.new { |h,k| h[k] = '' } |
| @signal_queue = Queue.new |
| @signal_pipes = IO.pipe |
| @signal_pipes[1].sync = true |
| @worker_queue = Queue.new |
| @shutdown_queue = Queue.new |
| end |
| |
| def add_connection(socket) |
| signal [:connection, socket] |
| end |
| |
| def spawn |
| @iom_thread = Thread.new do |
| @logger.debug "Starting #{self}" |
| run |
| end |
| end |
| |
| def shutdown(timeout = 0) |
| @logger.debug "#{self} is shutting down workers" |
| @worker_queue.clear |
| @num_threads.times { @worker_queue.push [:shutdown] } |
| signal [:shutdown, timeout] |
| @shutdown_queue.pop |
| @signal_pipes[0].close |
| @signal_pipes[1].close |
| @logger.debug "#{self} is shutting down, goodbye" |
| end |
| |
| def ensure_closed |
| kill_worker_threads if @worker_threads |
| @iom_thread.kill |
| end |
| |
| private |
| |
| def run |
| spin_worker_threads |
| |
| loop do |
| rd, = select([@signal_pipes[0], *@connections]) |
| if rd.delete @signal_pipes[0] |
| break if read_signals == :shutdown |
| end |
| rd.each do |fd| |
| if fd.handle.eof? |
| remove_connection fd |
| else |
| read_connection fd |
| end |
| end |
| end |
| join_worker_threads(@shutdown_timeout) |
| ensure |
| @shutdown_queue.push :shutdown |
| end |
| |
| def read_connection(fd) |
| @buffers[fd] << fd.read(DEFAULT_BUFFER) |
| frame = slice_frame!(@buffers[fd]) |
| if frame |
| @logger.debug "#{self} is processing a frame" |
| @worker_queue.push [:frame, fd, frame] |
| end |
| end |
| |
| def spin_worker_threads |
| @logger.debug "#{self} is spinning up worker threads" |
| @worker_threads = [] |
| @num_threads.times do |
| @worker_threads << spin_thread |
| end |
| end |
| |
| def spin_thread |
| Worker.new(@processor, @transport_factory, @protocol_factory, @logger, @worker_queue).spawn |
| end |
| |
| def signal(msg) |
| @signal_queue << msg |
| @signal_pipes[1].write " " |
| end |
| |
| def read_signals |
| # clear the signal pipe |
| # note that since read_nonblock is broken in jruby, |
| # we can only read up to a set number of signals at once |
| sigstr = @signal_pipes[0].readpartial(1024) |
| # now read the signals |
| begin |
| sigstr.length.times do |
| signal, obj = @signal_queue.pop(true) |
| case signal |
| when :connection |
| @connections << obj |
| when :shutdown |
| @shutdown_timeout = obj |
| return :shutdown |
| end |
| end |
| rescue ThreadError |
| # out of signals |
| # note that in a perfect world this would never happen, since we're |
| # only reading the number of signals pushed on the pipe, but given the lack |
| # of locks, in theory we could clear the pipe/queue while a new signal is being |
| # placed on the pipe, at which point our next read_signals would hit this error |
| end |
| end |
| |
| def remove_connection(fd) |
| # don't explicitly close it, a thread may still be writing to it |
| @connections.delete fd |
| @buffers.delete fd |
| end |
| |
| def join_worker_threads(shutdown_timeout) |
| start = Time.now |
| @worker_threads.each do |t| |
| if shutdown_timeout > 0 |
| timeout = (start + shutdown_timeout) - Time.now |
| break if timeout <= 0 |
| t.join(timeout) |
| else |
| t.join |
| end |
| end |
| kill_worker_threads |
| end |
| |
| def kill_worker_threads |
| @worker_threads.each do |t| |
| t.kill if t.status |
| end |
| @worker_threads.clear |
| end |
| |
| def slice_frame!(buf) |
| if buf.length >= 4 |
| size = buf.unpack('N').first |
| if buf.length >= size + 4 |
| buf.slice!(0, size + 4) |
| else |
| nil |
| end |
| else |
| nil |
| end |
| end |
| |
| class Worker # :nodoc: |
| def initialize(processor, transport_factory, protocol_factory, logger, queue) |
| @processor = processor |
| @transport_factory = transport_factory |
| @protocol_factory = protocol_factory |
| @logger = logger |
| @queue = queue |
| end |
| |
| def spawn |
| Thread.new do |
| @logger.debug "#{self} is spawning" |
| run |
| end |
| end |
| |
| private |
| |
| def run |
| loop do |
| cmd, *args = @queue.pop |
| case cmd |
| when :shutdown |
| @logger.debug "#{self} is shutting down, goodbye" |
| break |
| when :frame |
| fd, frame = args |
| begin |
| otrans = @transport_factory.get_transport(fd) |
| oprot = @protocol_factory.get_protocol(otrans) |
| membuf = MemoryBufferTransport.new(frame) |
| itrans = @transport_factory.get_transport(membuf) |
| iprot = @protocol_factory.get_protocol(itrans) |
| @processor.process(iprot, oprot) |
| rescue => e |
| @logger.error "#{Thread.current.inspect} raised error: #{e.inspect}\n#{e.backtrace.join("\n")}" |
| end |
| end |
| end |
| end |
| end |
| end |
| end |
| end |
