diff --git a/internal/lua/context.go b/internal/lua/context.go index ecb93e3..56af8cf 100644 --- a/internal/lua/context.go +++ b/internal/lua/context.go @@ -2,6 +2,8 @@ package lua import ( + "sync" + "github.com/valyala/fasthttp" ) @@ -21,14 +23,28 @@ type LuaContext struct { Exited bool } -// NewContext 创建请求上下文 +// luaContextPool LuaContext 对象池 +var luaContextPool = sync.Pool{ + New: func() interface{} { + return &LuaContext{ + Variables: make(map[string]string), + } + }, +} + +// AcquireContext 从池中获取 LuaContext +func AcquireContext(engine *LuaEngine, req *fasthttp.RequestCtx) *LuaContext { + lc := luaContextPool.Get().(*LuaContext) + lc.Engine = engine + lc.RequestCtx = req + lc.Phase = PhaseInit + // Variables 和 OutputBuffer 已在 Release 中重置 + return lc +} + +// NewContext 创建请求上下文(从池中获取) func NewContext(engine *LuaEngine, req *fasthttp.RequestCtx) *LuaContext { - return &LuaContext{ - Engine: engine, - RequestCtx: req, - Variables: make(map[string]string), - Phase: PhaseInit, - } + return AcquireContext(engine, req) } // InitCoroutine 初始化协程 @@ -99,14 +115,24 @@ func (c *LuaContext) Exit(code int) { c.RequestCtx.SetStatusCode(code) } -// Release 释放资源 +// Release 释放资源并放回池中 func (c *LuaContext) Release() { if c.Coroutine != nil { c.Coroutine.Close() c.Coroutine = nil } - c.Variables = nil - c.OutputBuffer = nil + + // 重置所有可变状态,防止请求间污染 + for k := range c.Variables { + delete(c.Variables, k) + } + c.OutputBuffer = c.OutputBuffer[:0] + c.Phase = PhaseInit + c.Exited = false + c.Engine = nil + c.RequestCtx = nil + + luaContextPool.Put(c) } // FlushOutput 刷新输出到响应 diff --git a/internal/lua/lua_bench_test.go b/internal/lua/lua_bench_test.go index 261a0f9..a91f59a 100644 --- a/internal/lua/lua_bench_test.go +++ b/internal/lua/lua_bench_test.go @@ -25,6 +25,24 @@ func BenchmarkCoroutineCreation(b *testing.B) { } } +// BenchmarkLuaContextPool 测试 LuaContext 池化开销 +func BenchmarkLuaContextPool(b *testing.B) { + engine, err := NewEngine(DefaultConfig()) + if err != nil { + b.Fatal(err) + } + defer engine.Close() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + ctx := NewContext(engine, nil) + ctx.SetVariable("key", "value") + ctx.Write([]byte("hello")) + ctx.SetPhase(PhaseContent) + ctx.Release() + } +} + // BenchmarkBytecodeCompilation 测试字节码编译开销 func BenchmarkBytecodeCompilation(b *testing.B) { engine, err := NewEngine(DefaultConfig()) diff --git a/internal/lua/lua_test.go b/internal/lua/lua_test.go index 1769dcf..0b22d61 100644 --- a/internal/lua/lua_test.go +++ b/internal/lua/lua_test.go @@ -94,6 +94,66 @@ func TestLuaContextFlushOutput(t *testing.T) { assert.NotNil(t, ctx.OutputBuffer) } +// TestLuaContextPoolStateIsolation 测试池化 context 请求间无状态污染 +func TestLuaContextPoolStateIsolation(t *testing.T) { + engine, err := NewEngine(DefaultConfig()) + require.NoError(t, err) + defer engine.Close() + + // 第一次使用:设置变量、输出、阶段、退出标记 + ctx1 := NewContext(engine, nil) + ctx1.SetVariable("key1", "value1") + ctx1.SetVariable("key2", "value2") + ctx1.Write([]byte("hello")) + ctx1.SetPhase(PhaseAccess) + ctx1.Exited = true + + // 释放回池 + ctx1.Release() + + // 第二次使用:从池中获取(可能是同一个对象) + ctx2 := NewContext(engine, nil) + + // 验证无状态污染 + assert.Equal(t, PhaseInit, ctx2.Phase, "Phase should be reset to PhaseInit") + assert.False(t, ctx2.Exited, "Exited should be reset to false") + assert.Empty(t, ctx2.OutputBuffer, "OutputBuffer should be empty") + assert.Empty(t, ctx2.Variables, "Variables map should be empty") + + // 验证旧的 key 不存在 + _, ok := ctx2.GetVariable("key1") + assert.False(t, ok, "key1 should not exist after release") + _, ok = ctx2.GetVariable("key2") + assert.False(t, ok, "key2 should not exist after release") + + ctx2.Release() +} + +// TestLuaContextPoolMultipleReuse 测试多次复用 +func TestLuaContextPoolMultipleReuse(t *testing.T) { + engine, err := NewEngine(DefaultConfig()) + require.NoError(t, err) + defer engine.Close() + + // 循环多次 release/acquire,验证状态始终正确 + for i := 0; i < 100; i++ { + ctx := NewContext(engine, nil) + ctx.SetVariable("iter", "val") + ctx.Write([]byte("data")) + ctx.SetPhase(PhaseLog) + ctx.Exited = true + ctx.Release() + } + + // 最后一次获取,验证状态干净 + ctx := NewContext(engine, nil) + assert.Equal(t, PhaseInit, ctx.Phase) + assert.False(t, ctx.Exited) + assert.Empty(t, ctx.OutputBuffer) + assert.Empty(t, ctx.Variables) + ctx.Release() +} + // TestLuaContextExecute 测试 Lua 执行 func TestLuaContextExecute(t *testing.T) { engine, err := NewEngine(DefaultConfig())