From 10e89a113f11fc5f6f034a248f7ebf689a4ae358 Mon Sep 17 00:00:00 2001
From: Roman Sandler <5535625+sandlerr@users.noreply.github.com>
Date: Mon, 26 Dec 2022 23:19:01 +1100
Subject: [PATCH] Initial Ractor support

---
 ext/sqlite3/extconf.rb          |  3 ++
 ext/sqlite3/sqlite3.c           |  5 ++
 lib/sqlite3/database.rb         |  9 +++-
 test/test_integration_ractor.rb | 84 +++++++++++++++++++++++++++++++++
 4 files changed, 100 insertions(+), 1 deletion(-)
 create mode 100644 test/test_integration_ractor.rb

diff --git a/ext/sqlite3/extconf.rb b/ext/sqlite3/extconf.rb
index c492d20d..36e40771 100644
--- a/ext/sqlite3/extconf.rb
+++ b/ext/sqlite3/extconf.rb
@@ -108,6 +108,9 @@ def configure_extension
         # Functions defined in 2.1 but not 2.0
         have_func('rb_integer_pack')
 
+        # Functions defined in 3.0 but not 2.7
+        have_func('rb_ext_ractor_safe')
+
         # These functions may not be defined
         have_func('sqlite3_initialize')
         have_func('sqlite3_backup_init')
diff --git a/ext/sqlite3/sqlite3.c b/ext/sqlite3/sqlite3.c
index 70612f1e..418c2050 100644
--- a/ext/sqlite3/sqlite3.c
+++ b/ext/sqlite3/sqlite3.c
@@ -84,6 +84,11 @@ static VALUE threadsafe_p(VALUE UNUSED(klass))
 
 void init_sqlite3_constants()
 {
+  #ifdef HAVE_RB_EXT_RACTOR_SAFE
+  if (sqlite3_threadsafe()) {
+    rb_ext_ractor_safe(true);
+  }
+  #endif
   VALUE mSqlite3Constants;
   VALUE mSqlite3Open;
 
diff --git a/lib/sqlite3/database.rb b/lib/sqlite3/database.rb
index 433cc642..6232909b 100644
--- a/lib/sqlite3/database.rb
+++ b/lib/sqlite3/database.rb
@@ -724,7 +724,14 @@ def translate_from_db types, row
 
     private
 
-    NULL_TRANSLATOR = lambda { |_, row| row }
+    # NULL_TRANSLATOR used to be a lambda, but a lambda can't be frozen (properly)
+    # and so can't work with ractors.
+    class NullTranslatorImplementation
+      def self.call(_, row)
+        row
+      end
+    end
+    NULL_TRANSLATOR = NullTranslatorImplementation
 
     def make_type_translator should_translate
       if should_translate
diff --git a/test/test_integration_ractor.rb b/test/test_integration_ractor.rb
new file mode 100644
index 00000000..1273de54
--- /dev/null
+++ b/test/test_integration_ractor.rb
@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+require 'helper'
+require 'fileutils'
+
+class TC_Integration_Ractor < SQLite3::TestCase
+  STRESS_DB_NAME = "stress.db"
+
+  def setup
+    teardown
+  end
+
+  def teardown
+    FileUtils.rm_rf(Dir.glob "#{STRESS_DB_NAME}*")
+  end
+
+  def test_ractor_share_database
+    skip('Requires Ruby with Ractors') if RUBY_VERSION < '3.2.0'
+    assert SQLite3.threadsafe?
+
+    db_receiver = Ractor.new do
+      db = Ractor.receive
+      Ractor.yield db.object_id
+      begin
+        db.execute("create table test_table ( b integer primary key)")
+        raise "Should have raised an exception in db.execute()"
+      rescue => e
+        Ractor.yield e
+      end
+    end
+    db_creator = Ractor.new(db_receiver) do |db_receiver|
+      db = SQLite3::Database.open(":memory:")
+      Ractor.yield db.object_id
+      db_receiver.send(db)
+      sleep 0.1
+      db.execute("create table test_table ( a integer primary key)")
+    end
+    first_oid = db_creator.take
+    second_oid = db_receiver.take
+    assert_not_equal first_oid, second_oid
+    ex = db_receiver.take
+    # For now, let's assert that you can't pass database connections around
+    # between different Ractors. Letting a live DB connection exist in two
+    # threads that are running concurrently might expose us to footguns and
+    # lead to data corruption, so we should avoid this possibility and wait
+    # until connections can be given away using `yield` or `send`.
+    assert_equal "prepare called on a closed database", ex.message
+  end
+
+  def test_ractor_stress
+    skip('Requires Ruby with Ractors') if RUBY_VERSION < '3.2.0'
+    assert SQLite3.threadsafe?
+
+    # Testing with a file instead of :memory: since it can be more realistic
+    # compared with real production use, and so discover problems that in-
+    # memory testing won't find. Trivial example: STRESS_DB_NAME needs to be
+    # frozen to pass into the Ractor, but :memory: might avoid that problem by
+    # using a literal string.
+    db = SQLite3::Database.open(STRESS_DB_NAME)
+    db.execute("PRAGMA journal_mode=WAL") # A little slow without this
+    db.execute("create table stress_test (a integer primary_key, b text)")
+    random = Random.new.freeze
+    ractors = (0..9).map do |ractor_number|
+      Ractor.new(random, ractor_number) do |random, ractor_number|
+        db_in_ractor = SQLite3::Database.open(STRESS_DB_NAME)
+        db_in_ractor.busy_handler do
+          sleep random.rand / 100 # Lots of busy errors happen with multiple concurrent writers
+          true
+        end
+        100.times do |i|
+          db_in_ractor.execute("insert into stress_test(a, b) values (#{ractor_number * 100 + i}, '#{random.rand}')")
+        end
+      end
+    end
+    ractors.each {|r| r.take}
+    final_check = Ractor.new do
+      db_in_ractor = SQLite3::Database.open(STRESS_DB_NAME)
+      res = db_in_ractor.execute("select count(*) from stress_test")
+      Ractor.yield res
+    end
+    res = final_check.take
+    assert_equal 1000, res[0][0]
+  end
+end