Change set_auth_data_hook to be implemented in C

This doesn't expose the global function PQsetAuthDataHook to ruby, but only the one per connection.
The conversion of "PGconn *" -> "PG::Connection object" is no longer done per WeakMap, but per st_table.
That should make it easier to get Ractor compatible.

Author: larskanis

ext/pg.h

@@ -113,6 +113,10 @@ typedef struct {
 	VALUE encoder_for_put_copy_data;
 	/* Kind of PG::Coder object for casting COPY rows to ruby values */
 	VALUE decoder_for_get_copy_data;
+#ifdef LIBPQ_HAS_PROMPT_OAUTH_DEVICE
+	/* Callback for retrieval of OAuth token */
+	VALUE auth_data_hook;
+#endif
 	/* Ruby encoding index of the client/internal encoding */
 	int enc_idx : PG_ENC_IDX_BITS;
 	/* flags controlling Symbol/String field names */

ext/pg_auth_hooks.c

@@ -8,49 +8,10 @@
 
 #ifdef LIBPQ_HAS_PROMPT_OAUTH_DEVICE
 
-#ifdef TRUFFLERUBY
-static VALUE auth_data_hook;
-#else
 /*
- * On Ruby verisons which support Ractors we store the global callback once
- * per Ractor.
+ * We store the pgconn pointers in a register to retrieve the PG::Connection VALUE in the oauth hook.
  */
-#include "ruby/ractor.h"
-static rb_ractor_local_key_t auth_data_hook_key;
-#endif
-
-static void
-auth_data_hook_init(void)
-{
-#ifdef TRUFFLERUBY
-	auth_data_hook = Qnil;
-	rb_gc_register_address(&auth_data_hook);
-#else
-	auth_data_hook_key = rb_ractor_local_storage_value_newkey();
-#endif
-}
-
-static VALUE
-auth_data_hook_get(void)
-{
-#ifdef TRUFFLERUBY
-	return auth_data_hook;
-#else
-	VALUE hook = Qnil;
-	rb_ractor_local_storage_value_lookup(auth_data_hook_key, &hook);
-	return hook;
-#endif
-}
-
-static void
-auth_data_hook_set(VALUE hook)
-{
-#ifdef TRUFFLERUBY
-	auth_data_hook = hook;
-#else
-	rb_ractor_local_storage_value_set(auth_data_hook_key, hook);
-#endif
-}
+struct st_table *pgconn2value;
 
 static VALUE rb_cPromptOAuthDevice;
 static VALUE rb_cOAuthBearerRequest;
@@ -298,16 +259,20 @@ oauth_bearer_request_hook_cleanup(VALUE self, VALUE ex)
  * currently-registered Ruby auth_data_hook object.
  */
 int
-auth_data_hook_proxy(PGauthData type, PGconn *conn, void *data)
+auth_data_hook_proxy(PGauthData type, PGconn *pgconn, void *data)
 {
-	VALUE proc = auth_data_hook_get(), ret = Qnil;
+	VALUE rb_conn = Qnil;
+	VALUE ret = Qnil;
+
+	if ( st_lookup(pgconn2value, (st_data_t)pgconn, (st_data_t*)&rb_conn) ) {
+		t_pg_connection *this = pg_get_connection( rb_conn );
+		VALUE proc = this->auth_data_hook;
 
-	if (proc != Qnil) {
 		if (type == PQAUTHDATA_PROMPT_OAUTH_DEVICE) {
 			t_pg_prompt_oauth_device *prompt;
 
 			VALUE v_prompt = TypedData_Make_Struct(rb_cPromptOAuthDevice, t_pg_prompt_oauth_device, &pg_prompt_oauth_device_type, prompt);
-			VALUE args[] = { proc, PTR2NUM(conn), v_prompt };
+			VALUE args[] = { proc, rb_conn, v_prompt };
 
 			prompt->prompt = data;
 
@@ -318,7 +283,7 @@ auth_data_hook_proxy(PGauthData type, PGconn *conn, void *data)
 			t_pg_oauth_bearer_request *request;
 
 			VALUE v_request = TypedData_Make_Struct(rb_cOAuthBearerRequest, t_pg_oauth_bearer_request, &pg_oauth_bearer_request_type, request);
-			VALUE args[] = { proc, PTR2NUM(conn), v_request };
+			VALUE args[] = { proc, rb_conn, v_request };
 
 			request->request = data;
 			request->request->cleanup = oauth_bearer_request_cleanup;
@@ -329,44 +294,30 @@ auth_data_hook_proxy(PGauthData type, PGconn *conn, void *data)
 		}
 	}
 
+  /* TODO: a hook can return 1, 0 or -1 */
 	return RTEST(ret);
 }
 
 /*
- * Document-method: PG.set_auth_data_hook
- *
  * call-seq:
- *   PG.set_auth_data_hook {|data| ... } -> Proc
- *
- * If you pass no arguments, it will reset the handler to the default.
+ *    PG.pgconn2value_size -> Integer
  */
 static VALUE
-pg_s_set_auth_data_hook(VALUE _self)
+pg_oauth_pgconn2value_size_get(VALUE self)
 {
-	PQsetAuthDataHook(gvl_auth_data_hook_proxy); // TODO: Add some safeguards?
-
-	VALUE old_proc = auth_data_hook_get(), proc;
-
-	if (rb_block_given_p()) {
-		proc = rb_block_proc();
-	} else {
-		/* if no block is given, set back to default */
-		proc = Qnil;
-	}
-
-	auth_data_hook_set(proc);
-
-	return old_proc;
+	return SIZET2NUM(rb_st_table_size(pgconn2value));
 }
 
+
 void
 init_pg_auth_hooks(void)
 {
-	auth_data_hook_init();
+	pgconn2value = st_init_numtable();
 
-	/* rb_mPG = rb_define_module("PG") */
+	PQsetAuthDataHook(gvl_auth_data_hook_proxy); // TODO: Add some safeguards?
 
-	rb_define_singleton_method(rb_mPG, "set_auth_data_hook", pg_s_set_auth_data_hook, 0);
+	/* rb_mPG = rb_define_module("PG") */
+	rb_define_private_method(rb_singleton_class(rb_mPG), "pgconn2value_size", pg_oauth_pgconn2value_size_get, 0);
 
 	rb_cPromptOAuthDevice = rb_define_class_under(rb_mPG, "PromptOAuthDevice", rb_cObject);
 	rb_undef_alloc_func(rb_cPromptOAuthDevice);

ext/pg_connection.c

@@ -22,6 +22,9 @@ static VALUE pgconn_wait_for_flush( VALUE self );
 static void pgconn_set_internal_encoding_index( VALUE );
 static const rb_data_type_t pg_connection_type;
 static VALUE pgconn_async_flush(VALUE self);
+#ifdef LIBPQ_HAS_PROMPT_OAUTH_DEVICE
+extern struct st_table *pgconn2value;
+#endif
 
 /*
  * Global functions
@@ -179,11 +182,17 @@ pgconn_gc_mark( void *_this )
 	rb_gc_mark_movable( this->trace_stream );
 	rb_gc_mark_movable( this->encoder_for_put_copy_data );
 	rb_gc_mark_movable( this->decoder_for_get_copy_data );
+#ifdef LIBPQ_HAS_PROMPT_OAUTH_DEVICE
+	rb_gc_mark_movable( this->auth_data_hook );
+#endif
 }
 
 static void
 pgconn_gc_compact( void *_this )
 {
+#ifdef LIBPQ_HAS_PROMPT_OAUTH_DEVICE
+	VALUE old_rb_conn;
+#endif
 	t_pg_connection *this = (t_pg_connection *)_this;
 	pg_gc_location( this->socket_io );
 	pg_gc_location( this->notice_receiver );
@@ -193,6 +202,15 @@ pgconn_gc_compact( void *_this )
 	pg_gc_location( this->trace_stream );
 	pg_gc_location( this->encoder_for_put_copy_data );
 	pg_gc_location( this->decoder_for_get_copy_data );
+
+#ifdef LIBPQ_HAS_PROMPT_OAUTH_DEVICE
+	pg_gc_location( this->auth_data_hook );
+	/* update the PG::Connection object which is maybe stored in pgconn2value */
+	if ( st_lookup(pgconn2value, (st_data_t)this->pgconn, (st_data_t*)&old_rb_conn) ) {
+		VALUE new_rb_conn = rb_gc_location(old_rb_conn);
+		st_insert( pgconn2value, (st_data_t)this->pgconn, (st_data_t)new_rb_conn );
+	}
+#endif
 }
 
 
@@ -210,9 +228,15 @@ pgconn_gc_free( void *_this )
 		}
 	}
 #endif
