Fix memory leak in fakeredis Lua script execution (#300)

## Summary

- Restructure Lua callback caching in the fakeredis monkeypatch to avoid
creating new `functools.partial` objects on every `eval()` call
- Add `collectgarbage()` after each script execution to clean up Lua
tables
- Fix flaky logging test that was broken by newer fakeredis DEBUG
logging

## Problem

The `memory://` backend was leaking memory when executing Lua scripts.
Every `eval()` created new partial objects for callbacks that were held
by the Lua runtime, preventing garbage collection.

Related: PrefectHQ/prefect#18605

## Testing

Created a test harness running a perpetual task at ~10 evals/second:
- **Before fix**: Memory grew several MB per minute  
- **After fix**: RSS stays flat at ~200KB above baseline after 1000+
evals

The same fix has been submitted upstream:
cunla/fakeredis-py#452

🤖 Generated with [Claude Code](https://claude.ai/code)

---------

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

Author: chrisguidry

src/docket/_redis.py

@@ -396,17 +396,20 @@ def patched_eval(
         sha1 = hashlib.sha1(script).hexdigest().encode()
         self._server.script_cache[sha1] = script
 
-        # Cache LuaRuntime and set_globals function on the server
+        # Cache LuaRuntime and all callbacks on the server
         if not hasattr(self._server, "_lua_runtime"):
             self._server._lua_runtime = LUA_MODULE.LuaRuntime(
                 encoding=None, unpack_returned_tuples=True
             )
+            lua_runtime = self._server._lua_runtime
             modules_import_str = "\n".join(
                 [f"{module} = require('{module}')" for module in self.load_lua_modules]
             )
-            self._server._lua_set_globals = self._server._lua_runtime.eval(
+
+            # Create set_globals for initial setup (sets callbacks once)
+            set_globals_init = lua_runtime.eval(
                 f"""
-                function(keys, argv, redis_call, redis_pcall, redis_log, cjson_encode, cjson_decode, cjson_null)
+                function(redis_call, redis_pcall, redis_log, cjson_encode, cjson_decode, cjson_null)
                     redis = {{}}
                     redis.call = redis_call
                     redis.pcall = redis_pcall
@@ -423,40 +426,94 @@ def patched_eval(
                     cjson.decode = cjson_decode
                     cjson.null = cjson_null
 
+                    KEYS = {{}}
+                    ARGV = {{}}
+                    {modules_import_str}
+                end
+                """
+            )
+
+            # Create set_keys_argv to update just KEYS/ARGV per call
+            self._server._lua_set_keys_argv = lua_runtime.eval(
+                """
+                function(keys, argv)
                     KEYS = keys
                     ARGV = argv
-                    {modules_import_str}
                 end
                 """
             )
-            # Capture expected globals once after first setup
-            self._server._lua_set_globals(
-                self._server._lua_runtime.table_from([]),
-                self._server._lua_runtime.table_from([]),
+
+            # Capture expected globals before setting up callbacks
+            set_globals_init(
                 lambda *args: None,
                 lambda *args: None,
                 lambda *args: None,
                 lambda *args: None,
                 lambda *args: None,
                 None,
             )
-            self._server._lua_expected_globals = set(
-                self._server._lua_runtime.globals().keys()
+            self._server._lua_expected_globals = set(lua_runtime.globals().keys())
+            expected_globals = self._server._lua_expected_globals
+
+            # Container to hold current socket - callbacks will look this up
+            self._server._lua_current_socket = [None]  # Use list for mutability
+
+            # Create wrapper callbacks that look up the current socket dynamically
+            def make_redis_call_wrapper() -> typing.Callable[..., typing.Any]:
+                server = self._server
+                lr = lua_runtime
+                eg = expected_globals
+
+                def wrapper(op: bytes, *args: typing.Any) -> typing.Any:
+                    socket = server._lua_current_socket[0]
+                    return socket._lua_redis_call(lr, eg, op, *args)
+
+                return wrapper
+
+            def make_redis_pcall_wrapper() -> typing.Callable[..., typing.Any]:
+                server = self._server
+                lr = lua_runtime
+                eg = expected_globals
+
+                def wrapper(op: bytes, *args: typing.Any) -> typing.Any:
+                    socket = server._lua_current_socket[0]
+                    return socket._lua_redis_pcall(lr, eg, op, *args)
+
+                return wrapper
+
+            # Cache the callback wrappers and static partials
+            self._server._lua_redis_call_wrapper = make_redis_call_wrapper()
+            self._server._lua_redis_pcall_wrapper = make_redis_pcall_wrapper()
+            self._server._lua_log_partial = functools.partial(
+                _lua_redis_log, lua_runtime, expected_globals
+            )
+            self._server._lua_cjson_encode_partial = functools.partial(
+                _lua_cjson_encode, lua_runtime, expected_globals
+            )
+            self._server._lua_cjson_decode_partial = functools.partial(
+                _lua_cjson_decode, lua_runtime, expected_globals
+            )
+
+            # Set up all callbacks once
+            set_globals_init(
+                self._server._lua_redis_call_wrapper,
+                self._server._lua_redis_pcall_wrapper,
+                self._server._lua_log_partial,
+                self._server._lua_cjson_encode_partial,
+                self._server._lua_cjson_decode_partial,
+                _lua_cjson_null,
             )
 
         lua_runtime = self._server._lua_runtime
-        set_globals = self._server._lua_set_globals
         expected_globals = self._server._lua_expected_globals
 
-        set_globals(
+        # Update the current socket so callbacks can find it
+        self._server._lua_current_socket[0] = self
+
+        # Only update KEYS and ARGV per call (callbacks are already set)
+        self._server._lua_set_keys_argv(
             lua_runtime.table_from(keys_and_args[:numkeys]),
             lua_runtime.table_from(keys_and_args[numkeys:]),
-            functools.partial(self._lua_redis_call, lua_runtime, expected_globals),
-            functools.partial(self._lua_redis_pcall, lua_runtime, expected_globals),
-            functools.partial(_lua_redis_log, lua_runtime, expected_globals),
-            functools.partial(_lua_cjson_encode, lua_runtime, expected_globals),
-            functools.partial(_lua_cjson_decode, lua_runtime, expected_globals),
-            _lua_cjson_null,
         )
 
         try:
@@ -479,6 +536,9 @@ def patched_eval(
 
         _check_for_lua_globals(lua_runtime, expected_globals)
 
+        # Clean up Lua tables (KEYS/ARGV) created for this script execution
+        lua_runtime.execute("collectgarbage()")
+
         return self._convert_lua_result(result, nested=False)
 
     ScriptingCommandsMixin.eval = patched_eval

tests/fundamentals/test_logging.py

@@ -28,9 +28,13 @@ async def the_task(
     with caplog.at_level(logging.INFO):
         await worker.run_until_finished()
 
-        assert "the_task('value-a', b=..., c='value-c', d=...)" in caplog.text
-        assert "value-b" not in caplog.text
-        assert "value-d" not in caplog.text
+        # Filter to only docket logs (exclude fakeredis DEBUG logs which contain raw pickle data)
+        docket_logs = "\n".join(
+            r.message for r in caplog.records if r.name.startswith("docket")
+        )
+        assert "the_task('value-a', b=..., c='value-c', d=...)" in docket_logs
+        assert "value-b" not in docket_logs
+        assert "value-d" not in docket_logs
 
 
 async def test_tasks_can_opt_into_logging_collection_lengths(