mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 01:47:47 +08:00
feat: 代理池均衡分发&批量添加代理节点
This commit is contained in:
wzw
@@ -0,0 +1,128 @@
|
||||
import { describe, expect, it } from 'vitest'
|
||||
|
||||
import {
|
||||
buildPoolProxyDistributionPlan,
|
||||
type PoolProxyDistributionKey,
|
||||
} from '@/features/pool/utils/poolProxyDistribution'
|
||||
|
||||
const nodes = [
|
||||
{ id: 'node-a', name: 'Node A' },
|
||||
{ id: 'node-b', name: 'Node B' },
|
||||
]
|
||||
|
||||
function key(id: string, nodeId?: string | null): PoolProxyDistributionKey {
|
||||
return {
|
||||
key_id: id,
|
||||
key_name: id,
|
||||
proxy: nodeId ? { node_id: nodeId, enabled: true } : null,
|
||||
}
|
||||
}
|
||||
|
||||
function fixedRng(): () => number {
|
||||
return () => 0
|
||||
}
|
||||
|
||||
function assignedIds(plan: ReturnType<typeof buildPoolProxyDistributionPlan>): string[] {
|
||||
return plan.assignments.flatMap(item => item.keys.map(key => key.key_id)).sort()
|
||||
}
|
||||
|
||||
describe('buildPoolProxyDistributionPlan', () => {
|
||||
it('keeps existing selected-node proxy bindings and fills empty capacity', () => {
|
||||
const plan = buildPoolProxyDistributionPlan({
|
||||
mode: 'fill',
|
||||
nodes,
|
||||
rng: fixedRng(),
|
||||
keys: [
|
||||
key('a-1', 'node-a'),
|
||||
key('a-2', 'node-a'),
|
||||
key('b-1', 'node-b'),
|
||||
key('new-1', null),
|
||||
],
|
||||
})
|
||||
|
||||
expect(plan.totalKeys).toBe(4)
|
||||
expect(plan.maxPerNode).toBe(2)
|
||||
expect(plan.retainedCount).toBe(3)
|
||||
expect(plan.changedCount).toBe(1)
|
||||
expect(assignedIds(plan)).toEqual(['a-1', 'a-2', 'b-1', 'new-1'])
|
||||
expect(plan.assignments.map(item => item.keys).map(keys => keys.length).sort()).toEqual([2, 2])
|
||||
})
|
||||
|
||||
it('moves overflowed existing bindings before final assignment', () => {
|
||||
const plan = buildPoolProxyDistributionPlan({
|
||||
mode: 'fill',
|
||||
nodes,
|
||||
rng: fixedRng(),
|
||||
keys: [
|
||||
key('a-1', 'node-a'),
|
||||
key('a-2', 'node-a'),
|
||||
key('a-3', 'node-a'),
|
||||
key('a-4', 'node-a'),
|
||||
key('b-1', 'node-b'),
|
||||
],
|
||||
})
|
||||
|
||||
const nodeA = plan.assignments.find(item => item.nodeId === 'node-a')!
|
||||
const nodeB = plan.assignments.find(item => item.nodeId === 'node-b')!
|
||||
|
||||
expect(plan.maxPerNode).toBe(3)
|
||||
expect(plan.overflowCount).toBe(1)
|
||||
expect(nodeA.keys).toHaveLength(3)
|
||||
expect(nodeB.keys).toHaveLength(2)
|
||||
expect(nodeB.changedKeys).toHaveLength(1)
|
||||
expect(assignedIds(plan)).toEqual(['a-1', 'a-2', 'a-3', 'a-4', 'b-1'])
|
||||
})
|
||||
|
||||
it('reassigns accounts bound to non-selected proxy nodes', () => {
|
||||
const plan = buildPoolProxyDistributionPlan({
|
||||
mode: 'fill',
|
||||
nodes,
|
||||
rng: fixedRng(),
|
||||
keys: [
|
||||
key('outside-1', 'node-c'),
|
||||
key('empty-1', null),
|
||||
],
|
||||
})
|
||||
|
||||
expect(plan.outsideSelectedProxyCount).toBe(1)
|
||||
expect(plan.changedCount).toBe(2)
|
||||
expect(plan.assignments.every(item => item.keys.length === 1)).toBe(true)
|
||||
})
|
||||
|
||||
it('force rewrites all accounts into balanced random targets', () => {
|
||||
const plan = buildPoolProxyDistributionPlan({
|
||||
mode: 'rewrite',
|
||||
nodes,
|
||||
rng: fixedRng(),
|
||||
keys: [
|
||||
key('k-1', 'node-a'),
|
||||
key('k-2', 'node-a'),
|
||||
key('k-3', 'node-b'),
|
||||
key('k-4', null),
|
||||
key('k-5', 'node-c'),
|
||||
],
|
||||
})
|
||||
|
||||
expect(plan.maxPerNode).toBe(3)
|
||||
expect(plan.retainedCount).toBe(0)
|
||||
expect(assignedIds(plan)).toEqual(['k-1', 'k-2', 'k-3', 'k-4', 'k-5'])
|
||||
expect(plan.assignments.map(item => item.keys).map(keys => keys.length).sort()).toEqual([2, 3])
|
||||
expect(plan.assignments.every(item => item.keys.length <= plan.maxPerNode)).toBe(true)
|
||||
})
|
||||
|
||||
it('supports fewer accounts than selected proxy nodes', () => {
|
||||
const plan = buildPoolProxyDistributionPlan({
|
||||
mode: 'rewrite',
|
||||
nodes: [
|
||||
...nodes,
|
||||
{ id: 'node-c', name: 'Node C' },
|
||||
],
|
||||
rng: fixedRng(),
|
||||
keys: [key('only-1', null)],
|
||||
})
|
||||
|
||||
expect(plan.maxPerNode).toBe(1)
|
||||
expect(plan.assignments.map(item => item.keys).map(keys => keys.length).sort()).toEqual([0, 0, 1])
|
||||
expect(assignedIds(plan)).toEqual(['only-1'])
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,266 @@
|
||||
export type PoolProxyDistributionMode = 'fill' | 'rewrite'
|
||||
|
||||
export interface PoolProxyDistributionKey {
|
||||
key_id: string
|
||||
key_name?: string | null
|
||||
proxy?: {
|
||||
node_id?: string | null
|
||||
enabled?: boolean
|
||||
} | null
|
||||
}
|
||||
|
||||
export interface PoolProxyDistributionNode {
|
||||
id: string
|
||||
name?: string | null
|
||||
}
|
||||
|
||||
export interface PoolProxyDistributionAssignment {
|
||||
nodeId: string
|
||||
targetCount: number
|
||||
retainedKeys: PoolProxyDistributionKey[]
|
||||
assignedKeys: PoolProxyDistributionKey[]
|
||||
changedKeys: PoolProxyDistributionKey[]
|
||||
keys: PoolProxyDistributionKey[]
|
||||
}
|
||||
|
||||
export interface PoolProxyDistributionPlan {
|
||||
mode: PoolProxyDistributionMode
|
||||
totalKeys: number
|
||||
nodeCount: number
|
||||
maxPerNode: number
|
||||
assignments: PoolProxyDistributionAssignment[]
|
||||
retainedCount: number
|
||||
changedCount: number
|
||||
outsideSelectedProxyCount: number
|
||||
overflowCount: number
|
||||
}
|
||||
|
||||
export interface PoolProxyDistributionOptions {
|
||||
mode: PoolProxyDistributionMode
|
||||
keys: PoolProxyDistributionKey[]
|
||||
nodes: PoolProxyDistributionNode[]
|
||||
rng?: () => number
|
||||
}
|
||||
|
||||
interface MutableAssignment {
|
||||
nodeId: string
|
||||
targetCount: number
|
||||
retainedKeys: PoolProxyDistributionKey[]
|
||||
assignedKeys: PoolProxyDistributionKey[]
|
||||
}
|
||||
|
||||
export function buildPoolProxyDistributionPlan(
|
||||
options: PoolProxyDistributionOptions,
|
||||
): PoolProxyDistributionPlan {
|
||||
const rng = options.rng ?? Math.random
|
||||
const keys = uniqueKeys(options.keys)
|
||||
const nodeIds = uniqueNodeIds(options.nodes)
|
||||
const selectedNodeSet = new Set(nodeIds)
|
||||
const totalKeys = keys.length
|
||||
const nodeCount = nodeIds.length
|
||||
const maxPerNode = nodeCount > 0 ? Math.ceil(totalKeys / nodeCount) : 0
|
||||
|
||||
if (totalKeys === 0 || nodeCount === 0) {
|
||||
return {
|
||||
mode: options.mode,
|
||||
totalKeys,
|
||||
nodeCount,
|
||||
maxPerNode,
|
||||
assignments: [],
|
||||
retainedCount: 0,
|
||||
changedCount: 0,
|
||||
outsideSelectedProxyCount: 0,
|
||||
overflowCount: 0,
|
||||
}
|
||||
}
|
||||
|
||||
const targetCounts = buildTargetCounts({
|
||||
keys,
|
||||
nodeIds,
|
||||
selectedNodeSet,
|
||||
mode: options.mode,
|
||||
rng,
|
||||
})
|
||||
|
||||
const mutableAssignments = new Map<string, MutableAssignment>()
|
||||
for (const nodeId of nodeIds) {
|
||||
mutableAssignments.set(nodeId, {
|
||||
nodeId,
|
||||
targetCount: targetCounts.get(nodeId) ?? 0,
|
||||
retainedKeys: [],
|
||||
assignedKeys: [],
|
||||
})
|
||||
}
|
||||
|
||||
const pendingKeys: PoolProxyDistributionKey[] = []
|
||||
let outsideSelectedProxyCount = 0
|
||||
let overflowCount = 0
|
||||
|
||||
if (options.mode === 'fill') {
|
||||
const keysByNode = new Map<string, PoolProxyDistributionKey[]>()
|
||||
for (const key of keys) {
|
||||
const nodeId = getKeyProxyNodeId(key)
|
||||
if (nodeId && selectedNodeSet.has(nodeId)) {
|
||||
const grouped = keysByNode.get(nodeId) ?? []
|
||||
grouped.push(key)
|
||||
keysByNode.set(nodeId, grouped)
|
||||
} else {
|
||||
if (nodeId) outsideSelectedProxyCount += 1
|
||||
pendingKeys.push(key)
|
||||
}
|
||||
}
|
||||
|
||||
for (const nodeId of nodeIds) {
|
||||
const assignment = mutableAssignments.get(nodeId)
|
||||
if (!assignment) continue
|
||||
const currentKeys = shuffle(keysByNode.get(nodeId) ?? [], rng)
|
||||
const retainedKeys = currentKeys.slice(0, assignment.targetCount)
|
||||
const overflowKeys = currentKeys.slice(assignment.targetCount)
|
||||
assignment.retainedKeys.push(...retainedKeys)
|
||||
pendingKeys.push(...overflowKeys)
|
||||
overflowCount += overflowKeys.length
|
||||
}
|
||||
} else {
|
||||
pendingKeys.push(...keys)
|
||||
}
|
||||
|
||||
const shuffledPendingKeys = shuffle(pendingKeys, rng)
|
||||
const slots = shuffle(buildOpenSlots(mutableAssignments), rng)
|
||||
for (let index = 0; index < shuffledPendingKeys.length; index += 1) {
|
||||
const nodeId = slots[index]
|
||||
if (!nodeId) break
|
||||
mutableAssignments.get(nodeId)?.assignedKeys.push(shuffledPendingKeys[index])
|
||||
}
|
||||
|
||||
const assignments = nodeIds.map((nodeId) => {
|
||||
const assignment = mutableAssignments.get(nodeId)!
|
||||
const nodeKeys = [...assignment.retainedKeys, ...assignment.assignedKeys]
|
||||
const changedKeys = nodeKeys.filter(key => getKeyProxyNodeId(key) !== nodeId)
|
||||
return {
|
||||
nodeId,
|
||||
targetCount: assignment.targetCount,
|
||||
retainedKeys: assignment.retainedKeys,
|
||||
assignedKeys: assignment.assignedKeys,
|
||||
changedKeys,
|
||||
keys: nodeKeys,
|
||||
}
|
||||
})
|
||||
|
||||
return {
|
||||
mode: options.mode,
|
||||
totalKeys,
|
||||
nodeCount,
|
||||
maxPerNode,
|
||||
assignments,
|
||||
retainedCount: assignments.reduce((sum, item) => sum + item.retainedKeys.length, 0),
|
||||
changedCount: assignments.reduce((sum, item) => sum + item.changedKeys.length, 0),
|
||||
outsideSelectedProxyCount,
|
||||
overflowCount,
|
||||
}
|
||||
}
|
||||
|
||||
function buildTargetCounts(options: {
|
||||
keys: PoolProxyDistributionKey[]
|
||||
nodeIds: string[]
|
||||
selectedNodeSet: Set<string>
|
||||
mode: PoolProxyDistributionMode
|
||||
rng: () => number
|
||||
}): Map<string, number> {
|
||||
const baseCount = Math.floor(options.keys.length / options.nodeIds.length)
|
||||
const extraCount = options.keys.length % options.nodeIds.length
|
||||
const existingCounts = new Map<string, number>()
|
||||
|
||||
if (options.mode === 'fill') {
|
||||
for (const key of options.keys) {
|
||||
const nodeId = getKeyProxyNodeId(key)
|
||||
if (nodeId && options.selectedNodeSet.has(nodeId)) {
|
||||
existingCounts.set(nodeId, (existingCounts.get(nodeId) ?? 0) + 1)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const extraNodeIds = new Set(
|
||||
options.nodeIds
|
||||
.map(nodeId => ({
|
||||
nodeId,
|
||||
existingCount: existingCounts.get(nodeId) ?? 0,
|
||||
rank: options.rng(),
|
||||
}))
|
||||
.sort((left, right) => {
|
||||
if (options.mode === 'fill' && left.existingCount !== right.existingCount) {
|
||||
return right.existingCount - left.existingCount
|
||||
}
|
||||
return left.rank - right.rank
|
||||
})
|
||||
.slice(0, extraCount)
|
||||
.map(item => item.nodeId),
|
||||
)
|
||||
|
||||
return new Map(
|
||||
options.nodeIds.map((nodeId) => [
|
||||
nodeId,
|
||||
baseCount + (extraNodeIds.has(nodeId) ? 1 : 0),
|
||||
]),
|
||||
)
|
||||
}
|
||||
|
||||
function buildOpenSlots(assignments: Map<string, MutableAssignment>): string[] {
|
||||
const slots: string[] = []
|
||||
for (const assignment of assignments.values()) {
|
||||
const openSlotCount = Math.max(
|
||||
assignment.targetCount - assignment.retainedKeys.length - assignment.assignedKeys.length,
|
||||
0,
|
||||
)
|
||||
for (let index = 0; index < openSlotCount; index += 1) {
|
||||
slots.push(assignment.nodeId)
|
||||
}
|
||||
}
|
||||
return slots
|
||||
}
|
||||
|
||||
function getKeyProxyNodeId(key: PoolProxyDistributionKey): string | null {
|
||||
const nodeId = key.proxy?.node_id?.trim()
|
||||
return nodeId || null
|
||||
}
|
||||
|
||||
function uniqueNodeIds(nodes: PoolProxyDistributionNode[]): string[] {
|
||||
const seen = new Set<string>()
|
||||
const ids: string[] = []
|
||||
for (const node of nodes) {
|
||||
const id = node.id.trim()
|
||||
if (!id || seen.has(id)) continue
|
||||
seen.add(id)
|
||||
ids.push(id)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
function uniqueKeys(keys: PoolProxyDistributionKey[]): PoolProxyDistributionKey[] {
|
||||
const seen = new Set<string>()
|
||||
const items: PoolProxyDistributionKey[] = []
|
||||
for (const key of keys) {
|
||||
const id = key.key_id.trim()
|
||||
if (!id || seen.has(id)) continue
|
||||
seen.add(id)
|
||||
items.push(key)
|
||||
}
|
||||
return items
|
||||
}
|
||||
|
||||
function shuffle<T>(items: T[], rng: () => number): T[] {
|
||||
const result = [...items]
|
||||
for (let index = result.length - 1; index > 0; index -= 1) {
|
||||
const swapIndex = Math.floor(clampRandom(rng()) * (index + 1))
|
||||
const current = result[index]
|
||||
result[index] = result[swapIndex]
|
||||
result[swapIndex] = current
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
function clampRandom(value: number): number {
|
||||
if (!Number.isFinite(value)) return 0
|
||||
if (value < 0) return 0
|
||||
if (value >= 1) return 0.999999999
|
||||
return value
|
||||
}
|
||||
Reference in New Issue
Block a user