zlib: Pool writers too

2 files changed, 50 insertions, 15 deletions

diff --git a/internal/zlib/reader.go b/internal/zlib/reader.go
index 5141b2f6..2234e7e0 100644
--- a/internal/zlib/reader.go
+++ b/internal/zlib/reader.go
@@ -7,13 +7,14 @@ Package zlib implements reading and writing of zlib format compressed data,
 as specified in RFC 1950.
 
 This package differs from the standard library's compress/zlib package
-in that it pools readers to reduce allocations. Writers are unchanged.
+in that it pools readers and writers to reduce allocations.
 
-Note that closing the reader causes it to be returned to a pool for
-reuse. Therefore, the caller must not retain references to the
-reader after closing it; in the standard library's compress/zlib package,
-it is legal to Reset a closed reader and continue using it; that is
-not allowed here, so there is simply no Resetter interface.
+Note that closing a reader or writer causes it to be returned to a pool
+for reuse. Therefore, the caller must not retain references to a
+reader or writer after closing it; in the standard library's
+compress/zlib package, it is legal to Reset a closed reader or writer
+and continue using it; that is not allowed here, so there is simply no
+Resetter interface.
 
 The implementation provides filters that uncompress during reading
 and compress during writing.  For example, to write compressed data
@@ -58,7 +59,7 @@ var (
 	ErrHeader = errors.New("zlib: invalid header")
 )
 
-var pool = sync.Pool{
+var readerPool = sync.Pool{
 	New: func() any {
 		r := new(reader)
 		return r
@@ -86,7 +87,7 @@ func NewReader(r io.Reader) (io.ReadCloser, error) {
 // NewReaderDict ignores the dictionary if the compressed data does not refer to it.
 // If the compressed data refers to a different dictionary, NewReaderDict returns [ErrDictionary].
 func NewReaderDict(r io.Reader, dict []byte) (io.ReadCloser, error) {
-	v := pool.Get()
+	v := readerPool.Get()
 	z, ok := v.(*reader)
 	if !ok {
 		panic("zlib: pool returned unexpected type")
@@ -140,7 +141,7 @@ func (z *reader) Close() error {
 		return z.err
 	}
 
-	pool.Put(z)
+	readerPool.Put(z)
 	return nil
 }
 
diff --git a/internal/zlib/writer.go b/internal/zlib/writer.go
index 14eca474..81c57f55 100644
--- a/internal/zlib/writer.go
+++ b/internal/zlib/writer.go
@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"hash"
 	"io"
+	"sync"
 
 	"codeberg.org/lindenii/furgit/internal/adler32"
 )
@@ -37,6 +38,12 @@ type Writer struct {
 	wroteHeader bool
 }
 
+var writerPool = sync.Pool{
+	New: func() any {
+		return new(Writer)
+	},
+}
+
 // NewWriter creates a new [Writer].
 // Writes to the returned Writer are compressed and written to w.
 //
@@ -66,11 +73,33 @@ func NewWriterLevelDict(w io.Writer, level int, dict []byte) (*Writer, error) {
 	if level < HuffmanOnly || level > BestCompression {
 		return nil, fmt.Errorf("zlib: invalid compression level: %d", level)
 	}
-	return &Writer{
-		w:     w,
-		level: level,
-		dict:  dict,
-	}, nil
+	v := writerPool.Get()
+	z, ok := v.(*Writer)
+	if !ok {
+		panic("zlib: pool returned unexpected type")
+	}
+
+	// flate.Writer can only be Reset with the same level/dictionary mode.
+	// Reuse it only when the configuration is unchanged and dictionary-free.
+	reuseCompressor := z.compressor != nil && z.level == level && z.dict == nil && dict == nil
+	if !reuseCompressor {
+		z.compressor = nil
+	}
+	if z.digest != nil {
+		z.digest.Reset()
+	}
+
+	*z = Writer{
+		w:          w,
+		level:      level,
+		dict:       dict,
+		compressor: z.compressor,
+		digest:     z.digest,
+	}
+	if z.compressor != nil {
+		z.compressor.Reset(w)
+	}
+	return z, nil
 }
 
 // Reset clears the state of the [Writer] z such that it is equivalent to its
@@ -190,5 +219,10 @@ func (z *Writer) Close() error {
 	// ZLIB (RFC 1950) is big-endian, unlike GZIP (RFC 1952).
 	binary.BigEndian.PutUint32(z.scratch[:], checksum)
 	_, z.err = z.w.Write(z.scratch[0:4])
-	return z.err
+	if z.err != nil {
+		return z.err
+	}
+
+	writerPool.Put(z)
+	return nil
 }