opt(memory): Use z.Calloc for allocating KVList (#1563)
KVs can take up a lot of memory in the stream framework. With this change, we allocate them using z.Allocator, and allow callers to KeyToList to use the allocator to generate KVs as well. After we call Send, we release them.
Also change Stream Framework to spit out StreamDone markers.
diff --git a/stream_test.go b/stream_test.go
index f3c4b1d..f0bd8bd 100644
--- a/stream_test.go
+++ b/stream_test.go
@@ -27,6 +27,8 @@
bpb "github.com/dgraph-io/badger/v2/pb"
"github.com/dgraph-io/badger/v2/y"
+ "github.com/dgraph-io/ristretto/z"
+ "github.com/golang/protobuf/proto"
"github.com/stretchr/testify/require"
)
@@ -50,7 +52,13 @@
}
func (c *collector) Send(list *bpb.KVList) error {
- c.kv = append(c.kv, list.Kv...)
+ for _, kv := range list.Kv {
+ if kv.StreamDone == true {
+ continue
+ }
+ cp := proto.Clone(kv).(*bpb.KV)
+ c.kv = append(c.kv, cp)
+ }
return nil
}
@@ -156,6 +164,7 @@
require.Equal(t, 50, count, "Count mismatch for pred: %s", pred)
}
require.NoError(t, db.Close())
+ require.Equal(t, int64(0), z.NumAllocBytes())
}
func TestStreamWithThreadId(t *testing.T) {
