92 lines
2.4 KiB
Go
92 lines
2.4 KiB
Go
package gateway
|
|
|
|
import (
|
|
"gpt-load/internal/affinity"
|
|
"gpt-load/internal/protocol"
|
|
"gpt-load/internal/scheduler"
|
|
"gpt-load/internal/state"
|
|
"gpt-load/internal/telemetry"
|
|
)
|
|
|
|
type requestAffinity struct {
|
|
key affinity.Key
|
|
observation affinity.Observation
|
|
preferredCredentialID uint
|
|
continuityKey string
|
|
kind string
|
|
}
|
|
|
|
func (handler *Handler) resolveRequestAffinity(
|
|
snapshot *state.ConfigSnapshot,
|
|
accessKeyID uint,
|
|
clientProtocol protocol.Protocol,
|
|
prefix []byte,
|
|
allowedCredentialRefs map[uint]state.CredentialRef,
|
|
promptCacheKey string,
|
|
) requestAffinity {
|
|
if handler == nil || snapshot == nil {
|
|
return requestAffinity{}
|
|
}
|
|
key := affinity.DeriveKey(
|
|
handler.encryption,
|
|
accessKeyID,
|
|
clientProtocol,
|
|
prefix,
|
|
)
|
|
// 执行层私有 replay scope 仍由提示词派生，不把客户端缓存分组当作会话身份。
|
|
result := requestAffinity{continuityKey: string(key), kind: telemetry.AffinityPromptPrefix}
|
|
if promptCacheKey != "" {
|
|
key = affinity.DerivePromptCacheKey(handler.encryption, accessKeyID, clientProtocol, promptCacheKey)
|
|
result.kind = telemetry.AffinityPromptCacheKey
|
|
}
|
|
if handler.affinityCache == nil ||
|
|
!handler.affinityCache.Configure(
|
|
snapshot.Revision,
|
|
snapshot.Settings.AffinityCapacity,
|
|
snapshot.Settings.AffinityTTL,
|
|
) {
|
|
return result
|
|
}
|
|
if !key.Valid() {
|
|
return result
|
|
}
|
|
result.key = key
|
|
observation := handler.affinityCache.Lookup(key)
|
|
resolved := result
|
|
resolved.observation = observation
|
|
if !observation.Found() {
|
|
return resolved
|
|
}
|
|
target := observation.Target
|
|
group, exists := snapshot.Groups[target.GroupID]
|
|
if !exists || !group.AffinityEnabled {
|
|
return resolved
|
|
}
|
|
ref, allowed := allowedCredentialRefs[target.CredentialID]
|
|
if !allowed || ref.GroupID != target.GroupID ||
|
|
ref.IdentityGeneration != target.IdentityGeneration {
|
|
return resolved
|
|
}
|
|
resolved.preferredCredentialID = target.CredentialID
|
|
return resolved
|
|
}
|
|
|
|
func (handler *Handler) recordAffinitySuccess(
|
|
request requestAffinity,
|
|
selection scheduler.Selection,
|
|
ref state.CredentialRef,
|
|
) {
|
|
if handler == nil || handler.affinityCache == nil || !request.key.Valid() ||
|
|
!selection.Group.AffinityEnabled {
|
|
return
|
|
}
|
|
handler.affinityCache.RecordSuccess(
|
|
request.key,
|
|
request.observation,
|
|
affinity.Target{
|
|
GroupID: selection.GroupID, CredentialID: selection.CredentialID,
|
|
IdentityGeneration: ref.IdentityGeneration,
|
|
},
|
|
)
|
|
}
|