mem: allow using io.WriterTo with a io.LimitedReader

[GiedriusS]

The caller of this function can wrap the io.Reader in a io.LimitedReader. This happens if some max message size is set. If so, this io.WriterTo check doesn't work anymore. Work around this by checking if it is maybe a io.LimitedReader.

Overall, the problem I'm trying to solve is that the constant

		buf := pool.Get(readAllBufSize)

32KiB is way too much in our use case. Messages are typically at max about only 1KiB in size so we always overallocate by ~31KiB in the best case scenario so we want to use the io.WriterTo branch so that we could appropriately size the buffer.

Is this OK? Any suggestions on how to make the code prettier? Also, maybe some suggestions on how to do something like io.LimitedReader on the io.WriterTo?

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (master@7472d57). Learn more about missing BASE report.
⚠️ Report is 4 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff            @@
##             master    #8697   +/-   ##
=========================================
  Coverage          ?   83.25%           
=========================================
  Files             ?      416           
  Lines             ?    32306           
  Branches          ?        0           
=========================================
  Hits              ?    26895           
  Misses            ?     4027           
  Partials          ?     1384           

Files with missing lines	Coverage Δ
mem/buffer_slice.go	96.57% <100.00%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[GiedriusS]

@ash2k any thoughts?

[easwars]

@arjan-bal : Can you please take a first look and see if this seems ok.

[ash2k]

I think this would circumvent the limit of the LimitedReader, wouldn't it? By directly accessing the underlying reader we are bypassing the limiter.

[GiedriusS]

Yeah but I'm not sure how to do a "LimitedWriter" if we could say it that way. Should we add something like https://github.com/nanmu42/limitio/blob/master/limitio.go to grpc-go code (as a library or our own impl)?

[arjan-bal]

Hi @GiedriusS, one solution I can think of is to create your own wrapper struct that wraps a io.LimitReader and also implements io.WriteTo. This would allow your reader to control the size of the temporary buffer being used. Here's a example implementation could work.

type LimitWriterTo struct {
    Reader io.Reader // The underlying io.LimitReader
}

func (l *LimitWriterTo) WriteTo(w io.Writer) (n int64, err error) {
    // Define your custom buffer size here (e.g., 64K, 128K)
    buffer := make([]byte, 1024) // You could get this from a buffer pool.
    
    // Use io.CopyBuffer internally with your custom buffer
    // or implement the read/write loop manually for ultimate control.
    return io.CopyBuffer(w, l.Reader, buffer) 
}

[GiedriusS]

The core issue is that grpc-go does an assertion (and it wraps io.Reader inside of a io.LimitedReader itself) whether it's a io.Reader and io.LimitedReader is not a io.Reader so I think this path would never be hit.

[arjan-bal]

Oh, I overlooked the line in the PR description: "This happens if some max message size is set."

gRPC controls the the reader type and not external code. Let me think about it a little more.

[arjan-bal]

How about we set the size of the copy buffer in the for loop below as min(readAllBufSize, LimitReader.N)?