From a25cba538036374676b746d9713eb665ed2178dc Mon Sep 17 00:00:00 2001
From: kingxt <kingxt4job@gmail.com>
Date: Wed, 3 Mar 2021 10:44:29 +0800
Subject: [PATCH] fix collection breaker (#537)

* fix collection breaker

* optimized

* optimized

* optimized
---
 core/stores/mongo/collection.go      | 4 ++--
 core/stores/mongo/collection_test.go | 2 +-
 core/stores/mongo/model.go           | 5 ++++-
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/core/stores/mongo/collection.go b/core/stores/mongo/collection.go
index 6fb199b2..1592dfa2 100644
--- a/core/stores/mongo/collection.go
+++ b/core/stores/mongo/collection.go
@@ -43,11 +43,11 @@ type (
 	}
 )
 
-func newCollection(collection *mgo.Collection) Collection {
+func newCollection(collection *mgo.Collection, brk breaker.Breaker) Collection {
 	return &decoratedCollection{
 		name:       collection.FullName,
 		collection: collection,
-		brk:        breaker.NewBreaker(),
+		brk:        brk,
 	}
 }
 
diff --git a/core/stores/mongo/collection_test.go b/core/stores/mongo/collection_test.go
index 2b6c7825..4723bf58 100644
--- a/core/stores/mongo/collection_test.go
+++ b/core/stores/mongo/collection_test.go
@@ -71,7 +71,7 @@ func TestNewCollection(t *testing.T) {
 		Database: nil,
 		Name:     "foo",
 		FullName: "bar",
-	})
+	}, breaker.GetBreaker("localhost"))
 	assert.Equal(t, "bar", col.(*decoratedCollection).name)
 }
 
diff --git a/core/stores/mongo/model.go b/core/stores/mongo/model.go
index f8c298ac..f42186d0 100644
--- a/core/stores/mongo/model.go
+++ b/core/stores/mongo/model.go
@@ -5,6 +5,7 @@ import (
 	"time"
 
 	"github.com/globalsign/mgo"
+	"github.com/tal-tech/go-zero/core/breaker"
 )
 
 type (
@@ -20,6 +21,7 @@ type (
 		session    *concurrentSession
 		db         *mgo.Database
 		collection string
+		brk        breaker.Breaker
 		opts       []Option
 	}
 )
@@ -46,6 +48,7 @@ func NewModel(url, collection string, opts ...Option) (*Model, error) {
 		// If name is empty, the database name provided in the dialed URL is used instead
 		db:         session.DB(""),
 		collection: collection,
+		brk:        breaker.GetBreaker(url),
 		opts:       opts,
 	}, nil
 }
@@ -66,7 +69,7 @@ func (mm *Model) FindId(id interface{}) (Query, error) {
 
 // GetCollection returns a Collection with given session.
 func (mm *Model) GetCollection(session *mgo.Session) Collection {
-	return newCollection(mm.db.C(mm.collection).With(session))
+	return newCollection(mm.db.C(mm.collection).With(session), mm.brk)
 }
 
 // Insert inserts docs into mm.