perf(volume): stream-count the gzip size when no Content-MD5 is set (#9433)

ParseUpload runs util.DecompressData on every gzipped multipart upload
just to record OriginalDataSize. The decompress materializes the full
uncompressed slice via bytes.Buffer.ReadFrom inside util.GunzipStream;
for a 64 MiB chunk that's a ~128 MiB heap spike per call (geometric
grow). On 6-way concurrent UploadPartCopy the spike dominated the
remaining heap profile after #9420/#9421/#9422/#9424/#9425.

When no Content-MD5 verification is requested the uncompressed bytes
aren't needed — only the length is. Stream the gunzip through
io.Discard and count: the pooled gzip.Reader's working set replaces
the materialized slice.

Unlike the previous attempt in #9426 the size still comes from the
real bytes, not from a client-set header.

  TotalAlloc per call, 4 MiB uncompressed body:
    materialize (was, still runs when MD5 is set):  ~16.8 MiB
    stream-count (no MD5):                             ~28 KiB

Refs #6541, #9426 (reverted in #9432).
This commit is contained in:
Chris Lu
2026-05-11 10:55:56 -07:00
committed by GitHub
parent a64483885c
commit 6001d65206
2 changed files with 135 additions and 4 deletions
+10 -4
View File
@@ -83,10 +83,16 @@ func ParseUpload(r *http.Request, sizeLimit int64, bytesBuffer *bytes.Buffer) (p
pu.UncompressedData = pu.Data
// println("received data", len(pu.Data), "isGzipped", pu.IsGzipped, "mime", pu.MimeType, "name", pu.FileName)
if pu.IsGzipped {
if unzipped, e := util.DecompressData(pu.Data); e == nil {
pu.OriginalDataSize = len(unzipped)
pu.UncompressedData = unzipped
// println("ungzipped data size", len(unzipped))
// MD5 check needs the uncompressed bytes; otherwise just count
// the gunzip stream — see #6541.
needMD5 := r.Header.Get("Content-MD5") != "" || pu.ContentMd5 != ""
if needMD5 {
if unzipped, err := util.DecompressData(pu.Data); err == nil {
pu.OriginalDataSize = len(unzipped)
pu.UncompressedData = unzipped
}
} else if n, err := util.GunzipStream(io.Discard, bytes.NewReader(pu.Data)); err == nil {
pu.OriginalDataSize = int(n)
}
} else {
ext := filepath.Base(pu.FileName)
@@ -0,0 +1,125 @@
package needle
import (
"bytes"
"compress/gzip"
"crypto/md5"
"encoding/base64"
"mime/multipart"
"net/http"
"net/http/httptest"
"net/textproto"
"runtime"
"testing"
)
// TestParseUpload_GzipStreamCount: gzipped uploads without Content-MD5
// must take the stream-count path and skip the uncompressed-slice
// allocation. See #6541.
func TestParseUpload_GzipStreamCount(t *testing.T) {
const uncompressedSize = 4 * 1024 * 1024 // 4 MiB
uncompressed := make([]byte, uncompressedSize)
for i := range uncompressed {
uncompressed[i] = byte(i*31 + 7)
}
gzipped := gzipBytes(t, uncompressed)
uncompressedMD5 := base64.StdEncoding.EncodeToString(md5sum(uncompressed))
cases := []struct {
name string
md5OnPart string
md5OnReq string
wantStream bool
}{
{name: "no MD5: stream-count", wantStream: true},
{name: "MD5 on part: materialize", md5OnPart: uncompressedMD5},
{name: "MD5 on request: materialize", md5OnReq: uncompressedMD5},
}
// Warm the gzip.Reader sync.Pool so per-test alloc is steady-state.
if _, err := ParseUpload(buildGzipReq(t, gzipped, "", ""), 256<<20, &bytes.Buffer{}); err != nil {
t.Fatalf("warmup ParseUpload: %v", err)
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
req := buildGzipReq(t, gzipped, c.md5OnPart, c.md5OnReq)
runtime.GC()
var before, after runtime.MemStats
runtime.ReadMemStats(&before)
pu, err := ParseUpload(req, 256<<20, &bytes.Buffer{})
runtime.ReadMemStats(&after)
allocated := after.TotalAlloc - before.TotalAlloc
if err != nil {
t.Fatalf("ParseUpload: %v", err)
}
if pu.OriginalDataSize != uncompressedSize {
t.Errorf("OriginalDataSize=%d, want %d", pu.OriginalDataSize, uncompressedSize)
}
if c.wantStream {
if allocated > uint64(uncompressedSize) {
t.Errorf("stream path allocated %d, bound %d", allocated, uncompressedSize)
}
} else {
if allocated < uint64(uncompressedSize) {
t.Errorf("materialize path allocated %d, want >= %d", allocated, uncompressedSize)
}
}
t.Logf("allocated=%d bytes (gzipped=%d, uncompressed=%d)",
allocated, len(gzipped), uncompressedSize)
})
}
}
func gzipBytes(t *testing.T, in []byte) []byte {
t.Helper()
var buf bytes.Buffer
gw := gzip.NewWriter(&buf)
if _, err := gw.Write(in); err != nil {
t.Fatalf("gzip write: %v", err)
}
if err := gw.Close(); err != nil {
t.Fatalf("gzip close: %v", err)
}
return buf.Bytes()
}
func md5sum(in []byte) []byte {
h := md5.New()
h.Write(in)
return h.Sum(nil)
}
func buildGzipReq(t *testing.T, body []byte, md5OnPart, md5OnReq string) *http.Request {
t.Helper()
var bodyBuf bytes.Buffer
mw := multipart.NewWriter(&bodyBuf)
h := make(textproto.MIMEHeader)
h.Set("Content-Disposition", `form-data; name="file"; filename="test.bin"`)
h.Set("Content-Encoding", "gzip")
if md5OnPart != "" {
h.Set("Content-MD5", md5OnPart)
}
fw, err := mw.CreatePart(h)
if err != nil {
t.Fatalf("CreatePart: %v", err)
}
if _, err := fw.Write(body); err != nil {
t.Fatalf("write part: %v", err)
}
if err := mw.Close(); err != nil {
t.Fatalf("close mw: %v", err)
}
req := httptest.NewRequest(http.MethodPost, "/upload", bytes.NewReader(bodyBuf.Bytes()))
req.Header.Set("Content-Type", mw.FormDataContentType())
req.ContentLength = int64(bodyBuf.Len())
if md5OnReq != "" {
req.Header.Set("Content-MD5", md5OnReq)
}
return req
}