luxfi/zapdb
1
0
Fork 0
mirror of https://github.com/luxfi/zapdb.git synced 2026-07-26 22:46:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

pb: convert protobuf to native binary encoding

Browse Source 
Replace protobuf serialization with native binary encoding using encoding/binary
package. The grpc build tag enables original protobuf behavior for backward
compatibility.

New files:
- pb/types_zap.go: Native type definitions with Marshal/Unmarshal/Size methods
- pb/marshal_zap.go: Interface wrappers for native encoding
- pb/marshal_grpc.go: Protobuf wrappers (build tag: grpc)
- pb/types_zap_test.go: Unit tests for binary encoding

Changes:
- Add //go:build grpc tag to pb/badgerpb4.pb.go
- Replace proto.Marshal/Unmarshal/Size with pb.Marshal/Unmarshal/Size
- Replace proto.Clone with native Clone method
- Remove google.golang.org/protobuf/proto import from non-grpc builds

Build modes:
- Default: Native binary encoding (no external dependencies)
- go build -tags=grpc: Use protobuf encoding

All tests pass in both modes.
This commit is contained in:
Zach Kelling
2026-01-26 22:11:24 -08:00
parent 28a4527f34
commit 1839d238f2
17 changed files with 1027 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backup.go
+3 -5
View File
@@ -13,8 +13,6 @@ import (
"fmt"
"io"
"google.golang.org/protobuf/proto"
"github.com/dgraph-io/badger/v4/pb"
"github.com/dgraph-io/badger/v4/y"
"github.com/dgraph-io/ristretto/v2/z"
@@ -137,10 +135,10 @@ func (stream *Stream) Backup(w io.Writer, since uint64) (uint64, error) {
}
func writeTo(list *pb.KVList, w io.Writer) error {
if err := binary.Write(w, binary.LittleEndian, uint64(proto.Size(list))); err != nil {
if err := binary.Write(w, binary.LittleEndian, uint64(pb.Size(list))); err != nil {
return err
}
buf, err := proto.Marshal(list)
buf, err := pb.Marshal(list)
if err != nil {
return err
}
@@ -253,7 +251,7 @@ func (db *DB) Load(r io.Reader, maxPendingWrites int) error {
}
list := &pb.KVList{}
if err := proto.Unmarshal(unmarshalBuf[:sz], list); err != nil {
if err := pb.Unmarshal(unmarshalBuf[:sz], list); err != nil {
return err
}
badger/cmd/bank.go
+1 -2
View File
@@ -21,7 +21,6 @@ import (
"time"
"github.com/spf13/cobra"
"google.golang.org/protobuf/proto"
"github.com/dgraph-io/badger/v4"
"github.com/dgraph-io/badger/v4/pb"
@@ -494,7 +493,7 @@ func runTest(cmd *cobra.Command, args []string) error {
stream.Send = func(buf *z.Buffer) error {
err := buf.SliceIterate(func(s []byte) error {
var kv pb.KV
if err := proto.Unmarshal(s, &kv); err != nil {
if err := pb.Unmarshal(s, &kv); err != nil {
return err
}
return batch.Set(kv.Key, kv.Value)
badger/cmd/read_bench.go
+1 -2
View File
@@ -16,7 +16,6 @@ import (
humanize "github.com/dustin/go-humanize"
"github.com/spf13/cobra"
"google.golang.org/protobuf/proto"
"github.com/dgraph-io/badger/v4"
"github.com/dgraph-io/badger/v4/pb"
@@ -204,7 +203,7 @@ func getSampleKeys(db *badger.DB, sampleSize int) ([][]byte, error) {
}
err := buf.SliceIterate(func(s []byte) error {
var kv pb.KV
if err := proto.Unmarshal(s, &kv); err != nil {
if err := pb.Unmarshal(s, &kv); err != nil {
return err
}
keys = append(keys, kv.Key)
batch.go
+1 -3
View File
@@ -11,8 +11,6 @@ import (
"sync"
"sync/atomic"
"google.golang.org/protobuf/proto"
"github.com/dgraph-io/badger/v4/pb"
"github.com/dgraph-io/badger/v4/y"
"github.com/dgraph-io/ristretto/v2/z"
@@ -104,7 +102,7 @@ func (wb *WriteBatch) Write(buf *z.Buffer) error {
err := buf.SliceIterate(func(s []byte) error {
kv := &pb.KV{}
if err := proto.Unmarshal(s, kv); err != nil {
if err := pb.Unmarshal(s, kv); err != nil {
return err
}
return wb.writeKV(kv)
key_registry.go
+2 -3
View File
@@ -19,7 +19,6 @@ import (
"github.com/dgraph-io/badger/v4/pb"
"github.com/dgraph-io/badger/v4/y"
"google.golang.org/protobuf/proto"
)
const (
@@ -185,7 +184,7 @@ func (kri *keyRegistryIterator) next() (*pb.DataKey, error) {
return nil, y.Wrapf(y.ErrChecksumMismatch, "Error while checking checksum for data key.")
}
dataKey := &pb.DataKey{}
if err = proto.Unmarshal(data, dataKey); err != nil {
if err = pb.Unmarshal(data, dataKey); err != nil {
return nil, y.Wrapf(err, "While unmarshal of datakey in keyRegistryIterator.next")
}
if len(kri.encryptionKey) > 0 {
@@ -394,7 +393,7 @@ func storeDataKey(buf *bytes.Buffer, storageKey []byte, k *pb.DataKey) error {
return y.Wrapf(err, "Error while encrypting datakey in storeDataKey")
}
var data []byte
if data, err = proto.Marshal(k); err != nil {
if data, err = pb.Marshal(k); err != nil {
err = y.Wrapf(err, "Error while marshaling datakey in storeDataKey")
var err2 error
// decrypting the datakey back.
manifest.go
+3 -5
View File
@@ -18,8 +18,6 @@ import (
"path/filepath"
"sync"
"google.golang.org/protobuf/proto"
"github.com/dgraph-io/badger/v4/options"
"github.com/dgraph-io/badger/v4/pb"
"github.com/dgraph-io/badger/v4/y"
@@ -199,7 +197,7 @@ func (mf *manifestFile) addChanges(changesParam []*pb.ManifestChange, opt Option
return nil
}
changes := pb.ManifestChangeSet{Changes: changesParam}
buf, err := proto.Marshal(&changes)
buf, err := pb.Marshal(&changes)
if err != nil {
return err
}
@@ -261,7 +259,7 @@ func helpRewrite(dir string, m *Manifest, extMagic uint16) (*os.File, int, error
changes := m.asChanges()
set := pb.ManifestChangeSet{Changes: changes}
changeBuf, err := proto.Marshal(&set)
changeBuf, err := pb.Marshal(&set)
if err != nil {
fp.Close()
return nil, 0, err
@@ -414,7 +412,7 @@ func ReplayManifestFile(fp *os.File, extMagic uint16, opt Options) (Manifest, in
}
var changeSet pb.ManifestChangeSet
if err := proto.Unmarshal(buf, &changeSet); err != nil {
if err := pb.Unmarshal(buf, &changeSet); err != nil {
return Manifest{}, 0, err
}
pb/badgerpb4.pb.go
+2
View File
@@ -1,3 +1,5 @@
//go:build grpc
//
// SPDX-FileCopyrightText: © 2017-2025 Istari Digital, Inc.
// SPDX-License-Identifier: Apache-2.0
pb/marshal_grpc.go
+28
View File
@@ -0,0 +1,28 @@
//go:build grpc
/*
* SPDX-FileCopyrightText: 2017-2025 Istari Digital, Inc.
* SPDX-License-Identifier: Apache-2.0
*/
package pb
import "google.golang.org/protobuf/proto"
// Marshal marshals a proto.Message to bytes.
func Marshal(m proto.Message) ([]byte, error) {
return proto.Marshal(m)
}
// Unmarshal unmarshals bytes into a proto.Message.
func Unmarshal(data []byte, m proto.Message) error {
return proto.Unmarshal(data, m)
}
// Size returns the encoded size of a proto.Message.
func Size(m proto.Message) int {
return proto.Size(m)
}
// MarshalOptions wraps proto.MarshalOptions for compatibility.
type MarshalOptions = proto.MarshalOptions
pb/marshal_zap.go
+50
View File
@@ -0,0 +1,50 @@
//go:build !grpc
/*
* SPDX-FileCopyrightText: 2017-2025 Istari Digital, Inc.
* SPDX-License-Identifier: Apache-2.0
*/
package pb
// Marshaler is the interface for types that can marshal themselves.
type Marshaler interface {
Marshal() ([]byte, error)
}
// Unmarshaler is the interface for types that can unmarshal themselves.
type Unmarshaler interface {
Unmarshal([]byte) error
}
// Sizer is the interface for types that can report their encoded size.
type Sizer interface {
Size() int
}
// Marshal marshals a Marshaler to bytes.
func Marshal(m Marshaler) ([]byte, error) {
return m.Marshal()
}
// Unmarshal unmarshals bytes into an Unmarshaler.
func Unmarshal(data []byte, m Unmarshaler) error {
return m.Unmarshal(data)
}
// Size returns the encoded size of a Sizer.
func Size(m Sizer) int {
return m.Size()
}
// MarshalOptions provides options for marshaling (compatibility with protobuf API).
type MarshalOptions struct{}
// MarshalAppend appends the marshaled form to the provided buffer.
func (MarshalOptions) MarshalAppend(b []byte, m Marshaler) ([]byte, error) {
data, err := m.Marshal()
if err != nil {
return nil, err
}
return append(b, data...), nil
}
pb/types_zap.go
+672
View File
@@ -0,0 +1,672 @@
//go:build !grpc
/*
* SPDX-FileCopyrightText: 2017-2025 Istari Digital, Inc.
* SPDX-License-Identifier: Apache-2.0
*/
// Package pb provides native binary encoding for badger protobuf types.
// This implementation replaces protobuf with direct binary encoding for
// reduced dependencies and faster serialization.
package pb
import (
"encoding/binary"
"errors"
"io"
)
// EncryptionAlgo defines encryption algorithm type.
type EncryptionAlgo int32
const (
EncryptionAlgo_aes EncryptionAlgo = 0
)
// ManifestChange_Operation defines manifest change operations.
type ManifestChange_Operation int32
const (
ManifestChange_CREATE ManifestChange_Operation = 0
ManifestChange_DELETE ManifestChange_Operation = 1
)
// Checksum_Algorithm defines checksum algorithm type.
type Checksum_Algorithm int32
const (
Checksum_CRC32C Checksum_Algorithm = 0
Checksum_XXHash64 Checksum_Algorithm = 1
)
var (
errBufferTooSmall = errors.New("buffer too small for unmarshaling")
errInvalidData = errors.New("invalid data format")
)
// KV represents a key-value pair.
type KV struct {
Key []byte
Value []byte
UserMeta []byte
Version uint64
ExpiresAt uint64
Meta []byte
StreamId uint32
StreamDone bool
}
func (k *KV) GetKey() []byte { return k.Key }
func (k *KV) GetValue() []byte { return k.Value }
func (k *KV) GetUserMeta() []byte { return k.UserMeta }
func (k *KV) GetVersion() uint64 { return k.Version }
func (k *KV) GetExpiresAt() uint64 { return k.ExpiresAt }
func (k *KV) GetMeta() []byte { return k.Meta }
func (k *KV) GetStreamId() uint32 { return k.StreamId }
func (k *KV) GetStreamDone() bool { return k.StreamDone }
func (k *KV) Reset() { *k = KV{} }
func (k *KV) String() string { return "KV{...}" }
// Size returns the encoded size of KV.
// Format: [keyLen:4][key][valueLen:4][value][userMetaLen:4][userMeta]
// [version:8][expiresAt:8][metaLen:4][meta][streamId:4][streamDone:1]
func (k *KV) Size() int {
return 4 + len(k.Key) +
4 + len(k.Value) +
4 + len(k.UserMeta) +
8 + // version
8 + // expiresAt
4 + len(k.Meta) +
4 + // streamId
1 // streamDone
}
// Marshal encodes KV to binary format.
func (k *KV) Marshal() ([]byte, error) {
buf := make([]byte, k.Size())
_, err := k.MarshalToSizedBuffer(buf)
return buf, err
}
// MarshalToSizedBuffer marshals KV to a pre-allocated buffer.
func (k *KV) MarshalToSizedBuffer(buf []byte) (int, error) {
if len(buf) < k.Size() {
return 0, io.ErrShortBuffer
}
offset := 0
// Key
binary.LittleEndian.PutUint32(buf[offset:], uint32(len(k.Key)))
offset += 4
copy(buf[offset:], k.Key)
offset += len(k.Key)
// Value
binary.LittleEndian.PutUint32(buf[offset:], uint32(len(k.Value)))
offset += 4
copy(buf[offset:], k.Value)
offset += len(k.Value)
// UserMeta
binary.LittleEndian.PutUint32(buf[offset:], uint32(len(k.UserMeta)))
offset += 4
copy(buf[offset:], k.UserMeta)
offset += len(k.UserMeta)
// Version
binary.LittleEndian.PutUint64(buf[offset:], k.Version)
offset += 8
// ExpiresAt
binary.LittleEndian.PutUint64(buf[offset:], k.ExpiresAt)
offset += 8
// Meta
binary.LittleEndian.PutUint32(buf[offset:], uint32(len(k.Meta)))
offset += 4
copy(buf[offset:], k.Meta)
offset += len(k.Meta)
// StreamId
binary.LittleEndian.PutUint32(buf[offset:], k.StreamId)
offset += 4
// StreamDone
if k.StreamDone {
buf[offset] = 1
} else {
buf[offset] = 0
}
offset++
return offset, nil
}
// Unmarshal decodes KV from binary format.
func (k *KV) Unmarshal(data []byte) error {
if len(data) < 37 { // minimum size: 4+0+4+0+4+0+8+8+4+0+4+1
return errBufferTooSmall
}
offset := 0
// Key
keyLen := int(binary.LittleEndian.Uint32(data[offset:]))
offset += 4
if offset+keyLen > len(data) {
return errBufferTooSmall
}
k.Key = make([]byte, keyLen)
copy(k.Key, data[offset:offset+keyLen])
offset += keyLen
// Value
if offset+4 > len(data) {
return errBufferTooSmall
}
valueLen := int(binary.LittleEndian.Uint32(data[offset:]))
offset += 4
if offset+valueLen > len(data) {
return errBufferTooSmall
}
k.Value = make([]byte, valueLen)
copy(k.Value, data[offset:offset+valueLen])
offset += valueLen
// UserMeta
if offset+4 > len(data) {
return errBufferTooSmall
}
userMetaLen := int(binary.LittleEndian.Uint32(data[offset:]))
offset += 4
if offset+userMetaLen > len(data) {
return errBufferTooSmall
}
k.UserMeta = make([]byte, userMetaLen)
copy(k.UserMeta, data[offset:offset+userMetaLen])
offset += userMetaLen
// Version
if offset+8 > len(data) {
return errBufferTooSmall
}
k.Version = binary.LittleEndian.Uint64(data[offset:])
offset += 8
// ExpiresAt
if offset+8 > len(data) {
return errBufferTooSmall
}
k.ExpiresAt = binary.LittleEndian.Uint64(data[offset:])
offset += 8
// Meta
if offset+4 > len(data) {
return errBufferTooSmall
}
metaLen := int(binary.LittleEndian.Uint32(data[offset:]))
offset += 4
if offset+metaLen > len(data) {
return errBufferTooSmall
}
k.Meta = make([]byte, metaLen)
copy(k.Meta, data[offset:offset+metaLen])
offset += metaLen
// StreamId
if offset+4 > len(data) {
return errBufferTooSmall
}
k.StreamId = binary.LittleEndian.Uint32(data[offset:])
offset += 4
// StreamDone
if offset+1 > len(data) {
return errBufferTooSmall
}
k.StreamDone = data[offset] != 0
return nil
}
// KVList represents a list of KV pairs.
type KVList struct {
Kv []*KV
AllocRef uint64
}
func (l *KVList) GetKv() []*KV { return l.Kv }
func (l *KVList) GetAllocRef() uint64 { return l.AllocRef }
func (l *KVList) Reset() { *l = KVList{} }
func (l *KVList) String() string { return "KVList{...}" }
// Size returns the encoded size of KVList.
func (l *KVList) Size() int {
size := 4 + 8 // count + allocRef
for _, kv := range l.Kv {
size += 4 + kv.Size() // length prefix + kv data
}
return size
}
// Marshal encodes KVList to binary format.
func (l *KVList) Marshal() ([]byte, error) {
buf := make([]byte, l.Size())
offset := 0
// Count
binary.LittleEndian.PutUint32(buf[offset:], uint32(len(l.Kv)))
offset += 4
// KVs
for _, kv := range l.Kv {
kvSize := kv.Size()
binary.LittleEndian.PutUint32(buf[offset:], uint32(kvSize))
offset += 4
kv.MarshalToSizedBuffer(buf[offset : offset+kvSize])
offset += kvSize
}
// AllocRef
binary.LittleEndian.PutUint64(buf[offset:], l.AllocRef)
return buf, nil
}
// Unmarshal decodes KVList from binary format.
func (l *KVList) Unmarshal(data []byte) error {
if len(data) < 12 { // minimum: count(4) + allocRef(8)
return errBufferTooSmall
}
offset := 0
// Count
count := int(binary.LittleEndian.Uint32(data[offset:]))
offset += 4
l.Kv = make([]*KV, count)
for i := 0; i < count; i++ {
if offset+4 > len(data) {
return errBufferTooSmall
}
kvSize := int(binary.LittleEndian.Uint32(data[offset:]))
offset += 4
if offset+kvSize > len(data) {
return errBufferTooSmall
}
l.Kv[i] = &KV{}
if err := l.Kv[i].Unmarshal(data[offset : offset+kvSize]); err != nil {
return err
}
offset += kvSize
}
// AllocRef
if offset+8 > len(data) {
return errBufferTooSmall
}
l.AllocRef = binary.LittleEndian.Uint64(data[offset:])
return nil
}
// ManifestChange represents a change to the manifest.
type ManifestChange struct {
Id uint64
Op ManifestChange_Operation
Level uint32
KeyId uint64
EncryptionAlgo EncryptionAlgo
Compression uint32
}
func (m *ManifestChange) GetId() uint64 { return m.Id }
func (m *ManifestChange) GetOp() ManifestChange_Operation { return m.Op }
func (m *ManifestChange) GetLevel() uint32 { return m.Level }
func (m *ManifestChange) GetKeyId() uint64 { return m.KeyId }
func (m *ManifestChange) GetEncryptionAlgo() EncryptionAlgo { return m.EncryptionAlgo }
func (m *ManifestChange) GetCompression() uint32 { return m.Compression }
func (m *ManifestChange) Reset() { *m = ManifestChange{} }
func (m *ManifestChange) String() string { return "ManifestChange{...}" }
// Size returns the encoded size of ManifestChange.
// Format: [id:8][op:4][level:4][keyId:8][encryptionAlgo:4][compression:4]
func (m *ManifestChange) Size() int {
return 8 + 4 + 4 + 8 + 4 + 4 // 32 bytes
}
// Marshal encodes ManifestChange to binary format.
func (m *ManifestChange) Marshal() ([]byte, error) {
buf := make([]byte, m.Size())
offset := 0
binary.LittleEndian.PutUint64(buf[offset:], m.Id)
offset += 8
binary.LittleEndian.PutUint32(buf[offset:], uint32(m.Op))
offset += 4
binary.LittleEndian.PutUint32(buf[offset:], m.Level)
offset += 4
binary.LittleEndian.PutUint64(buf[offset:], m.KeyId)
offset += 8
binary.LittleEndian.PutUint32(buf[offset:], uint32(m.EncryptionAlgo))
offset += 4
binary.LittleEndian.PutUint32(buf[offset:], m.Compression)
return buf, nil
}
// Unmarshal decodes ManifestChange from binary format.
func (m *ManifestChange) Unmarshal(data []byte) error {
if len(data) < 32 {
return errBufferTooSmall
}
offset := 0
m.Id = binary.LittleEndian.Uint64(data[offset:])
offset += 8
m.Op = ManifestChange_Operation(binary.LittleEndian.Uint32(data[offset:]))
offset += 4
m.Level = binary.LittleEndian.Uint32(data[offset:])
offset += 4
m.KeyId = binary.LittleEndian.Uint64(data[offset:])
offset += 8
m.EncryptionAlgo = EncryptionAlgo(binary.LittleEndian.Uint32(data[offset:]))
offset += 4
m.Compression = binary.LittleEndian.Uint32(data[offset:])
return nil
}
// ManifestChangeSet represents a set of manifest changes.
type ManifestChangeSet struct {
Changes []*ManifestChange
}
func (m *ManifestChangeSet) GetChanges() []*ManifestChange { return m.Changes }
func (m *ManifestChangeSet) Reset() { *m = ManifestChangeSet{} }
func (m *ManifestChangeSet) String() string { return "ManifestChangeSet{...}" }
// Size returns the encoded size of ManifestChangeSet.
func (m *ManifestChangeSet) Size() int {
size := 4 // count
for range m.Changes {
size += 4 + 32 // length prefix + ManifestChange size
}
return size
}
// Marshal encodes ManifestChangeSet to binary format.
func (m *ManifestChangeSet) Marshal() ([]byte, error) {
buf := make([]byte, m.Size())
offset := 0
binary.LittleEndian.PutUint32(buf[offset:], uint32(len(m.Changes)))
offset += 4
for _, change := range m.Changes {
changeData, err := change.Marshal()
if err != nil {
return nil, err
}
binary.LittleEndian.PutUint32(buf[offset:], uint32(len(changeData)))
offset += 4
copy(buf[offset:], changeData)
offset += len(changeData)
}
return buf, nil
}
// Unmarshal decodes ManifestChangeSet from binary format.
func (m *ManifestChangeSet) Unmarshal(data []byte) error {
if len(data) < 4 {
return errBufferTooSmall
}
offset := 0
count := int(binary.LittleEndian.Uint32(data[offset:]))
offset += 4
m.Changes = make([]*ManifestChange, count)
for i := 0; i < count; i++ {
if offset+4 > len(data) {
return errBufferTooSmall
}
changeSize := int(binary.LittleEndian.Uint32(data[offset:]))
offset += 4
if offset+changeSize > len(data) {
return errBufferTooSmall
}
m.Changes[i] = &ManifestChange{}
if err := m.Changes[i].Unmarshal(data[offset : offset+changeSize]); err != nil {
return err
}
offset += changeSize
}
return nil
}
// DataKey represents an encryption data key.
type DataKey struct {
KeyId uint64
Data []byte
Iv []byte
CreatedAt int64
}
func (d *DataKey) GetKeyId() uint64 { return d.KeyId }
func (d *DataKey) GetData() []byte { return d.Data }
func (d *DataKey) GetIv() []byte { return d.Iv }
func (d *DataKey) GetCreatedAt() int64 { return d.CreatedAt }
func (d *DataKey) Reset() { *d = DataKey{} }
func (d *DataKey) String() string { return "DataKey{...}" }
// Size returns the encoded size of DataKey.
// Format: [keyId:8][dataLen:4][data][ivLen:4][iv][createdAt:8]
func (d *DataKey) Size() int {
return 8 + 4 + len(d.Data) + 4 + len(d.Iv) + 8
}
// Marshal encodes DataKey to binary format.
func (d *DataKey) Marshal() ([]byte, error) {
buf := make([]byte, d.Size())
offset := 0
binary.LittleEndian.PutUint64(buf[offset:], d.KeyId)
offset += 8
binary.LittleEndian.PutUint32(buf[offset:], uint32(len(d.Data)))
offset += 4
copy(buf[offset:], d.Data)
offset += len(d.Data)
binary.LittleEndian.PutUint32(buf[offset:], uint32(len(d.Iv)))
offset += 4
copy(buf[offset:], d.Iv)
offset += len(d.Iv)
binary.LittleEndian.PutUint64(buf[offset:], uint64(d.CreatedAt))
return buf, nil
}
// Unmarshal decodes DataKey from binary format.
func (d *DataKey) Unmarshal(data []byte) error {
if len(data) < 24 { // minimum: keyId(8) + dataLen(4) + ivLen(4) + createdAt(8)
return errBufferTooSmall
}
offset := 0
d.KeyId = binary.LittleEndian.Uint64(data[offset:])
offset += 8
dataLen := int(binary.LittleEndian.Uint32(data[offset:]))
offset += 4
if offset+dataLen > len(data) {
return errBufferTooSmall
}
d.Data = make([]byte, dataLen)
copy(d.Data, data[offset:offset+dataLen])
offset += dataLen
if offset+4 > len(data) {
return errBufferTooSmall
}
ivLen := int(binary.LittleEndian.Uint32(data[offset:]))
offset += 4
if offset+ivLen > len(data) {
return errBufferTooSmall
}
d.Iv = make([]byte, ivLen)
copy(d.Iv, data[offset:offset+ivLen])
offset += ivLen
if offset+8 > len(data) {
return errBufferTooSmall
}
d.CreatedAt = int64(binary.LittleEndian.Uint64(data[offset:]))
return nil
}
// Checksum represents a checksum with algorithm.
type Checksum struct {
Algo Checksum_Algorithm
Sum uint64
}
func (c *Checksum) GetAlgo() Checksum_Algorithm { return c.Algo }
func (c *Checksum) GetSum() uint64 { return c.Sum }
func (c *Checksum) Reset() { *c = Checksum{} }
func (c *Checksum) String() string { return "Checksum{...}" }
// Size returns the encoded size of Checksum.
// Format: [algo:4][sum:8]
func (c *Checksum) Size() int {
return 4 + 8 // 12 bytes
}
// Marshal encodes Checksum to binary format.
func (c *Checksum) Marshal() ([]byte, error) {
buf := make([]byte, c.Size())
binary.LittleEndian.PutUint32(buf[0:], uint32(c.Algo))
binary.LittleEndian.PutUint64(buf[4:], c.Sum)
return buf, nil
}
// Unmarshal decodes Checksum from binary format.
func (c *Checksum) Unmarshal(data []byte) error {
if len(data) < 12 {
return errBufferTooSmall
}
c.Algo = Checksum_Algorithm(binary.LittleEndian.Uint32(data[0:]))
c.Sum = binary.LittleEndian.Uint64(data[4:])
return nil
}
// Match represents a match pattern.
type Match struct {
Prefix []byte
IgnoreBytes string
}
func (m *Match) GetPrefix() []byte { return m.Prefix }
func (m *Match) GetIgnoreBytes() string { return m.IgnoreBytes }
func (m *Match) Reset() { *m = Match{} }
func (m *Match) String() string { return "Match{...}" }
// Size returns the encoded size of Match.
// Format: [prefixLen:4][prefix][ignoreBytesLen:4][ignoreBytes]
func (m *Match) Size() int {
return 4 + len(m.Prefix) + 4 + len(m.IgnoreBytes)
}
// Marshal encodes Match to binary format.
func (m *Match) Marshal() ([]byte, error) {
buf := make([]byte, m.Size())
offset := 0
binary.LittleEndian.PutUint32(buf[offset:], uint32(len(m.Prefix)))
offset += 4
copy(buf[offset:], m.Prefix)
offset += len(m.Prefix)
binary.LittleEndian.PutUint32(buf[offset:], uint32(len(m.IgnoreBytes)))
offset += 4
copy(buf[offset:], m.IgnoreBytes)
return buf, nil
}
// Unmarshal decodes Match from binary format.
func (m *Match) Unmarshal(data []byte) error {
if len(data) < 8 { // minimum: prefixLen(4) + ignoreBytesLen(4)
return errBufferTooSmall
}
offset := 0
prefixLen := int(binary.LittleEndian.Uint32(data[offset:]))
offset += 4
if offset+prefixLen > len(data) {
return errBufferTooSmall
}
m.Prefix = make([]byte, prefixLen)
copy(m.Prefix, data[offset:offset+prefixLen])
offset += prefixLen
if offset+4 > len(data) {
return errBufferTooSmall
}
ignoreBytesLen := int(binary.LittleEndian.Uint32(data[offset:]))
offset += 4
if offset+ignoreBytesLen > len(data) {
return errBufferTooSmall
}
m.IgnoreBytes = string(data[offset : offset+ignoreBytesLen])
return nil
}
// Clone returns a deep copy of KV.
func (k *KV) Clone() *KV {
if k == nil {
return nil
}
clone := &KV{
Version: k.Version,
ExpiresAt: k.ExpiresAt,
StreamId: k.StreamId,
StreamDone: k.StreamDone,
}
if k.Key != nil {
clone.Key = make([]byte, len(k.Key))
copy(clone.Key, k.Key)
}
if k.Value != nil {
clone.Value = make([]byte, len(k.Value))
copy(clone.Value, k.Value)
}
if k.UserMeta != nil {
clone.UserMeta = make([]byte, len(k.UserMeta))
copy(clone.UserMeta, k.UserMeta)
}
if k.Meta != nil {
clone.Meta = make([]byte, len(k.Meta))
copy(clone.Meta, k.Meta)
}
return clone
}
pb/types_zap_test.go
+254
View File
@@ -0,0 +1,254 @@
//go:build !grpc
/*
* SPDX-FileCopyrightText: 2017-2025 Istari Digital, Inc.
* SPDX-License-Identifier: Apache-2.0
*/
package pb
import (
"testing"
)
func TestKVMarshalUnmarshal(t *testing.T) {
kv := &KV{
Key: []byte("test-key"),
Value: []byte("test-value"),
UserMeta: []byte{0x01},
Version: 12345,
ExpiresAt: 67890,
Meta: []byte{0x02},
StreamId: 42,
StreamDone: true,
}
data, err := kv.Marshal()
if err != nil {
t.Fatalf("Marshal failed: %v", err)
}
kv2 := &KV{}
if err := kv2.Unmarshal(data); err != nil {
t.Fatalf("Unmarshal failed: %v", err)
}
if string(kv2.Key) != string(kv.Key) {
t.Errorf("Key mismatch: got %s, want %s", kv2.Key, kv.Key)
}
if string(kv2.Value) != string(kv.Value) {
t.Errorf("Value mismatch: got %s, want %s", kv2.Value, kv.Value)
}
if kv2.Version != kv.Version {
t.Errorf("Version mismatch: got %d, want %d", kv2.Version, kv.Version)
}
if kv2.ExpiresAt != kv.ExpiresAt {
t.Errorf("ExpiresAt mismatch: got %d, want %d", kv2.ExpiresAt, kv.ExpiresAt)
}
if kv2.StreamId != kv.StreamId {
t.Errorf("StreamId mismatch: got %d, want %d", kv2.StreamId, kv.StreamId)
}
if kv2.StreamDone != kv.StreamDone {
t.Errorf("StreamDone mismatch: got %v, want %v", kv2.StreamDone, kv.StreamDone)
}
}
func TestKVListMarshalUnmarshal(t *testing.T) {
list := &KVList{
Kv: []*KV{
{Key: []byte("key1"), Value: []byte("value1"), Version: 1},
{Key: []byte("key2"), Value: []byte("value2"), Version: 2},
},
AllocRef: 999,
}
data, err := list.Marshal()
if err != nil {
t.Fatalf("Marshal failed: %v", err)
}
list2 := &KVList{}
if err := list2.Unmarshal(data); err != nil {
t.Fatalf("Unmarshal failed: %v", err)
}
if len(list2.Kv) != len(list.Kv) {
t.Fatalf("KV count mismatch: got %d, want %d", len(list2.Kv), len(list.Kv))
}
if list2.AllocRef != list.AllocRef {
t.Errorf("AllocRef mismatch: got %d, want %d", list2.AllocRef, list.AllocRef)
}
for i := range list.Kv {
if string(list2.Kv[i].Key) != string(list.Kv[i].Key) {
t.Errorf("KV[%d].Key mismatch", i)
}
}
}
func TestManifestChangeMarshalUnmarshal(t *testing.T) {
mc := &ManifestChange{
Id: 12345,
Op: ManifestChange_CREATE,
Level: 3,
KeyId: 67890,
EncryptionAlgo: EncryptionAlgo_aes,
Compression: 1,
}
data, err := mc.Marshal()
if err != nil {
t.Fatalf("Marshal failed: %v", err)
}
mc2 := &ManifestChange{}
if err := mc2.Unmarshal(data); err != nil {
t.Fatalf("Unmarshal failed: %v", err)
}
if mc2.Id != mc.Id {
t.Errorf("Id mismatch: got %d, want %d", mc2.Id, mc.Id)
}
if mc2.Op != mc.Op {
t.Errorf("Op mismatch: got %d, want %d", mc2.Op, mc.Op)
}
if mc2.Level != mc.Level {
t.Errorf("Level mismatch: got %d, want %d", mc2.Level, mc.Level)
}
}
func TestChecksumMarshalUnmarshal(t *testing.T) {
cs := &Checksum{
Algo: Checksum_XXHash64,
Sum: 123456789012345,
}
data, err := cs.Marshal()
if err != nil {
t.Fatalf("Marshal failed: %v", err)
}
cs2 := &Checksum{}
if err := cs2.Unmarshal(data); err != nil {
t.Fatalf("Unmarshal failed: %v", err)
}
if cs2.Algo != cs.Algo {
t.Errorf("Algo mismatch: got %d, want %d", cs2.Algo, cs.Algo)
}
if cs2.Sum != cs.Sum {
t.Errorf("Sum mismatch: got %d, want %d", cs2.Sum, cs.Sum)
}
}
func TestDataKeyMarshalUnmarshal(t *testing.T) {
dk := &DataKey{
KeyId: 123,
Data: []byte("encryption-key-data"),
Iv: []byte("initialization-vector"),
CreatedAt: 1609459200,
}
data, err := dk.Marshal()
if err != nil {
t.Fatalf("Marshal failed: %v", err)
}
dk2 := &DataKey{}
if err := dk2.Unmarshal(data); err != nil {
t.Fatalf("Unmarshal failed: %v", err)
}
if dk2.KeyId != dk.KeyId {
t.Errorf("KeyId mismatch: got %d, want %d", dk2.KeyId, dk.KeyId)
}
if string(dk2.Data) != string(dk.Data) {
t.Errorf("Data mismatch")
}
if dk2.CreatedAt != dk.CreatedAt {
t.Errorf("CreatedAt mismatch: got %d, want %d", dk2.CreatedAt, dk.CreatedAt)
}
}
func TestMarshalUnmarshalInterface(t *testing.T) {
kv := &KV{
Key: []byte("test"),
Value: []byte("value"),
Version: 100,
}
// Test via interface
data, err := Marshal(kv)
if err != nil {
t.Fatalf("Marshal failed: %v", err)
}
kv2 := &KV{}
if err := Unmarshal(data, kv2); err != nil {
t.Fatalf("Unmarshal failed: %v", err)
}
if kv2.Version != kv.Version {
t.Errorf("Version mismatch")
}
// Test Size
if Size(kv) != len(data) {
t.Errorf("Size mismatch: got %d, want %d", Size(kv), len(data))
}
}
func TestMarshalAppend(t *testing.T) {
kv := &KV{
Key: []byte("test"),
Value: []byte("value"),
Version: 100,
}
prefix := []byte("prefix-")
result, err := MarshalOptions{}.MarshalAppend(prefix, kv)
if err != nil {
t.Fatalf("MarshalAppend failed: %v", err)
}
if string(result[:7]) != "prefix-" {
t.Errorf("Prefix not preserved")
}
kv2 := &KV{}
if err := kv2.Unmarshal(result[7:]); err != nil {
t.Fatalf("Unmarshal failed: %v", err)
}
if kv2.Version != kv.Version {
t.Errorf("Version mismatch")
}
}
func TestKVClone(t *testing.T) {
kv := &KV{
Key: []byte("test-key"),
Value: []byte("test-value"),
UserMeta: []byte{0x01},
Version: 12345,
ExpiresAt: 67890,
Meta: []byte{0x02},
StreamId: 42,
StreamDone: true,
}
clone := kv.Clone()
// Verify values are equal
if string(clone.Key) != string(kv.Key) {
t.Errorf("Key mismatch")
}
if clone.Version != kv.Version {
t.Errorf("Version mismatch")
}
// Verify deep copy (modifying original doesn't affect clone)
kv.Key[0] = 'x'
if clone.Key[0] == 'x' {
t.Errorf("Clone shares memory with original")
}
}
stream.go
+3 -4
View File
@@ -14,7 +14,6 @@ import (
"time"
humanize "github.com/dustin/go-humanize"
"google.golang.org/protobuf/proto"
"github.com/dgraph-io/badger/v4/pb"
"github.com/dgraph-io/badger/v4/y"
@@ -507,7 +506,7 @@ func BufferToKVList(buf *z.Buffer) (*pb.KVList, error) {
var list pb.KVList
err := buf.SliceIterate(func(s []byte) error {
kv := new(pb.KV)
if err := proto.Unmarshal(s, kv); err != nil {
if err := pb.Unmarshal(s, kv); err != nil {
return err
}
list.Kv = append(list.Kv, kv)
@@ -517,7 +516,7 @@ func BufferToKVList(buf *z.Buffer) (*pb.KVList, error) {
}
func KVToBuffer(kv *pb.KV, buf *z.Buffer) {
in := buf.SliceAllocate(proto.Size(kv))[:0]
_, err := proto.MarshalOptions{}.MarshalAppend(in, kv)
in := buf.SliceAllocate(pb.Size(kv))[:0]
_, err := pb.MarshalOptions{}.MarshalAppend(in, kv)
y.AssertTrue(err == nil)
}
stream_test.go
+1 -2
View File
@@ -15,7 +15,6 @@ import (
"testing"
"github.com/stretchr/testify/require"
"google.golang.org/protobuf/proto"
bpb "github.com/dgraph-io/badger/v4/pb"
"github.com/dgraph-io/badger/v4/y"
@@ -50,7 +49,7 @@ func (c *collector) Send(buf *z.Buffer) error {
if kv.StreamDone == true {
return nil
}
cp := proto.Clone(kv).(*bpb.KV)
cp := kv.Clone()
c.kv = append(c.kv, cp)
}
return err
stream_writer.go
+1 -2
View File
@@ -11,7 +11,6 @@ import (
"sync"
"github.com/dustin/go-humanize"
"google.golang.org/protobuf/proto"
"github.com/dgraph-io/badger/v4/pb"
"github.com/dgraph-io/badger/v4/table"
@@ -141,7 +140,7 @@ func (sw *StreamWriter) Write(buf *z.Buffer) error {
err := buf.SliceIterate(func(s []byte) error {
var kv pb.KV
if err := proto.Unmarshal(s, &kv); err != nil {
if err := pb.Unmarshal(s, &kv); err != nil {
return err
}
if kv.StreamDone {
table/builder.go
+1 -2
View File
@@ -16,7 +16,6 @@ import (
fbs "github.com/google/flatbuffers/go"
"github.com/klauspost/compress/s2"
"google.golang.org/protobuf/proto"
"github.com/dgraph-io/badger/v4/fb"
"github.com/dgraph-io/badger/v4/options"
@@ -467,7 +466,7 @@ func (b *Builder) calculateChecksum(data []byte) []byte {
}
// Write checksum to the file.
chksum, err := proto.Marshal(&checksum)
chksum, err := pb.Marshal(&checksum)
y.Check(err)
// Write checksum size.
return chksum
table/table.go
+3 -4
View File
@@ -23,7 +23,6 @@ import (
"github.com/klauspost/compress/snappy"
"github.com/klauspost/compress/zstd"
"google.golang.org/protobuf/proto"
"github.com/dgraph-io/badger/v4/fb"
"github.com/dgraph-io/badger/v4/options"
@@ -236,7 +235,7 @@ func (b *Block) size() int64 {
func (b *Block) verifyCheckSum() error {
cs := &pb.Checksum{}
if err := proto.Unmarshal(b.checksum, cs); err != nil {
if err := pb.Unmarshal(b.checksum, cs); err != nil {
return y.Wrapf(err, "unable to unmarshal checksum for block")
}
return y.VerifyChecksum(b.data, cs)
@@ -369,7 +368,7 @@ func (t *Table) initBiggestAndSmallest() error {
checksum := &pb.Checksum{}
readPos -= checksumLen
buf = t.readNoFail(readPos, checksumLen)
_ = proto.Unmarshal(buf, checksum)
_ = pb.Unmarshal(buf, checksum)
fmt.Fprintf(&debugBuf, "checksum: %+v ", checksum)
// Read index size from the footer.
@@ -431,7 +430,7 @@ func (t *Table) initIndex() (*fb.BlockOffset, error) {
expectedChk := &pb.Checksum{}
readPos -= checksumLen
buf = t.readNoFail(readPos, checksumLen)
if err := proto.Unmarshal(buf, expectedChk); err != nil {
if err := pb.Unmarshal(buf, expectedChk); err != nil {
return nil, err
}
y/y_test.go
+1 -2
View File
@@ -15,7 +15,6 @@ import (
"time"
"github.com/stretchr/testify/require"
"google.golang.org/protobuf/proto"
"github.com/dgraph-io/badger/v4/pb"
"github.com/dgraph-io/ristretto/v2/z"
@@ -323,7 +322,7 @@ func TestAllocatorReuse(t *testing.T) {
kv.Version = uint64(sz)
list.Kv = append(list.Kv, kv)
}
_, err := proto.Marshal(&list)
_, err := pb.Marshal(&list)
require.NoError(t, err)
}
t.Logf("Allocator: %s\n", a)