selector.go

package registry

import (
	"errors"
	"log/slog"
	"sync"
	"time"

	"github.com/we7coreteam/registry-go-sdk/client"
	"github.com/we7coreteam/registry-go-sdk/client/auth"
	"github.com/we7coreteam/registry-go-sdk/types"
)

var DefaultServersCheckTimeout = 1500 * time.Millisecond
var DefaultChallengeCacheTTL = time.Hour

var globalChallengeCache = &challengeCache{
	entries: map[string]time.Time{},
}

type challengeCache struct {
	mu      sync.Mutex
	entries map[string]time.Time
}

func (c *challengeCache) block(host string, ttl time.Duration) {
	if host == "" || ttl <= 0 {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	c.entries[host] = time.Now().Add(ttl)
}

func (c *challengeCache) blocked(host string) bool {
	if host == "" {
		return false
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	expiresAt, ok := c.entries[host]
	if !ok {
		return false
	}
	if time.Now().After(expiresAt) {
		delete(c.entries, host)
		return false
	}
	return true
}

func getClientForCheck(server Server) client.Client {
	return client.NewClientWithAuthorizerAndTimeout(
		server.urlInfo.Scheme+"://"+server.urlInfo.Host,
		auth.NewAuthorizer(server.Credential.AccessKey, server.Credential.AccessSecret, true, DefaultServersCheckTimeout, defaultReqInterceptors...),
		true,
		DefaultServersCheckTimeout,
		defaultReqInterceptors...,
	)
}

func GetAvailableServersChan(servers []Server, clientFilters []func(client2 client.Client) (bool, error)) <-chan Server {
	resultsCh := make(chan Server, len(servers))
	if len(servers) == 0 {
		close(resultsCh)
		return resultsCh
	}

	filterNum := len(clientFilters)
	checkTime := time.Now()
	var wg sync.WaitGroup

	wg.Add(len(servers))

	for _, server := range servers {
		go func(tmpServer Server) {
			defer wg.Done()
			if globalChallengeCache.blocked(tmpServer.urlInfo.Host) {
				resultsCh <- tmpServer
				slog.Info("registry selector use challenge cached server", "server", tmpServer)
				return
			}

			cli := getClientForCheck(tmpServer)

			start := time.Now()
			var err error

			successFilterNum := 0
			for _, filter := range clientFilters {
				exists, err1 := filter(cli)
				if !exists {
					err = err1
					break
				}
				successFilterNum += 1
			}

			latency := time.Since(start)
			if err == nil && successFilterNum == filterNum {
				resultsCh <- tmpServer
			} else if errors.Is(err, types.ChallengeCodeErr) {
				globalChallengeCache.block(tmpServer.urlInfo.Host, DefaultChallengeCacheTTL)
				resultsCh <- tmpServer
			}

			slog.Info("registry selector check", "server", tmpServer, "checkTime", checkTime, "latency", latency, "filterNum", filterNum, "successFilterNum", successFilterNum, "err", err)
		}(server)
	}

	go func() {
		wg.Wait()
		close(resultsCh)
	}()

	return resultsCh
}