refactor: Use sync.Pool to reuse the []byte with specific length

internal/rdb/rdb_test.go

@@ -2,6 +2,8 @@ package rdb
 
 import (
 	"context"
+	"io"
+	"os"
 	"testing"
 
 	"RedisShake/internal/entry"
@@ -10,19 +12,37 @@ import (
 // BenchmarkParseRDB is a benchmark for ParseRDB
 // The baseline is "20	 350030327 ns/op	213804114 B/op	 1900715 allocs/op"
 func BenchmarkParseRDB(b *testing.B) {
+	if _, err := os.Stat("/tmp/dump.rdb"); err != nil && os.IsNotExist(err) {
+		sourcePath := "./dump.rdb"
+		sourceFile, err := os.Open(sourcePath)
+		if err != nil {
+			panic(err)
+		}
+		destPath := "/tmp/dump.rdb"
+		destFile, err := os.Create(destPath)
+		if err != nil {
+			panic(err)
+		}
+		_, err = io.Copy(destFile, sourceFile)
+		if err != nil {
+			panic(err)
+		}
+		destFile.Close()
+		sourceFile.Close()
+	}
+	b.ResetTimer()
 	b.ReportAllocs()
 	b.ResetTimer()
 	tempChan := make(chan *entry.Entry, 1024)
 	updateFunc := func(offset int64) {
 
 	}
-	b.N = 20
 
 	for i := 0; i < b.N; i++ {
-		loader := NewLoader("rdb", updateFunc, "./dump.rdb", tempChan)
+		loader := NewLoader("rdb", updateFunc, "/tmp/dump.rdb", tempChan)
 		go func() {
-			for temp := range tempChan {
-				print(temp.CmdName)
+			for _ = range tempChan {
+
 			}
 		}()
 		loader.ParseRDB(context.Background())

internal/rdb/structure/byte.go

@@ -2,13 +2,47 @@ package structure
 
 import (
 	"io"
+	"sync"
 
 	"RedisShake/internal/log"
 )
 
+var BytesPoolWithCap1 = sync.Pool{
+	New: func() interface{} {
+		tmp := make([]byte, 1)
+		return &tmp
+	},
+}
+
+var BytesPoolWithCap2 = sync.Pool{
+	New: func() interface{} {
+		tmp := make([]byte, 2)
+		return &tmp
+	},
+}
+
+var BytesPoolWithCap4 = sync.Pool{
+	New: func() interface{} {
+		tmp := make([]byte, 4)
+		return &tmp
+	},
+}
+
+var BytesPoolWithCap8 = sync.Pool{
+	New: func() interface{} {
+		tmp := make([]byte, 8)
+		return &tmp
+	},
+}
+
 func ReadByte(rd io.Reader) byte {
-	b := ReadBytes(rd, 1)[0]
-	return b
+	data := BytesPoolWithCap1.Get().(*[]byte)
+	if _, err := io.ReadFull(rd, *data); err != nil {
+		log.Panicf(err.Error())
+	}
+	result := (*data)[0]
+	BytesPoolWithCap1.Put(data)
+	return result
 }
 
 func ReadBytes(rd io.Reader, n int) []byte {

internal/rdb/structure/int.go

@@ -3,6 +3,8 @@ package structure
 import (
 	"encoding/binary"
 	"io"
+
+	"RedisShake/internal/log"
 )
 
 func ReadUint8(rd io.Reader) uint8 {
@@ -11,8 +13,13 @@ func ReadUint8(rd io.Reader) uint8 {
 }
 
 func ReadUint16(rd io.Reader) uint16 {
-	buf := ReadBytes(rd, 2)
-	return binary.LittleEndian.Uint16(buf)
+	data := BytesPoolWithCap2.Get().(*[]byte)
+	if _, err := io.ReadFull(rd, *data); err != nil {
+		log.Panicf(err.Error())
+	}
+	result := binary.LittleEndian.Uint16(*data)
+	BytesPoolWithCap2.Put(data)
+	return result
 }
 
 func ReadUint24(rd io.Reader) uint32 {
@@ -22,13 +29,23 @@ func ReadUint24(rd io.Reader) uint32 {
 }
 
 func ReadUint32(rd io.Reader) uint32 {
-	buf := ReadBytes(rd, 4)
-	return binary.LittleEndian.Uint32(buf)
+	data := BytesPoolWithCap4.Get().(*[]byte)
+	if _, err := io.ReadFull(rd, *data); err != nil {
+		log.Panicf(err.Error())
+	}
+	result := binary.LittleEndian.Uint32(*data)
+	BytesPoolWithCap4.Put(data)
+	return result
 }
 
 func ReadUint64(rd io.Reader) uint64 {
-	buf := ReadBytes(rd, 8)
-	return binary.LittleEndian.Uint64(buf)
+	data := BytesPoolWithCap8.Get().(*[]byte)
+	if _, err := io.ReadFull(rd, *data); err != nil {
+		log.Panicf(err.Error())
+	}
+	result := binary.LittleEndian.Uint64(*data)
+	BytesPoolWithCap8.Put(data)
+	return result
 }
 
 func ReadInt8(rd io.Reader) int8 {
@@ -37,8 +54,13 @@ func ReadInt8(rd io.Reader) int8 {
 }
 
 func ReadInt16(rd io.Reader) int16 {
-	buf := ReadBytes(rd, 2)
-	return int16(binary.LittleEndian.Uint16(buf))
+	data := BytesPoolWithCap2.Get().(*[]byte)
+	if _, err := io.ReadFull(rd, *data); err != nil {
+		log.Panicf(err.Error())
+	}
+	result := int16(binary.LittleEndian.Uint16(*data))
+	BytesPoolWithCap2.Put(data)
+	return result
 }
 
 func ReadInt24(rd io.Reader) int32 {
@@ -48,11 +70,21 @@ func ReadInt24(rd io.Reader) int32 {
 }
 
 func ReadInt32(rd io.Reader) int32 {
-	buf := ReadBytes(rd, 4)
-	return int32(binary.LittleEndian.Uint32(buf))
+	data := BytesPoolWithCap4.Get().(*[]byte)
+	if _, err := io.ReadFull(rd, *data); err != nil {
+		log.Panicf(err.Error())
+	}
+	result := int32(binary.LittleEndian.Uint32(*data))
+	BytesPoolWithCap4.Put(data)
+	return result
 }
 
 func ReadInt64(rd io.Reader) int64 {
-	buf := ReadBytes(rd, 8)
-	return int64(binary.LittleEndian.Uint64(buf))
+	data := BytesPoolWithCap8.Get().(*[]byte)
+	if _, err := io.ReadFull(rd, *data); err != nil {
+		log.Panicf(err.Error())
+	}
+	result := int64(binary.LittleEndian.Uint64(*data))
+	BytesPoolWithCap8.Put(data)
+	return result
 }

internal/rdb/structure/ziplist.go

@@ -80,8 +80,12 @@ func readZipListEntry(rd io.Reader, firstByte byte) string {
 		length := (int(firstByte&0x3f) << 8) | int(secondByte)
 		return string(ReadBytes(rd, length))
 	case zipStr32B:
-		lenBytes := ReadBytes(rd, 4)
-		length := binary.BigEndian.Uint32(lenBytes)
+		lenBytes := BytesPoolWithCap4.Get().(*[]byte)
+		if _, err := io.ReadFull(rd, *lenBytes); err != nil {
+			log.Panicf(err.Error())
+		}
+		length := binary.BigEndian.Uint32(*lenBytes)
+		BytesPoolWithCap4.Put(lenBytes)
 		return string(ReadBytes(rd, int(length)))
 	}
 	switch firstByte {