grok2api / backend /internal /application /gateway /selector_segmented.go
fromozuzhouzzz
Deploy grok2api v3.0.11 to HF Spaces
f1dd159
Raw
History Blame Contribute Delete
2.63 kB
package gateway
import (
"hash/fnv"
"sync/atomic"
"github.com/chenyme/grok2api/backend/internal/domain/account"
"github.com/chenyme/grok2api/backend/internal/pkg/perfmetrics"
)
const segmentedSelectorShards = 1024
type segmentedSelectorConfig struct {
enabled bool
minCandidates int
windowSize int
}
type segmentedSelectorState struct {
activeCursors [segmentedSelectorShards]atomic.Uint64
}
type segmentedSelectorCohort struct {
supportsModel bool
capabilityKnown bool
preferFreeBuild bool
tier int
priority int
billingFresh bool
}
type selectorLeaseObservation struct {
provider account.Provider
stage string
upstreamStarted atomic.Bool
completed atomic.Bool
}
func normalizeSegmentedSelectorConfig(value segmentedSelectorConfig) segmentedSelectorConfig {
if value.minCandidates < 100 {
value.minCandidates = 3000
}
if value.windowSize < 8 || value.windowSize > value.minCandidates {
value.windowSize = 64
}
return value
}
func segmentedSelectorShard(provider account.Provider, upstreamModel, quotaMode string) uint64 {
hash := fnv.New64a()
_, _ = hash.Write([]byte(provider))
_, _ = hash.Write([]byte{0})
_, _ = hash.Write([]byte(upstreamModel))
_, _ = hash.Write([]byte{0})
_, _ = hash.Write([]byte(quotaMode))
return hash.Sum64() % segmentedSelectorShards
}
func (o *selectorLeaseObservation) completeRelease() {
if o == nil {
return
}
if o.upstreamStarted.Load() {
o.record("failed")
return
}
o.record("skipped")
}
func (o *selectorLeaseObservation) complete(success bool) {
if o == nil {
return
}
outcome := "failed"
if success {
outcome = "success"
}
o.record(outcome)
}
func (o *selectorLeaseObservation) record(outcome string) {
if o == nil || !o.completed.CompareAndSwap(false, true) {
return
}
perfmetrics.Default.Inc("selector_segmented_active_upstream_total", perfmetrics.Labels{
Subsystem: "selector", Operation: "segmented_active_upstream", Provider: string(o.provider),
Stage: o.stage, Outcome: outcome,
})
}
func segmentedSelectorCohortBetter(left, right segmentedSelectorCohort) bool {
if left.supportsModel != right.supportsModel {
return left.supportsModel
}
if left.capabilityKnown != right.capabilityKnown {
return left.capabilityKnown
}
if left.preferFreeBuild != right.preferFreeBuild {
return left.preferFreeBuild
}
if left.tier != right.tier {
return left.tier < right.tier
}
if left.priority != right.priority {
return left.priority > right.priority
}
if left.billingFresh != right.billingFresh {
return left.billingFresh
}
return false
}