unix-ffi/re: Fix the PCRE2 memory leaks.

tools/ci.sh

@@ -78,6 +78,7 @@ function ci_package_tests_run {
         unix-ffi/gettext/test_gettext.py \
         unix-ffi/pwd/test_getpwnam.py \
         unix-ffi/re/test_re.py \
+        unix-ffi/re/test_re_leak.py \
         unix-ffi/sqlite3/test_sqlite3.py \
         unix-ffi/sqlite3/test_sqlite3_2.py \
         unix-ffi/sqlite3/test_sqlite3_3.py \

unix-ffi/re/manifest.py

@@ -1,4 +1,4 @@
-metadata(version="0.2.6")
+metadata(version="0.3.0")
 
 # Originally written by Paul Sokolovsky.
 

unix-ffi/re/re.py

@@ -2,6 +2,7 @@
 import ffilib
 import array
 import uctypes
+import weakref
 
 pcre2 = ffilib.open("libpcre2-8")
 
@@ -22,6 +23,12 @@
 #       PCRE2_SIZE *pcre2_get_ovector_pointer(pcre2_match_data *match_data);
 pcre2_get_ovector_pointer = pcre2.func("p", "pcre2_get_ovector_pointer_8", "p")
 
+#       void pcre2_code_free(pcre2_code *code);
+pcre2_code_free = pcre2.func("v", "pcre2_code_free_8", "p")
+
+#       void pcre2_match_data_free(pcre2_match_data *match_data);
+pcre2_match_data_free = pcre2.func("v", "pcre2_match_data_free_8", "p")
+
 #       pcre2_match_data *pcre2_match_data_create_from_pattern(const pcre2_code *code,
 #           pcre2_general_context *gcontext);
 pcre2_match_data_create_from_pattern = pcre2.func(
@@ -85,21 +92,29 @@ def span(self, n=0):
 class PCREPattern:
     def __init__(self, compiled_ptn):
         self.obj = compiled_ptn
+        # The compiled pattern lives in memory that PCRE2 owns and that the
+        # garbage collector knows nothing about, so release it once this object
+        # is collected.
+        weakref.finalize(self, pcre2_code_free, compiled_ptn)
 
     def search(self, s, pos=0, endpos=-1, _flags=0):
         assert endpos == -1, "pos: %d, endpos: %d" % (pos, endpos)
         buf = array.array("i", [0])
         pcre2_pattern_info(self.obj, PCRE2_INFO_CAPTURECOUNT, buf)
         cap_count = buf[0]
         match_data = pcre2_match_data_create_from_pattern(self.obj, None)
-        num = pcre2_match(self.obj, s, len(s), pos, _flags, match_data, None)
-        if num == -1:
-            # No match
-            return None
-        ov_ptr = pcre2_get_ovector_pointer(match_data)
-        # pcre2_get_ovector_pointer return PCRE2_SIZE
-        ov_buf = uctypes.bytearray_at(ov_ptr, PCRE2_SIZE_SIZE * (cap_count + 1) * 2)
-        ov = array.array(PCRE2_SIZE_TYPE, ov_buf)
+        try:
+            num = pcre2_match(self.obj, s, len(s), pos, _flags, match_data, None)
+            if num == -1:
+                # No match
+                return None
+            ov_ptr = pcre2_get_ovector_pointer(match_data)
+            # pcre2_get_ovector_pointer return PCRE2_SIZE.  The offsets are
+            # copied out here, because the match data is freed below.
+            ov_buf = uctypes.bytearray_at(ov_ptr, PCRE2_SIZE_SIZE * (cap_count + 1) * 2)
+            ov = array.array(PCRE2_SIZE_TYPE, ov_buf)
+        finally:
+            pcre2_match_data_free(match_data)
         # We don't care how many matching subexpressions we got, we
         # care only about total # of capturing ones (including empty)
         return PCREMatch(s, cap_count + 1, ov)
@@ -166,14 +181,41 @@ def findall(self, s):
             start = end
 
 
-def compile(pattern, flags=0):
+def _compile(pattern, flags):
+    # These are output arguments and must be of the size that pcre2_compile()
+    # writes: int for the error code, PCRE2_SIZE for the offset.
     errcode = bytes(4)
-    erroffset = bytes(4)
+    erroffset = bytes(PCRE2_SIZE_SIZE)
     regex = pcre2_compile(pattern, PCRE2_ZERO_TERMINATED, flags, errcode, erroffset, None)
-    assert regex
+    assert regex, "compile error %d at %d" % (
+        int.from_bytes(errcode, sys.byteorder),
+        int.from_bytes(erroffset, sys.byteorder),
+    )
     return PCREPattern(regex)
 
 
+# Compiled patterns are cached, the way CPython does it, so that using the same
+# pattern again does not compile it a second time.  compile() returns the
+# cached pattern, so re.compile(p) is re.compile(p), as in CPython.  A pattern
+# that is dropped from the cache is freed by the garbage collector once nothing
+# refers to it any more.
+_MAXCACHE = 32
+_cache = {}
+
+
+def compile(pattern, flags=0):
+    key = (pattern, flags)
+    r = _cache.get(key)
+    if r is None:
+        r = _compile(pattern, flags)
+        if len(_cache) >= _MAXCACHE:
+            # Drop the whole cache, the way CPython does, instead of keeping
+            # track of which entry was used last.
+            _cache.clear()
+        _cache[key] = r
+    return r
+
+
 def search(pattern, string, flags=0):
     r = compile(pattern, flags)
     return r.search(string)

unix-ffi/re/test_re_leak.py

@@ -0,0 +1,139 @@
+# Regression test for the memory that PCRE2 allocates behind this module: the
+# match data of every match, and every compiled pattern, have to be freed
+# again.  Otherwise each call leaks a few kilobytes.
+#
+# The match data is freed by the call that created it.  A compiled pattern is
+# freed by the garbage collector, through the weakref.finalize() that the
+# pattern registers for itself, so a pattern that is no longer reachable does
+# not have to be released by hand.
+#
+# The bounded cache that compile() keeps is covered as well: it must not grow
+# past its limit, and the patterns that it drops have to be freed.
+
+import gc
+import re
+
+
+def rss():
+    # Resident set size in KiB, from the second field of /proc/self/statm.
+    with open("/proc/self/statm") as f:
+        return int(f.read().split()[1]) * 4096 // 1024
+
+
+try:
+    rss()
+except OSError:
+    # No /proc, so memory use cannot be measured here.
+    print("SKIP")
+    raise SystemExit
+
+
+N = 4000
+LIMIT = 256  # KiB
+
+
+def check_no_leak(name, fn):
+    # Run the calls once to let the MicroPython heap grow to its steady state,
+    # so that only the memory allocated by PCRE2 is measured afterwards.
+    for _ in range(N):
+        fn()
+    gc.collect()
+    before = rss()
+    for _ in range(N):
+        fn()
+    gc.collect()
+    growth = rss() - before
+    assert growth < LIMIT, "%s leaks %d KiB per %d calls (%d bytes per call)" % (
+        name,
+        growth,
+        N,
+        growth * 1024 // N,
+    )
+
+
+text = "He was carefully disguised but captured quickly by police."
+p = re.compile("a(b)c")
+
+# Matching with a compiled pattern.
+check_no_leak("Pattern.search() with a match", lambda: p.search("xxabcxx"))
+check_no_leak("Pattern.search() without a match", lambda: p.search("xxxxxxx"))
+check_no_leak("Pattern.match()", lambda: p.match("abcxx"))
+check_no_leak("Pattern.sub()", lambda: p.sub("z", "xxabcxx"))
+check_no_leak("Pattern.split()", lambda: p.split("xxabcxx"))
+check_no_leak("Pattern.findall()", lambda: p.findall("xxabcxx abc"))
+
+# The module level functions, which compile a pattern of their own.
+check_no_leak("re.search()", lambda: re.search("a(b)c", "xxabcxx"))
+check_no_leak("re.match()", lambda: re.match("a(b)c", "abcxx"))
+check_no_leak("re.sub()", lambda: re.sub("a", "z", "caaab"))
+check_no_leak("re.split()", lambda: re.split(r"\W+", "Words, words, words."))
+check_no_leak("re.findall()", lambda: re.findall(r"(\w+)ly", text))
+
+
+# A pattern with several groups needs a larger match data block.
+def many_groups():
+    r = re.compile(r"(\w+)(\s+)(\w+)(\s+)(\w+)")
+    assert r.search("one two three").groups() == ("one", " ", "two", " ", "three")
+
+
+check_no_leak("pattern with several groups", many_groups)
+
+
+# The path that does not produce a usable pattern must not leak either.
+def failed_compile():
+    try:
+        re.compile("(")
+    except AssertionError:
+        pass
+
+
+check_no_leak("re.compile() of a bad pattern", failed_compile)
+
+
+# compile() returns the cached pattern, the way CPython does, so compiling the
+# same pattern again does not allocate.
+assert re.compile("a(b)c") is re.compile("a(b)c")
+check_no_leak("re.compile() with the same pattern", lambda: re.compile("a(b)c"))
+
+
+# The cache has to stay bounded, and a pattern that it drops has to be freed by
+# the garbage collector.
+counter = [0]
+
+
+def distinct_patterns():
+    counter[0] += 1
+    re.search("a%dc" % counter[0], "xxabcxx")
+
+
+# Push far more distinct patterns through the cache than it can hold: it has
+# to stop growing.
+for _ in range(re._MAXCACHE * 4):
+    distinct_patterns()
+assert len(re._cache) <= re._MAXCACHE, len(re._cache)
+
+check_no_leak("re.search() with distinct patterns", distinct_patterns)
+assert len(re._cache) <= re._MAXCACHE, len(re._cache)
+
+
+
+def compile_distinct():
+    counter[0] += 1
+    re.compile("b%dc" % counter[0])
+
+
+check_no_leak("re.compile() with distinct patterns", compile_distinct)
+assert len(re._cache) <= re._MAXCACHE, len(re._cache)
+
+
+# A replacement callback runs while sub() is still using its own pattern, and
+# may push further patterns through the cache.  The reference that sub() holds
+# has to keep that pattern alive across the callback.
+def reentrant_repl(m):
+    counter[0] += 1
+    re.search("z%dz" % counter[0], "nothing here")
+    return "z"
+
+
+check_no_leak("re.sub() with a reentrant callback", lambda: re.sub("a", reentrant_repl, "caaab"))
+assert len(re._cache) <= re._MAXCACHE, len(re._cache)