From 2026d4410bdaff0b8076521b172ccc7671e164f3 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Mon, 8 Jul 2024 23:41:02 +0800
Subject: [PATCH] fix: should not trigger breaker on duplicate key with mongodb
 (#4238)

---
 core/stores/mon/collection.go      | 26 +++++++++++----
 core/stores/mon/collection_test.go | 51 ++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 7 deletions(-)

diff --git a/core/stores/mon/collection.go b/core/stores/mon/collection.go
index e6f7e702..3509129b 100644
--- a/core/stores/mon/collection.go
+++ b/core/stores/mon/collection.go
@@ -2,6 +2,7 @@ package mon
 
 import (
 	"context"
+	"errors"
 	"time"
 
 	"github.com/zeromicro/go-zero/core/breaker"
@@ -15,7 +16,8 @@ import (
 const (
 	defaultSlowThreshold = time.Millisecond * 500
 	// spanName is the span name of the mongo calls.
-	spanName = "mongo"
+	spanName         = "mongo"
+	duplicateKeyCode = 11000
 
 	// mongodb method names
 	aggregate              = "Aggregate"
@@ -527,10 +529,20 @@ func (p keepablePromise) keep(err error) error {
 }
 
 func acceptable(err error) bool {
-	return err == nil || errorx.In(err, mongo.ErrNoDocuments, mongo.ErrNilValue,
-		mongo.ErrNilDocument, mongo.ErrNilCursor, mongo.ErrEmptySlice,
-		// session errors
-		session.ErrSessionEnded, session.ErrNoTransactStarted, session.ErrTransactInProgress,
-		session.ErrAbortAfterCommit, session.ErrAbortTwice, session.ErrCommitAfterAbort,
-		session.ErrUnackWCUnsupported, session.ErrSnapshotTransaction)
+	return err == nil || isDupKeyError(err) ||
+		errorx.In(err, mongo.ErrNoDocuments, mongo.ErrNilValue,
+			mongo.ErrNilDocument, mongo.ErrNilCursor, mongo.ErrEmptySlice,
+			// session errors
+			session.ErrSessionEnded, session.ErrNoTransactStarted, session.ErrTransactInProgress,
+			session.ErrAbortAfterCommit, session.ErrAbortTwice, session.ErrCommitAfterAbort,
+			session.ErrUnackWCUnsupported, session.ErrSnapshotTransaction)
+}
+
+func isDupKeyError(err error) bool {
+	var e mongo.WriteException
+	if !errors.As(err, &e) {
+		return false
+	}
+
+	return e.HasErrorCode(duplicateKeyCode)
 }
diff --git a/core/stores/mon/collection_test.go b/core/stores/mon/collection_test.go
index 4a3f3b56..575aa4fe 100644
--- a/core/stores/mon/collection_test.go
+++ b/core/stores/mon/collection_test.go
@@ -15,6 +15,7 @@ import (
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
 	mopt "go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/x/mongo/driver/session"
 )
 
 var errDummy = errors.New("dummy")
@@ -572,6 +573,56 @@ func TestDecoratedCollection_LogDuration(t *testing.T) {
 	assert.Contains(t, buf.String(), "slowcall")
 }
 
+func TestAcceptable(t *testing.T) {
+	tests := []struct {
+		name string
+		err  error
+		want bool
+	}{
+		{"NilError", nil, true},
+		{"NoDocuments", mongo.ErrNoDocuments, true},
+		{"NilValue", mongo.ErrNilValue, true},
+		{"NilDocument", mongo.ErrNilDocument, true},
+		{"NilCursor", mongo.ErrNilCursor, true},
+		{"EmptySlice", mongo.ErrEmptySlice, true},
+		{"SessionEnded", session.ErrSessionEnded, true},
+		{"NoTransactStarted", session.ErrNoTransactStarted, true},
+		{"TransactInProgress", session.ErrTransactInProgress, true},
+		{"AbortAfterCommit", session.ErrAbortAfterCommit, true},
+		{"AbortTwice", session.ErrAbortTwice, true},
+		{"CommitAfterAbort", session.ErrCommitAfterAbort, true},
+		{"UnackWCUnsupported", session.ErrUnackWCUnsupported, true},
+		{"SnapshotTransaction", session.ErrSnapshotTransaction, true},
+		{"DuplicateKeyError", mongo.WriteException{WriteErrors: []mongo.WriteError{{Code: duplicateKeyCode}}}, true},
+		{"OtherError", errors.New("other error"), false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, acceptable(tt.err))
+		})
+	}
+}
+
+func TestIsDupKeyError(t *testing.T) {
+	tests := []struct {
+		name string
+		err  error
+		want bool
+	}{
+		{"NilError", nil, false},
+		{"NonDupKeyError", errors.New("some other error"), false},
+		{"DupKeyError", mongo.WriteException{WriteErrors: []mongo.WriteError{{Code: duplicateKeyCode}}}, true},
+		{"OtherMongoError", mongo.WriteException{WriteErrors: []mongo.WriteError{{Code: 12345}}}, false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equal(t, tt.want, isDupKeyError(tt.err))
+		})
+	}
+}
+
 type mockPromise struct {
 	accepted bool
 	reason   string