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 15 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)?

[GiedriusS]

The default is quite huge

grpc-go/server.go

Line 59 in 1ed87ec

defaultServerMaxReceiveMessageSize = 1024 * 1024 * 4

so it wouldn't help in our case because the messages are really small (<1KB). Would you accept adding a LimitedWriter to grpc-go code?

[arjan-bal]

I see. Can you describe how you plan to use the new LimitedReader type to workaround this?

I have one approach in mind which involves having the LimitedReader conditionally implement io.WriterTo using a type assertion if the wrapped reader also implements io.WriterTo. I wanted to know if there's a simpler solution for this.

[GiedriusS]

I think you're correct - the decompressor through WriterTo will still be able to read as much as they want. I'm now actually thinking maybe really the best way would be to provide a way to use some specific bucketed []byte pool - that way the buckets would be created according to how big the users of grpc-go anticipate the inputs will be.

In the end there is probably no best decision and it depends on what do you think is an appropriate abstraction is here in the context of grpc-go.

[arjan-bal]

Here is the simplest solution I came up with:

1. Introduce a LimitWriter: Create a wrapper around io.Writer that restricts the number of bytes written. If a write exceeds the limit, it returns a specific sentinel error (e.g., ErrLimitExhausted).
2. Update ReadAll: In ReadAll, we check if the reader is an *io.LimitedReader and if the underlying io.Reader implements io.WriterTo. If so, we create a new writer using NewWriter(&result, pool), wrap it in our LimitWriter, and call WriteTo on the underlying reader. This effectively transfers the limit constraint from the reader to the writer.

To keep this optimization transparent to callers, ReadAll must trap the error returned by LimitWriter. It should translate that error into a successful return (nil error) and update the N field on the *io.LimitedReader to reflect the bytes actually consumed.

---

Here are some snippets to explain this:

The Helper Type

var ErrLimitExhausted = errors.New("limit exhausted")

type limitWriter struct {
	w io.Writer
	n int64
}

func (l *limitWriter) Write(p []byte) (n int, err error) {
	if l.n <= 0 {
		return 0, ErrLimitExhausted
	}
	// If p is larger than remaining limit, truncate it
	if int64(len(p)) > l.n {
		p = p[:l.n]
		err = ErrLimitExhausted // We will return this after the underlying write
	}
	n, wErr := l.w.Write(p)
	l.n -= int64(n)
	
	if wErr != nil {
		return n, wErr
	}
	return n, err
}

The Updated ReadAll Function

Here is how the logic fits into your existing context.

func ReadAll(r io.Reader, pool BufferPool) (BufferSlice, error) {
	var result BufferSlice

	// 1. Check for the specific optimization case: *io.LimitedReader wrapping an io.WriterTo
	if lr, ok := r.(*io.LimitedReader); ok {
		if wt, ok := lr.R.(io.WriterTo); ok {
			w := NewWriter(&result, pool)
			
			// Wrap the writer to enforce the reader's limit
			lw := &limitWriter{w: w, n: lr.N}
			
			// Delegate the heavy lifting to the underlying implementation
			n, err := wt.WriteTo(lw)
			
			// Update the LimitedReader state
			lr.N -= n
			
			// If we hit the limit, it's not a real error in this context; 
			// it just means we read everything we were allowed to.
			if err == ErrLimitExhausted {
				return result, nil
			}
			return result, err
		}
	}

	// 2. Standard optimization for direct io.WriterTo
	if wt, ok := r.(io.WriterTo); ok {
		w := NewWriter(&result, pool)
		_, err := wt.WriteTo(w)
		return result, err
	}

	// ... rest of fallback implementation ...
}

What do you think?

[github-actions]

This PR is labeled as requiring an update from the reporter, and no update has been received after 6 days. If no update is provided in the next 7 days, this issue will be automatically closed.