extend the write batch parser

write_batch.go

@@ -2,7 +2,10 @@ package gorocksdb
 
 // #include "rocksdb/c.h"
 import "C"
-import "io"
+import (
+	"errors"
+	"io"
+)
 
 // WriteBatch is a batching of Puts, Merges and Deletes.
 type WriteBatch struct {
@@ -102,14 +105,31 @@ type WriteBatchRecordType byte
 
 // Types of batch records.
 const (
-	WriteBatchRecordTypeDeletion WriteBatchRecordType = 0x0
-	WriteBatchRecordTypeValue    WriteBatchRecordType = 0x1
-	WriteBatchRecordTypeMerge    WriteBatchRecordType = 0x2
-	WriteBatchRecordTypeLogData  WriteBatchRecordType = 0x3
+	WriteBatchDeletionRecord                 WriteBatchRecordType = 0x0
+	WriteBatchValueRecord                    WriteBatchRecordType = 0x1
+	WriteBatchMergeRecord                    WriteBatchRecordType = 0x2
+	WriteBatchLogDataRecord                  WriteBatchRecordType = 0x3
+	WriteBatchCFDeletionRecord               WriteBatchRecordType = 0x4
+	WriteBatchCFValueRecord                  WriteBatchRecordType = 0x5
+	WriteBatchCFMergeRecord                  WriteBatchRecordType = 0x6
+	WriteBatchSingleDeletionRecord           WriteBatchRecordType = 0x7
+	WriteBatchCFSingleDeletionRecord         WriteBatchRecordType = 0x8
+	WriteBatchBeginPrepareXIDRecord          WriteBatchRecordType = 0x9
+	WriteBatchEndPrepareXIDRecord            WriteBatchRecordType = 0xA
+	WriteBatchCommitXIDRecord                WriteBatchRecordType = 0xB
+	WriteBatchRollbackXIDRecord              WriteBatchRecordType = 0xC
+	WriteBatchNoopRecord                     WriteBatchRecordType = 0xD
+	WriteBatchRangeDeletion                  WriteBatchRecordType = 0xF
+	WriteBatchCFRangeDeletion                WriteBatchRecordType = 0xE
+	WriteBatchCFBlobIndex                    WriteBatchRecordType = 0x10
+	WriteBatchBlobIndex                      WriteBatchRecordType = 0x11
+	WriteBatchBeginPersistedPrepareXIDRecord WriteBatchRecordType = 0x12
+	WriteBatchNotUsedRecord                  WriteBatchRecordType = 0x7F
 )
 
 // WriteBatchRecord represents a record inside a WriteBatch.
 type WriteBatchRecord struct {
+	CF    int
 	Key   []byte
 	Value []byte
 	Type  WriteBatchRecordType
@@ -129,36 +149,63 @@ func (iter *WriteBatchIterator) Next() bool {
 		return false
 	}
 	// reset the current record
+	iter.record.CF = 0
 	iter.record.Key = nil
 	iter.record.Value = nil
 
 	// parse the record type
-	recordType := WriteBatchRecordType(iter.data[0])
-	iter.record.Type = recordType
-	iter.data = iter.data[1:]
+	iter.record.Type = iter.decodeRecType()
 
-	// parse the key
-	x, n := iter.decodeVarint(iter.data)
-	if n == 0 {
-		iter.err = io.ErrShortBuffer
-		return false
-	}
-	k := n + int(x)
-	iter.record.Key = iter.data[n:k]
-	iter.data = iter.data[k:]
-
-	// parse the data
-	if recordType == WriteBatchRecordTypeValue || recordType == WriteBatchRecordTypeMerge {
-		x, n := iter.decodeVarint(iter.data)
-		if n == 0 {
-			iter.err = io.ErrShortBuffer
-			return false
+	switch iter.record.Type {
+	case
+		WriteBatchDeletionRecord,
+		WriteBatchSingleDeletionRecord:
+		iter.record.Key = iter.decodeSlice()
+	case
+		WriteBatchCFDeletionRecord,
+		WriteBatchCFSingleDeletionRecord:
+		iter.record.CF = int(iter.decodeVarint())
+		if iter.err == nil {
+			iter.record.Key = iter.decodeSlice()
+		}
+	case
+		WriteBatchValueRecord,
+		WriteBatchMergeRecord,
+		WriteBatchRangeDeletion,
+		WriteBatchBlobIndex:
+		iter.record.Key = iter.decodeSlice()
+		if iter.err == nil {
+			iter.record.Value = iter.decodeSlice()
 		}
-		k := n + int(x)
-		iter.record.Value = iter.data[n:k]
-		iter.data = iter.data[k:]
+	case
+		WriteBatchCFValueRecord,
+		WriteBatchCFRangeDeletion,
+		WriteBatchCFMergeRecord,
+		WriteBatchCFBlobIndex:
+		iter.record.CF = int(iter.decodeVarint())
+		if iter.err == nil {
+			iter.record.Key = iter.decodeSlice()
+		}
+		if iter.err == nil {
+			iter.record.Value = iter.decodeSlice()
+		}
+	case WriteBatchLogDataRecord:
+		iter.record.Value = iter.decodeSlice()
+	case
+		WriteBatchNoopRecord,
+		WriteBatchBeginPrepareXIDRecord,
+		WriteBatchBeginPersistedPrepareXIDRecord:
+	case
+		WriteBatchEndPrepareXIDRecord,
+		WriteBatchCommitXIDRecord,
+		WriteBatchRollbackXIDRecord:
+		iter.record.Value = iter.decodeSlice()
+	default:
+		iter.err = errors.New("unsupported wal record type")
 	}
-	return true
+
+	return iter.err == nil
+
 }
 
 // Record returns the current record.
@@ -171,19 +218,45 @@ func (iter *WriteBatchIterator) Error() error {
 	return iter.err
 }
 
-func (iter *WriteBatchIterator) decodeVarint(buf []byte) (x uint64, n int) {
-	// x, n already 0
-	for shift := uint(0); shift < 64; shift += 7 {
-		if n >= len(buf) {
-			return 0, 0
-		}
-		b := uint64(buf[n])
+func (iter *WriteBatchIterator) decodeSlice() []byte {
+	l := int(iter.decodeVarint())
+	if l > len(iter.data) {
+		iter.err = io.ErrShortBuffer
+	}
+	if iter.err != nil {
+		return []byte{}
+	}
+	ret := iter.data[:l]
+	iter.data = iter.data[l:]
+	return ret
+}
+
+func (iter *WriteBatchIterator) decodeRecType() WriteBatchRecordType {
+	if len(iter.data) == 0 {
+		iter.err = io.ErrShortBuffer
+		return WriteBatchNotUsedRecord
+	}
+	t := iter.data[0]
+	iter.data = iter.data[1:]
+	return WriteBatchRecordType(t)
+}
+
+func (iter *WriteBatchIterator) decodeVarint() uint64 {
+	var n int
+	var x uint64
+	for shift := uint(0); shift < 64 && n < len(iter.data); shift += 7 {
+		b := uint64(iter.data[n])
 		n++
 		x |= (b & 0x7F) << shift
 		if (b & 0x80) == 0 {
-			return x, n
+			iter.data = iter.data[n:]
+			return x
 		}
 	}
-	// The number is too large to represent in a 64-bit value.
-	return 0, 0
+	if n == len(iter.data) {
+		iter.err = io.ErrShortBuffer
+	} else {
+		iter.err = errors.New("malformed varint")
+	}
+	return 0
 }

write_batch_test.go

@@ -61,13 +61,13 @@ func TestWriteBatchIterator(t *testing.T) {
 	iter := wb.NewIterator()
 	ensure.True(t, iter.Next())
 	record := iter.Record()
-	ensure.DeepEqual(t, record.Type, WriteBatchRecordTypeValue)
+	ensure.DeepEqual(t, record.Type, WriteBatchValueRecord)
 	ensure.DeepEqual(t, record.Key, givenKey1)
 	ensure.DeepEqual(t, record.Value, givenVal1)
 
 	ensure.True(t, iter.Next())
 	record = iter.Record()
-	ensure.DeepEqual(t, record.Type, WriteBatchRecordTypeDeletion)
+	ensure.DeepEqual(t, record.Type, WriteBatchDeletionRecord)
 	ensure.DeepEqual(t, record.Key, givenKey2)
 
 	// there shouldn't be any left