test: rework and calibrate Lua DNS and routing benchmarks

2026-10-05 13:28:13 +03:00 · 2026-10-05 15:16:01 +08:00
parent 8d80e9924f
commit d86ad7c501
4 changed files with 341 additions and 131 deletions
@@ -0,0 +1,118 @@
+package dns
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/xtls/xray-core/common/net"
+	featureDNS "github.com/xtls/xray-core/features/dns"
+	lua "github.com/yuin/gopher-lua"
+)
+
+// BenchmarkLuaDNSHook isolates scalar argument bridging and a fixed return.
+// It excludes upstream queries, result decoding, and state pool management.
+func BenchmarkLuaDNSHook(b *testing.B) {
+	L := lua.NewState()
+	b.Cleanup(L.Close)
+	if err := L.DoString(`
+function HandleDNSQuery(domain, ipv4, ipv6, fake)
+    return true
+end
+`); err != nil {
+		b.Fatal(err)
+	}
+	L.SetContext(context.Background())
+	option := featureDNS.IPOption{IPv4Enable: true}
+	if err := callLuaQuery(L, "example.com", option); err != nil {
+		b.Fatal(err)
+	}
+	if L.Get(-3) != lua.LTrue {
+		b.Fatal("hook did not return true")
+	}
+	L.Pop(3)
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		if err := callLuaQuery(L, "example.com", option); err != nil {
+			b.Fatal(err)
+		}
+		L.Pop(3)
+	}
+}
+
+// BenchmarkLuaDNSQuery queries the same preselected, in-memory upstream.
+// client_query compares Client.QueryIP to a preloaded server:Query hook.
+// script_query additionally measures production pool and timeout management.
+// These cases do not measure DNS.LookupIP server selection or network latency.
+func BenchmarkLuaDNSQuery(b *testing.B) {
+	ctx := context.Background()
+	option := featureDNS.IPOption{IPv4Enable: true}
+	ip := net.ParseIP("127.0.0.1")
+	upstream := &benchmarkLuaNameServer{ips: []net.IP{ip}}
+	client := &Client{server: upstream, ipOption: &option, timeoutMs: time.Second}
+	server := &DNS{ctx: ctx, clients: []*Client{client}}
+	const script = `
+local server = require("xray.dns").Servers[1]
+function HandleDNSQuery(domain, ipv4, ipv6, fake)
+    return server:Query(domain, ipv4, ipv6, fake)
+end
+`
+	L := lua.NewState()
+	b.Cleanup(L.Close)
+	server.registerLua(L)
+	if err := L.DoString(script); err != nil {
+		b.Fatal(err)
+	}
+	L.SetContext(ctx)
+
+	path := filepath.Join(b.TempDir(), "query.lua")
+	if err := os.WriteFile(path, []byte(script), 0o600); err != nil {
+		b.Fatal(err)
+	}
+	engine, err := newScriptEngine(path, server)
+	if err != nil {
+		b.Fatal(err)
+	}
+	b.Cleanup(engine.close)
+	for _, bench := range []struct {
+		name  string
+		query func() ([]net.IP, uint32, error)
+	}{
+		{"client_query/native", func() ([]net.IP, uint32, error) {
+			return client.QueryIP(ctx, "example.com", option)
+		}},
+		{"client_query/lua", func() ([]net.IP, uint32, error) {
+			if err := callLuaQuery(L, "example.com", option); err != nil {
+				return nil, 0, err
+			}
+			ips, ttl, err := readLuaQueryResult(L)
+			L.Pop(3)
+			return ips, ttl, err
+		}},
+		{"script_query/lua", func() ([]net.IP, uint32, error) {
+			return engine.query("example.com", option)
+		}},
+	} {
+		b.Run(bench.name, func(b *testing.B) {
+			ips, ttl, err := bench.query()
+			if err != nil || ttl != 60 || len(ips) != 1 || !ips[0].Equal(ip) {
+				b.Fatalf("query() = %v, TTL %d, %v; want %v, TTL 60", ips, ttl, err, ip)
+			}
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				ips, ttl, err = bench.query()
+				if err != nil {
+					b.Fatal(err)
+				}
+			}
+			b.StopTimer()
+			if ttl != 60 || len(ips) != 1 || !ips[0].Equal(ip) {
+				b.Fatalf("query() = %v, TTL %d; want %v, TTL 60", ips, ttl, ip)
+			}
+		})
+	}
+}
@@ -270,59 +270,3 @@ func (*benchmarkLuaNameServer) IsDisableCache() bool { return true }
 func (s *benchmarkLuaNameServer) QueryIP(context.Context, string, featureDNS.IPOption) ([]net.IP, uint32, error) {
 	return s.ips, 60, nil
 }