-	if (this->pgconn != NULL)
+	if (this->pgconn != NULL) {
 		PQfinish( this->pgconn );
 
+#ifdef LIBPQ_HAS_PROMPT_OAUTH_DEVICE
+		/* Remove from auth hook callback table */
+		st_delete(pgconn2value, (st_data_t *)&this->pgconn, NULL );
+#endif
+	}
+
 	xfree(this);
 }
 
@@ -264,6 +288,9 @@ pgconn_s_allocate( VALUE klass )
 	RB_OBJ_WRITE(self, &this->type_map_for_results, pg_typemap_all_strings);
 	RB_OBJ_WRITE(self, &this->encoder_for_put_copy_data, Qnil);
 	RB_OBJ_WRITE(self, &this->decoder_for_get_copy_data, Qnil);
+#ifdef LIBPQ_HAS_PROMPT_OAUTH_DEVICE
+	RB_OBJ_WRITE(self, &this->auth_data_hook, Qnil);
+#endif
 	RB_OBJ_WRITE(self, &this->trace_stream, Qnil);
 	rb_ivar_set(self, rb_intern("@calls_to_put_copy_data"), INT2FIX(0));
 	rb_ivar_set(self, rb_intern("@iopts_for_reset"), Qnil);
@@ -623,22 +650,45 @@ pgconn_reset_poll(VALUE self)
 	return INT2FIX((int)status);
 }
 
