-
Notifications
You must be signed in to change notification settings - Fork 0
/
client.go
123 lines (105 loc) · 2.34 KB
/
client.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
package mongo
import (
"context"
"errors"
"time"
"go.mongodb.org/mongo-driver/mongo"
"github.com/ehwjh2010/viper/client/enums"
"github.com/ehwjh2010/viper/client/settings"
"github.com/ehwjh2010/viper/component/routine"
"github.com/ehwjh2010/viper/log"
)
type Client struct {
db *mongo.Database
cli *mongo.Client
rawConfig settings.Mongo // 数据库配置配置
pCount int // 心跳连续失败次数
rCount int // 重连连续失败次数
}
func NewClient(cli *mongo.Client, db *mongo.Database, rawConfig settings.Mongo) *Client {
return &Client{db: db, rawConfig: rawConfig, cli: cli}
}
// Heartbeat ping连接
func (c *Client) Heartbeat() error {
return c.cli.Ping(context.TODO(), nil)
}
// WatchHeartbeat 监测心跳和重连
func (c *Client) WatchHeartbeat() {
// TODO 监测逻辑接口化
fn := func() {
waitFlag := true
for {
if waitFlag {
<-time.After(3 * time.Second)
}
// 重连失败次数大于0, 直接重连
if c.rCount > 0 {
if c.rCount >= 3 {
<-time.After(enums.OneSecD)
}
if ok, _ := c.replaceDB(); ok {
c.rCount = 0
c.pCount = 0
waitFlag = true
} else {
c.rCount++
c.pCount++
waitFlag = false
}
continue
}
if c.Heartbeat() != nil {
c.pCount++
// 心跳连续3次失败, 触发重连
if c.pCount >= 3 {
if ok, _ := c.replaceDB(); ok {
c.rCount = 0
c.pCount = 0
waitFlag = true
} else {
c.rCount++
waitFlag = false
}
}
} else {
c.rCount = 0
c.pCount = 0
waitFlag = true
}
}
}
// 优先使用协程池监听, 如果没有使用原生协程监听
err := routine.AddTask(fn)
if err != nil {
if errors.Is(err, routine.NoEnableRoutinePool) {
go fn()
} else {
log.Warn("watch heartbeat failed")
}
}
}
// Close 关闭连接
func (c *Client) Close() error {
return c.cli.Disconnect(context.TODO())
}
// replaceDB 替换内部连接
func (c *Client) replaceDB() (bool, error) {
cli, db, err := setup(c.rawConfig)
if err != nil {
log.Err("reconnect mongo failed", err)
return false, err
}
// 关闭之前的连接
c.Close()
c.db = db
c.cli = cli
return true, nil
}
func (c *Client) getDB() *mongo.Database {
db := c.db
return db
}
// GetDB 获取原生db
func (c *Client) GetDB() *mongo.Database {
return c.getDB()
}