231 lines
6.9 KiB
Go
231 lines
6.9 KiB
Go
/*
|
||
Package common 提供 LDAP 连接池管理功能
|
||
|
||
LDAP 连接池是系统与 OpenLDAP 服务器交互的核心组件，负责：
|
||
- 管理 LDAP 连接的创建、复用和销毁
|
||
- 提供线程安全的连接获取和归还机制
|
||
- 控制最大连接数，避免资源耗尽
|
||
- 支持连接的健康检查和自动重连
|
||
|
||
连接池的设计目标：
|
||
1. 提高 LDAP 操作的性能，避免频繁建立连接
|
||
2. 控制并发连接数，保护 LDAP 服务器资源
|
||
3. 提供稳定可靠的连接管理机制
|
||
4. 支持优雅的连接回收和错误处理
|
||
*/
|
||
package common
|
||
|
||
import (
|
||
"fmt"
|
||
"log"
|
||
"math/rand"
|
||
"net"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/eryajf/go-ldap-admin/config"
|
||
|
||
ldap "github.com/go-ldap/ldap/v3"
|
||
)
|
||
|
||
// LDAP 连接池相关全局变量
|
||
var (
|
||
ldapPool *LdapConnPool // LDAP 连接池实例
|
||
ldapInit = false // LDAP 初始化标志
|
||
ldapInitOne sync.Once // 确保 LDAP 只初始化一次
|
||
)
|
||
|
||
// InitLDAP 初始化 LDAP 连接池
|
||
// 该函数在系统启动时调用，负责：
|
||
// 1. 建立与 LDAP 服务器的初始连接
|
||
// 2. 验证管理员账户凭据
|
||
// 3. 初始化连接池结构
|
||
// 4. 设置连接池参数
|
||
func InitLDAP() {
|
||
// 防止重复初始化
|
||
if ldapInit {
|
||
return
|
||
}
|
||
|
||
// 使用 sync.Once 确保只初始化一次
|
||
ldapInitOne.Do(func() {
|
||
ldapInit = true
|
||
})
|
||
|
||
// ==================== 建立初始 LDAP 连接 ====================
|
||
|
||
// 创建 LDAP 连接，设置 5 秒连接超时
|
||
ldapConn, err := ldap.DialURL(config.Conf.Ldap.Url, ldap.DialWithDialer(&net.Dialer{Timeout: 5 * time.Second}))
|
||
if err != nil {
|
||
Log.Panicf("初始化 LDAP 连接异常: %v", err)
|
||
panic(fmt.Errorf("初始化 LDAP 连接异常: %v", err))
|
||
}
|
||
|
||
// 使用管理员账户绑定 LDAP 连接
|
||
err = ldapConn.Bind(config.Conf.Ldap.AdminDN, config.Conf.Ldap.AdminPass)
|
||
if err != nil {
|
||
Log.Panicf("绑定 LDAP 管理员账号异常: %v", err)
|
||
panic(fmt.Errorf("绑定 LDAP 管理员账号异常: %v", err))
|
||
}
|
||
|
||
// ==================== 初始化连接池 ====================
|
||
|
||
// 创建连接池实例
|
||
ldapPool = &LdapConnPool{
|
||
conns: make([]*ldap.Conn, 0), // 可用连接切片
|
||
reqConns: make(map[uint64]chan *ldap.Conn), // 等待连接的请求映射
|
||
openConn: 0, // 当前打开的连接数
|
||
maxOpen: config.Conf.Ldap.MaxConn, // 最大连接数
|
||
}
|
||
|
||
// 将初始连接放入连接池
|
||
PutLADPConn(ldapConn)
|
||
|
||
// ==================== 输出初始化信息 ====================
|
||
|
||
// 构建显示用的 DSN（隐藏密码）
|
||
showDsn := fmt.Sprintf(
|
||
"%s:******@tcp(%s)",
|
||
config.Conf.Ldap.AdminDN,
|
||
config.Conf.Ldap.Url,
|
||
)
|
||
|
||
Log.Info("🔗 LDAP 连接池初始化完成! DSN: ", showDsn)
|
||
}
|
||
|
||
// GetLDAPConn 从连接池获取 LDAP 连接
|
||
// 返回一个可用的 LDAP 连接，使用完毕后需要调用 PutLADPConn 归还
|
||
func GetLDAPConn() (*ldap.Conn, error) {
|
||
return ldapPool.GetConnection()
|
||
}
|
||
|
||
// PutLADPConn 将 LDAP 连接归还到连接池
|
||
// 连接使用完毕后必须调用此方法归还，以便其他请求复用
|
||
func PutLADPConn(conn *ldap.Conn) {
|
||
ldapPool.PutConnection(conn)
|
||
}
|
||
|
||
// LdapConnPool LDAP 连接池结构体
|
||
// 实现了线程安全的连接池管理机制
|
||
type LdapConnPool struct {
|
||
mu sync.Mutex // 互斥锁，保证线程安全
|
||
conns []*ldap.Conn // 可用连接切片
|
||
reqConns map[uint64]chan *ldap.Conn // 等待连接的请求映射表
|
||
openConn int // 当前打开的连接数
|
||
maxOpen int // 最大允许的连接数
|
||
}
|
||
|
||
// GetConnection 从连接池获取一个可用的 LDAP 连接
|
||
// 该方法实现了连接池的核心逻辑：
|
||
// 1. 优先从池中获取现有连接
|
||
// 2. 检查连接健康状态
|
||
// 3. 控制最大连接数限制
|
||
// 4. 支持连接等待队列
|
||
func (lcp *LdapConnPool) GetConnection() (*ldap.Conn, error) {
|
||
lcp.mu.Lock()
|
||
|
||
// ==================== 从连接池获取现有连接 ====================
|
||
|
||
// 检查连接池中是否有可用连接
|
||
connNum := len(lcp.conns)
|
||
if connNum > 0 {
|
||
lcp.openConn++ // 增加打开连接计数
|
||
conn := lcp.conns[0] // 获取第一个连接
|
||
copy(lcp.conns, lcp.conns[1:]) // 移除已获取的连接
|
||
lcp.conns = lcp.conns[:connNum-1] // 调整切片长度
|
||
|
||
lcp.mu.Unlock()
|
||
|
||
// 检查连接是否已关闭，如果已关闭则创建新连接
|
||
if conn.IsClosing() {
|
||
return initLDAPConn()
|
||
}
|
||
return conn, nil
|
||
}
|
||
|
||
// ==================== 连接数限制处理 ====================
|
||
|
||
// 当连接池为空且已达到最大连接数限制时
|
||
if lcp.maxOpen != 0 && lcp.openConn >= lcp.maxOpen {
|
||
// 创建等待队列，当有连接归还时会通知等待的请求
|
||
req := make(chan *ldap.Conn, 1)
|
||
reqKey := lcp.nextRequestKeyLocked() // 生成唯一请求键
|
||
lcp.reqConns[reqKey] = req // 将请求加入等待队列
|
||
lcp.mu.Unlock()
|
||
|
||
// 阻塞等待连接归还
|
||
return <-req, nil
|
||
} else {
|
||
// 未达到连接数限制，创建新连接
|
||
lcp.openConn++
|
||
lcp.mu.Unlock()
|
||
return initLDAPConn()
|
||
}
|
||
}
|
||
|
||
// PutConnection 将 LDAP 连接归还到连接池
|
||
// 该方法负责连接的回收和分配：
|
||
// 1. 优先满足等待队列中的请求
|
||
// 2. 将健康的连接放回连接池
|
||
// 3. 丢弃已关闭的连接
|
||
func (lcp *LdapConnPool) PutConnection(conn *ldap.Conn) {
|
||
log.Println("🔄 归还一个 LDAP 连接到连接池")
|
||
lcp.mu.Lock()
|
||
defer lcp.mu.Unlock()
|
||
|
||
// ==================== 优先处理等待队列 ====================
|
||
|
||
// 检查是否有等待连接的请求
|
||
if num := len(lcp.reqConns); num > 0 {
|
||
var req chan *ldap.Conn
|
||
var reqKey uint64
|
||
|
||
// 获取第一个等待请求
|
||
for reqKey, req = range lcp.reqConns {
|
||
break
|
||
}
|
||
|
||
// 从等待队列中移除该请求
|
||
delete(lcp.reqConns, reqKey)
|
||
|
||
// 将连接发送给等待的请求
|
||
req <- conn
|
||
return
|
||
}
|
||
|
||
// ==================== 连接回收处理 ====================
|
||
|
||
// 没有等待请求时，将连接放回连接池
|
||
lcp.openConn-- // 减少打开连接计数
|
||
|
||
// 只有健康的连接才放回连接池
|
||
if !conn.IsClosing() {
|
||
lcp.conns = append(lcp.conns, conn)
|
||
}
|
||
}
|
||
|
||
// nextRequestKeyLocked 生成下一个唯一的请求键
|
||
// 用于标识等待队列中的请求，确保每个等待请求都有唯一标识
|
||
// 注意：调用此方法时必须已经持有互斥锁
|
||
func (lcp *LdapConnPool) nextRequestKeyLocked() uint64 {
|
||
for {
|
||
reqKey := rand.Uint64() // 生成随机数作为请求键
|
||
if _, ok := lcp.reqConns[reqKey]; !ok { // 确保键的唯一性
|
||
return reqKey
|
||
}
|
||
}
|
||
}
|
||
|
||
// 获取 ladp 连接
|
||
func initLDAPConn() (*ldap.Conn, error) {
|
||
ldap, err := ldap.DialURL(config.Conf.Ldap.Url, ldap.DialWithDialer(&net.Dialer{Timeout: 5 * time.Second}))
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
err = ldap.Bind(config.Conf.Ldap.AdminDN, config.Conf.Ldap.AdminPass)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return ldap, err
|
||
}
|