Restore message logging and clear write state on exceptional exits

lib/net/protocol.rb

@@ -356,9 +356,9 @@ def writing
       @debug_output << '<- ' if @debug_output
       yield
       @debug_output << "\n" if @debug_output
-      bytes = @written_bytes
+      @written_bytes
+    ensure
       @written_bytes = nil
-      bytes
     end
 
     def write0(*strs)
@@ -426,12 +426,15 @@ def initialize(*, **)
     def each_message_chunk
       LOG 'reading message...'
       LOG_off()
-      read_bytes = 0
-      while (line = readuntil("\r\n")) != ".\r\n"
-        read_bytes += line.size
-        yield line.delete_prefix('.')
+      begin
+        read_bytes = 0
+        while (line = readuntil("\r\n")) != ".\r\n"
+          read_bytes += line.size
+          yield line.delete_prefix('.')
+        end
+      ensure
+        LOG_on()
       end
-      LOG_on()
       LOG "read message (#{read_bytes} bytes)"
     end
 
@@ -457,29 +460,35 @@ def write_message_0(src)
     def write_message(src)
       LOG "writing message from #{src.class}"
       LOG_off()
-      len = writing {
-        using_each_crlf_line {
-          write_message_0 src
+      begin
+        len = writing {
+          using_each_crlf_line {
+            write_message_0 src
+          }
         }
-      }
-      LOG_on()
+      ensure
+        LOG_on()
+      end
       LOG "wrote #{len} bytes"
       len
     end
 
     def write_message_by_block(&block)
       LOG 'writing message from block'
       LOG_off()
-      len = writing {
-        using_each_crlf_line {
-          begin
-            block.call(WriteAdapter.new(self.method(:write_message_0)))
-          rescue LocalJumpError
-            # allow `break' from writer block
-          end
+      begin
+        len = writing {
+          using_each_crlf_line {
+            begin
+              block.call(WriteAdapter.new(self.method(:write_message_0)))
+            rescue LocalJumpError
+              # allow `break' from writer block
+            end
+          }
         }
-      }
-      LOG_on()
+      ensure
+        LOG_on()
+      end
       LOG "wrote #{len} bytes"
       len
     end
@@ -499,6 +508,8 @@ def using_each_crlf_line
         write0 "\r\n"
       end
       write0 ".\r\n"
+      nil
+    ensure
       @wbuf = nil
     end