lua: refactor to simplify state creation and pooled script execution

app/dns/lua.go

@@ -157,7 +157,7 @@ func (s *DNS) CallLuaHook(L *lua.LState, ctx context.Context, domain string, opt
 	}()
 	fn := L.GetGlobal("HandleDNSQuery")
 	if fn.Type() != lua.LTFunction {
-		return nil, 0, errors.New("DNS script must define HandleDNSQuery(domain, ipv4, ipv6, fake)")
+		return nil, 0, errors.New("DNS script must define HandleDNSQuery(...)")
 	}
 	if err := L.CallByParam(lua.P{Fn: fn, NRet: 3, Protect: true},
 		lua.LString(strings.ToLower(domain)), lua.LBool(option.IPv4Enable),

app/dns/script.go

@@ -26,23 +26,19 @@ func newScriptEngine(path string, server *DNS) (*scriptEngine, error) {
 		return nil, err
 	}
 	e := &scriptEngine{dns: server}
-	e.pool, err = luamgr.NewPool(server.ctx, func(poolCtx context.Context) (*lua.LState, error) {
-		initCtx, cancel := context.WithTimeout(poolCtx, scriptExecutionTimeout)
-		defer cancel()
-		L, err := program.NewState(initCtx, func(L *lua.LState) {
+	e.pool, err = luamgr.NewPool(server.ctx, program.NewStateFactory(
+		scriptExecutionTimeout,
+		func(L *lua.LState) {
 			geodata.RegisterLua(L)
 			log.RegisterLua(L)
 			server.RegisterLua(L)
-		})
-		if err != nil {
-			return nil, err
-		}
-		if L.GetGlobal("HandleDNSQuery").Type() != lua.LTFunction {
-			L.Close()
-			return nil, errors.New("DNS script must define HandleDNSQuery(domain, ipv4, ipv6, fake)")
-		}
-		return L, nil
-	})
+		},
+		func(L *lua.LState) error {
+			if L.GetGlobal("HandleDNSQuery").Type() != lua.LTFunction {
+				return errors.New("DNS script must define HandleDNSQuery(...)")
+			}
+			return nil
+		}))
 	if err != nil {
 		return nil, err
 	}
@@ -54,20 +50,13 @@ func (e *scriptEngine) close() {
 	e.pool.Close()
 }
 
-func (e *scriptEngine) query(domain string, option dns.IPOption) ([]net.IP, uint32, error) {
-	L, err := e.pool.Acquire()
-	if err != nil {
-		return nil, 0, err
-	}
-	reusable := false
-	defer func() {
-		e.pool.Release(L, reusable)
-	}()
-	queryCtx, cancel := context.WithTimeout(e.pool.Context(), scriptExecutionTimeout)
-	defer cancel()
-	ips, ttl, err := e.dns.CallLuaHook(L, queryCtx, domain, option)
-	if err == nil {
-		reusable = true
-	}
-	return ips, ttl, err
+func (e *scriptEngine) query(domain string, option dns.IPOption) (ips []net.IP, ttl uint32, err error) {
+	err = e.pool.WithState(func(L *lua.LState) error {
+		luaCtx, cancel := context.WithTimeout(e.pool.Context(), scriptExecutionTimeout)
+		defer cancel()
+		var luaErr error
+		ips, ttl, luaErr = e.dns.CallLuaHook(L, luaCtx, domain, option)
+		return luaErr
+	})
+	return
 }

app/router/lua.go

@@ -121,9 +121,9 @@ func pushLuaError(L *lua.LState, err error) {
 }
 
 // CallLuaHook invokes HandleRoute in the supplied state.
-func (r *Router) CallLuaHook(L *lua.LState, ctx context.Context, routeCtx routing.Context) (string, string, error) {
+func (r *Router) CallLuaHook(L *lua.LState, luaCtx context.Context, routeCtx routing.Context) (string, string, error) {
 	previous, top := L.Context(), L.GetTop()
-	L.SetContext(ctx)
+	L.SetContext(luaCtx)
 	defer func() {
 		L.SetTop(top)
 		if previous == nil {

app/router/script.go

@@ -27,24 +27,20 @@ func newScriptEngine(path string, router *Router) (*scriptEngine, error) {
 		return nil, err
 	}
 	e := &scriptEngine{router: router}
-	e.pool, err = luamgr.NewPool(router.ctx, func(poolCtx context.Context) (*lua.LState, error) {
-		initCtx, cancel := context.WithTimeout(poolCtx, scriptExecutionTimeout)
-		defer cancel()
-		L, err := program.NewState(initCtx, func(L *lua.LState) {
+	e.pool, err = luamgr.NewPool(router.ctx, program.NewStateFactory(
+		scriptExecutionTimeout,
+		func(L *lua.LState) {
 			geodata.RegisterLua(L)
 			log.RegisterLua(L)
 			router.RegisterLua(L)
 			dns.RegisterLua(L, router.dns)
-		})
-		if err != nil {
-			return nil, err
-		}
-		if L.GetGlobal("HandleRoute").Type() != lua.LTFunction {
-			L.Close()
-			return nil, errors.New("routing script must define HandleRoute(...)")
-		}
-		return L, nil
-	})
+		},
+		func(L *lua.LState) error {
+			if L.GetGlobal("HandleRoute").Type() != lua.LTFunction {
+				return errors.New("routing script must define HandleRoute(...)")
+			}
+			return nil
+		}))
 	if err != nil {
 		return nil, err
 	}
@@ -57,21 +53,17 @@ func (e *scriptEngine) close() {
 }
 
 func (e *scriptEngine) pickRoute(ctx routing.Context) (routing.Route, error) {
-	L, err := e.pool.Acquire()
+	var tag, ruleTag string
+	err := e.pool.WithState(func(L *lua.LState) error {
+		luaCtx, cancel := context.WithTimeout(e.pool.Context(), scriptExecutionTimeout)
+		defer cancel()
+		var luaErr error
+		tag, ruleTag, luaErr = e.router.CallLuaHook(L, luaCtx, ctx)
+		return luaErr
+	})
 	if err != nil {
 		return nil, err
 	}
-	reusable := false
-	defer func() {
-		e.pool.Release(L, reusable)
-	}()
-	callCtx, cancel := context.WithTimeout(e.pool.Context(), scriptExecutionTimeout)
-	defer cancel()
-	tag, ruleTag, err := e.router.CallLuaHook(L, callCtx, ctx)
-	if err != nil {
-		return nil, err
-	}
-	reusable = true
 	if tag == "" {
 		return nil, common.ErrNoClue
 	}