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/re.py

@@ -22,6 +22,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 +91,39 @@ def span(self, n=0):
 class PCREPattern:
     def __init__(self, compiled_ptn):
         self.obj = compiled_ptn
+        self.key = None  # set while this pattern is held by the cache
+
+    def _free(self):
+        # MicroPython does not run __del__ on instances of Python classes, so
+        # the compiled pattern cannot be released by the garbage collector and
+        # has to be freed explicitly.
+        if self.obj is not None:
+            if self.key is not None:
+                # Drop the pattern from the cache first, so that nothing hands
+                # out a pointer that is about to become invalid.
+                del _cache[self.key]
+                self.key = None
+            pcre2_code_free(self.obj)
+            self.obj = None
 
     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,37 +190,92 @@ def findall(self, s):
             start = end
 
 
-def compile(pattern, flags=0):
-    errcode = bytes(4)
-    erroffset = bytes(4)
+def _compile(pattern, flags):
+    # These are output arguments and must be writable and of the size that
+    # pcre2_compile() writes: int for the error code, PCRE2_SIZE for the offset.
+    errcode = array.array("i", [0])
+    erroffset = array.array(PCRE2_SIZE_TYPE, [0])
     regex = pcre2_compile(pattern, PCRE2_ZERO_TERMINATED, flags, errcode, erroffset, None)
-    assert regex
+    assert regex, "error %d compiling regex at offset %d" % (errcode[0], erroffset[0])
     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.
+#
+# The cache owns the patterns it holds and never evicts them.  A pattern that
+# is still being used, either by the caller or by a call further up the stack,
+# must not be freed underneath it; a replacement callback passed to sub() can
+# otherwise trigger exactly that.  The cache is bounded instead: once it is
+# full, further patterns are compiled and, where this module owns them, freed
+# again after use.
+_MAXCACHE = 32
+_cache = {}
+
+
+def _cached(pattern, flags):
+    # Return the compiled pattern, and whether the caller has to free it.
+    key = (pattern, flags)
+    r = _cache.get(key)
+    if r is not None:
+        return r, False
+    r = _compile(pattern, flags)
+    if len(_cache) < _MAXCACHE:
+        _cache[key] = r
+        r.key = key
+        return r, False
+    return r, True
+
+
+def compile(pattern, flags=0):
+    # The pattern belongs to the caller, so it is never freed here.
+    return _cached(pattern, flags)[0]
+
+
 def search(pattern, string, flags=0):
-    r = compile(pattern, flags)
-    return r.search(string)
+    r, owned = _cached(pattern, flags)
+    try:
+        return r.search(string)
+    finally:
+        if owned:
+            r._free()
 
 
 def match(pattern, string, flags=0):
-    r = compile(pattern, flags | PCRE2_ANCHORED)
-    return r.search(string)
+    r, owned = _cached(pattern, flags | PCRE2_ANCHORED)
+    try:
+        return r.search(string)
+    finally:
+        if owned:
+            r._free()
 
 
 def sub(pattern, repl, s, count=0, flags=0):
-    r = compile(pattern, flags)
-    return r.sub(repl, s, count)
+    r, owned = _cached(pattern, flags)
+    try:
+        return r.sub(repl, s, count)
+    finally:
+        if owned:
+            r._free()
 
 
 def split(pattern, s, maxsplit=0, flags=0):
-    r = compile(pattern, flags)
-    return r.split(s, maxsplit)
+    r, owned = _cached(pattern, flags)
+    try:
+        return r.split(s, maxsplit)
+    finally:
+        if owned:
+            r._free()
 
 
 def findall(pattern, s, flags=0):
-    r = compile(pattern, flags)
-    return r.findall(s)
+    r, owned = _cached(pattern, flags)
+    try:
+        return r.findall(s)
+    finally:
+        if owned:
+            r._free()
 
 
 def escape(s):

unix-ffi/re/test_re_leak.py

@@ -0,0 +1,149 @@
+# Regression test for the memory that PCRE2 allocates behind this module: the
+# match data of every match, and every pattern compiled by the module level
+# functions, have to be freed again.  Otherwise each call leaks a few
+# kilobytes.
+#
+# A pattern returned by re.compile() and kept by the caller is not covered
+# here.  MicroPython does not run __del__ on instances of Python classes, so
+# such a pattern can only be released explicitly.
+#
+# The bounded cache that the module level functions keep is covered: it must
+# not grow past its limit, and the patterns that do not fit into it must be
+# freed again.
+
+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))
+
+
+# Compiling, including the path that does not produce a usable pattern.
+def compile_and_free():
+    re.compile("a(b)c")._free()
+
+
+def free_twice():
+    r = re.compile("a(b)c")
+    r._free()
+    r._free()
+
+
+def failed_compile():
+    try:
+        re.compile("(")
+    except AssertionError:
+        pass
+
+
+check_no_leak("re.compile() and _free()", compile_and_free)
+check_no_leak("_free() called twice", free_twice)
+check_no_leak("re.compile() of a bad pattern", failed_compile)
+
+
+# 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")
+    r._free()
+
+
+check_no_leak("pattern with several groups", many_groups)
+
+
+# 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"))
+
+# _free() drops the pattern from the cache, so that nothing afterwards hands
+# out a pointer to memory that has been released.
+r = re.compile("zz(y)")
+r._free()
+assert re.search("zz(y)", "xxzzyxx").group(0) == "zzy"
+
+
+# The module level functions cache the patterns they compile.  That cache must
+# stay bounded, and a pattern that does not fit into it has to be freed again.
+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)
+
+
+# A replacement callback runs while sub() is still using its own pattern, and
+# may push further patterns through the cache.  The pattern that is in use must
+# survive that.
+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)