From 6001d65206bc00c7481c10d90cefecb7ad89a78f Mon Sep 17 00:00:00 2001 From: Chris Lu Date: Mon, 11 May 2026 10:55:56 -0700 Subject: [PATCH] perf(volume): stream-count the gzip size when no Content-MD5 is set (#9433) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- weed/storage/needle/needle_parse_upload.go | 14 +- .../needle/needle_parse_upload_stream_test.go | 125 ++++++++++++++++++ 2 files changed, 135 insertions(+), 4 deletions(-) create mode 100644 weed/storage/needle/needle_parse_upload_stream_test.go diff --git a/weed/storage/needle/needle_parse_upload.go b/weed/storage/needle/needle_parse_upload.go index aa54c3462..4b80897d7 100644 --- a/weed/storage/needle/needle_parse_upload.go +++ b/weed/storage/needle/needle_parse_upload.go @@ -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) diff --git a/weed/storage/needle/needle_parse_upload_stream_test.go b/weed/storage/needle/needle_parse_upload_stream_test.go new file mode 100644 index 000000000..80648ea73 --- /dev/null +++ b/weed/storage/needle/needle_parse_upload_stream_test.go @@ -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 +}