proxy_pass: rewrite to be async, using rack.hijack

This allows our reverse proxy to avoid having an innefficient 1:1
relationship between threads and upstream connections, reducing
memory usage when there are many upstream connections (possibly to
multiple backend machines).

1 files changed, 117 insertions, 3 deletions

diff --git a/test/test_proxy_pass.rb b/test/test_proxy_pass.rb
index 5398b29..5bf8722 100644
--- a/test/test_proxy_pass.rb
+++ b/test/test_proxy_pass.rb
@@ -6,13 +6,45 @@ require 'json'
 class TestProxyPass < Testcase
   ENV["N"].to_i > 1 and parallelize_me!
   include ServerHelper
+  OMFG = 'a' * (1024 * 1024 * 32)
 
   class ProxiedApp
     def call(env)
       h = [ %w(Content-Length 3), %w(Content-Type text/plain) ]
       case env['REQUEST_METHOD']
       when 'GET'
-        [ 200, h, [ "hi\n"] ]
+        case env['PATH_INFO']
+        when '/giant-body'
+          h = [ %W(Content-Length #{OMFG.size}), %w(Content-Type text/plain) ]
+          [ 200, h, [ OMFG ] ]
+        when %r{\A/slow-headers-(\d+(?:\.\d+)?)\z}
+          delay = $1.to_f
+          io = env['rack.hijack'].call
+          [ "HTTP/1.1 200 OK\r\n",
+            "Content-Length: 7\r\n",
+            "Content-Type: text/PAIN\r\n",
+            "connection: close\r\n\r\n",
+            "HIHIHI!"
+          ].each do |l|
+            io.write(l)
+            sleep delay
+          end
+          io.close
+        when %r{\A/chunky-slow-(\d+(?:\.\d+)?)\z}
+          delay = $1.to_f
+          chunky = Object.new
+          chunky.instance_variable_set(:@delay, delay)
+          def chunky.each
+            sleep @delay
+            yield "3\r\nHI!\r\n"
+            sleep @delay
+            yield "0\r\n\r\n"
+          end
+          h = [ %w(Content-Type text/pain), %w(Transfer-Encoding chunked) ]
+          [ 200, h, chunky ]
+        else
+          [ 200, h, [ "hi\n"] ]
+        end
       when 'HEAD'
         [ 200, h, [] ]
       when 'PUT'
@@ -72,6 +104,7 @@ class TestProxyPass < Testcase
         stderr_path err.path
       end
     end
+
     Net::HTTP.start(host, port) do |http|
       res = http.request(Net::HTTP::Get.new('/f00'))
       assert_equal 200, res.code.to_i
@@ -109,7 +142,7 @@ class TestProxyPass < Testcase
       require 'yahns/proxy_pass'
       @srv2.close
       cfg.instance_eval do
-        app(:rack, Yahns::ProxyPass.new("http://#{host2}:#{port2}/")) do
+        app(:rack, Yahns::ProxyPass.new("http://#{host2}:#{port2}")) do
           listen "#{host}:#{port}"
         end
         stderr_path err.path
@@ -126,6 +159,8 @@ class TestProxyPass < Testcase
       end
     end
 
+    check_pipelining(host, port)
+
     gplv3 = File.open('COPYING')
 
     Net::HTTP.start(host, port) do |http|
@@ -138,7 +173,7 @@ class TestProxyPass < Testcase
       assert_equal n, res['Content-Length'].to_i
       assert_nil res.body
 
-      # chunked encoding
+      # chunked encoding (PUT)
       req = Net::HTTP::Put.new('/')
       req.body_stream = gplv3
       req.content_type = 'application/octet-stream'
@@ -148,6 +183,18 @@ class TestProxyPass < Testcase
       assert_equal gplv3.read, res.body
       assert_equal 201, res.code.to_i
 
+      # chunked encoding (GET)
+      res = http.request(Net::HTTP::Get.new('/chunky-slow-0.1'))
+      assert_equal 200, res.code.to_i
+      assert_equal 'chunked', res['Transfer-encoding']
+      assert_equal "HI!", res.body
+
+      # slow headers (GET)
+      res = http.request(Net::HTTP::Get.new('/slow-headers-0.01'))
+      assert_equal 200, res.code.to_i
+      assert_equal 'text/PAIN', res['Content-Type']
+      assert_equal 'HIHIHI!', res.body
+
       # normal content-length
       gplv3.rewind
       req = Net::HTTP::Put.new('/')
@@ -158,10 +205,77 @@ class TestProxyPass < Testcase
       gplv3.rewind
       assert_equal gplv3.read, res.body
       assert_equal 201, res.code.to_i
+
+      # giant body
+      res = http.request(Net::HTTP::Get.new('/giant-body'))
+      assert_equal 200, res.code.to_i
+      assert_equal OMFG, res.body
+    end
+
+    # ensure we do not chunk responses back to an HTTP/1.0 client even if
+    # the proxy <-> upstream connection is chunky
+    %w(0 0.1).each do |delay|
+      begin
+        h10 = TCPSocket.new(host, port)
+        h10.write "GET /chunky-slow-#{delay} HTTP/1.0\r\n\r\n"
+        res = Timeout.timeout(60) { h10.read }
+        assert_match %r{^Connection: close\r\n}, res
+        assert_match %r{^Content-Type: text/pain\r\n}, res
+        assert_match %r{\r\n\r\nHI!\z}, res
+        refute_match %r{^Transfer-Encoding:}, res
+        refute_match %r{\r0\r\n}, res
+      ensure
+        h10.close
+      end
     end
   ensure
     gplv3.close if gplv3
     quit_wait pid
     quit_wait pid2
   end
+
+  def check_pipelining(host, port)
+    pl = TCPSocket.new(host, port)
+    r1 = ''
+    r2 = ''
+    r3 = ''
+    Timeout.timeout(60) do
+      pl.write "GET / HTTP/1.1\r\nHost: example.com\r\n\r\nGET /"
+      until r1 =~ /hi\n/
+        r1 << pl.readpartial(666)
+      end
+
+      pl.write "chunky-slow-0.1 HTTP/1.1\r\nHost: example.com\r\n\r\nP"
+      until r2 =~ /\r\n3\r\nHI!\r\n0\r\n\r\n/
+        r2 << pl.readpartial(666)
+      end
+
+      if false
+        pl.write "ET / HTTP/1.1\r\nHost: example.com\r\n\r\n"
+        until r3 =~ /hi\n/
+          r3 << pl.readpartial(666)
+        end
+      else
+        pl.write "UT / HTTP/1.1\r\nHost: example.com\r\n"
+        pl.write "Transfer-Encoding: chunked\r\n\r\n"
+        pl.write "6\r\nchunky\r\n"
+        pl.write "0\r\n\r\n"
+
+        until r3 =~ /chunky/
+          r3 << pl.readpartial(666)
+        end
+      end
+    end
+    r1 = r1.split("\r\n").reject { |x| x =~ /^Date: / }
+    r2 = r2.split("\r\n").reject { |x| x =~ /^Date: / }
+    assert_equal 'HTTP/1.1 200 OK', r1[0]
+    assert_equal 'HTTP/1.1 200 OK', r2[0]
+    assert_match %r{\r\n\r\nchunky\z}, r3
+    assert_match %r{\AHTTP/1\.1 201 Created\r\n}, r3
+  rescue => e
+    warn [ e.class, e.message ].inspect
+    warn e.backtrace.join("\n")
+  ensure
+    pl.close
+  end
 end