PCRE2: use thread local for jit stack and match data (#16175)

We can reuse the same JIT stack and matchdata for all instances. There's no need for a specific match data per `Regex`: We just allocate one and make sure it's not used by other threads.

The drawback is that we must allocate each matchdata with a maximum number of ovectors (65535). That might increase memory usage, though I failed to notice it in practice.
Maybe not allocating separate memory for every regular expression is helping?

Author: ysbaddaden

src/crystal/value_with_finalizer.cr

@@ -0,0 +1,11 @@
+# :nodoc:
+class Crystal::ValueWithFinalizer(T)
+  getter value : T
+
+  def initialize(@value : T, @finalizer : T ->)
+  end
+
+  def finalize
+    @finalizer.call(@value)
+  end
+end

src/regex/lib_pcre2.cr

@@ -264,10 +264,12 @@ lib LibPCRE2
 
   fun jit_stack_create = pcre2_jit_stack_create_8(startsize : LibC::SizeT, maxsize : LibC::SizeT, gcontext : GeneralContext*) : JITStack*
   fun jit_stack_assign = pcre2_jit_stack_assign_8(mcontext : MatchContext*, callable_function : Void* -> JITStack*, callable_data : Void*) : Void
+  fun jit_stack_free = pcre2_jit_stack_free_8(jit_stack : JITStack*) : Void
 
   fun pattern_info = pcre2_pattern_info_8(code : Code*, what : UInt32, where : Void*) : Int
 
   fun match = pcre2_match_8(code : Code*, subject : UInt8*, length : LibC::SizeT, startoffset : LibC::SizeT, options : UInt32, match_data : MatchData*, mcontext : MatchContext*) : Int
+  fun match_data_create = pcre2_match_data_create_8(ovecsize : UInt32, gcontext : GeneralContext*) : MatchData*
   fun match_data_create_from_pattern = pcre2_match_data_create_from_pattern_8(code : Code*, gcontext : GeneralContext*) : MatchData*
   fun match_data_free = pcre2_match_data_free_8(match_data : MatchData*) : Void
 

src/regex/pcre2.cr

@@ -1,5 +1,5 @@
 require "./lib_pcre2"
-require "crystal/thread_local_value"
+require "crystal/value_with_finalizer"
 
 # :nodoc:
 module Regex::PCRE2
@@ -207,12 +207,14 @@ module Regex::PCRE2
   private def match_impl(str, byte_index, options)
     match_data = match_data(str, byte_index, options) || return
 
-    ovector_count = LibPCRE2.get_ovector_count(match_data)
+    # We reuse the same `match_data` allocation, so we must reimplement the
+    # behavior of pcre2_match_data_create_from_pattern (get_ovector_count always
+    # returns 65535, aka the maximum).
+    ovector_count = capture_count_impl &+ 1
     ovector = Slice.new(LibPCRE2.get_ovector_pointer(match_data), ovector_count &* 2)
 
     # We need to dup the ovector because `match_data` is re-used for subsequent
-    # matches (see `@match_data`).
-    # Dup brings the ovector data into the realm of the GC.
+    # matches. We only dup the match data (not everything).
     ovector = ovector.dup
 
     ::Regex::MatchData.new(self, @re, str, byte_index, ovector.to_unsafe, ovector_count.to_i32 &- 1)
@@ -228,43 +230,43 @@ module Regex::PCRE2
 
   class_getter match_context : LibPCRE2::MatchContext* do
     match_context = LibPCRE2.match_context_create(nil)
-    LibPCRE2.jit_stack_assign(match_context, ->(_data) { Regex::PCRE2.jit_stack }, nil)
+    LibPCRE2.jit_stack_assign(match_context, ->(_data) { current_jit_stack.value }, nil)
     match_context
   end
 
-  # Returns a JIT stack that's shared in the current thread.
+  # JIT stack is unique per thread.
   #
-  # Only a single `match` function can run per thread at any given time, so there
-  # can't be any concurrent access to the JIT stack.
-  @@jit_stack = Crystal::ThreadLocalValue(LibPCRE2::JITStack*).new
-
-  def self.jit_stack
-    @@jit_stack.get do
-      LibPCRE2.jit_stack_create(32_768, 1_048_576, nil) || raise "Error allocating JIT stack"
-    end
+  # Only a single `match` function can run per thread at any given time, so
+  # there can't be any concurrent access to the JIT stack.
+  thread_local(current_jit_stack : ::Crystal::ValueWithFinalizer(::LibPCRE2::JITStack*)) do
+    ptr = LibPCRE2.jit_stack_create(32_768, 1_048_576, nil)
+    raise RuntimeError.new("Error allocating JIT stack") if ptr.null?
+    ::Crystal::ValueWithFinalizer.new(ptr, ->(value : ::LibPCRE2::JITStack*) { LibPCRE2.jit_stack_free(value) })
   end
 
-  # Match data is shared per instance and thread.
+  # Match data is unique per thread.
   #
-  # Match data contains a buffer for backtracking when matching in interpreted mode (non-JIT).
-  # This buffer is heap-allocated and should be re-used for subsequent matches.
-  @match_data = Crystal::ThreadLocalValue(LibPCRE2::MatchData*).new
-
-  private def match_data
-    @match_data.get do
-      LibPCRE2.match_data_create_from_pattern(@re, nil)
-    end
+  # Match data contains a buffer for backtracking when matching in interpreted
+  # mode (non-JIT). This buffer is heap-allocated and should be re-used for
+  # subsequent matches.
+  #
+  # Only a single `match` function can run per thread at any given time, so
+  # there can't be any concurrent access to the match data buffer.
+  thread_local(current_match_data : ::Crystal::ValueWithFinalizer(::LibPCRE2::MatchData*)) do
+    # The ovector size is clamped to 65535 pairs; we declare the maximum because
+    # we allocate the match data buffer once for the thread and need to adapt to
+    # any regular expression.
+    ptr = LibPCRE2.match_data_create(65_535, nil)
+    raise RuntimeError.new("Error allocating match data") if ptr.null?
+    ::Crystal::ValueWithFinalizer.new(ptr, ->(value : LibPCRE2::MatchData*) { LibPCRE2.match_data_free(value) })
   end
 
   def finalize
-    @match_data.consume_each do |match_data|
-      LibPCRE2.match_data_free(match_data)
-    end
     LibPCRE2.code_free @re
   end
 
   private def match_data(str, byte_index, options)
-    match_data = self.match_data
+    match_data = Regex::PCRE2.current_match_data.value
     match_count = LibPCRE2.match(@re, str, str.bytesize, byte_index, pcre2_match_options(options), match_data, PCRE2.match_context)
 
     if match_count < 0