Skip to content

phsl / new-api

Go to a project
Commit 2f14d89f by CaIon
Options

✨ feat(channel): implement multi-key mode handling and improve channel update logic

parent 8e11b67e
Showing with 153 additions and 47 deletions
controller/channel.go
...@@ -718,8 +718,13 @@ func DeleteChannelBatch(c *gin.Context) { ...@@ -718,8 +718,13 @@ func DeleteChannelBatch(c *gin.Context) {
return return
} }
type PatchChannel struct {
model.Channel
MultiKeyMode *string `json:"multi_key_mode"`
}
func UpdateChannel(c *gin.Context) { func UpdateChannel(c *gin.Context) {
channel := model.Channel{} channel := PatchChannel{}
err := c.ShouldBindJSON(&channel) err := c.ShouldBindJSON(&channel)
if err != nil { if err != nil {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
...@@ -761,6 +766,19 @@ func UpdateChannel(c *gin.Context) { ...@@ -761,6 +766,19 @@ func UpdateChannel(c *gin.Context) {
} }
} }
} }
if channel.MultiKeyMode != nil && *channel.MultiKeyMode != "" {
originChannel, err := model.GetChannelById(channel.Id, false)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
}
if originChannel.ChannelInfo.IsMultiKey {
channel.ChannelInfo = originChannel.ChannelInfo
channel.ChannelInfo.MultiKeyMode = constant.MultiKeyMode(*channel.MultiKeyMode)
}
}
err = channel.Update() err = channel.Update()
if err != nil { if err != nil {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
......
model/channel.go
...@@ -117,7 +117,15 @@ func (channel *Channel) GetNextEnabledKey() (string, *types.NewAPIError) { ...@@ -117,7 +117,15 @@ func (channel *Channel) GetNextEnabledKey() (string, *types.NewAPIError) {
// Randomly pick one enabled key // Randomly pick one enabled key
return keys[enabledIdx[rand.Intn(len(enabledIdx))]], nil return keys[enabledIdx[rand.Intn(len(enabledIdx))]], nil
case constant.MultiKeyModePolling: case constant.MultiKeyModePolling:
defer func() {
if !common.MemoryCacheEnabled {
_ = channel.Save()
} else {
CacheUpdateChannel(channel)
}
}()
// Start from the saved polling index and look for the next enabled key // Start from the saved polling index and look for the next enabled key
println(channel.ChannelInfo.MultiKeyPollingIndex)
start := channel.ChannelInfo.MultiKeyPollingIndex start := channel.ChannelInfo.MultiKeyPollingIndex
if start < 0 || start >= len(keys) { if start < 0 || start >= len(keys) {
start = 0 start = 0
...@@ -127,6 +135,7 @@ func (channel *Channel) GetNextEnabledKey() (string, *types.NewAPIError) { ...@@ -127,6 +135,7 @@ func (channel *Channel) GetNextEnabledKey() (string, *types.NewAPIError) {
if getStatus(idx) == common.ChannelStatusEnabled { if getStatus(idx) == common.ChannelStatusEnabled {
// update polling index for next call (point to the next position) // update polling index for next call (point to the next position)
channel.ChannelInfo.MultiKeyPollingIndex = (idx + 1) % len(keys) channel.ChannelInfo.MultiKeyPollingIndex = (idx + 1) % len(keys)
println(channel.ChannelInfo.MultiKeyPollingIndex)
return keys[idx], nil return keys[idx], nil
} }
} }
...@@ -273,14 +282,20 @@ func SearchChannels(keyword string, group string, model string, idSort bool) ([] ...@@ -273,14 +282,20 @@ func SearchChannels(keyword string, group string, model string, idSort bool) ([]
} }
func GetChannelById(id int, selectAll bool) (*Channel, error) { func GetChannelById(id int, selectAll bool) (*Channel, error) {
channel := Channel{Id: id} channel := &Channel{Id: id}
var err error = nil var err error = nil
if selectAll { if selectAll {
err = DB.First(&channel, "id = ?", id).Error err = DB.First(channel, "id = ?", id).Error
} else { } else {
err = DB.Omit("key").First(&channel, "id = ?", id).Error err = DB.Omit("key").First(channel, "id = ?", id).Error
}
if err != nil {
return nil, err
}
if channel == nil {
return nil, errors.New("channel not found")
} }
return &channel, err return channel, nil
} }
func BatchInsertChannels(channels []Channel) error { func BatchInsertChannels(channels []Channel) error {
......
model/cache.go → model/channel_cache.go
...@@ -14,7 +14,7 @@ import ( ...@@ -14,7 +14,7 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
) )
var group2model2channels map[string]map[string][]*Channel var group2model2channels map[string]map[string][]int
var channelsIDM map[int]*Channel var channelsIDM map[int]*Channel
var channelSyncLock sync.RWMutex var channelSyncLock sync.RWMutex
...@@ -34,10 +34,10 @@ func InitChannelCache() { ...@@ -34,10 +34,10 @@ func InitChannelCache() {
for _, ability := range abilities { for _, ability := range abilities {
groups[ability.Group] = true groups[ability.Group] = true
} }
newGroup2model2channels := make(map[string]map[string][]*Channel) newGroup2model2channels := make(map[string]map[string][]int)
newChannelsIDM := make(map[int]*Channel) newChannelsIDM := make(map[int]*Channel)
for group := range groups { for group := range groups {
newGroup2model2channels[group] = make(map[string][]*Channel) newGroup2model2channels[group] = make(map[string][]int)
} }
for _, channel := range channels { for _, channel := range channels {
newChannelsIDM[channel.Id] = channel newChannelsIDM[channel.Id] = channel
...@@ -46,9 +46,9 @@ func InitChannelCache() { ...@@ -46,9 +46,9 @@ func InitChannelCache() {
models := strings.Split(channel.Models, ",") models := strings.Split(channel.Models, ",")
for _, model := range models { for _, model := range models {
if _, ok := newGroup2model2channels[group][model]; !ok { if _, ok := newGroup2model2channels[group][model]; !ok {
newGroup2model2channels[group][model] = make([]*Channel, 0) newGroup2model2channels[group][model] = make([]int, 0)
} }
newGroup2model2channels[group][model] = append(newGroup2model2channels[group][model], channel) newGroup2model2channels[group][model] = append(newGroup2model2channels[group][model], channel.Id)
} }
} }
} }
...@@ -57,7 +57,7 @@ func InitChannelCache() { ...@@ -57,7 +57,7 @@ func InitChannelCache() {
for group, model2channels := range newGroup2model2channels { for group, model2channels := range newGroup2model2channels {
for model, channels := range model2channels { for model, channels := range model2channels {
sort.Slice(channels, func(i, j int) bool { sort.Slice(channels, func(i, j int) bool {
return channels[i].GetPriority() > channels[j].GetPriority() return newChannelsIDM[channels[i]].GetPriority() > newChannelsIDM[channels[j]].GetPriority()
}) })
newGroup2model2channels[group][model] = channels newGroup2model2channels[group][model] = channels
} }
...@@ -136,8 +136,12 @@ func getRandomSatisfiedChannel(group string, model string, retry int) (*Channel, ...@@ -136,8 +136,12 @@ func getRandomSatisfiedChannel(group string, model string, retry int) (*Channel,
} }
uniquePriorities := make(map[int]bool) uniquePriorities := make(map[int]bool)
for _, channel := range channels { for _, channelId := range channels {
if channel, ok := channelsIDM[channelId]; ok {
uniquePriorities[int(channel.GetPriority())] = true uniquePriorities[int(channel.GetPriority())] = true
} else {
return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelId)
}
} }
var sortedUniquePriorities []int var sortedUniquePriorities []int
for priority := range uniquePriorities { for priority := range uniquePriorities {
...@@ -152,10 +156,14 @@ func getRandomSatisfiedChannel(group string, model string, retry int) (*Channel, ...@@ -152,10 +156,14 @@ func getRandomSatisfiedChannel(group string, model string, retry int) (*Channel,
// get the priority for the given retry number // get the priority for the given retry number
var targetChannels []*Channel var targetChannels []*Channel
for _, channel := range channels { for _, channelId := range channels {
if channel, ok := channelsIDM[channelId]; ok {
if channel.GetPriority() == targetPriority { if channel.GetPriority() == targetPriority {
targetChannels = append(targetChannels, channel) targetChannels = append(targetChannels, channel)
} }
} else {
return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelId)
}
} }
// 平滑系数 // 平滑系数
...@@ -210,9 +218,11 @@ func CacheUpdateChannel(channel *Channel) { ...@@ -210,9 +218,11 @@ func CacheUpdateChannel(channel *Channel) {
} }
channelSyncLock.Lock() channelSyncLock.Lock()
defer channelSyncLock.Unlock() defer channelSyncLock.Unlock()
if channel == nil { if channel == nil {
return return
} }
println("CacheUpdateChannel:", channel.Id, channel.Name, channel.Status, channel.ChannelInfo.MultiKeyPollingIndex)
channelsIDM[channel.Id] = channel channelsIDM[channel.Id] = channel
} }
web/src/components/table/ChannelsTable.js
...@@ -42,19 +42,20 @@ import { ...@@ -42,19 +42,20 @@ import {
IconTreeTriangleDown, IconTreeTriangleDown,
IconSearch, IconSearch,
IconMore, IconMore,
IconList IconList, IconDescend2
} from '@douyinfe/semi-icons'; } from '@douyinfe/semi-icons';
import { loadChannelModels, isMobile, copy } from '../../helpers'; import { loadChannelModels, isMobile, copy } from '../../helpers';
import EditTagModal from '../../pages/Channel/EditTagModal.js'; import EditTagModal from '../../pages/Channel/EditTagModal.js';
import { useTranslation } from 'react-i18next'; import { useTranslation } from 'react-i18next';
import { useTableCompactMode } from '../../hooks/useTableCompactMode'; import { useTableCompactMode } from '../../hooks/useTableCompactMode';
import { FaRandom } from 'react-icons/fa';
const ChannelsTable = () => { const ChannelsTable = () => {
const { t } = useTranslation(); const { t } = useTranslation();
let type2label = undefined; let type2label = undefined;
const renderType = (type, multiKey = false) => { const renderType = (type, channelInfo = undefined) => {
if (!type2label) { if (!type2label) {
type2label = new Map(); type2label = new Map();
for (let i = 0; i < CHANNEL_OPTIONS.length; i++) { for (let i = 0; i < CHANNEL_OPTIONS.length; i++) {
...@@ -65,13 +66,20 @@ const ChannelsTable = () => { ...@@ -65,13 +66,20 @@ const ChannelsTable = () => {
let icon = getChannelIcon(type); let icon = getChannelIcon(type);
if (multiKey) { if (channelInfo?.is_multi_key) {
icon = ( icon = (
channelInfo?.multi_key_mode === 'random' ? (
<div className="flex items-center gap-1"> <div className="flex items-center gap-1">
<IconList className="text-blue-500" /> <FaRandom className="text-blue-500" />
{icon} {icon}
</div> </div>
); ) : (
<div className="flex items-center gap-1">
<IconDescend2 className="text-blue-500" />
{icon}
</div>
)
)
} }
return ( return (
...@@ -587,7 +595,12 @@ const ChannelsTable = () => { ...@@ -587,7 +595,12 @@ const ChannelsTable = () => {
/> />
</SplitButtonGroup> </SplitButtonGroup>
{record.status === 1 ? ( {record.channel_info?.is_multi_key ? (
<SplitButtonGroup
aria-label={t('多密钥渠道操作项目组')}
>
{
record.status === 1 ? (
<Button <Button
theme='light' theme='light'
type='warning' type='warning'
...@@ -605,6 +618,47 @@ const ChannelsTable = () => { ...@@ -605,6 +618,47 @@ const ChannelsTable = () => {
> >
{t('启用')} {t('启用')}
</Button> </Button>
)
}
<Dropdown
trigger='click'
position='bottomRight'
menu={[
{
node: 'item',
name: t('启用全部密钥'),
onClick: () => manageChannel(record.id, 'enable_all', record),
}
]}
>
<Button
theme='light'
type='secondary'
size="small"
icon={<IconTreeTriangleDown />}
/>
</Dropdown>
</SplitButtonGroup>
) : (
record.status === 1 ? (
<Button
theme='light'
type='warning'
size="small"
onClick={() => manageChannel(record.id, 'disable', record)}
>
{t('禁用')}
</Button>
) : (
<Button
theme='light'
type='secondary'
size="small"
onClick={() => manageChannel(record.id, 'enable', record)}
>
{t('启用')}
</Button>
)
)} )}
<Button <Button
...@@ -1014,6 +1068,11 @@ const ChannelsTable = () => { ...@@ -1014,6 +1068,11 @@ const ChannelsTable = () => {
} }
res = await API.put('/api/channel/', data); res = await API.put('/api/channel/', data);
break; break;
case 'enable_all':
data.channel_info = record.channel_info;
data.channel_info.multi_key_status_list = {};
res = await API.put('/api/channel/', data);
break;
} }
const { success, message } = res.data; const { success, message } = res.data;
if (success) { if (success) {
......
web/src/pages/Channel/EditChannel.js
...@@ -435,7 +435,7 @@ const EditChannel = (props) => { ...@@ -435,7 +435,7 @@ const EditChannel = (props) => {
const formValues = formApiRef.current ? formApiRef.current.getValues() : {}; const formValues = formApiRef.current ? formApiRef.current.getValues() : {};
let localInputs = { ...formValues }; let localInputs = { ...formValues };
if (localInputs.type === 41 && batch) { if (localInputs.type === 41) {
let keys = vertexKeys; let keys = vertexKeys;
if (keys.length === 0) { if (keys.length === 0) {
// 确保提交时也能解析，避免因异步延迟导致 keys 为空 // 确保提交时也能解析，避免因异步延迟导致 keys 为空
...@@ -460,7 +460,11 @@ const EditChannel = (props) => { ...@@ -460,7 +460,11 @@ const EditChannel = (props) => {
return; return;
} }
if (batch) {
localInputs.key = JSON.stringify(keys); localInputs.key = JSON.stringify(keys);
} else {
localInputs.key = JSON.stringify(keys[0]);
}
} }
delete localInputs.vertex_files; delete localInputs.vertex_files;
...@@ -561,7 +565,7 @@ const EditChannel = (props) => { ...@@ -561,7 +565,7 @@ const EditChannel = (props) => {
const batchAllowed = !isEdit || isMultiKeyChannel; const batchAllowed = !isEdit || isMultiKeyChannel;
const batchExtra = batchAllowed ? ( const batchExtra = batchAllowed ? (
<Space> <Space>
<Checkbox checked={batch} onChange={() => { <Checkbox disabled={isEdit} checked={batch} onChange={() => {
setBatch(!batch); setBatch(!batch);
if (batch) { if (batch) {
setMultiToSingle(false); setMultiToSingle(false);
...@@ -569,7 +573,7 @@ const EditChannel = (props) => { ...@@ -569,7 +573,7 @@ const EditChannel = (props) => {
} }
}}>{t('批量创建')}</Checkbox> }}>{t('批量创建')}</Checkbox>
{batch && ( {batch && (
<Checkbox checked={multiToSingle} onChange={() => { <Checkbox disabled={isEdit} checked={multiToSingle} onChange={() => {
setMultiToSingle(prev => !prev); setMultiToSingle(prev => !prev);
setInputs(prev => { setInputs(prev => {
const newInputs = { ...prev }; const newInputs = { ...prev };
...@@ -702,35 +706,26 @@ const EditChannel = (props) => { ...@@ -702,35 +706,26 @@ const EditChannel = (props) => {
) : ( ) : (
<> <>
{inputs.type === 41 ? ( {inputs.type === 41 ? (
<Form.TextArea <Form.Upload
field='key' field='vertex_files'
label={t('密钥')} label={t('密钥文件 (.json)')}
placeholder={ accept='.json'
'{\n' + draggable
' "type": "service_account",\n' + dragIcon={<IconBolt />}
' "project_id": "abc-bcd-123-456",\n' + dragMainText={t('点击上传文件或拖拽文件到这里')}
' "private_key_id": "123xxxxx456",\n' + dragSubText={t('仅支持 JSON 文件')}
' "private_key": "-----BEGIN PRIVATE KEY-----xxxx\n' + style={{ marginTop: 10 }}
' "client_email": "xxx@developer.gserviceaccount.com",\n' + uploadTrigger='custom'
' "client_id": "111222333",\n' + beforeUpload={() => false}
' "auth_uri": "https://accounts.google.com/o/oauth2/auth",\n' + onChange={handleVertexUploadChange}
' "token_uri": "https://oauth2.googleapis.com/token",\n' + fileList={vertexFileList}
' "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",\n' + rules={isEdit ? [] : [{ required: true, message: t('请上传密钥文件') }]}
' "client_x509_cert_url": "https://xxxxx.gserviceaccount.com",\n' +
' "universe_domain": "googleapis.com"\n' +
'}'
}
rules={isEdit ? [] : [{ required: true, message: t('请输入密钥') }]}
autosize
autoComplete='new-password'
onChange={(value) => handleInputChange('key', value)}
extraText={batchExtra} extraText={batchExtra}
showClear
/> />
) : ( ) : (
<Form.Input <Form.Input
field='key' field='key'
label={t('密钥')} label={isEdit ? t('密钥（编辑模式下，保存的密钥不会显示）') : t('密钥')}
placeholder={t(type2secretPrompt(inputs.type))} placeholder={t(type2secretPrompt(inputs.type))}
rules={isEdit ? [] : [{ required: true, message: t('请输入密钥') }]} rules={isEdit ? [] : [{ required: true, message: t('请输入密钥') }]}
autoComplete='new-password' autoComplete='new-password'
...@@ -743,6 +738,7 @@ const EditChannel = (props) => { ...@@ -743,6 +738,7 @@ const EditChannel = (props) => {
)} )}
{batch && multiToSingle && ( {batch && multiToSingle && (
<>
<Form.Select <Form.Select
field='multi_key_mode' field='multi_key_mode'
label={t('密钥聚合模式')} label={t('密钥聚合模式')}
...@@ -758,6 +754,14 @@ const EditChannel = (props) => { ...@@ -758,6 +754,14 @@ const EditChannel = (props) => {
handleInputChange('multi_key_mode', value); handleInputChange('multi_key_mode', value);
}} }}
/> />
{inputs.multi_key_mode === 'polling' && (
<Banner
type='warning'
description={t('轮询模式必须搭配Redis和内存缓存功能使用，否则性能将大幅降低，并且无法实现轮询功能')}
className='!rounded-lg mt-2'
/>
)}
</>
)} )}
{inputs.type === 18 && ( {inputs.type === 18 && (
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment