diff --git a/internal/rdb/rdb_test.go b/internal/rdb/rdb_test.go index 955fbce8..7361a793 100644 --- a/internal/rdb/rdb_test.go +++ b/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()) diff --git a/internal/rdb/structure/byte.go b/internal/rdb/structure/byte.go index f52ae0f5..9373ca73 100644 --- a/internal/rdb/structure/byte.go +++ b/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 { diff --git a/internal/rdb/structure/int.go b/internal/rdb/structure/int.go index 8ce5f6f0..44c015bc 100644 --- a/internal/rdb/structure/int.go +++ b/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 } diff --git a/internal/rdb/structure/ziplist.go b/internal/rdb/structure/ziplist.go index a1c47b54..5900fa60 100644 --- a/internal/rdb/structure/ziplist.go +++ b/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 {