Add OpenTelemetryHook

internal/proto/writer.go

@@ -1,7 +1,6 @@
 package proto
 
 import (
-	"bufio"
 	"encoding"
 	"fmt"
 	"io"
@@ -11,16 +10,22 @@ import (
 	"github.com/go-redis/redis/v8/internal/util"
 )
 
+type writer interface {
+	io.Writer
+	io.ByteWriter
+	io.StringWriter
+}
+
 type Writer struct {
-	wr *bufio.Writer
+	writer
 
 	lenBuf []byte
 	numBuf []byte
 }
 
-func NewWriter(wr io.Writer) *Writer {
+func NewWriter(wr writer) *Writer {
 	return &Writer{
-		wr: bufio.NewWriter(wr),
+		writer: wr,
 
 		lenBuf: make([]byte, 64),
 		numBuf: make([]byte, 64),
@@ -28,19 +33,16 @@ func NewWriter(wr io.Writer) *Writer {
 }
 
 func (w *Writer) WriteArgs(args []interface{}) error {
-	err := w.wr.WriteByte(ArrayReply)
-	if err != nil {
+	if err := w.WriteByte(ArrayReply); err != nil {
 		return err
 	}
 
-	err = w.writeLen(len(args))
-	if err != nil {
+	if err := w.writeLen(len(args)); err != nil {
 		return err
 	}
 
 	for _, arg := range args {
-		err := w.writeArg(arg)
-		if err != nil {
+		if err := w.WriteArg(arg); err != nil {
 			return err
 		}
 	}
@@ -51,11 +53,11 @@ func (w *Writer) WriteArgs(args []interface{}) error {
 func (w *Writer) writeLen(n int) error {
 	w.lenBuf = strconv.AppendUint(w.lenBuf[:0], uint64(n), 10)
 	w.lenBuf = append(w.lenBuf, '\r', '\n')
-	_, err := w.wr.Write(w.lenBuf)
+	_, err := w.Write(w.lenBuf)
 	return err
 }
 
-func (w *Writer) writeArg(v interface{}) error {
+func (w *Writer) WriteArg(v interface{}) error {
 	switch v := v.(type) {
 	case nil:
 		return w.string("")
@@ -108,18 +110,15 @@ func (w *Writer) writeArg(v interface{}) error {
 }
 
 func (w *Writer) bytes(b []byte) error {
-	err := w.wr.WriteByte(StringReply)
-	if err != nil {
+	if err := w.WriteByte(StringReply); err != nil {
 		return err
 	}
 
-	err = w.writeLen(len(b))
-	if err != nil {
+	if err := w.writeLen(len(b)); err != nil {
 		return err
 	}
 
-	_, err = w.wr.Write(b)
-	if err != nil {
+	if _, err := w.Write(b); err != nil {
 		return err
 	}
 
@@ -146,21 +145,8 @@ func (w *Writer) float(f float64) error {
 }
 
 func (w *Writer) crlf() error {
-	err := w.wr.WriteByte('\r')
-	if err != nil {
+	if err := w.WriteByte('\r'); err != nil {
 		return err
 	}
-	return w.wr.WriteByte('\n')
-}
-
-func (w *Writer) Buffered() int {
-	return w.wr.Buffered()
-}
-
-func (w *Writer) Reset(wr io.Writer) {
-	w.wr.Reset(wr)
-}
-
-func (w *Writer) Flush() error {
-	return w.wr.Flush()
+	return w.WriteByte('\n')
 }

internal/proto/writer_test.go

@@ -3,7 +3,6 @@ package proto_test
 import (
 	"bytes"
 	"encoding"
-	"io/ioutil"
 	"testing"
 	"time"
 
@@ -41,9 +40,6 @@ var _ = Describe("WriteBuffer", func() {
 		})
 		Expect(err).NotTo(HaveOccurred())
 
-		err = wr.Flush()
-		Expect(err).NotTo(HaveOccurred())
-
 		Expect(buf.Bytes()).To(Equal([]byte("*6\r\n" +
 			"$6\r\nstring\r\n" +
 			"$2\r\n12\r\n" +
@@ -59,9 +55,6 @@ var _ = Describe("WriteBuffer", func() {
 		err := wr.WriteArgs([]interface{}{tm})
 		Expect(err).NotTo(HaveOccurred())
 
-		err = wr.Flush()
-		Expect(err).NotTo(HaveOccurred())
-
 		Expect(buf.Len()).To(Equal(41))
 	})
 
@@ -69,26 +62,32 @@ var _ = Describe("WriteBuffer", func() {
 		err := wr.WriteArgs([]interface{}{&MyType{}})
 		Expect(err).NotTo(HaveOccurred())
 
-		err = wr.Flush()
-		Expect(err).NotTo(HaveOccurred())
-
 		Expect(buf.Len()).To(Equal(15))
 	})
 })
 
+type discard struct{}
+
+func (discard) Write(b []byte) (int, error) {
+	return len(b), nil
+}
+
+func (discard) WriteString(s string) (int, error) {
+	return len(s), nil
+}
+
+func (discard) WriteByte(c byte) error {
+	return nil
+}
+
 func BenchmarkWriteBuffer_Append(b *testing.B) {
-	buf := proto.NewWriter(ioutil.Discard)
+	buf := proto.NewWriter(discard{})
 	args := []interface{}{"hello", "world", "foo", "bar"}
 
 	for i := 0; i < b.N; i++ {
 		err := buf.WriteArgs(args)
 		if err != nil {
-			panic(err)
-		}
-
-		err = buf.Flush()
-		if err != nil {
-			panic(err)
+			b.Fatal(err)
 		}
 	}
 }