image/png: reduce memory allocs encoding images by reusing buffers

This change allows greatly reducing memory allocations with a slightly
performance improvement as well.

Instances of (*png).Encoder can have a optional BufferPool attached to
them. This allows reusing temporary buffers used when encoding a new
image. This buffers include instances to zlib.Writer and bufio.Writer.

Also, buffers for current and previous rows are saved in the encoder
instance and reused as long as their cap() is enough to fit the current
image row.

A new benchmark was added to demonstrate the performance improvement
when setting a BufferPool to an Encoder instance:

$ go test -bench BenchmarkEncodeGray -benchmem
BenchmarkEncodeGray-4                 	    1000	   2349584 ns/op	 130.75 MB/s	  852230 B/op	      32 allocs/op
BenchmarkEncodeGrayWithBufferPool-4   	    1000	   2241650 ns/op	 137.04 MB/s	     900 B/op	       3 allocs/op

Change-Id: I4488201ae53cb2ad010c68c1e0118ee12beae14e
Reviewed-on: https://go-review.googlesource.com/34150
Reviewed-by: Nigel Tao <[email protected]>
Run-TryBot: Nigel Tao <[email protected]>
TryBot-Result: Gobot Gobot <[email protected]>

Author: cezarsa

src/image/png/writer.go

@@ -17,17 +17,37 @@ import (
 // Encoder configures encoding PNG images.
 type Encoder struct {
 	CompressionLevel CompressionLevel
+
+	// BufferPool optionally specifies a buffer pool to get temporary
+	// EncoderBuffers when encoding an image.
+	BufferPool EncoderBufferPool
+}
+
+// EncoderBufferPool is an interface for getting and returning temporary
+// instances of the EncoderBuffer struct. This can be used to reuse buffers
+// when encoding multiple images.
+type EncoderBufferPool interface {
+	Get() *EncoderBuffer
+	Put(*EncoderBuffer)
 }
 
+// EncoderBuffer holds the buffers used for encoding PNG images.
+type EncoderBuffer encoder
+
 type encoder struct {
-	enc    *Encoder
-	w      io.Writer
-	m      image.Image
-	cb     int
-	err    error
-	header [8]byte
-	footer [4]byte
-	tmp    [4 * 256]byte
+	enc     *Encoder
+	w       io.Writer
+	m       image.Image
+	cb      int
+	err     error
+	header  [8]byte
+	footer  [4]byte
+	tmp     [4 * 256]byte
+	cr      [nFilter][]uint8
+	pr      []uint8
+	zw      *zlib.Writer
+	zwLevel int
+	bw      *bufio.Writer
 }
 
 type CompressionLevel int
@@ -273,12 +293,24 @@ func filter(cr *[nFilter][]byte, pr []byte, bpp int) int {
 	return filter
 }
 
-func writeImage(w io.Writer, m image.Image, cb int, level int) error {
-	zw, err := zlib.NewWriterLevel(w, level)
-	if err != nil {
-		return err
+func zeroMemory(v []uint8) {
+	for i := range v {
+		v[i] = 0
+	}
+}
+
+func (e *encoder) writeImage(w io.Writer, m image.Image, cb int, level int) error {
+	if e.zw == nil || e.zwLevel != level {
+		zw, err := zlib.NewWriterLevel(w, level)
+		if err != nil {
+			return err
+		}
+		e.zw = zw
+		e.zwLevel = level
+	} else {
+		e.zw.Reset(w)
 	}
-	defer zw.Close()
+	defer e.zw.Close()
 
 	bpp := 0 // Bytes per pixel.
 
@@ -304,12 +336,23 @@ func writeImage(w io.Writer, m image.Image, cb int, level int) error {
 	// other PNG filter types. These buffers are allocated once and re-used for each row.
 	// The +1 is for the per-row filter type, which is at cr[*][0].
 	b := m.Bounds()
-	var cr [nFilter][]uint8
-	for i := range cr {
-		cr[i] = make([]uint8, 1+bpp*b.Dx())
-		cr[i][0] = uint8(i)
+	sz := 1 + bpp*b.Dx()
+	for i := range e.cr {
+		if cap(e.cr[i]) < sz {
+			e.cr[i] = make([]uint8, sz)
+		} else {
+			e.cr[i] = e.cr[i][:sz]
+		}
+		e.cr[i][0] = uint8(i)
+	}
+	cr := e.cr
+	if cap(e.pr) < sz {
+		e.pr = make([]uint8, sz)
+	} else {
+		e.pr = e.pr[:sz]
+		zeroMemory(e.pr)
 	}
-	pr := make([]uint8, 1+bpp*b.Dx())
+	pr := e.pr
 
 	gray, _ := m.(*image.Gray)
 	rgba, _ := m.(*image.RGBA)
@@ -429,7 +472,7 @@ func writeImage(w io.Writer, m image.Image, cb int, level int) error {
 		}
 
 		// Write the compressed bytes.
-		if _, err := zw.Write(cr[f]); err != nil {
+		if _, err := e.zw.Write(cr[f]); err != nil {
 			return err
 		}
 
@@ -444,13 +487,16 @@ func (e *encoder) writeIDATs() {
 	if e.err != nil {
 		return
 	}
-	var bw *bufio.Writer
-	bw = bufio.NewWriterSize(e, 1<<15)
-	e.err = writeImage(bw, e.m, e.cb, levelToZlib(e.enc.CompressionLevel))
+	if e.bw == nil {
+		e.bw = bufio.NewWriterSize(e, 1<<15)
+	} else {
+		e.bw.Reset(e)
+	}
+	e.err = e.writeImage(e.bw, e.m, e.cb, levelToZlib(e.enc.CompressionLevel))
 	if e.err != nil {
 		return
 	}
-	e.err = bw.Flush()
+	e.err = e.bw.Flush()
 }
 
 // This function is required because we want the zero value of
@@ -489,7 +535,19 @@ func (enc *Encoder) Encode(w io.Writer, m image.Image) error {
 		return FormatError("invalid image size: " + strconv.FormatInt(mw, 10) + "x" + strconv.FormatInt(mh, 10))
 	}
 
-	var e encoder
+	var e *encoder
+	if enc.BufferPool != nil {
+		buffer := enc.BufferPool.Get()
+		e = (*encoder)(buffer)
+
+	}
+	if e == nil {
+		e = &encoder{}
+	}
+	if enc.BufferPool != nil {
+		defer enc.BufferPool.Put((*EncoderBuffer)(e))
+	}
+
 	e.enc = enc
 	e.w = w
 	e.m = m

src/image/png/writer_test.go

@@ -130,6 +130,31 @@ func BenchmarkEncodeGray(b *testing.B) {
 	}
 }
 
+type pool struct {
+	b *EncoderBuffer
+}
+
+func (p *pool) Get() *EncoderBuffer {
+	return p.b
+}
+
+func (p *pool) Put(b *EncoderBuffer) {
+	p.b = b
+}
+
+func BenchmarkEncodeGrayWithBufferPool(b *testing.B) {
+	b.StopTimer()
+	img := image.NewGray(image.Rect(0, 0, 640, 480))
+	e := Encoder{
+		BufferPool: &pool{},
+	}
+	b.SetBytes(640 * 480 * 1)
+	b.StartTimer()
+	for i := 0; i < b.N; i++ {
+		e.Encode(ioutil.Discard, img)
+	}
+}
+
 func BenchmarkEncodeNRGBOpaque(b *testing.B) {
 	b.StopTimer()
 	img := image.NewNRGBA(image.Rect(0, 0, 640, 480))