fix: Merge retry logic

events/dag_sync.go

@@ -14,7 +14,6 @@ import (
 	"sync"
 
 	"github.com/ipfs/go-cid"
-
 	"github.com/sourcenetwork/immutable"
 )
 
@@ -23,6 +22,8 @@ type DAGMergeChannel = immutable.Option[Channel[DAGMerge]]
 
 // DAGMerge is a notification that a merge can be performed up to the provided CID.
 type DAGMerge struct {
+	// DocID is the unique identifier for the document being merged.
+	DocID string
 	// Cid is the id of the composite commit that formed this update in the DAG.
 	Cid cid.Cid
 	// SchemaRoot is the root identifier of the schema that defined the shape of the document that was updated.

go.mod

@@ -185,10 +185,12 @@ require (
 	github.com/ipfs/go-ipfs-delay v0.0.1 // indirect
 	github.com/ipfs/go-ipfs-pq v0.0.3 // indirect
 	github.com/ipfs/go-ipfs-util v0.0.3 // indirect
+	github.com/ipfs/go-ipld-legacy v0.2.1 // indirect
 	github.com/ipfs/go-log v1.0.5 // indirect
 	github.com/ipfs/go-metrics-interface v0.0.1 // indirect
 	github.com/ipfs/go-peertaskqueue v0.8.1 // indirect
 	github.com/ipfs/kubo v0.25.0 // indirect
+	github.com/ipld/go-codec-dagpb v1.6.0 // indirect
 	github.com/jackpal/go-nat-pmp v1.0.2 // indirect
 	github.com/jbenet/go-temp-err-catcher v0.1.0 // indirect
 	github.com/jmhodges/levigo v1.0.0 // indirect

go.sum

@@ -1166,6 +1166,8 @@ github.com/viant/assertly v0.4.8/go.mod h1:aGifi++jvCrUaklKEKT0BU95igDNaqkvz+49u
 github.com/viant/toolbox v0.24.0/go.mod h1:OxMCG57V0PXuIP2HNQrtJf2CjqdmbrOx5EkMILuUhzM=
 github.com/vito/go-sse v1.0.0 h1:e6/iTrrvy8BRrOwJwmQmlndlil+TLdxXvHi55ZDzH6M=
 github.com/vito/go-sse v1.0.0/go.mod h1:2wkcaQ+jtlZ94Uve8gYZjFpL68luAjssTINA2hpgcZs=
+github.com/warpfork/go-testmark v0.12.1 h1:rMgCpJfwy1sJ50x0M0NgyphxYYPMOODIJHhsXyEHU0s=
+github.com/warpfork/go-testmark v0.12.1/go.mod h1:kHwy7wfvGSPh1rQJYKayD4AbtNaeyZdcGi9tNJTaa5Y=
 github.com/warpfork/go-wish v0.0.0-20220906213052-39a1cc7a02d0 h1:GDDkbFiaK8jsSDJfjId/PEGEShv6ugrt4kYsC5UIDaQ=
 github.com/warpfork/go-wish v0.0.0-20220906213052-39a1cc7a02d0/go.mod h1:x6AKhvSSexNrVSrViXSHUEbICjmGXhtgABaHIySUSGw=
 github.com/wasmerio/wasmer-go v1.0.4 h1:MnqHoOGfiQ8MMq2RF6wyCeebKOe84G88h5yv+vmxJgs=

internal/db/merge.go

@@ -13,6 +13,7 @@
 import (
 	"container/list"
 	"context"
+	"sync"
 
 	"github.com/ipfs/go-cid"
 	"github.com/ipld/go-ipld-prime/linking"
@@ -34,6 +35,7 @@
 )
 
 func (db *db) handleMerges(ctx context.Context, merges events.Subscription[events.DAGMerge]) {
+	queue := newMergeQueue()
 	for {
 		select {
 		case <-ctx.Done():
@@ -43,28 +45,37 @@
 				return
 			}
 			go func() {
-				err := db.executeMerge(ctx, merge)
+				// ensure only one merge per docID
+				queue.add(merge.DocID)
+				defer queue.done(merge.DocID)
+
+				var err error
+				// retry merge up to max txn retries
+				for i := 0; i < db.MaxTxnRetries(); i++ {
+					err = db.executeMerge(ctx, merge)
+					if errors.Is(err, badger.ErrTxnConflict) {
+						continue // retry merge
+					}
+					break // merge success or error
+				}
+
 				if err != nil {
 					log.ErrorContextE(
 						ctx,
 						"Failed to execute merge",
 						err,
-						corelog.String("CID", merge.Cid.String()),
-						corelog.String("Error", err.Error()),
-					)
+						corelog.Any("Error", err),
+						corelog.Any("Event", merge))
+				}
+				if merge.Wg != nil {
+					merge.Wg.Done()
 				}
 			}()
 		}
 	}
 }
 
 func (db *db) executeMerge(ctx context.Context, dagMerge events.DAGMerge) error {
-	defer func() {
-		// Notify the caller that the merge is complete.
-		if dagMerge.Wg != nil {
-			dagMerge.Wg.Done()
-		}
-	}()
 	ctx, txn, err := ensureContextTxn(ctx, db, false)
 	if err != nil {
 		return err
@@ -79,7 +90,7 @@
 	ls := cidlink.DefaultLinkSystem()
 	ls.SetReadStorage(txn.DAGstore().AsIPLDStorage())
 
-	docID, err := getDocIDFromBlock(ctx, ls, dagMerge.Cid)
+	docID, err := client.NewDocIDFromString(dagMerge.DocID)
 	if err != nil {
 		return err
 	}
@@ -100,35 +111,57 @@
 		return err
 	}
 
-	for retry := 0; retry < db.MaxTxnRetries(); retry++ {
-		err := mp.mergeComposites(ctx)
-		if err != nil {
-			return err
-		}
-		err = syncIndexedDoc(ctx, docID, col)
-		if err != nil {
-			return err
-		}
-		err = txn.Commit(ctx)
-		if err != nil {
-			if errors.Is(err, badger.ErrTxnConflict) {
-				txn, err = db.NewTxn(ctx, false)
-				if err != nil {
-					return err
-				}
-				ctx = SetContextTxn(ctx, txn)
-				mp.txn = txn
-				mp.lsys.SetReadStorage(txn.DAGstore().AsIPLDStorage())
-				// Reset the CRDTs to avoid reusing the old transaction.
-				mp.mCRDTs = make(map[string]merklecrdt.MerkleCRDT)
-				continue
-			}
-			return err
-		}
-		break
+	err = mp.mergeComposites(ctx)
+	if err != nil {
+		return err
 	}
 
-	return nil
+	err = syncIndexedDoc(ctx, docID, col)
+	if err != nil {
+		return err
+	}
+
+	return txn.Commit(ctx)
+}
+
+// mergeQueue is synchronization source to ensure that concurrent
+// document merges do not cause transaction conflicts.
+type mergeQueue struct {
+	docs map[string]chan struct{}
+	mu   sync.Mutex
+}
+
+func newMergeQueue() *mergeQueue {
+	return &mergeQueue{
+		docs: make(map[string]chan struct{}),
+	}
+}
+
+// add adds a docID to the queue. If the docID is already in the queue, it will
+// wait for the docID to be removed from the queue. For every add call, done must
+// be called to remove the docID from the queue. Otherwise, subsequent add calls will
+// block forever.
+func (dq *mergeQueue) add(docID string) {
+	dq.mu.Lock()
+	done, ok := dq.docs[docID]
+	if !ok {
+		dq.docs[docID] = make(chan struct{})
+	}
+	dq.mu.Unlock()
+	if ok {
+		<-done
+		dq.add(docID)
+	}
+}
+
+func (dq *mergeQueue) done(docID string) {
+	dq.mu.Lock()
+	defer dq.mu.Unlock()
+	done, ok := dq.docs[docID]
+	if ok {
+		delete(dq.docs, docID)
+		close(done)
+	}
 }
 
 type mergeProcessor struct {
@@ -333,18 +366,6 @@
 	return mcrdt, nil
 }
 
-func getDocIDFromBlock(ctx context.Context, ls linking.LinkSystem, cid cid.Cid) (client.DocID, error) {
-	nd, err := ls.Load(linking.LinkContext{Ctx: ctx}, cidlink.Link{Cid: cid}, coreblock.SchemaPrototype)
-	if err != nil {
-		return client.DocID{}, err
-	}
-	block, err := coreblock.GetFromNode(nd)
-	if err != nil {
-		return client.DocID{}, err
-	}
-	return client.NewDocIDFromString(string(block.Delta.GetDocID()))
-}
-
 func getCollectionFromRootSchema(ctx context.Context, db *db, rootSchema string) (*collection, error) {
 	cols, err := db.getCollections(
 		ctx,

internal/db/merge_test.go

@@ -13,6 +13,7 @@ package db
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/fxamacker/cbor/v2"
 	"github.com/ipld/go-ipld-prime"
@@ -59,6 +60,7 @@ func TestMerge_SingleBranch_NoError(t *testing.T) {
 	require.NoError(t, err)
 
 	err = db.executeMerge(ctx, events.DAGMerge{
+		DocID:      docID.String(),
 		Cid:        compInfo2.link.Cid,
 		SchemaRoot: col.SchemaRoot(),
 	})
@@ -103,6 +105,7 @@ func TestMerge_DualBranch_NoError(t *testing.T) {
 	require.NoError(t, err)
 
 	err = db.executeMerge(ctx, events.DAGMerge{
+		DocID:      docID.String(),
 		Cid:        compInfo2.link.Cid,
 		SchemaRoot: col.SchemaRoot(),
 	})
@@ -112,6 +115,7 @@ func TestMerge_DualBranch_NoError(t *testing.T) {
 	require.NoError(t, err)
 
 	err = db.executeMerge(ctx, events.DAGMerge{
+		DocID:      docID.String(),
 		Cid:        compInfo3.link.Cid,
 		SchemaRoot: col.SchemaRoot(),
 	})
@@ -159,6 +163,7 @@ func TestMerge_DualBranchWithOneIncomplete_CouldNotFindCID(t *testing.T) {
 	require.NoError(t, err)
 
 	err = db.executeMerge(ctx, events.DAGMerge{
+		DocID:      docID.String(),
 		Cid:        compInfo2.link.Cid,
 		SchemaRoot: col.SchemaRoot(),
 	})
@@ -177,6 +182,7 @@ func TestMerge_DualBranchWithOneIncomplete_CouldNotFindCID(t *testing.T) {
 	require.NoError(t, err)
 
 	err = db.executeMerge(ctx, events.DAGMerge{
+		DocID:      docID.String(),
 		Cid:        compInfo3.link.Cid,
 		SchemaRoot: col.SchemaRoot(),
 	})
@@ -292,3 +298,25 @@ func encodeValue(val any) []byte {
 	}
 	return b
 }
+
+func TestMergeQueue(t *testing.T) {
+	q := newMergeQueue()
+
+	testDocID := "test"
+
+	q.add(testDocID)
+	go q.add(testDocID)
+	// give time for the goroutine to block
+	time.Sleep(10 * time.Millisecond)
+	require.Len(t, q.docs, 1)
+	q.done(testDocID)
+	// give time for the goroutine to add the docID
+	time.Sleep(10 * time.Millisecond)
+	q.mu.Lock()
+	require.Len(t, q.docs, 1)
+	q.mu.Unlock()
+	q.done(testDocID)
+	q.mu.Lock()
+	require.Len(t, q.docs, 0)
+	q.mu.Unlock()
+}