TRKJ/mayfly-go
2
0
Fork 0
mirror of https://gitee.com/dromara/mayfly-go synced 2025-11-02 23:40:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: 移除隧道连接时检测是否正在使用

Browse Source
This commit is contained in:
meilin.huang
2025-05-26 22:33:51 +08:00
parent d6280ea280
commit e0c01d4561
19 changed files with 314 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
frontend/src/views/ops/db/DbList.vue
View File

@@ -153,7 +153,7 @@
:close-on-click-modal="false"
v-model="sqlExecLogDialog.visible"
:destroy-on-close="true"
body-class="h-250"
body-class="h-[65vh]"
>
<db-sql-exec-log :db-id="sqlExecLogDialog.dbId" :dbs="sqlExecLogDialog.dbs" />
</el-dialog>

69
frontend/src/views/ops/db/DbTransferFile.vue
View File

@@ -1,28 +1,25 @@
<template>
<div class="db-transfer-file">
<el-dialog :title="title" v-model="dialogVisible" :close-on-click-modal="false" :destroy-on-close="true" width="1000px">
<el-dialog
@open="search()"
:title="title"
v-model="dialogVisible"
:close-on-click-modal="false"
:destroy-on-close="true"
body-class="h-[65vh]"
width="1000px"
>
<page-table
ref="pageTableRef"
:data="state.tableData"
v-model:query-form="state.query"
:page-api="dbApi.dbTransferFileList"
:lazy="true"
:show-selection="true"
v-model:selection-data="state.selectionData"
:columns="columns"
@page-num-change="
(args: any) => {
state.query.pageNum = args.pageNum;
search();
}
"
@page-size-change="
(args: any) => {
state.query.pageSize = args.pageNum;
search();
}
"
>
<template #tableHeader>
<el-button v-auth="perms.del" :disabled="state.selectionData.length < 1" @click="del()" type="danger" icon="delete">
<el-button v-auth="perms.del" :disabled="state.selectionData.length < 1" @click="onDel()" type="danger" icon="delete">
{{ $t('common.delete') }}
</el-button>
</template>
@@ -41,17 +38,20 @@
<template #action="{ data }">
<el-button
v-if="actionBtns[perms.run] && data.status === DbTransferFileStatusEnum.Success.value"
@click="openRun(data)"
@click="onOpenRun(data)"
type="primary"
link
>{{ $t('db.run') }}</el-button
>
<el-button v-if="data.logId" @click="openLog(data)" type="success" link>{{ $t('db.log') }}</el-button>
{{ $t('db.run') }}
</el-button>
<el-button v-if="data.logId" @click="onOpenLog(data)" type="success" link>{{ $t('db.log') }}</el-button>
</template>
</page-table>
<TerminalLog v-model:log-id="state.logsDialog.logId" v-model:visible="state.logsDialog.visible" :title="state.logsDialog.title" />
</el-dialog>
<TerminalLog v-model:log-id="state.logsDialog.logId" v-model:visible="state.logsDialog.visible" :title="state.logsDialog.title" />
<el-dialog :title="state.runDialog.title" v-model="state.runDialog.visible" :destroy-on-close="true" width="600px">
<el-form :model="state.runDialog.runForm" ref="runFormRef" label-width="auto" :rules="state.runDialog.formRules">
<el-form-item :label="$t('db.dbFileType')" prop="dbType">
@@ -70,17 +70,15 @@
</el-form>
<template #footer>
<div>
<el-button @click="state.runDialog.cancel()">{{ $t('common.cancel') }}</el-button>
<el-button type="primary" :loading="state.runDialog.loading" @click="state.runDialog.btnOk">{{ $t('common.confirm') }}</el-button>
</div>
<el-button @click="state.runDialog.onCancel()">{{ $t('common.cancel') }}</el-button>
<el-button type="primary" :loading="state.runDialog.loading" @click="state.runDialog.onConfirm">{{ $t('common.confirm') }}</el-button>
</template>
</el-dialog>
</div>
</template>
<script lang="ts" setup>
import { onMounted, reactive, Ref, ref, watch } from 'vue';
import { onMounted, reactive, Ref, ref, useTemplateRef, watch } from 'vue';
import { dbApi } from '@/views/ops/db/api';
import { getDbDialect } from '@/views/ops/db/dialect';
import PageTable from '@/components/pagetable/PageTable.vue';
@@ -109,6 +107,8 @@ const props = defineProps({
const dialogVisible = defineModel<boolean>('visible', { default: false });
const pageTableRef: Ref<any> = useTemplateRef('pageTableRef');
const columns = ref([
TableColumn.new('fileKey', 'db.file').setMinWidth(280).isSlot(),
TableColumn.new('createTime', 'db.execTime').setMinWidth(180).isTime(),
@@ -168,11 +168,11 @@ const state = reactive({
targetDbType: '',
},
loading: false,
cancel: function () {
onCancel: function () {
state.runDialog.visible = false;
state.runDialog.runForm = {} as any;
},
btnOk: async function () {
onConfirm: async function () {
await useI18nFormValidate(runFormRef);
if (state.runDialog.runForm.targetDbType !== state.runDialog.runForm.dbType) {
ElMessage.warning(t('db.targetDbTypeSelectError', { dbType: state.runDialog.runForm.dbType }));
@@ -181,7 +181,7 @@ const state = reactive({
state.runDialog.runForm.clientId = getClientId();
await dbApi.dbTransferFileRun.request(state.runDialog.runForm);
useI18nOperateSuccessMsg();
state.runDialog.cancel();
state.runDialog.onCancel();
await search();
},
onSelectRunTargetDb: function (param: any) {
@@ -195,14 +195,13 @@ const state = reactive({
});
const search = async () => {
const { total, list } = await dbApi.dbTransferFileList.request(state.query);
state.tableData = list;
pageTableRef.value.total = total;
pageTableRef.value?.search();
// const { total, list } = await dbApi.dbTransferFileList.request(state.query);
// state.tableData = list;
// pageTableRef.value.total = total;
};
const pageTableRef: Ref<any> = ref(null);
const del = async function () {
const onDel = async function () {
try {
await useI18nDeleteConfirm(state.selectionData.map((x: any) => x.fileKey).join('、'));
await dbApi.dbTransferFileDel.request({ fileId: state.selectionData.map((x: any) => x.id).join(',') });
@@ -213,7 +212,7 @@ const del = async function () {
}
};
const openLog = function (data: any) {
const onOpenLog = function (data: any) {
state.logsDialog.logId = data.logId;
state.logsDialog.visible = true;
state.logsDialog.title = t('db.log');
@@ -221,7 +220,7 @@ const openLog = function (data: any) {
};
// 运行sql，弹出选择需要运行的库，默认运行当前数据库，需要保证数据库类型与sql文件一致
const openRun = function (data: any) {
const onOpenRun = function (data: any) {
state.runDialog.runForm = { id: data.id, dbType: data.fileDbType } as any;
state.runDialog.visible = true;
};

17
frontend/src/views/ops/db/component/table/DbTableDataForm.vue
View File

@@ -1,10 +1,9 @@
<template>
<el-dialog v-model="visible" :title="title" :destroy-on-close="true" width="600px">
<el-form ref="dataForm" :model="modelValue" :show-message="false" label-width="auto" size="small">
<el-dialog v-model="visible" :title="title" :destroy-on-close="true" width="600px" body-class="h-[65vh] overflow-auto">
<el-form ref="dataForm" :model="modelValue" scroll-to-error :show-message="false" label-width="auto" size="small">
<el-form-item
v-for="column in columns"
:key="column.columnName"
class="mb-1 w-full"
:prop="column.columnName"
:required="props.tableName != '' && !column.nullable && !column.isPrimaryKey && !column.autoIncrement"
>
@@ -24,10 +23,8 @@
</el-form-item>
</el-form>
<template #footer v-if="props.tableName">
<span class="dialog-footer">
<el-button @click="closeDialog">{{ $t('common.cancel') }}</el-button>
<el-button type="primary" @click="confirm">{{ $t('common.confirm') }}</el-button>
</span>
<el-button @click="onCloseDialog">{{ $t('common.cancel') }}</el-button>
<el-button type="primary" @click="onConfirm">{{ $t('common.confirm') }}</el-button>
</template>
</el-dialog>
</template>
@@ -79,12 +76,12 @@ const setOldValue = () => {
}
};
const closeDialog = () => {
const onCloseDialog = () => {
visible.value = false;
modelValue.value = {};
};
const confirm = async () => {
const onConfirm = async () => {
await useI18nFormValidate(dataForm);
const dbInst = props.dbInst;
@@ -107,7 +104,7 @@ const confirm = async () => {
}
dbInst.promptExeSql(db, sql, null, () => {
closeDialog();
onCloseDialog();
emit('submitSuccess');
});
};

2
frontend/src/views/ops/machine/ScriptManage.vue
View File

@@ -68,7 +68,7 @@
:close-on-click-modal="false"
:modal="false"
@close="closeTermnial"
body-class="h-[560px]"
body-class="h-[65vh]"
draggable
append-to-body
>

2
frontend/src/views/ops/machine/cronjob/CronJobExecList.vue
View File

@@ -9,7 +9,7 @@
:show-close="true"
:destroy-on-close="true"
width="65%"
body-class="h-200"
body-class="h-[65vh]"
>
<page-table
ref="pageTableRef"

34
frontend/src/views/ops/machine/file/MachineFileContent.vue
View File

@@ -15,14 +15,14 @@
<template #footer>
<el-button @click="handleClose">{{ $t('common.cancel') }}</el-button>
<el-button v-auth="'machine:file:write'" type="primary" @click="updateContent">{{ $t('common.save') }}</el-button>
<el-button v-loading="saveing" v-auth="'machine:file:write'" type="primary" @click="updateContent">{{ $t('common.save') }}</el-button>
</template>
</el-dialog>
</div>
</template>
<script lang="ts" setup>
import { computed, reactive, toRefs, watch } from 'vue';
import { computed, reactive, Ref, ref, toRefs, watch } from 'vue';
import { machineApi } from '../api';
import MonacoEditor from '@/components/monaco/MonacoEditor.vue';
import { useI18nSaveSuccessMsg } from '@/hooks/useI18n';
@@ -42,9 +42,10 @@ const emit = defineEmits(['cancel', 'update:machineId']);
const updateFileContent = machineApi.updateFileContent;
const saveing: Ref<any> = ref(false);
const state = reactive({
loadingContent: false,
content: '',
fileType: '',
});
@@ -83,17 +84,22 @@ const handleClose = () => {
};
const updateContent = async () => {
await updateFileContent.request({
content: state.content,
id: props.fileId,
path: props.path,
machineId: props.machineId,
authCertName: props.authCertName,
protocol: props.protocol,
});
useI18nSaveSuccessMsg();
handleClose();
state.content = '';
try {
saveing.value = true;
await updateFileContent.request({
content: fileContent.value,
id: props.fileId,
path: props.path,
machineId: props.machineId,
authCertName: props.authCertName,
protocol: props.protocol,
});
useI18nSaveSuccessMsg();
handleClose();
fileContent.value = '';
} finally {
saveing.value = false;
}
};
const getFileType = (path: string) => {

6
frontend/src/views/ops/redis/Info.vue
View File

@@ -22,9 +22,9 @@
<el-row :gutter="20">
<el-col :lg="12" :md="12">
<el-descriptions class="redis-info info-cluster" :title="$t('redis.node')" :column="3" size="small" border>
<el-descriptions-item :label="$t('redis.clusterEnable')">{{ info.Cluster.cluster_enabled }}</el-descriptions-item>
<el-descriptions-item label="DB">{{ info.Cluster.databases }}</el-descriptions-item>
<el-descriptions-item :label="$t('redis.nodeCount')">{{ info.Cluster.nodecount }}</el-descriptions-item>
<el-descriptions-item :label="$t('redis.clusterEnable')">{{ info.Cluster?.cluster_enabled }}</el-descriptions-item>
<el-descriptions-item label="DB">{{ info.Cluster?.databases }}</el-descriptions-item>
<el-descriptions-item :label="$t('redis.nodeCount')">{{ info.Cluster?.nodecount }}</el-descriptions-item>
</el-descriptions>
</el-col>

2
server/go.mod
View File

@@ -39,7 +39,7 @@ require (
gopkg.in/yaml.v3 v3.0.1
// gorm
gorm.io/driver/mysql v1.5.7
gorm.io/gorm v1.26.1
gorm.io/gorm v1.30.0
)
require (

20
server/internal/db/dbm/dbm.go
View File

@@ -9,6 +9,7 @@ import (
_ "mayfly-go/internal/db/dbm/oracle"
_ "mayfly-go/internal/db/dbm/postgres"
_ "mayfly-go/internal/db/dbm/sqlite"
"mayfly-go/internal/machine/mcm"
"mayfly-go/pkg/logx"
"mayfly-go/pkg/pool"
)
@@ -17,6 +18,25 @@ var (
poolGroup = pool.NewPoolGroup[*dbi.DbConn]()
)
func init() {
mcm.AddCheckSshTunnelMachineUseFunc(func(machineId int) bool {
items := poolGroup.AllPool()
for _, v := range items {
if v.Stats().TotalConns == 0 {
continue // 连接池中没有连接，跳过
}
conn, err := v.Get(context.Background())
if err != nil {
continue // 获取连接失败，跳过
}
if conn.Info.SshTunnelMachineId == machineId {
return true
}
}
return false
})
}
// GetDbConn 从连接池中获取连接信息
func GetDbConn(ctx context.Context, dbId uint64, database string, getDbInfo func() (*dbi.DbInfo, error)) (*dbi.DbConn, error) {
connId := dbi.GetDbConnId(dbId, database)

22
server/internal/es/application/es_instance.go
View File

@@ -8,6 +8,7 @@ import (
"mayfly-go/internal/es/domain/repository"
"mayfly-go/internal/es/esm/esi"
"mayfly-go/internal/es/imsg"
"mayfly-go/internal/machine/mcm"
"mayfly-go/internal/pkg/consts"
tagapp "mayfly-go/internal/tag/application"
tagdto "mayfly-go/internal/tag/application/dto"
@@ -40,6 +41,25 @@ var _ Instance = &instanceAppImpl{}
var poolGroup = pool.NewPoolGroup[*esi.EsConn]()
func init() {
mcm.AddCheckSshTunnelMachineUseFunc(func(machineId int) bool {
items := poolGroup.AllPool()
for _, v := range items {
if v.Stats().TotalConns == 0 {
continue // 连接池中没有连接，跳过
}
conn, err := v.Get(context.Background())
if err != nil {
continue // 获取连接失败，跳过
}
if conn.Info.SshTunnelMachineId == machineId {
return true
}
}
return false
})
}
type instanceAppImpl struct {
base.AppImpl[*entity.EsInstance, repository.EsInstance]
@@ -234,6 +254,8 @@ func (app *instanceAppImpl) Delete(ctx context.Context, instanceId uint64) error
return errorx.NewBiz("db instnace not found")
}
poolGroup.Close(fmt.Sprintf("es-%d", instanceId))
return app.Tx(ctx, func(ctx context.Context) error {
// 删除该实例
return app.DeleteById(ctx, instanceId)

23
server/internal/machine/mcm/client_cache.go
View File

@@ -9,6 +9,27 @@ var (
poolGroup = pool.NewPoolGroup[*Cli]()
)
func init() {
AddCheckSshTunnelMachineUseFunc(func(machineId int) bool {
// 遍历所有redis连接实例，若存在redis实例使用该ssh隧道机器，则返回true，表示还在使用中...
items := poolGroup.AllPool()
for _, v := range items {
if v.Stats().TotalConns == 0 {
continue // 连接池中没有连接，跳过
}
cli, err := v.Get(context.Background())
if err != nil {
continue // 获取连接失败，跳过
}
sshTunnelMachine := cli.Info.SshTunnelMachine
if sshTunnelMachine != nil && sshTunnelMachine.Id == uint64(machineId) {
return true
}
}
return false
})
}
// 从缓存中获取客户端信息，不存在则回调获取机器信息函数，并新建。
// @param 机器的授权凭证名
func GetMachineCli(ctx context.Context, authCertName string, getMachine func(string) (*MachineInfo, error)) (*Cli, error) {
@@ -39,7 +60,7 @@ func DeleteCli(id uint64) {
continue
}
if conn.Info.Id == id {
pool.Close()
poolGroup.Close(conn.Info.AuthCertName)
}
}
}

45
server/internal/machine/mcm/sshtunnel.go
View File

@@ -10,10 +10,15 @@ import (
"mayfly-go/pkg/utils/netx"
"net"
"sync"
"time"
"golang.org/x/crypto/ssh"
)
// type SshTunnelAble interface {
// GetSshTunnelMachineId() int
// }
var (
// 所有检测ssh隧道机器是否被使用的函数
checkSshTunnelMachineHasUseFuncs []CheckSshTunnelMachineHasUseFunc
@@ -132,7 +137,20 @@ func GetSshTunnelMachine(ctx context.Context, machineId int, getMachine func(uin
logx.Infof("connect to the ssh tunnel machine for the first time[%d][%s:%d]", machineId, mi.Ip, mi.Port)
return stm, err
})
}, pool.WithIdleTimeout[*SshTunnelMachine](50*time.Minute), pool.WithOnConnClose(func(conn *SshTunnelMachine) error {
mid := int(conn.mi.Id)
logx.Debugf("periodically check if the ssh tunnel machine [%d] is still in use...", mid)
for _, checkUseFunc := range checkSshTunnelMachineHasUseFuncs {
// 如果一个在使用则返回不关闭，不继续后续检查
if checkUseFunc(mid) {
return fmt.Errorf("ssh tunnel machine [%s] is still in use", conn.mi.Name)
}
}
return nil
}))
if err != nil {
return nil, err
}
@@ -142,18 +160,19 @@ func GetSshTunnelMachine(ctx context.Context, machineId int, getMachine func(uin
// 关闭ssh隧道机器的指定隧道
func CloseSshTunnelMachine(machineId uint64, tunnelId string) {
//sshTunnelMachine := mcIdPool[machineId]
//if sshTunnelMachine == nil {
// return
//}
//
//sshTunnelMachine.mutex.Lock()
//defer sshTunnelMachine.mutex.Unlock()
//t := sshTunnelMachine.tunnels[tunnelId]
//if t != nil {
// t.Close()
// delete(sshTunnelMachine.tunnels, tunnelId)
//}
sshTunnelMachinePool, ok := tunnelPoolGroup.Get(fmt.Sprintf("machine-tunnel-%d", machineId))
if !ok {
return
}
sshTunnelMachine, err := sshTunnelMachinePool.Get(context.Background())
if err != nil {
return
}
t := sshTunnelMachine.tunnels[tunnelId]
if t != nil {
t.Close()
delete(sshTunnelMachine.tunnels, tunnelId)
}
}
type Tunnel struct {

20
server/internal/mongo/mgm/conn_cache.go
View File

@@ -2,6 +2,7 @@ package mgm
import (
"context"
"mayfly-go/internal/machine/mcm"
"mayfly-go/pkg/pool"
)
@@ -9,6 +10,25 @@ var (
poolGroup = pool.NewPoolGroup[*MongoConn]()
)
func init() {
mcm.AddCheckSshTunnelMachineUseFunc(func(machineId int) bool {
items := poolGroup.AllPool()
for _, v := range items {
if v.Stats().TotalConns == 0 {
continue // 连接池中没有连接，跳过
}
conn, err := v.Get(context.Background())
if err != nil {
continue // 获取连接失败，跳过
}
if conn.Info.SshTunnelMachineId == machineId {
return true
}
}
return false
})
}
// 从缓存中获取mongo连接信息, 若缓存中不存在则会使用回调函数获取mongoInfo进行连接并缓存
func GetMongoConn(ctx context.Context, mongoId uint64, getMongoInfo func() (*MongoInfo, error)) (*MongoConn, error) {
pool, err := poolGroup.GetCachePool(getConnId(mongoId), func() (*MongoConn, error) {

24
server/internal/redis/rdm/conn_cache.go
View File

@@ -2,16 +2,34 @@ package rdm
import (
"context"
"mayfly-go/internal/machine/mcm"
"mayfly-go/pkg/pool"
)
func init() {
}
var (
poolGroup = pool.NewPoolGroup[*RedisConn]()
)
func init() {
mcm.AddCheckSshTunnelMachineUseFunc(func(machineId int) bool {
// 遍历所有redis连接实例，若存在redis实例使用该ssh隧道机器，则返回true，表示还在使用中...
items := poolGroup.AllPool()
for _, v := range items {
if v.Stats().TotalConns == 0 {
continue // 连接池中没有连接，跳过
}
rc, err := v.Get(context.Background())
if err != nil {
continue // 获取连接失败，跳过
}
if rc.Info.SshTunnelMachineId == machineId {
return true
}
}
return false
})
}
// 从缓存中获取redis连接信息, 若缓存中不存在则会使用回调函数获取redisInfo进行连接并缓存
func GetRedisConn(ctx context.Context, redisId uint64, db int, getRedisInfo func() (*RedisInfo, error)) (*RedisConn, error) {
p, err := poolGroup.GetCachePool(getConnId(redisId, db), func() (*RedisConn, error) {

60
server/pkg/pool/cache_pool.go
View File

@@ -8,35 +8,27 @@ import (
"time"
)
var CachePoolDefaultConfig = PoolConfig{
MaxConns: 1,
IdleTimeout: 60 * time.Minute,
WaitTimeout: 10 * time.Second,
HealthCheckInterval: 10 * time.Minute,
}
type cacheEntry[T Conn] struct {
conn T
lastActive time.Time
}
func (e *cacheEntry[T]) Close() {
if err := e.conn.Close(); err != nil {
logx.Errorf("cache pool - closing connection error: %v", err)
}
}
type CachePool[T Conn] struct {
factory func() (T, error)
mu sync.RWMutex
cache map[string]*cacheEntry[T] // 使用字符串键的缓存
config PoolConfig
config PoolConfig[T]
closeCh chan struct{}
closed bool
}
func NewCachePool[T Conn](factory func() (T, error), opts ...Option) *CachePool[T] {
config := CachePoolDefaultConfig
func NewCachePool[T Conn](factory func() (T, error), opts ...Option[T]) *CachePool[T] {
config := PoolConfig[T]{
MaxConns: 1,
IdleTimeout: 60 * time.Minute,
WaitTimeout: 10 * time.Second,
HealthCheckInterval: 10 * time.Minute,
}
for _, opt := range opts {
opt(&config)
}
@@ -81,8 +73,9 @@ func (p *CachePool[T]) Get(ctx context.Context) (T, error) {
return entry.conn, nil
}
// 清理超时连接
entry.Close()
delete(p.cache, key)
if !p.closeConn(key, entry, false) {
return entry.conn, nil
}
}
// 创建新连接
@@ -151,8 +144,9 @@ func (p *CachePool[T]) Close() {
p.closed = true
close(p.closeCh)
for _, entry := range p.cache {
entry.Close()
for key, entry := range p.cache {
// 强制关闭连接
p.closeConn(key, entry, true)
}
// 触发关闭回调
@@ -212,13 +206,33 @@ func (p *CachePool[T]) cleanupIdle() {
cutoff := time.Now().Add(-p.config.IdleTimeout)
for key, entry := range p.cache {
if entry.lastActive.Before(cutoff) {
entry.Close()
delete(p.cache, key)
if entry.lastActive.Before(cutoff) || entry.conn.Ping() != nil {
logx.Infof("cache pool - cleaning up idle connection, key: %s", key)
// 如果连接超时或不可用，则关闭连接
p.closeConn(key, entry, false)
}
}
}
func (p *CachePool[T]) closeConn(key string, entry *cacheEntry[T], force bool) bool {
if !force {
// 如果不是强制关闭且有连接关闭回调，则调用回调
// 如果回调返回错误，则不关闭连接
if onConnClose := p.config.OnConnClose; onConnClose != nil {
if err := onConnClose(entry.conn); err != nil {
logx.Infof("cache pool - connection close callback returned error, skip closing connection:: %v", err)
return false
}
}
}
if err := entry.conn.Close(); err != nil {
logx.Errorf("cache pool - closing connection error: %v", err)
}
delete(p.cache, key)
return true
}
// 生成缓存键
func generateCacheKey() string {
return stringx.RandUUID()

18
server/pkg/pool/chan_pool.go
View File

@@ -10,13 +10,6 @@ import (
"time"
)
var ChanPoolDefaultConfig = PoolConfig{
MaxConns: 5,
IdleTimeout: 60 * time.Minute,
WaitTimeout: 10 * time.Second,
HealthCheckInterval: 10 * time.Minute,
}
// chanConn 封装连接及其元数据
type chanConn[T Conn] struct {
conn T
@@ -41,7 +34,7 @@ type ChanPool[T Conn] struct {
mu sync.RWMutex
factory func() (T, error)
idleConns chan *chanConn[T]
config PoolConfig
config PoolConfig[T]
currentConns int32
stats PoolStats
closeChan chan struct{} // 用于关闭健康检查 goroutine
@@ -56,9 +49,14 @@ type PoolStats struct {
WaitCount int64 // 等待连接次数
}
func NewChannelPool[T Conn](factory func() (T, error), opts ...Option) *ChanPool[T] {
func NewChannelPool[T Conn](factory func() (T, error), opts ...Option[T]) *ChanPool[T] {
// 1. 初始化配置（使用默认值 + Option 覆盖）
config := ChanPoolDefaultConfig
config := PoolConfig[T]{
MaxConns: 5,
IdleTimeout: 60 * time.Minute,
WaitTimeout: 10 * time.Second,
HealthCheckInterval: 10 * time.Minute,
}
for _, opt := range opts {
opt(&config)
}

33
server/pkg/pool/config.go
View File

@@ -8,48 +8,57 @@ import (
var ErrPoolClosed = errors.New("pool is closed")
// PoolConfig 连接池配置
type PoolConfig struct {
type PoolConfig[T Conn] struct {
MaxConns int // 最大连接数
IdleTimeout time.Duration // 空闲连接超时时间
WaitTimeout time.Duration // 获取连接超时时间
HealthCheckInterval time.Duration // 健康检查间隔
OnPoolClose func() error // 连接池关闭时的回调
OnConnClose func(conn T) error // 连接关闭时的回调,若err != nil则不关闭连接
}
// Option 函数类型，用于配置 Pool
type Option func(*PoolConfig)
type Option[T Conn] func(*PoolConfig[T])
// WithMaxConns 设置最大连接数
func WithMaxConns(maxConns int) Option {
return func(c *PoolConfig) {
func WithMaxConns[T Conn](maxConns int) Option[T] {
return func(c *PoolConfig[T]) {
c.MaxConns = maxConns
}
}
// WithIdleTimeout 设置空闲超时
func WithIdleTimeout(timeout time.Duration) Option {
return func(c *PoolConfig) {
func WithIdleTimeout[T Conn](timeout time.Duration) Option[T] {
return func(c *PoolConfig[T]) {
c.IdleTimeout = timeout
}
}
// WithWaitTimeout 设置等待超时
func WithWaitTimeout(timeout time.Duration) Option {
return func(c *PoolConfig) {
func WithWaitTimeout[T Conn](timeout time.Duration) Option[T] {
return func(c *PoolConfig[T]) {
c.WaitTimeout = timeout
}
}
// WithHealthCheckInterval 设置健康检查间隔
func WithHealthCheckInterval(interval time.Duration) Option {
return func(c *PoolConfig) {
func WithHealthCheckInterval[T Conn](interval time.Duration) Option[T] {
return func(c *PoolConfig[T]) {
c.HealthCheckInterval = interval
}
}
// WithOnPoolClose 设置连接池关闭回调
func WithOnPoolClose(fn func() error) Option {
return func(c *PoolConfig) {
func WithOnPoolClose[T Conn](fn func() error) Option[T] {
return func(c *PoolConfig[T]) {
c.OnPoolClose = fn
}
}
// WithOnConnClose 设置连接关闭回调, 若返回的错误不为nil，则不关闭连接
func WithOnConnClose[T Conn](fn func(conn T) error) Option[T] {
return func(c *PoolConfig[T]) {
c.OnConnClose = fn
}
}

16
server/pkg/pool/group.go
View File

@@ -23,7 +23,7 @@ func NewPoolGroup[T Conn]() *PoolGroup[T] {
func (pg *PoolGroup[T]) GetOrCreate(
key string,
poolFactory func() Pool[T],
opts ...Option,
opts ...Option[T],
) (Pool[T], error) {
// 先尝试读锁获取
pg.mu.RLock()
@@ -63,19 +63,29 @@ func (pg *PoolGroup[T]) GetOrCreate(
}
// GetChanPool 获取或创建 ChannelPool 类型连接池
func (pg *PoolGroup[T]) GetChanPool(key string, factory func() (T, error), opts ...Option) (Pool[T], error) {
func (pg *PoolGroup[T]) GetChanPool(key string, factory func() (T, error), opts ...Option[T]) (Pool[T], error) {
return pg.GetOrCreate(key, func() Pool[T] {
return NewChannelPool(factory, opts...)
}, opts...)
}
// GetCachePool 获取或创建 CachePool 类型连接池
func (pg *PoolGroup[T]) GetCachePool(key string, factory func() (T, error), opts ...Option) (Pool[T], error) {
func (pg *PoolGroup[T]) GetCachePool(key string, factory func() (T, error), opts ...Option[T]) (Pool[T], error) {
return pg.GetOrCreate(key, func() Pool[T] {
return NewCachePool(factory, opts...)
}, opts...)
}
// Get 获取指定 key 的连接池
func (pg *PoolGroup[T]) Get(key string) (Pool[T], bool) {
pg.mu.RLock()
defer pg.mu.RUnlock()
if p, ok := pg.poolGroup[key]; ok {
return p, true
}
return nil, false
}
func (pg *PoolGroup[T]) Close(key string) error {
pg.mu.Lock()
defer pg.mu.Unlock()

60
server/pkg/pool/pool_test.go
View File

@@ -88,10 +88,10 @@ func newMockConn(id int) *mockConn {
func TestChanPool_Basic(t *testing.T) {
var idGen int
pool := NewChannelPool(func() (Conn, error) {
pool := NewChannelPool(func() (*mockConn, error) {
idGen++
return newMockConn(idGen), nil
}, WithMaxConns(2), WithIdleTimeout(time.Second))
}, WithMaxConns[*mockConn](2), WithIdleTimeout[*mockConn](time.Second))
ctx := context.Background()
conn1, _ := pool.Get(ctx)
@@ -112,9 +112,9 @@ func TestChanPool_Basic(t *testing.T) {
}
func TestChanPool_WaitTimeout(t *testing.T) {
pool := NewChannelPool(func() (Conn, error) {
pool := NewChannelPool(func() (*mockConn, error) {
return newMockConn(1), nil
}, WithMaxConns(1), WithWaitTimeout(100*time.Millisecond))
}, WithMaxConns[*mockConn](1), WithWaitTimeout[*mockConn](100*time.Millisecond))
ctx := context.Background()
conn1, _ := pool.Get(ctx)
@@ -132,9 +132,9 @@ func TestChanPool_WaitTimeout(t *testing.T) {
}
func TestChanPool_ContextCancel(t *testing.T) {
pool := NewChannelPool(func() (Conn, error) {
pool := NewChannelPool(func() (*mockConn, error) {
return newMockConn(1), nil
}, WithMaxConns(1))
}, WithMaxConns[*mockConn](1))
ctx, cancel := context.WithCancel(context.Background())
conn, _ := pool.Get(ctx)
@@ -145,9 +145,9 @@ func TestChanPool_ContextCancel(t *testing.T) {
}
func TestChanPool_Resize(t *testing.T) {
pool := NewChannelPool(func() (Conn, error) {
pool := NewChannelPool(func() (*mockConn, error) {
return newMockConn(1), nil
}, WithMaxConns(2))
}, WithMaxConns[*mockConn](2))
ctx := context.Background()
conn1, _ := pool.Get(ctx)
conn2, _ := pool.Get(ctx)
@@ -158,9 +158,9 @@ func TestChanPool_Resize(t *testing.T) {
}
func TestChanPool_HealthCheck(t *testing.T) {
pool := NewChannelPool(func() (Conn, error) {
pool := NewChannelPool(func() (*mockConn, error) {
return newMockConn(1), nil
}, WithMaxConns(1), WithIdleTimeout(10*time.Millisecond), WithHealthCheckInterval(10*time.Millisecond))
}, WithMaxConns[*mockConn](1), WithIdleTimeout[*mockConn](10*time.Millisecond), WithHealthCheckInterval[*mockConn](10*time.Millisecond))
ctx := context.Background()
conn, _ := pool.Get(ctx)
_ = pool.Put(conn)
@@ -176,10 +176,10 @@ func TestChanPool_HealthCheck(t *testing.T) {
func TestCachePool_Basic(t *testing.T) {
var idGen int
pool := NewCachePool(func() (Conn, error) {
pool := NewCachePool(func() (*mockConn, error) {
idGen++
return newMockConn(idGen), nil
}, WithMaxConns(2), WithIdleTimeout(time.Second))
}, WithMaxConns[*mockConn](2), WithIdleTimeout[*mockConn](time.Second))
ctx := context.Background()
conn1, _ := pool.Get(ctx)
@@ -193,9 +193,9 @@ func TestCachePool_Basic(t *testing.T) {
}
func TestCachePool_TimeoutCleanup(t *testing.T) {
pool := NewCachePool(func() (Conn, error) {
pool := NewCachePool(func() (*mockConn, error) {
return newMockConn(1), nil
}, WithMaxConns(1), WithIdleTimeout(10*time.Millisecond), WithHealthCheckInterval(10*time.Millisecond))
}, WithMaxConns[*mockConn](1), WithIdleTimeout[*mockConn](10*time.Millisecond), WithHealthCheckInterval[*mockConn](10*time.Millisecond))
ctx := context.Background()
conn, _ := pool.Get(ctx)
_ = pool.Put(conn)
@@ -209,10 +209,10 @@ func TestCachePool_TimeoutCleanup(t *testing.T) {
func TestCachePool_OverMaxConns(t *testing.T) {
var idGen int
pool := NewCachePool(func() (Conn, error) {
pool := NewCachePool(func() (*mockConn, error) {
idGen++
return newMockConn(idGen), nil
}, WithMaxConns(1))
}, WithMaxConns[*mockConn](1))
ctx := context.Background()
conn1, _ := pool.Get(ctx)
_ = pool.Put(conn1)
@@ -231,9 +231,9 @@ func TestCachePool_OverMaxConns(t *testing.T) {
}
func TestCachePool_Resize(t *testing.T) {
pool := NewCachePool(func() (Conn, error) {
pool := NewCachePool(func() (*mockConn, error) {
return newMockConn(1), nil
}, WithMaxConns(2))
}, WithMaxConns[*mockConn](2))
ctx := context.Background()
conn1, _ := pool.Get(ctx)
_ = pool.Put(conn1)
@@ -288,9 +288,9 @@ func TestPoolGroup_Concurrent(t *testing.T) {
// ========== 压力测试 ==========
func BenchmarkChanPool_Concurrent(b *testing.B) {
pool := NewChannelPool(func() (Conn, error) {
pool := NewChannelPool(func() (*mockConn, error) {
return newMockConn(1), nil
}, WithMaxConns(100))
}, WithMaxConns[*mockConn](100))
b.ResetTimer()
b.RunParallel(func(pb *testing.PB) {
@@ -307,9 +307,9 @@ func BenchmarkChanPool_Concurrent(b *testing.B) {
}
func BenchmarkCachePool_Concurrent(b *testing.B) {
pool := NewCachePool(func() (Conn, error) {
pool := NewCachePool(func() (*mockConn, error) {
return newMockConn(1), nil
}, WithMaxConns(100))
}, WithMaxConns[*mockConn](100))
b.ResetTimer()
b.RunParallel(func(pb *testing.PB) {
@@ -332,9 +332,9 @@ func TestChanPool_Stress(t *testing.T) {
iterations = 1000
)
pool := NewChannelPool(func() (Conn, error) {
pool := NewChannelPool(func() (*mockConn, error) {
return newMockConn(1), nil
}, WithMaxConns(20), WithWaitTimeout(time.Second))
}, WithMaxConns[*mockConn](20), WithWaitTimeout[*mockConn](time.Second))
var wg sync.WaitGroup
var errCount int32
@@ -389,9 +389,9 @@ func TestCachePool_Stress(t *testing.T) {
iterations = 1000
)
pool := NewCachePool(func() (Conn, error) {
pool := NewCachePool(func() (*mockConn, error) {
return newMockConn(1), nil
}, WithMaxConns(20), WithIdleTimeout(time.Minute))
}, WithMaxConns[*mockConn](20), WithIdleTimeout[*mockConn](time.Minute))
var wg sync.WaitGroup
var errCount int32
@@ -430,11 +430,11 @@ func TestCachePool_Stress(t *testing.T) {
// 测试连接池在连接失效时的行为
func TestChanPool_InvalidConn(t *testing.T) {
pool := NewChannelPool(func() (Conn, error) {
pool := NewChannelPool(func() (*mockConn, error) {
conn := newMockConn(1)
conn.pingErr = errors.New("connection invalid")
return conn, nil
}, WithMaxConns(1), WithHealthCheckInterval(10*time.Millisecond))
}, WithMaxConns[*mockConn](1), WithHealthCheckInterval[*mockConn](10*time.Millisecond))
ctx := context.Background()
conn, _ := pool.Get(ctx)
@@ -458,9 +458,9 @@ func TestChanPool_InvalidConn(t *testing.T) {
// 测试连接池在并发关闭时的行为
func TestChanPool_ConcurrentClose(t *testing.T) {
pool := NewChannelPool(func() (Conn, error) {
pool := NewChannelPool(func() (*mockConn, error) {
return newMockConn(1), nil
}, WithMaxConns(10))
}, WithMaxConns[*mockConn](10))
var wg sync.WaitGroup
const goroutines = 10