luxfi/zapdb
1
0
Fork 0
mirror of https://github.com/luxfi/zapdb.git synced 2026-07-27 06:54:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix(OOM): Use z.Allocator for Stream (#1611)

Browse Source 
* Let's use allocator again
* Switch to z.NumAllocBytes
* Make stream work with both another Badger DB or with a file to backup.
* Add a test for Allocator
* Use allocator for Backup
* Bring in latest Ristretto

Co-authored-by: Daniel Mai <daniel@dgraph.io>
This commit is contained in:
Manish R Jain
2020-12-04 12:59:40 -08:00
committed by GitHub
co-authored by Daniel Mai
parent 3bd1f12565
commit b744d51732
8 changed files with 89 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backup.go
+9 -5
View File
@@ -62,6 +62,7 @@ func (db *DB) Backup(w io.Writer, since uint64) (uint64, error) {
func (stream *Stream) Backup(w io.Writer, since uint64) (uint64, error) {
stream.KeyToList = func(key []byte, itr *Iterator) (*pb.KVList, error) {
list := &pb.KVList{}
a := itr.alloc
for ; itr.Valid(); itr.Next() {
item := itr.Item()
if !bytes.Equal(item.Key(), key) {
@@ -77,7 +78,10 @@ func (stream *Stream) Backup(w io.Writer, since uint64) (uint64, error) {
if !item.IsDeletedOrExpired() {
// No need to copy value, if item is deleted or expired.
var err error
valCopy, err = item.ValueCopy(nil)
err = item.Value(func(val []byte) error {
valCopy = a.Copy(val)
return nil
})
if err != nil {
stream.db.opt.Errorf("Key [%x, %d]. Error while fetching value [%v]\n",
item.Key(), item.Version(), err)
@@ -87,14 +91,14 @@ func (stream *Stream) Backup(w io.Writer, since uint64) (uint64, error) {
// clear txn bits
meta := item.meta &^ (bitTxn | bitFinTxn)
kv := itr.NewKV()
kv := y.NewKV(a)
*kv = pb.KV{
Key: item.KeyCopy(nil),
Key: a.Copy(item.Key()),
Value: valCopy,
UserMeta: []byte{item.UserMeta()},
UserMeta: a.Copy([]byte{item.UserMeta()}),
Version: item.Version(),
ExpiresAt: item.ExpiresAt(),
Meta: []byte{meta},
Meta: a.Copy([]byte{meta}),
}
list.Kv = append(list.Kv, kv)
badger/cmd/stream.go
+35 -22
View File
@@ -39,6 +39,7 @@ This command streams the contents of this DB into another DB with the given opti
}
var outDir string
var outFile string
var compressionType uint32
func init() {
@@ -46,6 +47,8 @@ func init() {
RootCmd.AddCommand(streamCmd)
streamCmd.Flags().StringVarP(&outDir, "out", "o", "",
"Path to output DB. The directory should be empty.")
streamCmd.Flags().StringVarP(&outFile, "", "f", "",
"Run a backup to this file.")
streamCmd.Flags().BoolVarP(&readOnly, "read_only", "", true,
"Option to open input DB in read-only mode")
streamCmd.Flags().IntVarP(&numVersions, "num_versions", "", 0,
@@ -59,20 +62,6 @@ func init() {
}
func stream(cmd *cobra.Command, args []string) error {
// Check that outDir doesn't exist or is empty.
if _, err := os.Stat(outDir); err == nil {
f, err := os.Open(outDir)
if err != nil {
return err
}
defer f.Close()
_, err = f.Readdirnames(1)
if err != io.EOF {
return errors.Errorf("cannot run stream tool on non-empty output directory %s", outDir)
}
}
// Options for input DB.
if numVersions <= 0 {
numVersions = math.MaxInt32
@@ -94,20 +83,44 @@ func stream(cmd *cobra.Command, args []string) error {
return errors.Errorf(
"compression value must be one of 0 (disabled), 1 (Snappy), or 2 (ZSTD)")
}
outOpt := inOpt.
WithDir(outDir).
WithValueDir(outDir).
WithNumVersionsToKeep(numVersions).
WithCompression(options.CompressionType(compressionType)).
WithReadOnly(false)
inDB, err := badger.OpenManaged(inOpt)
if err != nil {
return y.Wrapf(err, "cannot open DB at %s", sstDir)
}
defer inDB.Close()
err = inDB.StreamDB(outOpt)
stream := inDB.NewStreamAt(math.MaxUint64)
if len(outDir) > 0 {
if _, err := os.Stat(outDir); err == nil {
f, err := os.Open(outDir)
if err != nil {
return err
}
defer f.Close()
_, err = f.Readdirnames(1)
if err != io.EOF {
return errors.Errorf(
"cannot run stream tool on non-empty output directory %s", outDir)
}
}
stream.LogPrefix = "DB.Stream"
outOpt := inOpt.
WithDir(outDir).
WithValueDir(outDir).
WithNumVersionsToKeep(numVersions).
WithCompression(options.CompressionType(compressionType)).
WithReadOnly(false)
err = inDB.StreamDB(outOpt)
} else if len(outFile) > 0 {
stream.LogPrefix = "DB.Backup"
f, err := os.OpenFile(outFile, os.O_RDWR|os.O_CREATE, 0666)
y.Check(err)
_, err = stream.Backup(f, 0)
}
fmt.Println("Done.")
return err
}
go.mod
+1 -1
View File
@@ -7,7 +7,7 @@ go 1.12
require (
github.com/DataDog/zstd v1.4.1
github.com/cespare/xxhash v1.1.0
github.com/dgraph-io/ristretto v0.0.4-0.20201125174811-766bca5e9938
github.com/dgraph-io/ristretto v0.0.4-0.20201204194510-b7ca2e90f544
github.com/dustin/go-humanize v1.0.0
github.com/golang/protobuf v1.3.1
github.com/golang/snappy v0.0.1
go.sum
+2
View File
@@ -19,6 +19,8 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/dgraph-io/ristretto v0.0.4-0.20201125174811-766bca5e9938 h1:FdSJif9oUVeH+MpsScsrL6OAbdW0pUYvXmkdhDSWWcQ=
github.com/dgraph-io/ristretto v0.0.4-0.20201125174811-766bca5e9938/go.mod h1:tv2ec8nA7vRpSYX7/MbP52ihrUMXIHit54CQMq8npXQ=
github.com/dgraph-io/ristretto v0.0.4-0.20201204194510-b7ca2e90f544 h1:6vntPuznvHo+vxTe3KZYzJeorSUt5wkY+1ICtn1GEj0=
github.com/dgraph-io/ristretto v0.0.4-0.20201204194510-b7ca2e90f544/go.mod h1:tv2ec8nA7vRpSYX7/MbP52ihrUMXIHit54CQMq8npXQ=
github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2 h1:tdlZCpZ/P9DhczCTSixgIKmwPv6+wP5DGjqLYw5SUiA=
github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw=
github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo=
iterator.go
+2 -16
View File
@@ -26,8 +26,8 @@ import (
"sync/atomic"
"time"
"github.com/dgraph-io/badger/v2/pb"
"github.com/dgraph-io/badger/v2/table"
"github.com/dgraph-io/ristretto/z"
"github.com/dgraph-io/badger/v2/y"
)
@@ -436,7 +436,7 @@ type Iterator struct {
// iterators created by the stream interface
ThreadId int
reuse []*pb.KV
alloc *z.Allocator
}
// NewIterator returns a new iterator. Depending upon the options, either only keys, or both
@@ -496,20 +496,6 @@ func (txn *Txn) NewKeyIterator(key []byte, opt IteratorOptions) *Iterator {
return txn.NewIterator(opt)
}
// NewKV must be called serially. It is NOT thread-safe.
func (it *Iterator) NewKV() *pb.KV {
if len(it.reuse) == 0 {
return &pb.KV{}
}
kv := it.reuse[len(it.reuse)-1]
it.reuse = it.reuse[:len(it.reuse)-1]
if kv == nil {
kv = &pb.KV{}
}
kv.Reset()
return kv
}
func (it *Iterator) newItem() *Item {
item := it.waste.pop()
if item == nil {
stream.go
+11 -10
View File
@@ -97,7 +97,8 @@ func (st *Stream) SendDoneMarkers(done bool) {
// ToList is a default implementation of KeyToList. It picks up all valid versions of the key,
// skipping over deleted or expired keys.
func (st *Stream) ToList(key []byte, itr *Iterator) (*pb.KVList, error) {
ka := y.Copy(key)
a := itr.alloc
ka := a.Copy(key)
list := &pb.KVList{}
for ; itr.Valid(); itr.Next() {
@@ -110,11 +111,11 @@ func (st *Stream) ToList(key []byte, itr *Iterator) (*pb.KVList, error) {
break
}
kv := itr.NewKV()
kv := y.NewKV(a)
kv.Key = ka
if err := item.Value(func(val []byte) error {
kv.Value = y.Copy(val)
kv.Value = a.Copy(val)
return nil
}); err != nil {
@@ -122,7 +123,7 @@ func (st *Stream) ToList(key []byte, itr *Iterator) (*pb.KVList, error) {
}
kv.Version = item.Version()
kv.ExpiresAt = item.ExpiresAt()
kv.UserMeta = y.Copy([]byte{item.UserMeta()})
kv.UserMeta = a.Copy([]byte{item.UserMeta()})
list.Kv = append(list.Kv, kv)
if st.db.opt.NumVersionsToKeep == 1 {
@@ -195,6 +196,10 @@ func (st *Stream) produceKVs(ctx context.Context, threadId int) error {
itr.ThreadId = threadId
defer itr.Close()
itr.alloc = z.NewAllocator(1 << 20)
itr.alloc.Tag = "Stream.Iterate"
defer itr.alloc.Release()
// This unique stream id is used to identify all the keys from this iteration.
streamId := atomic.AddUint32(&st.nextStreamId, 1)
@@ -228,6 +233,7 @@ func (st *Stream) produceKVs(ctx context.Context, threadId int) error {
}
// Now convert to key value.
itr.alloc.Reset()
list, err := st.KeyToList(item.KeyCopy(nil), itr)
if err != nil {
return err
@@ -245,9 +251,6 @@ func (st *Stream) produceKVs(ctx context.Context, threadId int) error {
return err
}
}
if len(itr.reuse) < 100 {
itr.reuse = append(itr.reuse, list.Kv...)
}
}
// Mark the stream as done.
if st.doneMarkers {
@@ -337,11 +340,9 @@ outer:
}
speed := bytesSent / durSec
var ms z.MemStats
z.ReadMemStats(&ms)
st.db.opt.Infof("%s Time elapsed: %s, bytes sent: %s, speed: %s/sec, jemalloc: %s\n",
st.LogPrefix, y.FixedDuration(dur), humanize.IBytes(bytesSent),
humanize.IBytes(speed), humanize.IBytes(ms.Active))
humanize.IBytes(speed), humanize.IBytes(uint64(z.NumAllocBytes())))
case kvs, ok := <-st.kvChan:
if !ok {
stream_test.go
+1 -2
View File
@@ -300,8 +300,7 @@ func TestStreamCustomKeyToList(t *testing.T) {
if err != nil {
return nil, err
}
kv := itr.NewKV()
*kv = pb.KV{
kv := &pb.KV{
Key: y.Copy(item.Key()),
Value: val,
}
y/y_test.go
+28
View File
@@ -9,6 +9,8 @@ import (
"testing"
"time"
"github.com/dgraph-io/badger/v2/pb"
"github.com/dgraph-io/ristretto/z"
"github.com/stretchr/testify/require"
)
@@ -276,3 +278,29 @@ func TestEncodedSize(t *testing.T) {
require.Equal(t, valBufSize+uint32(2)+expVarintSize, valStruct.EncodedSize())
}
func TestAllocatorReuse(t *testing.T) {
a := z.NewAllocator(1024)
defer a.Release()
N := 1024
buf := make([]byte, 4096)
rand.Read(buf)
for i := 0; i < N; i++ {
a.Reset()
var list pb.KVList
for j := 0; j < N; j++ {
kv := NewKV(a)
sz := rand.Intn(1024)
kv.Key = a.Copy(buf[:sz])
kv.Value = a.Copy(buf[:4*sz])
kv.Meta = a.Copy([]byte{1})
kv.Version = uint64(sz)
list.Kv = append(list.Kv, kv)
}
_, err := list.Marshal()
require.NoError(t, err)
}
t.Logf("Allocator: %s\n", a)
}