-
-// BenchmarkLuaDNSQuery measures a preloaded DNS script using server:Query.
-// The direct case measures the same DNS client without Lua.
-func BenchmarkLuaDNSQuery(b *testing.B) {
-	option := featureDNS.IPOption{IPv4Enable: true}
-	ip := net.ParseIP("127.0.0.1")
-	upstream := &benchmarkLuaNameServer{ips: []net.IP{ip}}
-	client := &Client{server: upstream, ipOption: &option, timeoutMs: time.Second}
-	server := &DNS{clients: []*Client{client}}
-	L := lua.NewState()
-	defer L.Close()
-	server.registerLua(L)
-	if err := L.DoString(`
-local server = require("xray.dns").Servers[1]
-function HandleDNSQuery(domain, ipv4, ipv6, fake)
-    return server:Query(domain, ipv4, ipv6, fake)
-end
-`); err != nil {
-		b.Fatal(err)
-	}
-
-	ctx := context.Background()
-	L.SetContext(ctx)
-	for _, bench := range []struct {
-		name  string
-		query func() ([]net.IP, uint32, error)
-	}{
-		{"direct", func() ([]net.IP, uint32, error) { return client.QueryIP(ctx, "example.com", option) }},
-		{"lua_script", func() ([]net.IP, uint32, error) {
-			if err := callLuaQuery(L, "example.com", option); err != nil {
-				return nil, 0, err
-			}
-			ips, ttl, err := readLuaQueryResult(L)
-			L.Pop(3)
-			return ips, ttl, err
-		}},
-	} {
-		b.Run(bench.name, func(b *testing.B) {
-			b.ReportAllocs()
-			b.ResetTimer()
-			var ips []net.IP
-			var ttl uint32
-			var err error
-			for i := 0; i < b.N; i++ {
-				ips, ttl, err = bench.query()
-				if err != nil {
-					b.Fatal(err)
-				}
-			}
-			b.StopTimer()
-			if ttl != 60 || len(ips) != 1 || !ips[0].Equal(ip) {
-				b.Fatalf("query() = %v, TTL %d; want %v, TTL 60", ips, ttl, ip)
-			}
-		})
-	}
-}
@@ -0,0 +1,223 @@
+package router
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/xtls/xray-core/common/geodata"
+	"github.com/xtls/xray-core/common/net"
+	"github.com/xtls/xray-core/common/session"
+	"github.com/xtls/xray-core/features/routing"
+	routing_session "github.com/xtls/xray-core/features/routing/session"
+	lua "github.com/yuin/gopher-lua"
+)
+
+func benchmarkRouteContext(target net.Destination) *routing_session.Context {
+	// Use the production context: its IP getters construct a slice per call.
+	// The cached IP slices in luaRouteTestContext would undercount this cost.
+	return &routing_session.Context{
+		Inbound: &session.Inbound{
+			Tag:    "in",
+			Source: net.TCPDestination(net.LocalHostIP, 1234),
+			Local:  net.TCPDestination(net.LocalHostIP, 5678),
+		},
+		Outbound: &session.Outbound{Target: target},
+		Content:  &session.Content{Protocol: "tls"},
+	}
+}
+
+func benchmarkRouteState(b *testing.B, r *Router, script string) *lua.LState {
+	b.Helper()
+	L := lua.NewState()
+	b.Cleanup(L.Close)
+	r.RegisterLua(L)
+	geodata.RegisterLua(L)
+	if err := L.DoString(script); err != nil {
+		b.Fatal(err)
+	}
+	L.SetContext(context.Background())
+	return L
+}
+
+// BenchmarkLuaRouteHook isolates argument bridging and a fixed-return hook.
+// It excludes rules, result decoding, the state pool, and Route construction.
+func BenchmarkLuaRouteHook(b *testing.B) {
+	L := benchmarkRouteState(b, new(Router), `
+function HandleRoute(ctx, inboundTag, sourcePort, targetPort, localPort,
+    targetDomain, network, protocol, user, vlessRoute, skipDNSResolve)
+    return "out", "rule"
+end
+`)
+	ctx := benchmarkRouteContext(net.TCPDestination(net.LocalHostIP, 443))
+	if err := callLuaRoute(L, ctx); err != nil {
+		b.Fatal(err)
+	}
+	outboundTag, ruleTag, err := readLuaRouteResult(L)
+	L.Pop(3)
+	if err != nil || outboundTag != "out" || ruleTag != "rule" {
+		b.Fatalf("hook() = %q, %q, %v", outboundTag, ruleTag, err)
+	}
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		if err := callLuaRoute(L, ctx); err != nil {
+			b.Fatal(err)
+		}
+		L.Pop(3)
+	}
+}
+
+// BenchmarkLuaRoute compares equivalent ordered rules on the same session.
+// rules returns tags only; pick_route uses Router.PickRoute on both sides.
+// All compilation, matcher construction, and pool startup are outside timing.
+func BenchmarkLuaRoute(b *testing.B) {
+	for _, name := range []string{"scalar", "ip", "domain", "domain_32_last"} {
+		b.Run(name, func(b *testing.B) {
+			config, script, ctx, wantTag, wantRule := benchmarkRouteFixture(b, name)
+			native := new(Router)
+			if err := native.Init(context.Background(), config, nil, nil, nil); err != nil {
+				b.Fatal(err)
+			}
+			L := benchmarkRouteState(b, native, script)
+
+			path := filepath.Join(b.TempDir(), "route.lua")
+			if err := os.WriteFile(path, []byte(script), 0o600); err != nil {
+				b.Fatal(err)
+			}
+			scripted := new(Router)
+			if err := scripted.Init(context.Background(), &Config{Script: path}, nil, nil, nil); err != nil {
+				b.Fatal(err)
+			}
+			if err := scripted.Start(); err != nil {
+				b.Fatal(err)
+			}
+			b.Cleanup(func() {
+				if err := scripted.Close(); err != nil {
+					b.Error(err)
+				}
+			})
+
+			for _, bench := range []struct {
+				name  string
+				route func() (string, string, error)
+			}{
+				{"rules/native", func() (string, string, error) {
+					rule, _, err := native.pickRouteInternal(ctx)
+					if err != nil {
+						return "", "", err
+					}
+					tag, err := rule.GetTag()
+					return tag, rule.RuleTag, err
+				}},
+				{"rules/lua", func() (string, string, error) {
+					if err := callLuaRoute(L, ctx); err != nil {
+						return "", "", err
+					}
+					tag, ruleTag, err := readLuaRouteResult(L)
+					L.Pop(3)
+					return tag, ruleTag, err
+				}},
+				{"pick_route/native", func() (string, string, error) {
+					return benchmarkPickRoute(native, ctx)
+				}},
+				{"pick_route/lua", func() (string, string, error) {
+					return benchmarkPickRoute(scripted, ctx)
+				}},
+			} {
+				b.Run(bench.name, func(b *testing.B) {
+					// Validate and warm both paths before measuring steady state.
+					tag, ruleTag, err := bench.route()
+					if err != nil || tag != wantTag || ruleTag != wantRule {
+						b.Fatalf("route() = %q, %q, %v; want %q, %q", tag, ruleTag, err, wantTag, wantRule)
+					}
+					b.ReportAllocs()
+					b.ResetTimer()
+					for i := 0; i < b.N; i++ {
+						tag, ruleTag, err = bench.route()
+						if err != nil {
+							b.Fatal(err)
+						}
+					}
+					b.StopTimer()
+					if tag != wantTag || ruleTag != wantRule {
+						b.Fatalf("route() = %q, %q; want %q, %q", tag, ruleTag, wantTag, wantRule)
+					}
+				})
+			}
+		})
+	}
+}
+
+func benchmarkPickRoute(r *Router, ctx routing.Context) (string, string, error) {
+	route, err := r.PickRoute(ctx)
+	if err != nil {
+		return "", "", err
+	}
+	return route.GetOutboundTag(), route.GetRuleTag(), nil
+}
+
+func benchmarkRouteFixture(b *testing.B, name string) (*Config, string, routing.Context, string, string) {
+	b.Helper()
+	config := new(Config)
+	ctx := benchmarkRouteContext(net.TCPDestination(net.LocalHostIP, 443))
+	prelude := `local router = require("xray.router")
+local geodata = require("xray.geodata")
+`
+	body := `if inboundTag == "in" and network == router.NetworkTCP then return "out", "rule" end`
+	wantTag, wantRule := "out", "rule"
+	if name == "scalar" || name == "ip" {
+		rule := &RoutingRule{
+			TargetTag:  &RoutingRule_Tag{Tag: wantTag},
+			RuleTag:    wantRule,
+			InboundTag: []string{"in"},
+			Networks:   []net.Network{net.Network_TCP},
+		}
+		if name == "ip" {
+			var err error
+			rule.Ip, err = geodata.ParseIPRules([]string{"127.0.0.0/8"})
+			if err != nil {
+				b.Fatal(err)
+			}
+			prelude += `local matcher = geodata.BuildIPMatcher("127.0.0.0/8")` + "\n"
+			body = `if inboundTag == "in" and network == router.NetworkTCP and matcher:AnyMatch(ctx:GetTargetIPs()) then return "out", "rule" end`
+		}
+		config.Rule = []*RoutingRule{rule}
+	} else {
+		count := 1
+		if name == "domain_32_last" {
+			count = 32
+		}
+		var rules strings.Builder
+		rules.WriteString("local rules = {\n")
+		for i := 0; i < count; i++ {
+			domain := fmt.Sprintf("route-%d.example.com", i)
+			tag, ruleTag := fmt.Sprintf("out-%d", i), fmt.Sprintf("rule-%d", i)
+			domains, err := geodata.ParseDomainRules([]string{"full:" + domain}, geodata.Domain_Domain)
+			if err != nil {
+				b.Fatal(err)
+			}
+			config.Rule = append(config.Rule, &RoutingRule{
+				TargetTag: &RoutingRule_Tag{Tag: tag}, RuleTag: ruleTag, Domain: domains,
+			})
+			fmt.Fprintf(&rules, "{geodata.BuildDomainMatcher(%q), %q, %q},\n", "full:"+domain, tag, ruleTag)
+			if i == count-1 {
+				ctx.Outbound.Target = net.TCPDestination(net.DomainAddress(domain), 443)
+				wantTag, wantRule = tag, ruleTag
+			}
+		}
+		rules.WriteString("}\n")
+		prelude += rules.String()
+		body = `for i = 1, #rules do
+        local rule = rules[i]
+        if rule[1]:MatchAny(targetDomain) then return rule[2], rule[3] end
+    end`
+	}
+	script := prelude + `function HandleRoute(ctx, inboundTag, sourcePort, targetPort, localPort,
+    targetDomain, network, protocol, user, vlessRoute, skipDNSResolve)
+    ` + body + "\nend\n"
+	return config, script, ctx, wantTag, wantRule
+}
@@ -271,79 +271,4 @@ func TestFindProcess(t *testing.T) {
 	}
 }

