125 lines
2.4 KiB
Go
125 lines
2.4 KiB
Go
package core
|
|
|
|
import (
|
|
"database/sql"
|
|
"sync"
|
|
"sync/atomic"
|
|
)
|
|
|
|
// ReadWriteDB 读写分离数据库连接
|
|
type ReadWriteDB struct {
|
|
master *sql.DB // 主库（写）
|
|
slaves []*sql.DB // 从库列表（读）
|
|
policy ReadPolicy // 读负载均衡策略
|
|
counter uint64 // 轮询计数器
|
|
mu sync.RWMutex // 读写锁
|
|
}
|
|
|
|
// NewReadWriteDB 创建读写分离数据库连接
|
|
func NewReadWriteDB(master *sql.DB, slaves []*sql.DB, policy ReadPolicy) *ReadWriteDB {
|
|
return &ReadWriteDB{
|
|
master: master,
|
|
slaves: slaves,
|
|
policy: policy,
|
|
}
|
|
}
|
|
|
|
// GetMaster 获取主库连接（用于写操作）
|
|
func (rw *ReadWriteDB) GetMaster() *sql.DB {
|
|
return rw.master
|
|
}
|
|
|
|
// GetSlave 获取从库连接（用于读操作）
|
|
func (rw *ReadWriteDB) GetSlave() *sql.DB {
|
|
rw.mu.RLock()
|
|
defer rw.mu.RUnlock()
|
|
|
|
if len(rw.slaves) == 0 {
|
|
// 没有从库，使用主库
|
|
return rw.master
|
|
}
|
|
|
|
switch rw.policy {
|
|
case Random:
|
|
// 随机选择一个从库
|
|
idx := int(atomic.LoadUint64(&rw.counter)) % len(rw.slaves)
|
|
return rw.slaves[idx]
|
|
|
|
case RoundRobin:
|
|
// 轮询选择从库
|
|
idx := int(atomic.AddUint64(&rw.counter, 1)) % len(rw.slaves)
|
|
return rw.slaves[idx]
|
|
|
|
case LeastConn:
|
|
// 选择连接数最少的从库（简化实现）
|
|
return rw.selectLeastConn()
|
|
|
|
default:
|
|
return rw.slaves[0]
|
|
}
|
|
}
|
|
|
|
// selectLeastConn 选择连接数最少的从库
|
|
func (rw *ReadWriteDB) selectLeastConn() *sql.DB {
|
|
if len(rw.slaves) == 0 {
|
|
return rw.master
|
|
}
|
|
|
|
minConn := -1
|
|
selected := rw.slaves[0]
|
|
|
|
for _, slave := range rw.slaves {
|
|
stats := slave.Stats()
|
|
openConnections := stats.OpenConnections
|
|
|
|
if minConn == -1 || openConnections < minConn {
|
|
minConn = openConnections
|
|
selected = slave
|
|
}
|
|
}
|
|
|
|
return selected
|
|
}
|
|
|
|
// AddSlave 添加从库
|
|
func (rw *ReadWriteDB) AddSlave(slave *sql.DB) {
|
|
rw.mu.Lock()
|
|
defer rw.mu.Unlock()
|
|
rw.slaves = append(rw.slaves, slave)
|
|
}
|
|
|
|
// RemoveSlave 移除从库
|
|
func (rw *ReadWriteDB) RemoveSlave(slave *sql.DB) {
|
|
rw.mu.Lock()
|
|
defer rw.mu.Unlock()
|
|
|
|
for i, s := range rw.slaves {
|
|
if s == slave {
|
|
rw.slaves = append(rw.slaves[:i], rw.slaves[i+1:]...)
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
// Close 关闭所有连接
|
|
func (rw *ReadWriteDB) Close() error {
|
|
rw.mu.Lock()
|
|
defer rw.mu.Unlock()
|
|
|
|
// 关闭主库
|
|
if rw.master != nil {
|
|
if err := rw.master.Close(); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
// 关闭所有从库
|
|
for _, slave := range rw.slaves {
|
|
if err := slave.Close(); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|