Switch to non-blocking reads, set ZMQ::LINGER option and close all so…

…ckets before closing the context

lib/mongrel2/connection.rb

@@ -12,16 +12,25 @@ def initialize(uuid, sub, pub)
       # Connect to receive requests
       @reqs = CTX.socket(ZMQ::PULL)
       @reqs.connect(sub)
+      @reqs.setsockopt(ZMQ::LINGER, 0)
 
       # Connect to send responses
       @resp = CTX.socket(ZMQ::PUB)
       @resp.connect(pub)
       @resp.setsockopt(ZMQ::IDENTITY, uuid)
+      @resp.setsockopt(ZMQ::LINGER, 0)
     end
 
     def recv
-      msg = @reqs.recv(0)
-      msg.nil? ? nil : Request.parse(msg)
+      msg = nil
+      ready_sockets = ZMQ.select([@reqs], nil, nil, 30)
+      if !ready_sockets.nil?
+        ready_sockets[0].each do | socket |
+          msg = socket.recv(ZMQ::NOBLOCK)
+          msg = Request.parse(msg) unless msg.nil?
+        end
+      end
+      msg
     end
 
     def reply(req, body, status = 200, headers = {})
@@ -31,8 +40,9 @@ def reply(req, body, status = 200, headers = {})
     end
 
     def close
-      # I think I should be able to just close the context
-      CTX.close rescue nil
+      @resp.close
+      @reqs.close
+      CTX.close
     end
   end
 end

lib/rack/handler/mongrel2.rb

@@ -4,68 +4,62 @@
 module Rack
   module Handler
     class Mongrel2
-      class << self
-        def run(app, options = {})
-          options = {
-            :recv => ENV['RACK_MONGREL2_RECV'],
-            :send => ENV['RACK_MONGREL2_SEND'],
-            :uuid => ENV['RACK_MONGREL2_UUID']
-          }.merge(options)
-
-          raise ArgumentError.new('Must specify an :recv or set RACK_MONGREL2_RECV') if options[:recv].nil?
-          raise ArgumentError.new('Must specify an :recv or set RACK_MONGREL2_SEND') if options[:send].nil?
-          raise ArgumentError.new('Must specify an :uuid or set RACK_MONGREL2_UUID') if options[:uuid].nil?
+
+      def self.run(app, options = {})
+
+        raise ArgumentError.new('Must specify :recv') if options[:recv].nil?
+        raise ArgumentError.new('Must specify :recv') if options[:send].nil?
+        raise ArgumentError.new('Must specify :uuid') if options[:uuid].nil?
 
-          conn = ::Mongrel2::Connection.new(options[:uuid], options[:recv], options[:send])
+        conn = ::Mongrel2::Connection.new(options[:uuid], options[:recv], options[:send])
 
-          running = true
+        running = true
 
-          # This doesn't work at all until zmq fixes their shit (in 2.1.x I think), but trap it now anyway.
-          %w(INT TERM KILL).each do |sig|
-            trap(sig) do
-              conn.close
-              running = false
-            end
+        %w(INT TERM KILL).each do | sig |
+          Signal.trap(sig) do
+            running = false
           end
+        end
 
-          while running
-            req = conn.recv rescue nil
-            next if req.nil? || req.disconnect?
-            break if !running
+        while running do
+          req = conn.recv
+          next if req.nil? || req.disconnect?
+          break if !running
 
-            script_name = ENV['RACK_RELATIVE_URL_ROOT'] || req.headers['PATTERN'].split('(', 2).first.gsub(/\/$/, '')
-            env = {
-              'rack.version' => Rack::VERSION,
-              'rack.url_scheme' => 'http', # Only HTTP for now
-              'rack.input' => StringIO.new(req.body),
-              'rack.errors' => $stderr,
-              'rack.multithread' => true,
-              'rack.multiprocess' => true,
-              'rack.run_once' => false,
-              'mongrel2.pattern' => req.headers['PATTERN'],
-              'REQUEST_METHOD' => req.headers['METHOD'],
-              'SCRIPT_NAME' => script_name,
-              'PATH_INFO' => req.headers['PATH'].gsub(script_name, ''),
-              'QUERY_STRING' => req.headers['QUERY'] || ''
-            }
+          script_name = ENV['RACK_RELATIVE_URL_ROOT'] || req.headers['PATTERN'].split('(', 2).first.gsub(/\/$/, '')
+          env = {
+            'rack.version' => Rack::VERSION,
+            'rack.url_scheme' => 'http', # Only HTTP for now
+            'rack.input' => StringIO.new(req.body),
+            'rack.errors' => $stderr,
+            'rack.multithread' => true,
+            'rack.multiprocess' => true,
+            'rack.run_once' => false,
+            'mongrel2.pattern' => req.headers['PATTERN'],
+            'REQUEST_METHOD' => req.headers['METHOD'],
+            'SCRIPT_NAME' => script_name,
+            'PATH_INFO' => req.headers['PATH'].gsub(script_name, ''),
+            'QUERY_STRING' => req.headers['QUERY'] || ''
+          }
 
-            env['SERVER_NAME'], env['SERVER_PORT'] = req.headers['host'].split(':', 2)
-            req.headers.each do |key, val|
-              unless key =~ /content_(type|length)/i
-                key = "HTTP_#{key.upcase.gsub('-', '_')}"
-              end
-              env[key] = val
+          env['SERVER_NAME'], env['SERVER_PORT'] = req.headers['host'].split(':', 2)
+          req.headers.each do |key, val|
+            unless key =~ /content_(type|length)/i
+              key = "HTTP_#{key.upcase.gsub('-', '_')}"
             end
-
-            status, headers, rack_response = app.call(env)
-            body = ''
-            rack_response.each { |b| body << b }
-            conn.reply(req, body, status, headers)
+            env[key] = val
           end
-        ensure
-          conn.close if conn.respond_to?(:close)
+
+          status, headers, rack_response = app.call(env)
+          body = ''
+          rack_response.each { |b| body << b }
+          conn.reply(req, body, status, headers)
         end
-      end
-    end
-  end
-end
+
+        conn.close
+        Process.exit!
+      end #def self.run
+
+    end #class Mongrel2
+  end #module Handler
+end #module Rack