From c0cd58c19201f5a7dbb1dbfe7a3e11b90997b4a4 Mon Sep 17 00:00:00 2001
From: Pavel Safronov <pavel@dolthub.com>
Date: Tue, 7 Feb 2023 17:03:59 -0800
Subject: [PATCH] Add child refs to memtable only when a chunk was added.

Fixes https://github.com/dolthub/dolt/issues/5301
---
 go/store/nbs/mem_table.go | 21 ++++++++++++++++-----
 go/store/nbs/store.go     | 10 +++++-----
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/go/store/nbs/mem_table.go b/go/store/nbs/mem_table.go
index aea022b7ac..4476979ec8 100644
--- a/go/store/nbs/mem_table.go
+++ b/go/store/nbs/mem_table.go
@@ -33,6 +33,14 @@ import (
 	"github.com/dolthub/dolt/go/store/hash"
 )
 
+type addChunkResult int
+
+const (
+	chunkExists addChunkResult = iota
+	chunkAdded
+	chunkNotAdded
+)
+
 func WriteChunks(chunks []chunks.Chunk) (string, []byte, error) {
 	var size uint64
 	for _, chunk := range chunks {
@@ -46,7 +54,8 @@ func WriteChunks(chunks []chunks.Chunk) (string, []byte, error) {
 
 func writeChunksToMT(mt *memTable, chunks []chunks.Chunk) (string, []byte, error) {
 	for _, chunk := range chunks {
-		if !mt.addChunk(addr(chunk.Hash()), chunk.Data()) {
+		res := mt.addChunk(addr(chunk.Hash()), chunk.Data())
+		if res == chunkNotAdded {
 			return "", nil, errors.New("didn't create this memory table with enough space to add all the chunks")
 		}
 	}
@@ -78,17 +87,19 @@ func newMemTable(memTableSize uint64) *memTable {
 	return &memTable{chunks: map[addr][]byte{}, maxData: memTableSize}
 }
 
-func (mt *memTable) addChunk(h addr, data []byte) bool {
+func (mt *memTable) addChunk(h addr, data []byte) addChunkResult {
 	if len(data) == 0 {
 		panic("NBS blocks cannot be zero length")
 	}
 	if _, ok := mt.chunks[h]; ok {
-		return true
+		return chunkExists
 	}
+
 	dataLen := uint64(len(data))
 	if mt.totalData+dataLen > mt.maxData {
-		return false
+		return chunkNotAdded
 	}
+
 	mt.totalData += dataLen
 	mt.chunks[h] = data
 	mt.order = append(mt.order, hasRecord{
@@ -97,7 +108,7 @@ func (mt *memTable) addChunk(h addr, data []byte) bool {
 		len(mt.order),
 		false,
 	})
-	return true
+	return chunkAdded
 }
 
 func (mt *memTable) addChildRefs(addrs hash.HashSet) {
diff --git a/go/store/nbs/store.go b/go/store/nbs/store.go
index 69a05016ed..ada2f18223 100644
--- a/go/store/nbs/store.go
+++ b/go/store/nbs/store.go
@@ -623,8 +623,8 @@ func (nbs *NomsBlockStore) addChunk(ctx context.Context, ch chunks.Chunk, addrs
 	}
 	a := addr(ch.Hash())
 
-	ok := nbs.mt.addChunk(a, ch.Data())
-	if !ok {
+	addChunkRes := nbs.mt.addChunk(a, ch.Data())
+	if addChunkRes == chunkNotAdded {
 		ts, err := nbs.tables.append(ctx, nbs.mt, checker, nbs.stats)
 		if err != nil {
 			if errors.Is(err, ErrDanglingRef) {
@@ -634,12 +634,12 @@ func (nbs *NomsBlockStore) addChunk(ctx context.Context, ch chunks.Chunk, addrs
 		}
 		nbs.tables = ts
 		nbs.mt = newMemTable(nbs.mtSize)
-		ok = nbs.mt.addChunk(a, ch.Data())
+		addChunkRes = nbs.mt.addChunk(a, ch.Data())
 	}
-	if ok {
+	if addChunkRes == chunkAdded {
 		nbs.mt.addChildRefs(addrs)
 	}
-	return ok, nil
+	return true, nil
 }
 
 // refCheck checks that no dangling references are being committed.