From f2cc89edc030e964ff29f1dcf58b76cb05f3e6d7 Mon Sep 17 00:00:00 2001
From: Andy Arthur <andy@dolthub.com>
Date: Mon, 23 Jan 2023 17:44:56 -0800
Subject: [PATCH] go/store/nbs: fix reentrant deadlock in ref checker

---
 go/store/nbs/generational_chunk_store.go | 19 ++++++-------------
 go/store/nbs/store.go                    | 11 +++++------
 2 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/go/store/nbs/generational_chunk_store.go b/go/store/nbs/generational_chunk_store.go
index fee6c53b11..5fc91b6311 100644
--- a/go/store/nbs/generational_chunk_store.go
+++ b/go/store/nbs/generational_chunk_store.go
@@ -136,6 +136,8 @@ func (gcs *GenerationalNBS) Has(ctx context.Context, h hash.Hash) (bool, error)
 
 // HasMany returns a new HashSet containing any members of |hashes| that are absent from the store.
 func (gcs *GenerationalNBS) HasMany(ctx context.Context, hashes hash.HashSet) (absent hash.HashSet, err error) {
+	gcs.newGen.mu.RLock()
+	defer gcs.newGen.mu.RUnlock()
 	return gcs.hasMany(toHasRecords(hashes))
 }
 
@@ -146,6 +148,9 @@ func (gcs *GenerationalNBS) hasMany(recs []hasRecord) (absent hash.HashSet, err
 	} else if len(absent) == 0 {
 		return absent, nil
 	}
+
+	gcs.oldGen.mu.RLock()
+	defer gcs.oldGen.mu.RUnlock()
 	return gcs.oldGen.hasMany(recs)
 }
 
@@ -166,19 +171,7 @@ func (gcs *GenerationalNBS) errorIfDangling(ctx context.Context, addrs hash.Hash
 // to Flush(). Put may be called concurrently with other calls to Put(),
 // Get(), GetMany(), Has() and HasMany().
 func (gcs *GenerationalNBS) Put(ctx context.Context, c chunks.Chunk, getAddrs chunks.GetAddrsCb) error {
-	addrs, err := getAddrs(ctx, c)
-	if err != nil {
-		return err
-	}
-
-	err = gcs.errorIfDangling(ctx, addrs)
-	if err != nil {
-		return err
-	}
-
-	return gcs.newGen.Put(ctx, c, func(ctx context.Context, c chunks.Chunk) (hash.HashSet, error) {
-		return nil, nil
-	})
+	return gcs.newGen.Put(ctx, c, getAddrs)
 }
 
 // Returns the NomsVersion with which this ChunkSource is compatible.
diff --git a/go/store/nbs/store.go b/go/store/nbs/store.go
index 3988e77580..2024c9dbfa 100644
--- a/go/store/nbs/store.go
+++ b/go/store/nbs/store.go
@@ -623,12 +623,11 @@ func (nbs *NomsBlockStore) addChunk(ctx context.Context, ch chunks.Chunk, addrs
 type refCheck func(reqs []hasRecord) (hash.HashSet, error)
 
 func (nbs *NomsBlockStore) errorIfDangling(checker refCheck) error {
-	nbs.mu.RLock()
-	defer nbs.mu.RUnlock()
 	if (nbs.mt == nil || nbs.mt.pendingRefs == nil) && (len(nbs.tables.novel) == 0) {
 		return nil // no refs to check
 	}
 
+	sort.Sort(hasRecordByPrefix(nbs.mt.pendingRefs))
 	absent, err := checker(nbs.mt.pendingRefs)
 	if err != nil {
 		return err
@@ -860,15 +859,13 @@ func (nbs *NomsBlockStore) HasMany(ctx context.Context, hashes hash.HashSet) (ha
 	defer nbs.stats.HasLatency.SampleTimeSince(t1)
 	nbs.stats.AddressesPerHas.SampleLen(hashes.Size())
 
+	nbs.mu.RLock()
+	defer nbs.mu.RUnlock()
 	return nbs.hasMany(toHasRecords(hashes))
 }
 
 func (nbs *NomsBlockStore) hasMany(reqs []hasRecord) (hash.HashSet, error) {
-	sort.Sort(hasRecordByPrefix(reqs))
-
 	tables, remaining, err := func() (tables chunkReader, remaining bool, err error) {
-		nbs.mu.RLock()
-		defer nbs.mu.RUnlock()
 		tables = nbs.tables
 
 		remaining = true
@@ -983,9 +980,11 @@ func (nbs *NomsBlockStore) commit(ctx context.Context, current, last hash.Hash,
 	}
 
 	// check for dangling references in |nbs.mt|
+	nbs.mu.Lock()
 	if err = nbs.errorIfDangling(checker); err != nil {
 		return false, err
 	}
+	nbs.mu.Unlock()
 
 	err = func() error {
 		// This is unfortunate. We want to serialize commits to the same store