提升通过域名查找服务的性能，轻松支持海量域名

pkg/serverconfigs/server_address_group.go

@@ -1,19 +1,71 @@
 package serverconfigs
 
-import "strings"
+import (
+	"github.com/TeaOSLab/EdgeCommon/pkg/configutils"
+	"strings"
+	"sync"
+)
 
 type ServerAddressGroup struct {
 	fullAddr string
-	Servers  []*ServerConfig
+	servers  []*ServerConfig
+
+	// 域名和服务映射
+	strictDomainMap map[string]map[string]*ServerConfig // domain[:2] => {domain => *ServerConfig}
+	fuzzyDomainMap  map[string]*ServerConfig            // special domain => *ServerConfig
+
+	cacheLocker       sync.RWMutex
+	cacheDomainMap    map[string]map[string]*ServerConfig // domain[:2] => {domain => *ServerConfig}
+	countCacheDomains int
+
+	// 支持CNAME的服务
+	cnameDomainMap map[string]map[string]*ServerConfig // domain[:2] => {domain => *ServerConfig}
+
+	// 第一个TLS Server
+	firstTLSServer *ServerConfig
 }
 
 func NewServerAddressGroup(fullAddr string) *ServerAddressGroup {
-	return &ServerAddressGroup{fullAddr: fullAddr}
+	return &ServerAddressGroup{
+		fullAddr:        fullAddr,
+		strictDomainMap: map[string]map[string]*ServerConfig{},
+		fuzzyDomainMap:  map[string]*ServerConfig{},
+		cacheDomainMap:  map[string]map[string]*ServerConfig{},
+		cnameDomainMap:  map[string]map[string]*ServerConfig{},
+	}
 }
 
 // Add 添加服务
 func (this *ServerAddressGroup) Add(server *ServerConfig) {
-	this.Servers = append(this.Servers, server)
+	for _, serverName := range server.AllStrictNames() {
+		var prefix = this.domainPrefix(serverName)
+		domainsMap, ok := this.strictDomainMap[prefix]
+		if ok {
+			domainsMap[serverName] = server
+		} else {
+			this.strictDomainMap[prefix] = map[string]*ServerConfig{serverName: server}
+		}
+
+		// CNAME
+		if server.SupportCNAME {
+			cnameDomainsMap, ok := this.cnameDomainMap[prefix]
+			if ok {
+				cnameDomainsMap[serverName] = server
+			} else {
+				this.cnameDomainMap[prefix] = map[string]*ServerConfig{serverName: server}
+			}
+		}
+	}
+	for _, serverName := range server.AllFuzzyNames() {
+		this.fuzzyDomainMap[serverName] = server
+	}
+
+	this.servers = append(this.servers, server)
+
+	// 第一个TLS Server
+	if this.firstTLSServer == nil && server.SSLPolicy() != nil && server.SSLPolicy().IsOn {
+		this.firstTLSServer = server
+	}
 }
 
 // FullAddr 获取完整的地址
@@ -40,6 +92,11 @@ func (this *ServerAddressGroup) Addr() string {
 	return strings.TrimPrefix(this.fullAddr, protocol.String()+"://")
 }
 
+// Servers 读取所有服务
+func (this *ServerAddressGroup) Servers() []*ServerConfig {
+	return this.servers
+}
+
 // IsHTTP 判断当前分组是否为HTTP
 func (this *ServerAddressGroup) IsHTTP() bool {
 	p := this.Protocol()
@@ -78,8 +135,82 @@ func (this *ServerAddressGroup) IsUDP() bool {
 
 // FirstServer 获取第一个Server
 func (this *ServerAddressGroup) FirstServer() *ServerConfig {
-	if len(this.Servers) > 0 {
-		return this.Servers[0]
+	if len(this.servers) > 0 {
+		return this.servers[0]
 	}
 	return nil
 }
+
+// FirstTLSServer 获取第一个TLS Server
+func (this *ServerAddressGroup) FirstTLSServer() *ServerConfig {
+	return this.firstTLSServer
+}
+
+// MatchServerName 使用域名查找服务
+func (this *ServerAddressGroup) MatchServerName(serverName string) *ServerConfig {
+	var prefix = this.domainPrefix(serverName)
+
+	// 试图从缓存中读取
+	this.cacheLocker.RLock()
+	if len(this.cacheDomainMap) > 0 {
+		domainMap, ok := this.cacheDomainMap[prefix]
+		if ok {
+			server, ok := domainMap[serverName]
+			if ok {
+				return server
+			}
+		}
+	}
+	this.cacheLocker.RUnlock()
+
+	domainMap, ok := this.strictDomainMap[prefix]
+	if ok {
+		server, ok := domainMap[serverName]
+		if ok {
+			return server
+		}
+	}
+	for pattern, server := range this.fuzzyDomainMap {
+		if configutils.MatchDomain(pattern, serverName) {
+			// 加入到缓存
+			this.cacheLocker.Lock()
+
+			// 限制缓存的最大尺寸，防止内存耗尽
+			if this.countCacheDomains < 1_000_000 {
+				domainMap, ok := this.cacheDomainMap[prefix]
+				if ok {
+					domainMap[serverName] = server
+				} else {
+					this.cacheDomainMap[prefix] = map[string]*ServerConfig{serverName: server}
+				}
+				this.countCacheDomains++
+			}
+			this.cacheLocker.Unlock()
+
+			return server
+		}
+	}
+	return nil
+}
+
+// MatchServerCNAME 使用CNAME查找服务
+func (this *ServerAddressGroup) MatchServerCNAME(serverName string) *ServerConfig {
+	var prefix = this.domainPrefix(serverName)
+
+	domainMap, ok := this.cnameDomainMap[prefix]
+	if ok {
+		server, ok := domainMap[serverName]
+		if ok {
+			return server
+		}
+	}
+
+	return nil
+}
+
+func (this *ServerAddressGroup) domainPrefix(domain string) string {
+	if len(domain) < 2 {
+		return domain
+	}
+	return domain[:2]
+}