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 <nigeltao@golang.org>
Run-TryBot: Nigel Tao <nigeltao@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

diff --git a/src/image/png/writer.go b/src/image/png/writer.go
index dd87d81629196002f12a6a4a83476cf9ee0c7c09..49f1ad2e7fa1fb77b39e44f54b15119b0cfba94a 100644 (file)
--- a/src/image/png/writer.go
+++ b/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

diff --git a/src/image/png/writer_test.go b/src/image/png/writer_test.go
index d67a815698fd0e75358b92b0257e88a45f0ee574..b1f97b1d7bf9798cd7c70bcfc8af579fee8a8039 100644 (file)
--- a/src/image/png/writer_test.go
+++ b/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))