Add ImageFactory and Buffer to Options

decoder/decoder.go

@@ -29,6 +29,7 @@ package decoder
 */
 import "C"
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"image"
@@ -47,22 +48,26 @@ type Decoder struct {
 	sPtr    C.size_t
 }
 
-// NewDecoder return new decoder instance
 func NewDecoder(r io.Reader, options *Options) (d *Decoder, err error) {
-	var data []byte
+	if options == nil {
+		options = &Options{}
+	}
+
+	if options.ImageFactory == nil {
+		options.ImageFactory = &DefaultImageFactory{}
+	}
 
-	if data, err = io.ReadAll(r); err != nil {
+	buf := bytes.NewBuffer(options.Buffer)
+
+	if _, err = io.Copy(buf, r); err != nil {
 		return nil, err
 	}
 
-	if len(data) == 0 {
+	if len(buf.Bytes()) == 0 {
 		return nil, errors.New("data is empty")
 	}
 
-	if options == nil {
-		options = &Options{}
-	}
-	d = &Decoder{data: data, options: options}
+	d = &Decoder{data: buf.Bytes(), options: options}
 
 	if d.config, err = d.options.GetConfig(); err != nil {
 		return nil, err
@@ -87,10 +92,7 @@ func (d *Decoder) Decode() (image.Image, error) {
 	d.config.output.colorspace = C.MODE_RGBA
 	d.config.output.is_external_memory = 1
 
-	img := image.NewNRGBA(image.Rectangle{Max: image.Point{
-		X: int(d.config.output.width),
-		Y: int(d.config.output.height),
-	}})
+	img := d.options.ImageFactory.Get(int(d.config.output.width), int(d.config.output.height))
 
 	buff := (*C.WebPRGBABuffer)(unsafe.Pointer(&d.config.output.u[0]))
 	buff.stride = C.int(img.Stride)

decoder/decoder_benchmark_test.go

@@ -0,0 +1,60 @@
+package decoder
+
+import (
+	"bytes"
+	"image"
+	"os"
+	"testing"
+)
+
+func loadImage(b *testing.B) []byte {
+	filename := "../test_data/images/100x150_lossless.webp"
+	data, err := os.ReadFile(filename)
+	if err != nil {
+		b.Fatal(err)
+	}
+	return data
+}
+
+func BenchmarkDecodePooled(b *testing.B) {
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	data := loadImage(b)
+
+	imagePool := NewImagePool()
+	bufferPool := NewBufferPool()
+
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			buf := bufferPool.Get()
+			decoder, err := NewDecoder(bytes.NewReader(data), &Options{ImageFactory: imagePool, Buffer: buf})
+			img, err := decoder.Decode()
+			if err != nil {
+				b.Fatal(err)
+			}
+
+			// put everything back
+			imagePool.Put(img.(*image.NRGBA))
+			bufferPool.Put(buf)
+		}
+	})
+}
+
+func BenchmarkDecodeUnPooled(b *testing.B) {
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	data := loadImage(b)
+
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			decoder, err := NewDecoder(bytes.NewReader(data), &Options{})
+			img, err := decoder.Decode()
+			if err != nil {
+				b.Fatal(err)
+			}
+			_ = img
+		}
+	})
+}

decoder/decoder_test.go

@@ -22,11 +22,12 @@
 package decoder
 
 import (
-	"github.com/kolesa-team/go-webp/utils"
-	"github.com/stretchr/testify/require"
 	"image"
 	"os"
 	"testing"
+
+	"github.com/kolesa-team/go-webp/utils"
+	"github.com/stretchr/testify/require"
 )
 
 func TestNewDecoder(t *testing.T) {

decoder/options.go

@@ -41,6 +41,16 @@ type Options struct {
 	Flip                   bool
 	DitheringStrength      int
 	AlphaDitheringStrength int
+
+	// These two are optimizations that require a little extra work on the caller side.
+
+	// if nil, DefaultImageFactory will be used.  If non-nil, decode will return an image that must be put back into the pool
+	// when you're done with it
+	ImageFactory ImageFactory
+	// if nil, a default buffer will be used.  If non-nil, decode will use this buffer to store data from the reader.
+	// The idea is that this buffer be reused, so either pass this back in next time you call decode, or put it back into
+	// a pool when you're done with it.
+	Buffer []byte
 }
 
 // GetConfig build WebPDecoderConfig for libwebp
@@ -89,3 +99,13 @@ func (o *Options) GetConfig() (*C.WebPDecoderConfig, error) {
 
 	return &config, nil
 }
+
+type ImageFactory interface {
+	Get(width, height int) *image.NRGBA
+}
+
+type DefaultImageFactory struct{}
+
+func (d *DefaultImageFactory) Get(width, height int) *image.NRGBA {
+	return image.NewNRGBA(image.Rect(0, 0, width, height))
+}

decoder/pool.go

@@ -0,0 +1,75 @@
+package decoder
+
+import (
+	"image"
+	"sync"
+	"sync/atomic"
+)
+
+type ImagePool struct {
+	poolMap map[int]*sync.Pool
+	lock    *sync.Mutex
+	Count   int64
+}
+
+func NewImagePool() *ImagePool {
+	return &ImagePool{
+		poolMap: make(map[int]*sync.Pool),
+		lock:    &sync.Mutex{},
+	}
+}
+
+func (n *ImagePool) Get(width, height int) *image.NRGBA {
+	dimPool := n.getPool(width, height)
+
+	img := dimPool.Get().(*image.NRGBA)
+	img.Rect.Max.X = width
+	img.Rect.Max.Y = height
+	return img
+}
+
+func (n *ImagePool) getPool(width int, height int) *sync.Pool {
+	dim := width * height
+
+	n.lock.Lock()
+	dimPool, ok := n.poolMap[dim]
+	if !ok {
+		atomic.AddInt64(&n.Count, 1)
+		dimPool = &sync.Pool{
+			New: func() interface{} {
+				return image.NewNRGBA(image.Rect(0, 0, width, height))
+			},
+		}
+		n.poolMap[dim] = dimPool
+	}
+	n.lock.Unlock()
+	return dimPool
+}
+
+func (n *ImagePool) Put(img *image.NRGBA) {
+	dimPool := n.getPool(img.Rect.Dx(), img.Rect.Dy())
+	dimPool.Put(img)
+}
+
+type BufferPool struct {
+	pool *sync.Pool // pointer because noCopy
+}
+
+func NewBufferPool() *BufferPool {
+	return &BufferPool{
+		pool: &sync.Pool{
+			New: func() interface{} {
+				return make([]byte, 0, 1024)
+			},
+		},
+	}
+}
+
+func (b *BufferPool) Get() []byte {
+	return b.pool.Get().([]byte)
+}
+
+func (b *BufferPool) Put(buf []byte) {
+	buf = buf[:0]
+	b.pool.Put(buf)
+}