DomainMatcher: Reduce runtime memory usage and startup peak memory on iOS (#5924)

https://github.com/XTLS/Xray-core/pull/5814#issuecomment-4231071433

Closes https://github.com/XTLS/Xray-core/issues/4422

common/geodata/domain_matcher.go

@@ -2,7 +2,9 @@ package geodata
 
 import (
 	"context"
+	"runtime"
 	"strings"
+	"sync"
 
 	"github.com/xtls/xray-core/common/errors"
 	"github.com/xtls/xray-core/common/geodata/strmatcher"
@@ -13,7 +15,14 @@ type DomainMatcher interface {
 	MatchAny(input string) bool
 }
 
-func buildDomainMatcher(rules []*DomainRule) (DomainMatcher, error) {
+type DomainMatcherFactory interface {
+	BuildMatcher(rules []*DomainRule) (DomainMatcher, error)
+}
+
+type MphDomainMatcherFactory struct{}
+
+// BuildMatcher implements DomainMatcherFactory.
+func (f *MphDomainMatcherFactory) BuildMatcher(rules []*DomainRule) (DomainMatcher, error) {
 	g := strmatcher.NewMphValueMatcher()
 	for i, r := range rules {
 		switch v := r.Value.(type) {
@@ -47,6 +56,105 @@ func buildDomainMatcher(rules []*DomainRule) (DomainMatcher, error) {
 	return g, nil
 }
 
+type CompactDomainMatcherFactory struct {
+	sync.Mutex
+	shared map[string]strmatcher.MatcherGroup // TODO: cleanup
+}
+
+func (f *CompactDomainMatcherFactory) getOrCreateFrom(rule *GeoSiteRule) (strmatcher.MatcherGroup, error) {
+	key := rule.File + ":" + rule.Code + "@" + rule.Attrs
+
+	f.Lock()
+	defer f.Unlock()
+
+	if m := f.shared[key]; m != nil {
+		return m, nil
+	}
+
+	g := strmatcher.NewLinearValueMatcher()
+	domains, err := loadSiteWithAttrs(rule.File, rule.Code, rule.Attrs)
+	if err != nil {
+		return nil, err
+	}
+	for i, d := range domains {
+		domains[i] = nil // peak mem
+		m, err := parseDomain(d)
+		if err != nil {
+			errors.LogError(context.Background(), "ignore invalid geosite entry in ", rule.File, ":", rule.Code, " at index ", i, ", ", err)
+			continue
+		}
+		g.Add(m, 0)
+	}
+	f.shared[key] = g
+	return g, err
+}
+
+// BuildMatcher implements DomainMatcherFactory.
+func (f *CompactDomainMatcherFactory) BuildMatcher(rules []*DomainRule) (DomainMatcher, error) {
+	compact := &CompactDomainMatcher{
+		matchers: make([]strmatcher.MatcherGroup, 0, len(rules)),
+		values:   make([]uint32, 0, len(rules)),
+	}
+	custom := strmatcher.NewLinearValueMatcher()
+	var idx uint32
+	for _, r := range rules {
+		switch v := r.Value.(type) {
+		case *DomainRule_Custom:
+			m, err := parseDomain(v.Custom)
+			if err != nil {
+				return nil, err
+			}
+			custom.Add(m, 0)
+		case *DomainRule_Geosite:
+			m, err := f.getOrCreateFrom(v.Geosite)
+			if err != nil {
+				return nil, err
+			}
+			compact.matchers = append(compact.matchers, m)
+			compact.values = append(compact.values, idx)
+			idx++
+		default:
+			panic("unknown domain rule type")
+		}
+	}
+	if len(compact.matchers) != len(rules) {
+		compact.matchers = append(compact.matchers, custom)
+		compact.values = append(compact.values, idx+1)
+	}
+	return compact, nil
+}
+
+type CompactDomainMatcher struct {
+	matchers []strmatcher.MatcherGroup
+	values   []uint32
+}
+
+func (c *CompactDomainMatcher) Add(matcher strmatcher.MatcherGroup, value uint32) {
+	c.matchers = append(c.matchers, matcher)
+	c.values = append(c.values, value)
+}
+
+// Match implements DomainMatcher.
+func (c *CompactDomainMatcher) Match(input string) []uint32 {
+	result := make([]uint32, 0)
+	for i, m := range c.matchers {
+		if m.MatchAny(input) {
+			result = append(result, c.values[i])
+		}
+	}
+	return result
+}
+
+// MatchAny implements DomainMatcher.
+func (c *CompactDomainMatcher) MatchAny(input string) bool {
+	for _, m := range c.matchers {
+		if m.MatchAny(input) {
+			return true
+		}
+	}
+	return false
+}
+
 func parseDomain(d *Domain) (strmatcher.Matcher, error) {
 	if d == nil {
 		return nil, errors.New("domain must not be nil")
@@ -64,3 +172,12 @@ func parseDomain(d *Domain) (strmatcher.Matcher, error) {
 		return nil, errors.New("unknown domain type: ", d.Type)
 	}
 }
+
+func newDomainMatcherFactory() DomainMatcherFactory {
+	switch runtime.GOOS {
+	case "ios":
+		return &CompactDomainMatcherFactory{shared: make(map[string]strmatcher.MatcherGroup)}
+	default:
+		return &MphDomainMatcherFactory{}
+	}
+}

common/geodata/domain_registry.go

@@ -1,13 +1,17 @@
 package geodata
 
-type DomainRegistry struct{}
+type DomainRegistry struct {
+	factory DomainMatcherFactory
+}
 
 func (r *DomainRegistry) BuildDomainMatcher(rules []*DomainRule) (DomainMatcher, error) {
-	return buildDomainMatcher(rules)
+	return r.factory.BuildMatcher(rules)
 }
 
 func newDomainRegistry() *DomainRegistry {
-	return &DomainRegistry{}
+	return &DomainRegistry{
+		factory: newDomainMatcherFactory(),
+	}
 }
 
 var DomainReg = newDomainRegistry()