-// BenchmarkLuaRoute measures a preloaded routing script using its context bridge.
-// The direct case runs an equivalent native routing rule.
-func BenchmarkLuaRoute(b *testing.B) {
-	r := new(Router)
-	if err := r.Init(context.Background(), &Config{Rule: []*RoutingRule{{
-		TargetTag:  &RoutingRule_Tag{Tag: "out"},
-		RuleTag:    "rule",
-		InboundTag: []string{"in"},
-		Networks:   []net.Network{net.Network_TCP},
-		Ip: []*geodata.IPRule{{
-			Value: &geodata.IPRule_Custom{Custom: &geodata.CIDRRule{
-				Cidr: &geodata.CIDR{Ip: []byte{127, 0, 0, 0}, Prefix: 8},
-			}},
-		}},
-	}}}, nil, nil, nil); err != nil {
-		b.Fatal(err)
-	}
-	L := lua.NewState()
-	defer L.Close()
-	r.RegisterLua(L)
-	geodata.RegisterLua(L)
-	if err := L.DoString(`
-local router = require("xray.router")
-local matcher = require("xray.geodata").BuildIPMatcher("127.0.0.0/8")
-function HandleRoute(ctx, inboundTag, sourcePort, targetPort, localPort,
-    targetDomain, network, protocol, user, vlessRoute, skipDNSResolve)
-    if inboundTag == "in" and network == router.NetworkTCP and matcher:AnyMatch(ctx:GetTargetIPs()) then
-        return "out", "rule"
-    end
-end
-`); err != nil {
-		b.Fatal(err)
-	}
-
-	L.SetContext(context.Background())
-	ctx := newLuaRouteTestContext()
-	for _, benchmark := range []struct {
-		name  string
-		route func() (string, string, error)
-	}{
-		{"direct", func() (string, string, error) {
-			route, err := r.PickRoute(ctx)
-			if err != nil {
-				return "", "", err
-			}
-			return route.GetOutboundTag(), route.GetRuleTag(), nil
-		}},
-		{"lua_script", func() (string, string, error) {
-			if err := callLuaRoute(L, ctx); err != nil {
-				return "", "", err
-			}
-			outboundTag, ruleTag, err := readLuaRouteResult(L)
-			L.Pop(3)
-			return outboundTag, ruleTag, err
-		}},
-	} {
-		b.Run(benchmark.name, func(b *testing.B) {
-			b.ReportAllocs()
-			b.ResetTimer()
-			var outboundTag, ruleTag string
-			var err error
-			for i := 0; i < b.N; i++ {
-				outboundTag, ruleTag, err = benchmark.route()
-				if err != nil {
-					b.Fatal(err)
-				}
-			}
-			b.StopTimer()
-			if outboundTag != "out" || ruleTag != "rule" {
-				b.Fatalf("route() = %q, %q; want out, rule", outboundTag, ruleTag)
-			}
-		})
-	}
-}
-
 var _ routing.Context = (*luaRouteTestContext)(nil)