elmprodvpn/selective-vpn-api/app/resolver/live_batch_select.go

package resolver

import "strings"

func ClassifyLiveBatchHost(
	host string,
	cache DomainCacheState,
	cacheSourceForHost func(string) DomainCacheSource,
	wildcards WildcardMatcher,
) (priority int, nxHeavy bool) {
	h := strings.TrimSpace(strings.ToLower(host))
	if h == "" {
		return 2, false
	}
	if wildcards.IsExact(h) {
		return 1, false
	}
	source := cacheSourceForHost(h)
	rec, ok := cache.Domains[h]
	if !ok {
		return 2, false
	}
	entry := GetCacheEntryBySource(rec, source)
	if entry == nil {
		return 2, false
	}
	stored := NormalizeCacheIPs(entry.IPs)
	state := NormalizeDomainState(entry.State, entry.Score)
	errKind, hasErr := NormalizeCacheErrorKind(entry.LastErrorKind)
	nxHeavy = hasErr && errKind == DNSErrorNXDomain && (state == DomainStateQuarantine || state == DomainStateHardQuar || entry.Score <= -10)

	switch {
	case len(stored) > 0:
		return 1, false
	case state == DomainStateActive || state == DomainStateStable || state == DomainStateSuspect:
		return 1, false
	case nxHeavy:
		return 3, true
	default:
		return 2, false
	}
}

func SplitLiveBatchCandidates(
	candidates []string,
	cache DomainCacheState,
	cacheSourceForHost func(string) DomainCacheSource,
	wildcards WildcardMatcher,
) (p1, p2, p3 []string, nxHeavyTotal int) {
	for _, host := range candidates {
		h := strings.TrimSpace(strings.ToLower(host))
		if h == "" {
			continue
		}
		prio, nxHeavy := ClassifyLiveBatchHost(h, cache, cacheSourceForHost, wildcards)
		switch prio {
		case 1:
			p1 = append(p1, h)
		case 3:
			nxHeavyTotal++
			p3 = append(p3, h)
		case 2:
			p2 = append(p2, h)
		default:
			if nxHeavy {
				nxHeavyTotal++
				p3 = append(p3, h)
			} else {
				p2 = append(p2, h)
			}
		}
	}
	return p1, p2, p3, nxHeavyTotal
}

func PickAdaptiveLiveBatch(
	candidates []string,
	target int,
	now int,
	nxHeavyPct int,
	cache DomainCacheState,
	cacheSourceForHost func(string) DomainCacheSource,
	wildcards WildcardMatcher,
) ([]string, int, int, int, int, int) {
	if len(candidates) == 0 {
		return nil, 0, 0, 0, 0, 0
	}
	if target <= 0 {
		p1, p2, p3, nxTotal := SplitLiveBatchCandidates(candidates, cache, cacheSourceForHost, wildcards)
		return append([]string(nil), candidates...), len(p1), len(p2), len(p3), nxTotal, 0
	}
	if target > len(candidates) {
		target = len(candidates)
	}
	if nxHeavyPct < 0 {
		nxHeavyPct = 0
	}
	if nxHeavyPct > 100 {
		nxHeavyPct = 100
	}

	start := now % len(candidates)
	if start < 0 {
		start = 0
	}
	rotated := make([]string, 0, len(candidates))
	for i := 0; i < len(candidates); i++ {
		idx := (start + i) % len(candidates)
		rotated = append(rotated, candidates[idx])
	}
	p1, p2, p3, nxTotal := SplitLiveBatchCandidates(rotated, cache, cacheSourceForHost, wildcards)
	out := make([]string, 0, target)
	selectedP1 := 0
	selectedP2 := 0
	selectedP3 := 0

	take := func(src []string, n int) ([]string, int) {
		if n <= 0 || len(src) == 0 {
			return src, 0
		}
		if n > len(src) {
			n = len(src)
		}
		out = append(out, src[:n]...)
		return src[n:], n
	}

	remain := target
	var took int
	p1, took = take(p1, remain)
	selectedP1 += took
	remain = target - len(out)
	p2, took = take(p2, remain)
	selectedP2 += took
	remain = target - len(out)

	p3Cap := (target * nxHeavyPct) / 100
	if nxHeavyPct > 0 && p3Cap == 0 {
		p3Cap = 1
	}
	if len(out) == 0 && len(p3) > 0 && p3Cap == 0 {
		p3Cap = 1
	}
	if p3Cap > remain {
		p3Cap = remain
	}
	p3, took = take(p3, p3Cap)
	selectedP3 += took

	if len(out) == 0 && len(p3) > 0 && target > 0 {
		remain = target - len(out)
		p3, took = take(p3, remain)
		selectedP3 += took
	}

	nxSkipped := nxTotal - selectedP3
	if nxSkipped < 0 {
		nxSkipped = 0
	}
	return out, selectedP1, selectedP2, selectedP3, nxTotal, nxSkipped
}