diff --git a/src/pkg/compress/gzip/gzip.go b/src/pkg/compress/gzip/gzip.go index 45558b74289..fe32d6871ae 100644 --- a/src/pkg/compress/gzip/gzip.go +++ b/src/pkg/compress/gzip/gzip.go @@ -26,14 +26,15 @@ const ( // to its wrapped io.Writer. type Writer struct { Header - w io.Writer - level int - compressor *flate.Writer - digest hash.Hash32 - size uint32 - closed bool - buf [10]byte - err error + w io.Writer + level int + wroteHeader bool + compressor *flate.Writer + digest hash.Hash32 + size uint32 + closed bool + buf [10]byte + err error } // NewWriter creates a new Writer that satisfies writes by compressing data @@ -62,14 +63,39 @@ func NewWriterLevel(w io.Writer, level int) (*Writer, error) { if level < DefaultCompression || level > BestCompression { return nil, fmt.Errorf("gzip: invalid compression level: %d", level) } - return &Writer{ + z := new(Writer) + z.init(w, level) + return z, nil +} + +func (z *Writer) init(w io.Writer, level int) { + digest := z.digest + if digest != nil { + digest.Reset() + } else { + digest = crc32.NewIEEE() + } + compressor := z.compressor + if compressor != nil { + compressor.Reset(w) + } + *z = Writer{ Header: Header{ OS: 255, // unknown }, - w: w, - level: level, - digest: crc32.NewIEEE(), - }, nil + w: w, + level: level, + digest: digest, + compressor: compressor, + } +} + +// Reset discards the Writer z's state and makes it equivalent to the +// result of its original state from NewWriter or NewWriterLevel, but +// writing to w instead. This permits reusing a Writer rather than +// allocating a new one. +func (z *Writer) Reset(w io.Writer) { + z.init(w, z.level) } // GZIP (RFC 1952) is little-endian, unlike ZLIB (RFC 1950). @@ -138,7 +164,8 @@ func (z *Writer) Write(p []byte) (int, error) { } var n int // Write the GZIP header lazily. - if z.compressor == nil { + if !z.wroteHeader { + z.wroteHeader = true z.buf[0] = gzipID1 z.buf[1] = gzipID2 z.buf[2] = gzipDeflate @@ -183,7 +210,9 @@ func (z *Writer) Write(p []byte) (int, error) { return n, z.err } } - z.compressor, _ = flate.NewWriter(z.w, z.level) + if z.compressor == nil { + z.compressor, _ = flate.NewWriter(z.w, z.level) + } } z.size += uint32(len(p)) z.digest.Write(p) @@ -206,8 +235,11 @@ func (z *Writer) Flush() error { if z.closed { return nil } - if z.compressor == nil { + if !z.wroteHeader { z.Write(nil) + if z.err != nil { + return z.err + } } z.err = z.compressor.Flush() return z.err @@ -222,7 +254,7 @@ func (z *Writer) Close() error { return nil } z.closed = true - if z.compressor == nil { + if !z.wroteHeader { z.Write(nil) if z.err != nil { return z.err diff --git a/src/pkg/compress/gzip/gzip_test.go b/src/pkg/compress/gzip/gzip_test.go index 39afb331daf..119be2e135b 100644 --- a/src/pkg/compress/gzip/gzip_test.go +++ b/src/pkg/compress/gzip/gzip_test.go @@ -214,3 +214,18 @@ func TestConcat(t *testing.T) { t.Fatalf("ReadAll = %q, %v, want %q, nil", data, err, "hello world") } } + +func TestWriterReset(t *testing.T) { + buf := new(bytes.Buffer) + buf2 := new(bytes.Buffer) + z := NewWriter(buf) + msg := []byte("hello world") + z.Write(msg) + z.Close() + z.Reset(buf2) + z.Write(msg) + z.Close() + if buf.String() != buf2.String() { + t.Errorf("buf2 %q != original buf of %q", buf2.String(), buf.String()) + } +}