+#ifdef LIBPQ_HAS_PROMPT_OAUTH_DEVICE
+
 /*
  * call-seq:
- *    conn.pgconn_address()
- *
- * Returns the PGconn address.
+ *    conn.auth_data_hook(&block)
  *
- * It can be used to compare with the address received by the OAuth hook:
+ * Set a auth data hook.
+ */
+static VALUE
+pgconn_auth_data_hook_set(VALUE self, VALUE proc)
+{
+	t_pg_connection *this = pg_get_connection( self );
+
+	if (rb_obj_is_proc(proc)) {
+		/* set proc */
+		st_insert( pgconn2value, (st_data_t)this->pgconn, (st_data_t)self );
+	} else if (NIL_P(proc)) {
+		/* if nil is given, set back to default */
+		st_delete( pgconn2value, (st_data_t *)&this->pgconn, NULL );
+	} else {
+		rb_raise(rb_eArgError, "Proc object expected");
+	}
+	RB_OBJ_WRITE(self, &this->auth_data_hook, proc);
+	return proc;
+}
+
+/*
+ * call-seq:
+ *    conn.auth_data_hook()
  *
- *   PG.set_auth_data_hook do |pgconn_address, data|
+ * Returns the defined auth data hook.
  */
 static VALUE
-pgconn_pgconn_address(VALUE self)
+pgconn_auth_data_hook_get(VALUE self)
 {
-	return PTR2NUM(pg_get_pgconn(self));
+	return pg_get_connection(self)->auth_data_hook;
 }
 
+#endif
 
 /*
  * call-seq:
@@ -4766,7 +4816,10 @@ init_pg_connection(void)
 	rb_define_private_method(rb_cPGconn, "reset_start2", pgconn_reset_start2, 1);
 	rb_define_method(rb_cPGconn, "reset_poll", pgconn_reset_poll, 0);
 	rb_define_alias(rb_cPGconn, "close", "finish");
-	rb_define_private_method(rb_cPGconn, "pgconn_address", pgconn_pgconn_address, 0);
+#ifdef LIBPQ_HAS_PROMPT_OAUTH_DEVICE
+	rb_define_method(rb_cPGconn, "auth_data_hook=", pgconn_auth_data_hook_set, 1);
+	rb_define_method(rb_cPGconn, "auth_data_hook", pgconn_auth_data_hook_get, 0);
+#endif
 
 	/******     PG::Connection INSTANCE METHODS: Connection Status     ******/
 	rb_define_method(rb_cPGconn, "db", pgconn_db, 0);

lib/pg/connection.rb

@@ -943,11 +943,11 @@ def new(*args, **kwargs)
 			end
 			conn = self.connect_start(iopts) or
 										raise(PG::Error, "Unable to create a new connection")
-			if set_auth_data_hook
-				@@auth_mutex.synchronize do
-					@@pgconn_map[conn.send(:pgconn_address)] = conn
-				end
-				conn.instance_variable_set(:@auth_data_hook, set_auth_data_hook)
+
+			if conn.respond_to?(:auth_data_hook)
+				conn.auth_data_hook = set_auth_data_hook
+			elsif set_auth_data_hook
+				raise ArgumentError, "invalid option set_auth_data_hook"
 			end
 
 			raise PG::ConnectionBad, conn.error_message if conn.status == PG::CONNECTION_BAD
@@ -1096,23 +1096,5 @@ def async_api=(enable)
 		end
 	end
 
-	if PG.respond_to?(:set_auth_data_hook)
-		def call_auth_data_hook(data)
-			@auth_data_hook.call(self, data)
-		end
-
-		@@auth_mutex = Mutex.new
-		@@pgconn_map = ObjectSpace::WeakMap.new
-
-		PG.set_auth_data_hook do |conn_num, data|
-			@@auth_mutex.synchronize do
-				conn = @@pgconn_map[conn_num]
-				if conn
-					conn.call_auth_data_hook(data)
-				end
-			end
-		end
-	end
-
 	self.async_api = true
 end # class PG::Connection

spec/pg/connection_async_spec.rb

@@ -264,16 +264,17 @@ def interrupt_thread(exc=nil)
 				end
 			end
 
-			10.times do
+			before = PG.send(:pgconn2value_size)
+			20.times do
 				conn = PG.connect(host: "localhost", port: @port, dbname: "test", user: "testuseroauth", oauth_issuer: "http://localhost:#{@port + 3}", oauth_client_id: "foo", set_auth_data_hook: hook)
 				conn.exec("SELECT 1")
 			end
 
-			before = PG::Connection.class_variable_get(:@@pgconn_map).keys.size
 			GC.start
-			after = PG::Connection.class_variable_get(:@@pgconn_map).keys.size
+			after = PG.send(:pgconn2value_size)
 
-			expect(before - after).to be_between(2, 20)
+			# Number of GC'ed objects
+			expect(before + 20 - after).to be_between(1, 50)
 		end
 
 		# TODO: Is resetting the global hook still useful, when the hook is per connection?

spec/pg/gc_compact_spec.rb

@@ -24,6 +24,8 @@
 require_relative '../helpers'
 
 describe "GC.compact", if: GC.respond_to?(:compact) do
+	hook_called = false
+
 	before :all do
 		TM1 = Class.new(PG::TypeMapByClass) do
 			def conv_array(value)
@@ -57,6 +59,20 @@ def conv_array(value)
 			CANCON.socket_io
 		end
 
+		if PG::Connection.instance_methods.include?(:auth_data_hook)
+			build_oauth_validator
+			@old_env, ENV["PGOAUTHDEBUG"] = ENV["PGOAUTHDEBUG"], "UNSAFE"
+			HOOKED_CONN = PG::Connection.connect_start( "host=localhost port=#{@port} dbname=test user=testuseroauth oauth_issuer=http://localhost:#{@port + 3} oauth_client_id=foo" )
+			HOOKED_CONN.auth_data_hook = proc do |conn, data|
+				case data
+				when PG::OAuthBearerRequest
+					data.token = "yes"
+					hook_called = true
+					true
+				end
+			end
+		end
+
 		begin
 			# Use GC.verify_compaction_references instead of GC.compact .
 			# This has the advantage that all movable objects are actually moved.
@@ -111,6 +127,14 @@ def conv_array(value)
 		expect( CANCON.socket_io ).to be_kind_of( IO )
 	end
 
+	it "should compact PG::Connection in pgconn2value", :postgresql_18 do
+		wait_for_polling_ok(HOOKED_CONN)
+		expect( HOOKED_CONN.error_message ).to eq("")
+		HOOKED_CONN.finish
+		expect( hook_called ).to be_truthy
+		ENV["PGOAUTHDEBUG"] = @old_env
+	end
+
 	after :all do
 		CONN2.close
 	end