connection: split large frame in #send_immediately

6 files changed, 53 insertions, 44 deletions

diff --git a/lib/plum/binary_string.rb b/lib/plum/binary_string.rb
index 400c57b..2b50b38 100644
--- a/lib/plum/binary_string.rb
+++ b/lib/plum/binary_string.rb
@@ -70,6 +70,20 @@ module Plum
           # I want to write `enum_for(__method__, n)`!
         end
       end
+
+      # Splits this String into chunks.
+      # @param n [Integer] max chunk bytesize
+      # @return [Array<String>] the slices
+      def chunk(n)
+        res = []
+        pos = 0
+        lim = bytesize
+        while pos < lim
+          res << byteslice(pos, n)
+          pos += n
+        end
+        res
+      end
     end
   end
 end
diff --git a/lib/plum/connection.rb b/lib/plum/connection.rb
index f158560..73f6206 100644
--- a/lib/plum/connection.rb
+++ b/lib/plum/connection.rb
@@ -92,7 +92,14 @@ module Plum
 
     def send_immediately(frame)
       callback(:send_frame, frame)
-      @writer.call(frame.assemble)
+
+      if frame.length <= @remote_settings[:max_frame_size]
+        @writer.call(frame.assemble)
+      else
+        frame.split(@remote_settings[:max_frame_size]) { |splitted|
+          @writer.call(splitted.assemble)
+        }
+      end
     end
 
     def validate_received_frame(frame)
diff --git a/lib/plum/frame_utils.rb b/lib/plum/frame_utils.rb
index 20c79de..447c154 100644
--- a/lib/plum/frame_utils.rb
+++ b/lib/plum/frame_utils.rb
@@ -3,33 +3,28 @@ using Plum::BinaryString
 
 module Plum
   module FrameUtils
-    # Splits the DATA frame into multiple frames if the payload size exceeds max size.
+    # Splits this frame into multiple frames not to exceed MAX_FRAME_SIZE.
     # @param max [Integer] The maximum size of a frame payload.
-    # @return [Array<Frame>] The splitted frames.
-    def split_data(max)
-      return [self] if self.length <= max
-      raise "Frame type must be DATA" unless self.type == :data
-
-      fragments = self.payload.each_byteslice(max).to_a
-      frames = fragments.map {|fragment| Frame.new(type: :data, flags: [], stream_id: self.stream_id, payload: fragment) }
-      frames.first.flags = self.flags - [:end_stream]
-      frames.last.flags = self.flags & [:end_stream]
-      frames
-    end
-
-    # Splits the HEADERS or PUSH_PROMISE frame into multiple frames if the payload size exceeds max size.
-    # @param max [Integer] The maximum size of a frame payload.
-    # @return [Array<Frame>] The splitted frames.
-    def split_headers(max)
-      return [self] if self.length <= max
-      raise "Frame type must be HEADERS or PUSH_PROMISE" unless [:headers, :push_promise].include?(self.type)
-
-      fragments = self.payload.each_byteslice(max).to_a
-      frames = fragments.map {|fragment| Frame.new(type: :continuation, flags: [], stream_id: self.stream_id, payload: fragment) }
-      frames.first.type_value = self.type_value
-      frames.first.flags = self.flags - [:end_headers]
-      frames.last.flags = self.flags & [:end_headers]
-      frames
+    # @yield [Frame] The splitted frames.
+    def split(max)
+      return yield self if @length <= max
+      first, *mid, last = @payload.chunk(max)
+      case type
+      when :data
+        yield Frame.new(type_value: 0, stream_id: @stream_id, payload: first, flags_value: @flags_value & ~1)
+        mid.each { |slice|
+          yield Frame.new(type_value: 0, stream_id: @stream_id, payload: slice, flags_value: 0)
+        }
+        yield Frame.new(type_value: 0, stream_id: @stream_id, payload: last, flags_value: @flags_value & 1)
+      when :headers, :push_promise
+        yield Frame.new(type_value: @type_value, stream_id: @stream_id, payload: first, flags_value: @flags_value & ~4)
+        mid.each { |slice|
+          yield Frame.new(type: :continuation, stream_id: @stream_id, payload: slice, flags_value: 0)
+        }
+        yield Frame.new(type: :continuation, stream_id: @stream_id, payload: last, flags_value: @flags_value & 4)
+      else
+        raise NotImplementedError.new("frame split of frame with type #{type} is not supported")
+      end
     end
 
     # Parses SETTINGS frame payload. Ignores unknown settings type (see RFC7540 6.5.2).
