diff --git a/.gitignore b/.gitignore index 75d21f6..fe51a48 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ poetry-installer-error-*.log docs/_build .DS_Store dist/ +.idea/ diff --git a/extism/extism.py b/extism/extism.py index d66eb5c..9652294 100644 --- a/extism/extism.py +++ b/extism/extism.py @@ -270,10 +270,10 @@ def __init__(self, f: Function): _lib.extism_function_set_namespace(self.pointer, f.namespace.encode()) def __del__(self): - if not hasattr(self, "pointer"): + if not hasattr(self, "pointer") or self.pointer is None: return - if self.pointer is not None: - _lib.extism_function_free(self.pointer) + _lib.extism_function_free(self.pointer) + self.pointer = None def _map_arg(arg_name, xs) -> Tuple[ValType, Callable[[Any, Any], Any]]: @@ -507,7 +507,7 @@ def __init__( raise Error(msg.decode()) def __del__(self): - if not hasattr(self, "pointer"): + if not hasattr(self, "pointer") or self.pointer is None or self.pointer == -1: return _lib.extism_compiled_plugin_free(self.pointer) self.pointer = -1 @@ -551,7 +551,9 @@ def __init__( config: Optional[Any] = None, functions: Optional[List[Function]] = HOST_FN_REGISTRY, ): - if not isinstance(plugin, CompiledPlugin): + # Track if we created the CompiledPlugin (so we know to free it) + self._owns_compiled_plugin = not isinstance(plugin, CompiledPlugin) + if self._owns_compiled_plugin: plugin = CompiledPlugin(plugin, wasi, functions) self.compiled_plugin = plugin @@ -625,10 +627,14 @@ def call( return parse(buf) def __del__(self): - if not hasattr(self, "pointer"): + if not hasattr(self, "plugin") or self.plugin == -1: return _lib.extism_plugin_free(self.plugin) self.plugin = -1 + # Free the compiled plugin if we created it + if getattr(self, "_owns_compiled_plugin", False) and self.compiled_plugin is not None: + self.compiled_plugin.__del__() + self.compiled_plugin = None def __enter__(self): return self diff --git a/tests/test_extism.py b/tests/test_extism.py index c401a56..18fee5e 100644 --- a/tests/test_extism.py +++ b/tests/test_extism.py @@ -1,14 +1,17 @@ from collections import namedtuple -import unittest -import extism +import gc import hashlib import json +import pickle import time -from threading import Thread +import typing +import unittest from datetime import datetime, timedelta from os.path import join, dirname -import typing -import pickle +from threading import Thread + +import extism +from extism.extism import CompiledPlugin, _ExtismFunctionMetadata, TypeInferredFunction # A pickle-able object. @@ -47,6 +50,83 @@ def test_can_free_plugin(self): plugin = extism.Plugin(self._manifest()) del plugin + def test_plugin_del_frees_native_resources(self): + """Test that Plugin.__del__ properly frees native resources. + + This tests the fix for a bug where Plugin.__del__ checked for + 'self.pointer' instead of 'self.plugin', causing extism_plugin_free + to never be called and leading to memory leaks. + + This also tests that __del__ can be safely called multiple times + (via context manager exit and garbage collection) without causing + double-free errors. + """ + with extism.Plugin(self._manifest(), functions=[]) as plugin: + j = json.loads(plugin.call("count_vowels", "test")) + self.assertEqual(j["count"], 1) + # Plugin should own the compiled plugin it created + self.assertTrue(plugin._owns_compiled_plugin) + + # Verify plugin was freed after exiting context + self.assertEqual(plugin.plugin, -1, + "Expected plugin.plugin to be -1 after __del__, indicating extism_plugin_free was called") + # Verify compiled plugin was also freed (since Plugin owned it) + self.assertIsNone(plugin.compiled_plugin, + "Expected compiled_plugin to be None after __del__, indicating it was also freed") + + def test_compiled_plugin_del_frees_native_resources(self): + """Test that CompiledPlugin.__del__ properly frees native resources. + + Unlike Plugin, CompiledPlugin has no context manager so __del__ is only + called once by garbage collection. This also tests that __del__ can be + safely called multiple times without causing double-free errors. + """ + compiled = CompiledPlugin(self._manifest(), functions=[]) + # Verify pointer exists before deletion + self.assertTrue(hasattr(compiled, 'pointer')) + self.assertNotEqual(compiled.pointer, -1) + + # Create a plugin from compiled to ensure it works + plugin = extism.Plugin(compiled) + j = json.loads(plugin.call("count_vowels", "test")) + self.assertEqual(j["count"], 1) + + # Plugin should NOT own the compiled plugin (it was passed in) + self.assertFalse(plugin._owns_compiled_plugin) + + # Clean up plugin first + plugin.__del__() + self.assertEqual(plugin.plugin, -1) + + # Compiled plugin should NOT have been freed by Plugin.__del__ + self.assertNotEqual(compiled.pointer, -1, + "Expected compiled.pointer to NOT be -1 since Plugin didn't own it") + + # Now clean up compiled plugin manually + compiled.__del__() + + # Verify compiled plugin was freed + self.assertEqual(compiled.pointer, -1, + "Expected compiled.pointer to be -1 after __del__, indicating extism_compiled_plugin_free was called") + + def test_extism_function_metadata_del_frees_native_resources(self): + """Test that _ExtismFunctionMetadata.__del__ properly frees native resources.""" + def test_host_fn(inp: str) -> str: + return inp + + func = TypeInferredFunction(None, "test_func", test_host_fn, []) + metadata = _ExtismFunctionMetadata(func) + + # Verify pointer exists before deletion + self.assertTrue(hasattr(metadata, 'pointer')) + self.assertIsNotNone(metadata.pointer) + + metadata.__del__() + + # Verify function was freed (pointer set to None) + self.assertIsNone(metadata.pointer, + "Expected metadata.pointer to be None after __del__, indicating extism_function_free was called") + def test_errors_on_bad_manifest(self): self.assertRaises( extism.Error, lambda: extism.Plugin({"invalid_manifest": True})