[tree] return blob builders to pool after use

dolthub · May 21, 2024 · b024e60 · b024e60
1 parent 4a0f03d
commit b024e60
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 4 deletions.
diff --git a/go/libraries/doltcore/migrate/tuples.go b/go/libraries/doltcore/migrate/tuples.go
@@ -247,6 +247,8 @@ func translateStringField(ctx context.Context, ns tree.NodeStore, value types.St
 		// note: previously, TEXT fields were serialized as types.String
 		rd := strings.NewReader(string(value))
 		bb := ns.BlobBuilder()
+		defer ns.PutBlobBuilder(bb)
+
 		bb.Init(len(value))
 		_, addr, err := bb.Chunk(ctx, rd)
 		if err != nil {
@@ -316,6 +318,8 @@ func translateJSONField(ctx context.Context, ns tree.NodeStore, value types.JSON
 	buf := bytes.NewBuffer([]byte(s))
 
 	bb := ns.BlobBuilder()
+	defer ns.PutBlobBuilder(bb)
+
 	bb.Init(len(s))
 	_, addr, err := bb.Chunk(ctx, buf)
 	if err != nil {
@@ -346,6 +350,8 @@ func translateBlobField(ctx context.Context, ns tree.NodeStore, value types.Blob
 	}
 
 	bb := ns.BlobBuilder()
+	defer ns.PutBlobBuilder(bb)
+
 	bb.Init(int(value.Len()))
 	_, addr, err := bb.Chunk(ctx, bytes.NewReader(buf))
 	if err != nil {

diff --git a/go/store/prolly/tree/node_store.go b/go/store/prolly/tree/node_store.go
@@ -48,6 +48,7 @@ type NodeStore interface {
 	Format() *types.NomsBinFormat
 
 	BlobBuilder() *BlobBuilder
+	PutBlobBuilder(*BlobBuilder)
 }
 
 type nodeStore struct {
@@ -182,6 +183,11 @@ func (ns nodeStore) BlobBuilder() *BlobBuilder {
 	return bb
 }
 
+// PutBlobBuilder implements NodeStore.
+func (ns nodeStore) PutBlobBuilder(bb *BlobBuilder) {
+	ns.bbp.Put(bb)
+}
+
 func (ns nodeStore) Format() *types.NomsBinFormat {
 	nbf, err := types.GetFormatForVersionString(ns.store.Version())
 	if err != nil {

diff --git a/go/store/prolly/tree/prolly_fields.go b/go/store/prolly/tree/prolly_fields.go
@@ -397,6 +397,7 @@ func serializeGeometry(v interface{}) []byte {
 
 func SerializeBytesToAddr(ctx context.Context, ns NodeStore, r io.Reader, dataSize int) (hash.Hash, error) {
 	bb := ns.BlobBuilder()
+	defer ns.PutBlobBuilder(bb)
 	bb.Init(dataSize)
 	_, addr, err := bb.Chunk(ctx, r)
 	if err != nil {

diff --git a/go/store/prolly/tree/testutils.go b/go/store/prolly/tree/testutils.go
@@ -264,12 +264,12 @@ func NewTestNodeStore() NodeStore {
 	ts := &chunks.TestStorage{}
 	ns := NewNodeStore(ts.NewViewWithFormat(types.Format_DOLT.VersionString()))
 	bb := &blobBuilderPool
-	return nodeStoreValidator{ns: ns, bb: bb}
+	return nodeStoreValidator{ns: ns, bbp: bb}
 }
 
 type nodeStoreValidator struct {
-	ns NodeStore
-	bb *sync.Pool
+	ns  NodeStore
+	bbp *sync.Pool
 }
 
 func (v nodeStoreValidator) Read(ctx context.Context, ref hash.Hash) (Node, error) {
@@ -320,13 +320,18 @@ func (v nodeStoreValidator) Pool() pool.BuffPool {
 }
 
 func (v nodeStoreValidator) BlobBuilder() *BlobBuilder {
-	bb := v.bb.Get().(*BlobBuilder)
+	bb := v.bbp.Get().(*BlobBuilder)
 	if bb.ns == nil {
 		bb.SetNodeStore(v)
 	}
 	return bb
 }
 
+// PutBlobBuilder implements NodeStore.
+func (v nodeStoreValidator) PutBlobBuilder(bb *BlobBuilder) {
+	v.bbp.Put(bb)
+}
+
 func (v nodeStoreValidator) Format() *types.NomsBinFormat {
 	return v.ns.Format()
 }