Skip to content

receive.handler pooling improvements#299

Open
gavins-db wants to merge 21 commits intodatabricks:db_mainfrom
gavins-db:gs/pooling-improvements
Open

receive.handler pooling improvements#299
gavins-db wants to merge 21 commits intodatabricks:db_mainfrom
gavins-db:gs/pooling-improvements

Conversation

@gavins-db
Copy link

@gavins-db gavins-db commented Feb 19, 2026

Changes

This PR has three primary changes related to existing pooling in the receive handler that should improve memory usage and performance. Those changes are detailed blow. In addition to the changes above, I added a syncutil.Pool utility that is a generic wrapper for pooling any type. This is a pattern I've found extremely useful in past lives to reduce boilerplate, improve readability, and make sure we're actually using pools correctly. And I also added flags to configure the pooling behavior (enable, max capacity, etc.).

  1. The first change removes the copyBufPool pool which was used as an intermediate buffer for compressed request bodies. This was an unnecessary intermediate buffer; the destination buffer (a bytes.Buffer) has a ReadFrom method that can more efficiently copy over the contents of the source reader using its internal byte slice.

    Note: This means that we no longer validate rate limiting at the intermediate buffer level when no content-length header was provided, but rather at the final buffer level. This is OKAY. We have reduced the overall heap size by removing this unnecessary intermediate buffer. And this pattern ensures we fully drain the request body before returning a response which can help avoid tcp connection reset that occur if there is > 256KB of data remaining in the request body (see net/http/server.go#1089). There actually may be other places we should be discarding remaining r.Body data before returning a response as well, but that was not my focus in this PR.

  2. The second change fixes an issue that was likely rendering the uncompressed slice pool ineffective due to how the s2.Decode function works internally. Specifically, s2.Decode will allocate a new slice if the provided dst byte slice does not have enough capacity to hold the decoded data. This means that any uncompressed body larger than 128KB (the initial capacity of the byte slice) would be allocated and then thrown to GC. Now, we ensure the buffer is at least the capacity of the decoded data before passing it to the s2.Decode function so that we guarantee the buffer is large enough to be used in all cases.

  3. The third change was to remove the writeRequestPool which was using the .Reset() of a proto object. This is not a good idea as it will replace the reference with a new struct, so we would effectively be pooling a pointer. I hope we'll be able to add pooling back to this in a future pr.

Testing

  • Tests were added for the syncutil.Pool implementation.
  • go test -race -v -count=10 -timeout 15m ./pkg/receive/.... I'm relying on existing tests for the receive handler. But I verified that the tests break if the pools were improperly configured by temporarily skipping reset steps.

Comment on lines 97 to 99
defaultCompressedBufCap = 32 * 1024
maxPooledCompressedCap = 1 << 20 // 1MB
maxPooledDecompressedCap = 4 << 20 // 4MB
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm curious about these numbers- where did they come from originally? i.e. do we have empirical data / histograms on incoming request sizes? Do we know what contributes to such large request sizes?


// Default / max capacities for pooled buffers. These caps prevent "pool ballooning"
// where a single large request permanently inflates process RSS.
defaultCompressedBufCap = 32 * 1024
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

note: I removed this defaultCompressedBufCap variable for 3 reasons:

  1. We don't want to over allocate from the beginning. As buffers in the pool are reused, they will grow to reach the steady state of a usual request.
  2. I wasn't able to find any empirical evidence that this 32KB actually is an expected average size.
  3. When pooling is deactivated, I want it to operate as close to the previous implementation as possible (before there was any pooling). The previous implementation stared with a default size of a) content-length (we still do this a la buffer.Grow(content-length)) b) 512b if no content-length header is set (which actually is what the buffer will do internally).

@yuchen-db yuchen-db self-requested a review February 23, 2026 20:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant