Cache short text numerics as frozen BigDecimals (2.5–4.2× faster numeric decoding)

ext/extconf.rb

@@ -309,6 +309,7 @@ module PG
 				"or upgrade your database to at least PostgreSQL-10."
 # optional headers/functions
 have_func 'PQresultMemorySize', 'libpq-fe.h' # since PostgreSQL-12
+have_func 'rb_ractor_shareable_p', 'ruby/ractor.h' # since ruby-3.0, but not in Truffleruby
 have_func 'rb_hash_new_capa' # since ruby-3.2
 
 have_header 'inttypes.h'

ext/pg.h

@@ -195,6 +195,7 @@ typedef VALUE (* t_pg_typecast_copy_get)( t_typemap *, VALUE, int, int, int );
 #define PG_CODER_FORMAT_ERROR_TO_RAISE 0x4
 #define PG_CODER_FORMAT_ERROR_TO_STRING 0x8
 #define PG_CODER_FORMAT_ERROR_TO_PARTIAL 0xc
+#define PG_CODER_DISABLE_CACHE 0x10
 
 struct pg_coder {
 	t_pg_coder_enc_func enc_func;

ext/pg_coder.c

@@ -627,6 +627,7 @@ init_pg_coder(void)
 	rb_define_const( rb_cPG_Coder, "FORMAT_ERROR_TO_RAISE", INT2NUM(PG_CODER_FORMAT_ERROR_TO_RAISE));
 	rb_define_const( rb_cPG_Coder, "FORMAT_ERROR_TO_STRING", INT2NUM(PG_CODER_FORMAT_ERROR_TO_STRING));
 	rb_define_const( rb_cPG_Coder, "FORMAT_ERROR_TO_PARTIAL", INT2NUM(PG_CODER_FORMAT_ERROR_TO_PARTIAL));
+	rb_define_const( rb_cPG_Coder, "DISABLE_CACHE", INT2NUM(PG_CODER_DISABLE_CACHE));
 
 	/*
 	 * Name of the coder or the corresponding data type.

ext/pg_text_decoder.c

@@ -39,6 +39,11 @@
 #include <sys/socket.h>
 #endif
 #include <string.h>
+#include "ruby/ractor.h"
+#include "ruby/thread_native.h"
+#if !defined(_WIN32)
+#include <pthread.h>
+#endif
 
 VALUE rb_mPG_TextDecoder;
 static ID s_id_Rational;
@@ -157,6 +162,57 @@ pg_text_dec_integer(t_pg_coder *conv, const char *val, int len, int tuple, int f
 	return rb_cstr2inum(val, 10);
 }
 
+#define NUMERIC_CACHE_SLOT_BITS 12
+
+typedef struct {
+	uint64_t text;
+	VALUE value;
+} t_numeric_cache_slot;
+
+static t_numeric_cache_slot numeric_cache[1 << NUMERIC_CACHE_SLOT_BITS];
+static rb_nativethread_lock_t numeric_cache_lock;
+static VALUE numeric_cache_marker;
+
+static void
+numeric_cache_mark(void *_unused)
+{
+	size_t i;
+	for( i = 0; i < sizeof(numeric_cache) / sizeof(*numeric_cache); i++ ){
+		rb_gc_mark_movable(numeric_cache[i].value);
+	}
+}
+
+static void
+numeric_cache_compact(void *_unused)
+{
+	size_t i;
+	for( i = 0; i < sizeof(numeric_cache) / sizeof(*numeric_cache); i++ ){
+		pg_gc_location(numeric_cache[i].value);
+	}
+}
+
+static const rb_data_type_t numeric_cache_type = {
+	"PG::TextDecoder::Numeric cache",
+	{ numeric_cache_mark, NULL, NULL, numeric_cache_compact },
+	0, 0, RUBY_TYPED_WB_PROTECTED,
+};
+
+#if !defined(_WIN32)
+/* Another Ractor may have held the lock or been writing a slot when the process forked. */
+static void
+numeric_cache_reset_in_child(void)
+{
+	rb_nativethread_lock_initialize(&numeric_cache_lock);
+	memset(numeric_cache, 0, sizeof(numeric_cache));
+}
+#endif
+
+static VALUE
+pg_text_dec_numeric_uncached(const char *val, int len)
+{
+	return rb_funcall(rb_cObject, s_id_BigDecimal, 1, rb_str_new(val, len));
+}
+
 /*
  * Document-class: PG::TextDecoder::Numeric < PG::SimpleDecoder
  *
@@ -165,11 +221,58 @@ pg_text_dec_integer(t_pg_coder *conv, const char *val, int len, int tuple, int f
  *
  * As soon as this class is used, it requires the 'bigdecimal' gem.
  *
+ * Short values like +0.00+ or +19.90+ repeat a lot, so their frozen BigDecimal objects are
+ * cached process-wide once seen twice in a row, and the same object is returned for the same text.
+ * Set the flag PG::Coder::DISABLE_CACHE to get a new object for every value.
+ *
  */
 static VALUE
 pg_text_dec_numeric(t_pg_coder *conv, const char *val, int len, int tuple, int field, int enc_idx)
 {
-	return rb_funcall(rb_cObject, s_id_BigDecimal, 1, rb_str_new(val, len));
+	uint64_t text = 0;
+	t_numeric_cache_slot *slot;
+	VALUE value;
+
+	/* NaN and Infinity don't end with a digit. They stay uncached to honor BigDecimal.mode exceptions.
+	 * Up to 8 characters covers values like 0.00, 19.90 or 1299.00, and fits the key into one uint64_t. */
+	if( (conv->flags & PG_CODER_DISABLE_CACHE) || len < 1 || len > (int)sizeof(text) || !rb_isdigit(val[len - 1]) ){
+		return pg_text_dec_numeric_uncached(val, len);
+	}
+
+	memcpy(&text, val, len);
+	slot = &numeric_cache[(text * UINT64_C(0x9E3779B97F4A7C15)) >> (64 - NUMERIC_CACHE_SLOT_BITS)];
+
+	/* The lock guards against Ractors running in parallel. It is never held across BigDecimal(), which can raise. */
+	/* A text is cached the second time it's seen in a row in its slot. Texts seen only once only leave their key,
+	 * which spares them rb_ractor_shareable_p() and a promotion to the old generation. */
+	rb_nativethread_lock_lock(&numeric_cache_lock);
+	if( slot->text == text ){
+		value = slot->value;
+	} else {
+		value = Qundef;
+		slot->text = text;
+		slot->value = Qnil;
+	}
+	rb_nativethread_lock_unlock(&numeric_cache_lock);
+	if( value != Qundef && value != Qnil ) return value;
+	if( value == Qundef ) return pg_text_dec_numeric_uncached(val, len);
+
+	value = pg_text_dec_numeric_uncached(val, len);
+
+#ifdef HAVE_RB_RACTOR_SHAREABLE_P
+	/* Write cache only if value can be shared across ractors. */
+	if( rb_ractor_shareable_p(value) ){
+#else
+	/* Truffleruby doesn't support ractor.
+	 * Cache the frozen value. */
+	rb_obj_freeze(value);
+	{
+#endif
+		rb_nativethread_lock_lock(&numeric_cache_lock);
+		if( slot->text == text ) RB_OBJ_WRITE(numeric_cache_marker, &slot->value, value);
+		rb_nativethread_lock_unlock(&numeric_cache_lock);
+	}
+	return value;
 }
 
 /* called per autoload when TextDecoder::Numeric is used */
@@ -1021,6 +1124,14 @@ init_pg_text_decoder(void)
 	s_pos_inf = rb_eval_string("1.0/0.0");
 	s_neg_inf = rb_eval_string("-1.0/0.0");
 
+	rb_nativethread_lock_initialize(&numeric_cache_lock);
+#if !defined(_WIN32)
+	pthread_atfork(NULL, NULL, numeric_cache_reset_in_child);
+#endif
+	rb_gc_register_address(&numeric_cache_marker);
+	/* The GC skips dmark for a NULL data pointer, so it must point somewhere. */
+	numeric_cache_marker = TypedData_Wrap_Struct(0, &numeric_cache_type, numeric_cache);
+
 	/* This module encapsulates all decoder classes with text input format */
 	rb_mPG_TextDecoder = rb_define_module_under( rb_mPG, "TextDecoder" );
 	rb_define_private_method(rb_singleton_class(rb_mPG_TextDecoder), "init_date", init_pg_text_decoder_date, 0);

spec/pg/gc_compact_spec.rb

@@ -63,6 +63,14 @@ def conv_array(value)
 			CANCON.socket_io
 		end
 
+		begin
+			TEXTDEC_NUMERIC = PG::TextDecoder::Numeric.new
+		rescue LoadError
+		else
+			TEXTDEC_NUMERIC.decode('1299.00') # cached the second time it's seen
+			BIGDECIMAL = TEXTDEC_NUMERIC.decode('1299.00')
+		end
+
 		begin
 			# Use GC.verify_compaction_references instead of GC.compact .
 			# This has the advantage that all movable objects are actually moved.
@@ -127,6 +135,11 @@ def conv_array(value)
 		expect( CANCON.socket_io ).to be_kind_of( IO )
 	end
 
+	it "should compact PG::TextDecoder::Numeric cache" do
+		skip "bigdecimal not available" unless defined?(BIGDECIMAL)
+		expect( BIGDECIMAL ).to equal( TEXTDEC_NUMERIC.decode('1299.00') )
+	end
+
 	after :all do
 		CONN2.close
 		CONN3.close

spec/pg/type_spec.rb

@@ -250,6 +250,64 @@ def textdec_timestamptz_decode_should_fail(str)
 				end
 			end
 
+			context 'numerics', :bigdecimal do
+				let(:textdec_numeric) { PG::TextDecoder::Numeric.new }
+
+				it 'decodes like BigDecimal()' do
+					['0', '0.00', '-0', '1.200', '1299.00', '12345678', '123456789', 'NaN', 'Infinity', '-Infinity'].each do |value|
+						expect( textdec_numeric.decode(value)._dump ).to eq( BigDecimal(value)._dump )
+					end
+				end
+				it 'returns the same frozen object for equal short values' do
+					expect( textdec_numeric.decode('1299.00') ).to be_frozen
+					expect( textdec_numeric.decode('1299.00') ).to equal( textdec_numeric.decode('1299.00') )
+					expect( textdec_numeric.decode('1299.00') ).not_to equal( textdec_numeric.decode('1299.0') )
+				end
+				it 'returns a new object for a value seen only once' do
+					first = textdec_numeric.decode('31.4159')
+					cached = textdec_numeric.decode('31.4159')
+					expect( cached ).not_to equal( first )
+					expect( textdec_numeric.decode('31.4159') ).to equal( cached )
+				end
+				it 'keeps cached objects alive across GC' do
+					textdec_numeric.decode('7.77')
+					cached_object_id = textdec_numeric.decode('7.77').object_id
+					GC.start
+					expect( textdec_numeric.decode('7.77').object_id ).to eq( cached_object_id )
+				end
+				it 'raises on NaN and Infinity in BigDecimal exception mode' do
+					textdec_numeric.decode('NaN')
+					textdec_numeric.decode('Infinity')
+					BigDecimal.save_exception_mode do
+						BigDecimal.mode(BigDecimal::EXCEPTION_NaN | BigDecimal::EXCEPTION_INFINITY, true)
+						expect{ textdec_numeric.decode('NaN') }.to raise_error( FloatDomainError )
+						expect{ textdec_numeric.decode('Infinity') }.to raise_error( FloatDomainError )
+					end
+				end
+				it 'returns a new object for values longer than 8 characters' do
+					expect( textdec_numeric.decode('123456789') ).not_to equal( textdec_numeric.decode('123456789') )
+				end
+				it 'returns a new object with DISABLE_CACHE' do
+					textdec_numeric.flags = PG::Coder::DISABLE_CACHE
+					expect( textdec_numeric.decode('0.00') ).not_to equal( textdec_numeric.decode('0.00') )
+				end
+				it 'raises on invalid input like BigDecimal()' do
+					['', '-', '1.5x', "1\0002"].each do |value|
+						2.times { expect{ textdec_numeric.decode(value) }.to raise_error( ArgumentError ) }
+					end
+				end
+				it 'decodes in parallel Ractors', :ractor do
+					PG::TextDecoder::Numeric # autoload from the main Ractor
+					ractors = 4.times.map do
+						Ractor.new do
+							decoder = PG::TextDecoder::Numeric.new
+							10_000.times.all? { |i| decoder.decode("#{i % 2000}.00") == i % 2000 }
+						end
+					end
+					expect( ractors.map(&:value) ).to all( be true )
+				end
+			end
+
 			context 'dates' do
 				let(:textdec_date) { PG::TextDecoder::Date.new }