Add pool optimization to frame decoders

pkg/frame/buffer.go

@@ -0,0 +1,8 @@
+package frame
+
+import "image"
+
+type buffer struct {
+	image image.Image
+	raw   []uint8
+}

pkg/frame/compressed.go

@@ -6,7 +6,9 @@ import (
 	"image/jpeg"
 )
 
-func decodeMJPEG(frame []byte, width, height int) (image.Image, func(), error) {
-	img, err := jpeg.Decode(bytes.NewReader(frame))
-	return img, func() {}, err
+func decodeMJPEG() decoderFunc {
+	return func(frame []byte, width, height int) (image.Image, func(), error) {
+		img, err := jpeg.Decode(bytes.NewReader(frame))
+		return img, func() {}, err
+	}
 }

pkg/frame/decode.go

@@ -5,22 +5,22 @@ import (
 )
 
 func NewDecoder(f Format) (Decoder, error) {
-	var decoder decoderFunc
+	var buildDecoder func() decoderFunc
 
 	switch f {
 	case FormatI420:
-		decoder = decodeI420
+		buildDecoder = decodeI420
 	case FormatNV21:
-		decoder = decodeNV21
+		buildDecoder = decodeNV21
 	case FormatYUY2:
-		decoder = decodeYUY2
+		buildDecoder = decodeYUY2
 	case FormatUYVY:
-		decoder = decodeUYVY
+		buildDecoder = decodeUYVY
 	case FormatMJPEG:
-		decoder = decodeMJPEG
+		buildDecoder = decodeMJPEG
 	default:
 		return nil, fmt.Errorf("%s is not supported", f)
 	}
 
-	return decoder, nil
+	return buildDecoder(), nil
 }

pkg/frame/yuv.go

@@ -5,47 +5,51 @@ import (
 	"image"
 )
 
-func decodeI420(frame []byte, width, height int) (image.Image, func(), error) {
-	yi := width * height
-	cbi := yi + width*height/4
-	cri := cbi + width*height/4
+func decodeI420() decoderFunc {
+	return func(frame []byte, width, height int) (image.Image, func(), error) {
+		yi := width * height
+		cbi := yi + width*height/4
+		cri := cbi + width*height/4
 
-	if cri > len(frame) {
-		return nil, func() {}, fmt.Errorf("frame length (%d) less than expected (%d)", len(frame), cri)
+		if cri > len(frame) {
+			return nil, func() {}, fmt.Errorf("frame length (%d) less than expected (%d)", len(frame), cri)
+		}
+
+		return &image.YCbCr{
+			Y:              frame[:yi],
+			YStride:        width,
+			Cb:             frame[yi:cbi],
+			Cr:             frame[cbi:cri],
+			CStride:        width / 2,
+			SubsampleRatio: image.YCbCrSubsampleRatio420,
+			Rect:           image.Rect(0, 0, width, height),
+		}, func() {}, nil
 	}
-
-	return &image.YCbCr{
-		Y:              frame[:yi],
-		YStride:        width,
-		Cb:             frame[yi:cbi],
-		Cr:             frame[cbi:cri],
-		CStride:        width / 2,
-		SubsampleRatio: image.YCbCrSubsampleRatio420,
-		Rect:           image.Rect(0, 0, width, height),
-	}, func() {}, nil
 }
 
-func decodeNV21(frame []byte, width, height int) (image.Image, func(), error) {
-	yi := width * height
-	ci := yi + width*height/2
+func decodeNV21() decoderFunc {
+	return func(frame []byte, width, height int) (image.Image, func(), error) {
+		yi := width * height
+		ci := yi + width*height/2
 
-	if ci > len(frame) {
-		return nil, func() {}, fmt.Errorf("frame length (%d) less than expected (%d)", len(frame), ci)
+		if ci > len(frame) {
+			return nil, func() {}, fmt.Errorf("frame length (%d) less than expected (%d)", len(frame), ci)
+		}
+
+		var cb, cr []byte
+		for i := yi; i < ci; i += 2 {
+			cb = append(cb, frame[i])
+			cr = append(cr, frame[i+1])
+		}
+
+		return &image.YCbCr{
+			Y:              frame[:yi],
+			YStride:        width,
+			Cb:             cb,
+			Cr:             cr,
+			CStride:        width / 2,
+			SubsampleRatio: image.YCbCrSubsampleRatio420,
+			Rect:           image.Rect(0, 0, width, height),
+		}, func() {}, nil
 	}
-
-	var cb, cr []byte
-	for i := yi; i < ci; i += 2 {
-		cb = append(cb, frame[i])
-		cr = append(cr, frame[i+1])
-	}
-
-	return &image.YCbCr{
-		Y:              frame[:yi],
-		YStride:        width,
-		Cb:             cb,
-		Cr:             cr,
-		CStride:        width / 2,
-		SubsampleRatio: image.YCbCrSubsampleRatio420,
-		Rect:           image.Rect(0, 0, width, height),
-	}, func() {}, nil
 }

pkg/frame/yuv_cgo.go

@@ -12,66 +12,70 @@ import (
 // void decodeUYVYCGO(uint8_t* y, uint8_t* cb, uint8_t* cr, uint8_t* uyvy, int width, int height);
 import "C"
 
-func decodeYUY2(frame []byte, width, height int) (image.Image, func(), error) {
-	yi := width * height
-	ci := yi / 2
-	fi := yi + 2*ci
+func decodeYUY2() decoderFunc {
+	return func(frame []byte, width, height int) (image.Image, func(), error) {
+		yi := width * height
+		ci := yi / 2
+		fi := yi + 2*ci
 
-	if len(frame) != fi {
-		return nil, func() {}, fmt.Errorf("frame length (%d) less than expected (%d)", len(frame), fi)
+		if len(frame) != fi {
+			return nil, func() {}, fmt.Errorf("frame length (%d) less than expected (%d)", len(frame), fi)
+		}
+
+		y := make([]byte, yi)
+		cb := make([]byte, ci)
+		cr := make([]byte, ci)
+
+		C.decodeYUY2CGO(
+			(*C.uchar)(&y[0]),
+			(*C.uchar)(&cb[0]),
+			(*C.uchar)(&cr[0]),
+			(*C.uchar)(&frame[0]),
+			C.int(width), C.int(height),
+		)
+
+		return &image.YCbCr{
+			Y:              y,
+			YStride:        width,
+			Cb:             cb,
+			Cr:             cr,
+			CStride:        width / 2,
+			SubsampleRatio: image.YCbCrSubsampleRatio422,
+			Rect:           image.Rect(0, 0, width, height),
+		}, func() {}, nil
 	}
-
-	y := make([]byte, yi)
-	cb := make([]byte, ci)
-	cr := make([]byte, ci)
-
-	C.decodeYUY2CGO(
-		(*C.uchar)(&y[0]),
-		(*C.uchar)(&cb[0]),
-		(*C.uchar)(&cr[0]),
-		(*C.uchar)(&frame[0]),
-		C.int(width), C.int(height),
-	)
-
-	return &image.YCbCr{
-		Y:              y,
-		YStride:        width,
-		Cb:             cb,
-		Cr:             cr,
-		CStride:        width / 2,
-		SubsampleRatio: image.YCbCrSubsampleRatio422,
-		Rect:           image.Rect(0, 0, width, height),
-	}, func() {}, nil
 }
 
-func decodeUYVY(frame []byte, width, height int) (image.Image, func(), error) {
-	yi := width * height
-	ci := yi / 2
-	fi := yi + 2*ci
+func decodeUYVY() decoderFunc {
+	return func(frame []byte, width, height int) (image.Image, func(), error) {
+		yi := width * height
+		ci := yi / 2
+		fi := yi + 2*ci
 
-	if len(frame) != fi {
-		return nil, func() {}, fmt.Errorf("frame length (%d) less than expected (%d)", len(frame), fi)
+		if len(frame) != fi {
+			return nil, func() {}, fmt.Errorf("frame length (%d) less than expected (%d)", len(frame), fi)
+		}
+
+		y := make([]byte, yi)
+		cb := make([]byte, ci)
+		cr := make([]byte, ci)
+
+		C.decodeUYVYCGO(
+			(*C.uchar)(&y[0]),
+			(*C.uchar)(&cb[0]),
+			(*C.uchar)(&cr[0]),
+			(*C.uchar)(&frame[0]),
+			C.int(width), C.int(height),
+		)
+
+		return &image.YCbCr{
+			Y:              y,
+			YStride:        width,
+			Cb:             cb,
+			Cr:             cr,
+			CStride:        width / 2,
+			SubsampleRatio: image.YCbCrSubsampleRatio422,
+			Rect:           image.Rect(0, 0, width, height),
+		}, func() {}, nil
 	}
-
-	y := make([]byte, yi)
-	cb := make([]byte, ci)
-	cr := make([]byte, ci)
-
-	C.decodeUYVYCGO(
-		(*C.uchar)(&y[0]),
-		(*C.uchar)(&cb[0]),
-		(*C.uchar)(&cr[0]),
-		(*C.uchar)(&frame[0]),
-		C.int(width), C.int(height),
-	)
-
-	return &image.YCbCr{
-		Y:              y,
-		YStride:        width,
-		Cb:             cb,
-		Cr:             cr,
-		CStride:        width / 2,
-		SubsampleRatio: image.YCbCrSubsampleRatio422,
-		Rect:           image.Rect(0, 0, width, height),
-	}, func() {}, nil
 }

pkg/frame/yuv_nocgo.go

@@ -5,74 +5,97 @@ package frame
 import (
 	"fmt"
 	"image"
+	"sync"
 )
 
-func decodeYUY2(frame []byte, width, height int) (image.Image, func(), error) {
-	yi := width * height
-	ci := yi / 2
-	fi := yi + 2*ci
-
-	if len(frame) != fi {
-		return nil, func() {}, fmt.Errorf("frame length (%d) less than expected (%d)", len(frame), fi)
+func decodeYUY2() decoderFunc {
+	pool := sync.Pool{
+		New: func() interface{} {
+			return &buffer{
+				image: &image.YCbCr{},
+			}
+		},
 	}
 
-	y := make([]byte, yi)
-	cb := make([]byte, ci)
-	cr := make([]byte, ci)
+	return func(frame []byte, width, height int) (image.Image, func(), error) {
+		buff := pool.Get().(*buffer)
 
-	fast := 0
-	slow := 0
-	for i := 0; i < fi; i += 4 {
-		y[fast] = frame[i]
-		cb[slow] = frame[i+1]
-		y[fast+1] = frame[i+2]
-		cr[slow] = frame[i+3]
-		fast += 2
-		slow++
+		yi := width * height
+		ci := yi / 2
+		fi := yi + 2*ci
+
+		if len(frame) != fi {
+			pool.Put(buff)
+			return nil, func() {}, fmt.Errorf("frame length (%d) less than expected (%d)", len(frame), fi)
+		}
+
+		if len(buff.raw) < fi {
+			need := fi - len(buff.raw)
+			buff.raw = append(buff.raw, make([]uint8, need)...)
+		}
+		y := buff.raw[:yi:yi]
+		cb := buff.raw[yi : yi+ci : yi+ci]
+		cr := buff.raw[yi+ci : fi : fi]
+
+		fast := 0
+		slow := 0
+		for i := 0; i < fi; i += 4 {
+			y[fast] = frame[i]
+			cb[slow] = frame[i+1]
+			y[fast+1] = frame[i+2]
+			cr[slow] = frame[i+3]
+			fast += 2
+			slow++
+		}
+
+		img := buff.image.(*image.YCbCr)
+		img.Y = y
+		img.YStride = width
+		img.Cb = cb
+		img.Cr = cr
+		img.CStride = width / 2
+		img.SubsampleRatio = image.YCbCrSubsampleRatio422
+		img.Rect.Min.X = 0
+		img.Rect.Min.Y = 0
+		img.Rect.Max.X = width
+		img.Rect.Max.Y = height
+		return img, func() { pool.Put(buff) }, nil
 	}
-
-	return &image.YCbCr{
-		Y:              y,
-		YStride:        width,
-		Cb:             cb,
-		Cr:             cr,
-		CStride:        width / 2,
-		SubsampleRatio: image.YCbCrSubsampleRatio422,
-		Rect:           image.Rect(0, 0, width, height),
-	}, func() {}, nil
 }
 
-func decodeUYVY(frame []byte, width, height int) (image.Image, func(), error) {
-	yi := width * height
-	ci := yi / 2
-	fi := yi + 2*ci
+func decodeUYVY() decoderFunc {
+	return func(frame []byte, width, height int) (image.Image, func(), error) {
+		yi := width * height
+		ci := yi / 2
+		fi := yi + 2*ci
 
-	if len(frame) != fi {
-		return nil, func() {}, fmt.Errorf("frame length (%d) less than expected (%d)", len(frame), fi)
+		if len(frame) != fi {
+			return nil, func() {}, fmt.Errorf("frame length (%d) less than expected (%d)", len(frame), fi)
+		}
+
+		y := make([]byte, yi)
+		cb := make([]byte, ci)
+		cr := make([]byte, ci)
+
+		fast := 0
+		slow := 0
+		for i := 0; i < fi; i += 4 {
+			cb[slow] = frame[i]
+			y[fast] = frame[i+1]
+			cr[slow] = frame[i+2]
+			y[fast+1] = frame[i+3]
+			fast += 2
+			slow++
+		}
+
+		return &image.YCbCr{
+			Y:              y,
+			YStride:        width,
+			Cb:             cb,
+			Cr:             cr,
+			CStride:        width / 2,
+			SubsampleRatio: image.YCbCrSubsampleRatio422,
+			Rect:           image.Rect(0, 0, width, height),
+		}, func() {}, nil
 	}
-
-	y := make([]byte, yi)
-	cb := make([]byte, ci)
-	cr := make([]byte, ci)
-
-	fast := 0
-	slow := 0
-	for i := 0; i < fi; i += 4 {
-		cb[slow] = frame[i]
-		y[fast] = frame[i+1]
-		cr[slow] = frame[i+2]
-		y[fast+1] = frame[i+3]
-		fast += 2
-		slow++
-	}
-
-	return &image.YCbCr{
-		Y:              y,
-		YStride:        width,
-		Cb:             cb,
-		Cr:             cr,
-		CStride:        width / 2,
-		SubsampleRatio: image.YCbCrSubsampleRatio422,
-		Rect:           image.Rect(0, 0, width, height),
-	}, func() {}, nil
 }

pkg/frame/yuv_test.go

@@ -27,7 +27,7 @@ func TestDecodeYUY2(t *testing.T) {
 		Rect:           image.Rect(0, 0, width, height),
 	}
 
-	img, _, err := decodeYUY2(input, width, height)
+	img, _, err := decodeYUY2()(input, width, height)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -56,7 +56,7 @@ func TestDecodeUYVY(t *testing.T) {
 		Rect:           image.Rect(0, 0, width, height),
 	}
 
-	img, _, err := decodeUYVY(input, width, height)
+	img, _, err := decodeUYVY()(input, width, height)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -76,11 +76,13 @@ func BenchmarkDecodeYUY2(b *testing.B) {
 		sz := sz
 		b.Run(fmt.Sprintf("%dx%d", sz.width, sz.height), func(b *testing.B) {
 			input := make([]byte, sz.width*sz.height*2)
+			decode := decodeYUY2()
 			for i := 0; i < b.N; i++ {
-				_, _, err := decodeYUY2(input, sz.width, sz.height)
+				_, release, err := decode(input, sz.width, sz.height)
 				if err != nil {
 					b.Fatal(err)
 				}
+				release()
 			}
 		})
 	}