when wrapping the body, body.close writes the log

We no longer write the log out at the end of the body.each call.
This is a behavioral change, but fortunately all Rack servers
I've seen call body.close inside an ensure.

This allows us to later pass along the "to_path" method
and not rely on "each" to write the log.

3 files changed, 21 insertions, 15 deletions

diff --git a/ext/clogger_ext/clogger.c b/ext/clogger_ext/clogger.c
index 415fe32..100392d 100644
--- a/ext/clogger_ext/clogger.c
+++ b/ext/clogger_ext/clogger.c
@@ -638,12 +638,11 @@ static VALUE body_iter_i(VALUE str, VALUE memop)
         return rb_yield(str);
 }
 
-static VALUE wrap_each(struct clogger *c)
+static VALUE wrap_close(struct clogger *c)
 {
-        c->body_bytes_sent = 0;
-        rb_iterate(rb_each, c->body, body_iter_i, (VALUE)&c->body_bytes_sent);
-
-        return c->body;
+        if (rb_respond_to(c->body, close_id))
+                return rb_funcall(c->body, close_id, 0);
+        return Qnil;
 }
 
 /**
@@ -659,8 +658,10 @@ static VALUE clogger_each(VALUE self)
         struct clogger *c = clogger_get(self);
 
         rb_need_block();
+        c->body_bytes_sent = 0;
+        rb_iterate(rb_each, c->body, body_iter_i, (VALUE)&c->body_bytes_sent);
 
-        return rb_ensure(wrap_each, (VALUE)c, cwrite, (VALUE)c);
+        return self;
 }
 
 /**
@@ -675,9 +676,7 @@ static VALUE clogger_close(VALUE self)
 {
         struct clogger *c = clogger_get(self);
 
-        if (rb_respond_to(c->body, close_id))
-                return rb_funcall(c->body, close_id, 0);
-        return Qnil;
+        return rb_ensure(wrap_close, (VALUE)c, cwrite, (VALUE)c);
 }
 
 /* :nodoc: */
diff --git a/lib/clogger/pure.rb b/lib/clogger/pure.rb
index e9a8e6a..50e4f6e 100644
--- a/lib/clogger/pure.rb
+++ b/lib/clogger/pure.rb
@@ -43,12 +43,13 @@ class Clogger
       @body_bytes_sent += Rack::Utils.bytesize(part)
       yield part
     end
-    ensure
-      log(@env, @status, @headers)
+    self
   end
 
   def close
     @body.close if @body.respond_to?(:close)
+    ensure
+      log(@env, @status, @headers)
   end
 
   def reentrant?
diff --git a/test/test_clogger.rb b/test/test_clogger.rb
index a20793c..1311017 100644
--- a/test/test_clogger.rb
+++ b/test/test_clogger.rb
@@ -55,6 +55,7 @@ class TestClogger < Test::Unit::TestCase
     assert_equal("302 Found", status)
     assert_equal({}, headers)
     body.each { |part| assert false }
+    body.close
     str = str.string
     r = %r{\Ahome - - \[[^\]]+\] "GET /hello\?goodbye=true HTTP/1.0" 302 -\n\z}
     assert_match r, str
@@ -134,7 +135,9 @@ class TestClogger < Test::Unit::TestCase
       'HTTP_COOKIE' => cookie,
     }
     req = @req.merge(req)
-    cl.call(req).last.each { |part| part }
+    body = cl.call(req).last
+    body.each { |part| part }
+    body.close
     str = str.string
     assert(str.size > 128)
     assert_match %r["echo and socat \\o/" "#{cookie}" \d+\.\d{3}], str
@@ -221,6 +224,7 @@ class TestClogger < Test::Unit::TestCase
     status, headers, body = cl.call(@req)
     tmp = []
     body.each { |s| tmp << s }
+    body.close
     assert_equal %w(a b c), tmp
     str = str.string
     assert_match %r[" 200 3 \d+\.\d{4}\n\z], str
@@ -279,6 +283,7 @@ class TestClogger < Test::Unit::TestCase
     cl = Clogger.new(app, :logger => str, :format => '$response_length')
     status, header, bodies = cl.call(@req)
     bodies.each { |part| part }
+    bodies.close
     assert_equal "-\n", str.string
   end
 
@@ -290,6 +295,7 @@ class TestClogger < Test::Unit::TestCase
     status, headers, body = cl.call(@req)
     tmp = []
     body.each { |s| tmp << s }
+    body.close
     assert_equal %w(a b c), tmp
     str = str.string
     assert_match %r[" 200 3 "-" "echo and socat \\o/"\n\z], str
@@ -402,7 +408,7 @@ class TestClogger < Test::Unit::TestCase
     req = Rack::Utils::HeaderHash.new(@req)
     app = lambda { |env| [302, [ %w(a) ], []] }
     cl = Clogger.new(app, :logger => str, :format => Rack_1_0)
-    assert_nothing_raised { cl.call(req).last.each {} }
+    assert_nothing_raised { cl.call(req).last.each {}.close }
     assert str.size > 0
   end
 
@@ -415,7 +421,7 @@ class TestClogger < Test::Unit::TestCase
     }
     app = lambda { |env| [302, [ %w(a) ], []] }
     cl = Clogger.new(app, :logger => str, :format => Rack_1_0)
-    assert_nothing_raised { cl.call(req).last.each {} }
+    assert_nothing_raised { cl.call(req).last.each {}.close }
     assert str.size > 0
   end
 
@@ -425,7 +431,7 @@ class TestClogger < Test::Unit::TestCase
     r = nil
     app = lambda { |env| [302, [ %w(a) ], [FooString.new(body)]] }
     cl = Clogger.new(app, :logger => str, :format => '$body_bytes_sent')
-    assert_nothing_raised { cl.call(@req).last.each { |x| r = x } }
+    assert_nothing_raised { cl.call(@req).last.each { |x| r = x }.close }
     assert str.size > 0
     assert_equal body.size.to_s << "\n", str.string
     assert_equal r, body