diff --git a/lib/plum/server/http_connection.rb b/lib/plum/server/http_connection.rb
index 74b7ad8..116fdcb 100644
--- a/lib/plum/server/http_connection.rb
+++ b/lib/plum/server/http_connection.rb
@@ -77,9 +77,8 @@ module Plum
                                   ":authority" => @_headers["host"] })
                          .reject {|n, v| ["connection", "http2-settings", "upgrade", "host"].include?(n) }
 
-      headers_s = Frame.headers(1, encoder.encode(headers), :end_headers).split_headers(max_frame_size) # stream ID is 1
-      data_s = Frame.data(1, @_body, :end_stream).split_data(max_frame_size)
-      (headers_s + data_s).each {|frag| stream.receive_frame(frag) }
+      stream.receive_frame Frame.headers(1, encoder.encode(headers), :end_headers)
+      stream.receive_frame Frame.data(1, @_body, :end_stream)
     end
   end
 end
diff --git a/lib/plum/stream.rb b/lib/plum/stream.rb
index 87e228e..04190fc 100644
--- a/lib/plum/stream.rb
+++ b/lib/plum/stream.rb
@@ -58,7 +58,6 @@ module Plum
     end
 
     # Closes this stream. Sends RST_STREAM frame to the peer.
-    # @param error_type [Symbol] The error type to be contained in the RST_STREAM frame.
     def close
       @state = :closed
       callback(:close)
diff --git a/lib/plum/stream_utils.rb b/lib/plum/stream_utils.rb
index e344528..ff14c7a 100644
--- a/lib/plum/stream_utils.rb
+++ b/lib/plum/stream_utils.rb
@@ -9,10 +9,8 @@ module Plum
     def promise(headers)
       stream = @connection.reserve_stream(weight: self.weight + 1, parent: self)
       encoded = @connection.hpack_encoder.encode(headers)
-      original = Frame.push_promise(id, stream.id, encoded, :end_headers)
-      original.split_headers(@connection.remote_settings[:max_frame_size]).each do |frame|
-        send frame
-      end
+      frame = Frame.push_promise(id, stream.id, encoded, :end_headers)
+      send frame
       stream
     end
 
@@ -22,10 +20,8 @@ module Plum
     def send_headers(headers, end_stream:)
       max = @connection.remote_settings[:max_frame_size]
       encoded = @connection.hpack_encoder.encode(headers)
-      original_frame = Frame.headers(id, encoded, :end_headers, (end_stream && :end_stream || nil))
-      original_frame.split_headers(max).each do |frame|
-        send frame
-      end
+      frame = Frame.headers(id, encoded, :end_headers, (end_stream && :end_stream || nil))
+      send frame
       @state = :half_closed_local if end_stream
     end
 
@@ -35,14 +31,13 @@ module Plum
     def send_data(data, end_stream: true)
       max = @connection.remote_settings[:max_frame_size]
       if data.is_a?(IO)
-        while !data.eof? && fragment = data.readpartial(max)
+        until data.eof?
+          fragment = data.readpartial(max)
           send Frame.data(id, fragment, (end_stream && data.eof? && :end_stream))
         end
       else
-        original = Frame.data(id, data, (end_stream && :end_stream))
-        original.split_data(max).each do |frame|
-          send frame
-        end
+        frame = Frame.data(id, data, (end_stream && :end_stream))
+        send frame
       end
       @state = :half_closed_local if end_stream
     end