compress/flate, archive/zip: reduce memory allocations

Author: subtle-byte

src/archive/zip/zip_test.go

@@ -24,13 +24,18 @@ func TestOver65kFiles(t *testing.T) {
 	if testing.Short() && testenv.Builder() == "" {
 		t.Skip("skipping in short mode")
 	}
+
+	memStatBeforeWriter := runtime.MemStats{}
+	runtime.GC() // to guarantee mem stat is up to date
+	runtime.ReadMemStats(&memStatBeforeWriter)
+
 	buf := new(strings.Builder)
 	w := NewWriter(buf)
 	const nFiles = (1 << 16) + 42
 	for i := 0; i < nFiles; i++ {
 		_, err := w.CreateHeader(&FileHeader{
 			Name:   fmt.Sprintf("%d.dat", i),
-			Method: Store, // avoid Issue 6136 and Issue 6138
+			Method: Deflate,
 		})
 		if err != nil {
 			t.Fatalf("creating file %d: %v", i, err)
@@ -39,6 +44,11 @@ func TestOver65kFiles(t *testing.T) {
 	if err := w.Close(); err != nil {
 		t.Fatalf("Writer.Close: %v", err)
 	}
+
+	memStatBeforeReader := runtime.MemStats{}
+	runtime.GC()
+	runtime.ReadMemStats(&memStatBeforeReader)
+
 	s := buf.String()
 	zr, err := NewReader(strings.NewReader(s), int64(len(s)))
 	if err != nil {
@@ -53,6 +63,26 @@ func TestOver65kFiles(t *testing.T) {
 			t.Fatalf("File(%d) = %q, want %q", i, zr.File[i].Name, want)
 		}
 	}
+	for i := 0; i < nFiles; i++ {
+		f, err := zr.File[i].Open()
+		if err != nil {
+			t.Fatalf("File(%d).Open: %v", i, err)
+		}
+		if err := f.Close(); err != nil {
+			t.Fatalf("File(%d).Open().Close: %v", i, err)
+		}
+	}
+
+	memStatAfter := runtime.MemStats{}
+	runtime.GC()
+	runtime.ReadMemStats(&memStatAfter)
+
+	if alloc := memStatBeforeReader.TotalAlloc - memStatBeforeWriter.TotalAlloc; alloc > 150<<20 {
+		t.Errorf("allocated too much (%v MB) creating zip file", alloc>>20)
+	}
+	if alloc := memStatAfter.TotalAlloc - memStatBeforeReader.TotalAlloc; alloc > 70<<20 {
+		t.Errorf("allocated too much (%v MB) reading zip file", alloc>>20)
+	}
 }
 
 func TestModTime(t *testing.T) {

src/compress/flate/inflate.go

@@ -267,6 +267,7 @@ type Reader interface {
 type decompressor struct {
 	// Input source.
 	r       Reader
+	bufR    *bufio.Reader // used if provided reader does not implement io.ByteReader
 	roffset int64
 
 	// Input bits, in top of b.
@@ -746,11 +747,18 @@ func (f *decompressor) huffSym(h *huffmanDecoder) (int, error) {
 	}
 }
 
-func makeReader(r io.Reader) Reader {
+func (f *decompressor) makeReader(r io.Reader) {
 	if rr, ok := r.(Reader); ok {
-		return rr
+		f.bufR = nil
+		f.r = rr
+		return
+	}
+	if f.bufR != nil {
+		f.bufR.Reset(r)
+	} else {
+		f.bufR = bufio.NewReader(r)
 	}
-	return bufio.NewReader(r)
+	f.r = f.bufR
 }
 
 func fixedHuffmanDecoderInit() {
@@ -775,12 +783,13 @@ func fixedHuffmanDecoderInit() {
 
 func (f *decompressor) Reset(r io.Reader, dict []byte) error {
 	*f = decompressor{
-		r:        makeReader(r),
+		bufR:     f.bufR,
 		bits:     f.bits,
 		codebits: f.codebits,
 		dict:     f.dict,
 		step:     (*decompressor).nextBlock,
 	}
+	f.makeReader(r)
 	f.dict.init(maxMatchOffset, dict)
 	return nil
 }
@@ -797,7 +806,7 @@ func NewReader(r io.Reader) io.ReadCloser {
 	fixedHuffmanDecoderInit()
 
 	var f decompressor
-	f.r = makeReader(r)
+	f.makeReader(r)
 	f.bits = new([maxNumLit + maxNumDist]int)
 	f.codebits = new([numCodes]int)
 	f.step = (*decompressor).nextBlock
@@ -816,7 +825,7 @@ func NewReaderDict(r io.Reader, dict []byte) io.ReadCloser {
 	fixedHuffmanDecoderInit()
 
 	var f decompressor
-	f.r = makeReader(r)
+	f.makeReader(r)
 	f.bits = new([maxNumLit + maxNumDist]int)
 	f.codebits = new([numCodes]int)
 	f.step = (*decompressor).nextBlock

src/compress/flate/inflate_test.go

@@ -5,6 +5,7 @@
 package flate
 
 import (
+	"bufio"
 	"bytes"
 	"io"
 	"strings"
@@ -95,3 +96,36 @@ func TestResetDict(t *testing.T) {
 		}
 	}
 }
+
+func TestReaderResetReusesReaderBuffer(t *testing.T) {
+	encodedBuf := &bytes.Buffer{}
+	w, _ := NewWriter(encodedBuf, 1)
+	w.Write([]byte("abc"))
+	w.Close()
+
+	encodedNotByteReader := struct{ io.Reader }{encodedBuf}
+
+	f := NewReader(encodedNotByteReader)
+	bufReader1 := f.(*decompressor).r.(*bufio.Reader)
+	f.(Resetter).Reset(encodedNotByteReader, nil)
+	bufReader2 := f.(*decompressor).r.(*bufio.Reader)
+	if bufReader1 != bufReader2 {
+		t.Fatalf("bufio.Reader was not reused")
+	}
+	f.(Resetter).Reset(encodedBuf, nil)
+	bufReader3 := f.(*decompressor).r
+	if bufReader3 != encodedBuf {
+		t.Fatalf("bufio.Reader was reused, but should be used the provided io.ByteReader directly instead")
+	}
+	f.(Resetter).Reset(encodedNotByteReader, nil)
+	_, ok := f.(*decompressor).r.(*bufio.Reader)
+	if !ok {
+		t.Fatalf("new bufio.Reader should be created")
+	}
+	f.Close()
+	f = NewReader(encodedBuf)
+	bufReader4 := f.(*decompressor).r
+	if bufReader4 != encodedBuf {
+		t.Fatalf("bufio.Reader implementing io.ByteReader was not used directly")
+	}
+}