Add IO::PrefixSuffixBuffer for capturing truncated streams

spec/std/io/prefix_suffix_buffer_spec.cr

@@ -0,0 +1,119 @@
+require "../spec_helper"
+require "io/prefix_suffix_buffer"
+
+describe IO::PrefixSuffixBuffer do
+  describe ".new(Int)" do
+    it "allocates buffer" do
+      IO::PrefixSuffixBuffer.new(12).capacity.should eq(24)
+    end
+
+    it "raises if negative capacity" do
+      expect_raises(ArgumentError, "Negative capacity") do
+        IO::PrefixSuffixBuffer.new(-1)
+      end
+    end
+  end
+
+  describe ".new(Bytes, Bytes)" do
+    it "receives buffers" do
+      IO::PrefixSuffixBuffer.new(Bytes.new(5), Bytes.new(7)).capacity.should eq(12)
+    end
+  end
+
+  it "#total_size" do
+    (IO::PrefixSuffixBuffer.new << "foo").total_size.should eq(3)
+  end
+
+  describe "#write" do
+    it "writes" do
+      io = IO::PrefixSuffixBuffer.new
+      io.total_size.should eq(0)
+      io.write Slice.new("hello".to_unsafe, 3)
+      io.total_size.should eq(3)
+      io.to_s.should eq("hel")
+    end
+
+    it "writes to capacity" do
+      s = "hi" * 100
+      io = IO::PrefixSuffixBuffer.new(100)
+      io.write Slice.new(s.to_unsafe, s.bytesize)
+      io.to_s.should eq(s)
+    end
+
+    it "writes single byte" do
+      io = IO::PrefixSuffixBuffer.new
+      io.write_byte 97_u8
+      io.to_s.should eq("a")
+    end
+
+    it "writes multiple times" do
+      io = IO::PrefixSuffixBuffer.new
+      io << "foo" << "bar"
+      io.to_s.should eq("foobar")
+    end
+
+    it "supports an empty prefix buffer" do
+      io = IO::PrefixSuffixBuffer.new(Bytes.empty, Bytes.new(3))
+      io << "abcdef"
+      io.to_s.should eq("\n...omitted 3 bytes...\ndef")
+    end
+
+    it "supports an empty suffix buffer" do
+      io = IO::PrefixSuffixBuffer.new(Bytes.new(3), Bytes.empty)
+      io << "abcdef"
+      io.to_s.should eq("abc\n...omitted 3 bytes...\n")
+    end
+  end
+
+  describe "#to_s" do
+    it "appends to another buffer" do
+      s1 = IO::PrefixSuffixBuffer.new
+      s1 << "hello"
+
+      s2 = IO::PrefixSuffixBuffer.new
+      s1.to_s(s2)
+      s2.to_s.should eq("hello")
+    end
+
+    it "appends to itself" do
+      io = IO::PrefixSuffixBuffer.new(33)
+      io << "-" * 33
+      io.to_s(io)
+      io.to_s.should eq "-" * 66
+    end
+
+    describe "truncation" do
+      it "basic" do
+        io = IO::PrefixSuffixBuffer.new(5)
+        io << "abcdefghijklmnopqrstuvwxyz"
+        io.to_s.should eq("abcde\n...omitted 16 bytes...\nvwxyz")
+      end
+
+      it "chunked" do
+        buffer = IO::PrefixSuffixBuffer.new(4)
+        buffer << "---" << "-X-" << "---"
+        buffer.to_s.should eq "----\n...omitted 1 bytes...\n----"
+      end
+
+      it "one" do
+        io = IO::PrefixSuffixBuffer.new(1)
+        io << "abc"
+        io.to_s.should eq("a\n...omitted 1 bytes...\nc")
+      end
+
+      it "longer" do
+        io = IO::PrefixSuffixBuffer.new(10)
+        io << "abcdefghijklmnopqrstuvwxyz"
+        io.to_s.should eq("abcdefghij\n...omitted 6 bytes...\nqrstuvwxyz")
+      end
+
+      it "chars" do
+        io = IO::PrefixSuffixBuffer.new(10)
+        "abcdefghijklmnopqrstuvwxyz".each_char do |char|
+          io << char
+        end
+        io.to_s.should eq("abcdefghij\n...omitted 6 bytes...\nqrstuvwxyz")
+      end
+    end
+  end
+end

spec/std/process_spec.cr

@@ -848,7 +848,7 @@ describe Process do
       result.error?.should be_nil
     end
 
-    pending "truncates error output", tags: %w[slow] do
+    it "truncates error output", tags: %w[slow] do
       dashes32 = "-" * (32 << 10)
       input = IO::Memory.new("#{dashes32}X#{dashes32}")
       result = Process.capture_result(to_ary(stdin_to_stderr_command), input: input)
@@ -942,7 +942,7 @@ describe Process do
       result.error?.should be_nil
     end
 
-    pending "truncates error output", tags: %w[slow] do
+    it "truncates error output", tags: %w[slow] do
       dashes32 = "-" * (32 << 10)
       input = IO::Memory.new("#{dashes32}X#{dashes32}")
       result = Process.capture_result?(to_ary(stdin_to_stderr_command), input: input).should be_a(Process::Result)

src/io/prefix_suffix_buffer.cr

@@ -0,0 +1,119 @@
+require "io/memory"
+
+# `PrefixSuffixBuffer` is an `IO` that retains the first bytes in the prefix
+# buffer and the last bytes in the suffix buffer.
+#
+# When writing more bytes than fit in the buffers, only the start and end bytes
+# are preserved.
+# `#to_s` renders the buffer contents with a message indicating how many
+# bytes were omitted.
+class IO::PrefixSuffixBuffer < IO
+  # Creates an instance with prefix and suffix buffers of the given size.
+  #
+  # The maximum size that this instance can consume without omitting data is
+  # `size * 2`.
+  def self.new(size : Int32 = 32) : self
+    buffer_size = size * 2
+    String.check_capacity_in_bounds(buffer_size)
+
+    buffer = Bytes.new(buffer_size)
+    new(buffer[0, size], buffer + size)
+  end
+
+  # Creates an instance using the given buffers.
+  #
+  # The maximum size that this instance can consume without omitting data is the
+  # sum of the buffer sizes.
+  def initialize(@prefix : Bytes, @suffix : Bytes)
+    @pos = 0
+  end
+
+  def read(slice : Bytes) : NoReturn
+    raise "Can't read from IO::PrefixSuffixBuffer"
+  end
+
+  def write(slice : Bytes) : Nil
+    check_open
+
+    return if slice.empty?
+
+    total = slice.size
+
+    # Fill the prefix linearly.
+    slice += fill(@prefix, slice, @pos)
+
+    # The suffix works as a ring buffer.
+    suffix = @suffix
+    if slice.size >= suffix.bytesize
+      ring_pos = slice.size
+      slice = slice[(-suffix.bytesize)..]
+    else
+      ring_pos = -suffix.bytesize
+    end
+
+    # The first chunk goes to the ring buffer after `ring_pos`
+    slice += fill(suffix, slice, ((@pos - @prefix.bytesize).clamp(0..) + ring_pos) % suffix.bytesize) if suffix.bytesize > 0
+
+    # The second chunk goes to the ring buffer before `ring_pos`
+    fill(suffix, slice, 0)
+
+    # We add the total size, not the number of actually written bytes because
+    # we may skip writing some of the middle bytes if `total` exceeds the buffer
+    # capacity.
+    @pos += total
+  end
+
+  private def fill(buffer, slice, pos)
+    max_size = buffer.bytesize - pos
+
+    return 0 if max_size <= 0
+    count = slice.size.clamp(..max_size)
+    slice[0, count].copy_to(buffer.to_unsafe + pos, count)
+    count
+  end
+
+  def capacity
+    @prefix.bytesize + @suffix.bytesize
+  end
+
+  def total_size
+    @pos
+  end
+
+  def to_s : String
+    capacity = total_size.clamp(0, @prefix.bytesize + @suffix.bytesize + 50)
+    String.build(capacity) do |io|
+      to_s(io)
+    end
+  end
+
+  # Appends the buffer to the given `IO`.
+  #
+  # When the total size of the consumed data exceeds the buffer size, the middle
+  # part is omitted and replaced by a message that indicates the number of
+  # skipped bytes.
+  def to_s(io : IO) : Nil
+    prefix = @prefix
+    suffix = @suffix
+    total = self.total_size
+
+    buffer_size = prefix.bytesize + suffix.bytesize
+
+    io.write prefix[0, total.clamp(..prefix.bytesize)]
+
+    ring_pos = total - prefix.bytesize
+
+    if ring_pos > suffix.bytesize
+      io << "\n...omitted " << (total - buffer_size) << " bytes...\n"
+
+      if suffix.bytesize > 0
+        ring_pos %= suffix.bytesize
+        io.write suffix + ring_pos
+      end
+    end
+
+    if suffix.bytesize > 0
+      io.write suffix[0, ring_pos.clamp(0..)]
+    end
+  end
+end

src/process/capture.cr

@@ -35,7 +35,8 @@ class Process
     #
     # If `error` was not captured, returns the empty string.
     #
-    # The captured error stream might be truncated.
+    # The captured error stream might be truncated. If the total output is larger
+    # than 64kB, only the first 32kB and the last 32kB are preserved.
     def error : String
       @error || ""
     end
@@ -44,7 +45,8 @@ class Process
     #
     # If `error` was not captured, returns `nil`.
     #
-    # The captured error stream might be truncated.
+    # The captured error stream might be truncated. If the total output is larger
+    # than 64kB, only the first 32kB and the last 32kB are preserved.
     def error? : String?
       @error
     end
@@ -93,7 +95,7 @@ class Process
 
   private def self.capture_result_impl(output, error, & : -> Process)
     if error == Redirect::Pipe
-      error = captured_error = IO::Memory.new
+      error = captured_error = IO::PrefixSuffixBuffer.new(32 << 10)
     end
 
     process = yield error