Stop sqlite calling into moved Ruby objects

ext/sqlite3/aggregator.c

@@ -254,10 +254,17 @@ rb_sqlite3_define_aggregator2(VALUE self, VALUE aggregator, VALUE ruby_name)
     CHECK(ctx->db, status);
 
     rb_ary_push(aggregators, aw);
+    RB_OBJ_WRITE(self, &ctx->aggregators, aggregators);
 
     return self;
 }
 
+void
+rb_sqlite3_aggregator_pin_instances(VALUE aw)
+{
+    rb_sqlite3_pin_array_and_contents(rb_iv_get(aw, "-instances"));
+}
+
 void
 rb_sqlite3_aggregator_init(void)
 {

ext/sqlite3/aggregator.h

@@ -7,4 +7,8 @@ VALUE rb_sqlite3_define_aggregator2(VALUE self, VALUE aggregator, VALUE ruby_nam
 
 void rb_sqlite3_aggregator_init(void);
 
+/* sqlite stores each live instance's VALUE in its aggregate context, so those
+ * must not move either. */
+void rb_sqlite3_aggregator_pin_instances(VALUE aw);
+
 #endif

ext/sqlite3/database.c

@@ -70,11 +70,64 @@ close_or_discard_db(sqlite3RubyPtr ctx)
 }
 
 
+void
+rb_sqlite3_pin_array_and_contents(VALUE ary)
+{
+    long i;
+
+    if (NIL_P(ary) || !ary) { return; }
+
+    rb_gc_mark(ary);
+    for (i = 0; i < RARRAY_LEN(ary); i++) {
+        rb_gc_mark(RARRAY_AREF(ary, i));
+    }
+}
+
+static int
+pin_hash_value(VALUE UNUSED(key), VALUE value, VALUE UNUSED(arg))
+{
+    rb_gc_mark(value);
+    return ST_CONTINUE;
+}
+
+static void
+pin_hash_and_contents(VALUE hash)
+{
+    if (NIL_P(hash) || !hash) { return; }
+
+    rb_gc_mark(hash);
+    rb_hash_foreach(hash, pin_hash_value, 0);
+}
+
+/* Each wrapper also owns live aggregate instances, whose VALUEs sqlite keeps in
+ * its own aggregate contexts. */
+static void
+pin_aggregators(VALUE aggregators)
+{
+    long i;
+
+    rb_sqlite3_pin_array_and_contents(aggregators);
+
+    if (NIL_P(aggregators) || !aggregators) { return; }
+
+    for (i = 0; i < RARRAY_LEN(aggregators); i++) {
+        rb_sqlite3_aggregator_pin_instances(RARRAY_AREF(aggregators, i));
+    }
+}
+
 static void
 database_mark(void *ctx)
 {
     sqlite3RubyPtr c = (sqlite3RubyPtr)ctx;
+
+    /* sqlite holds raw pointers to these, so they must not move. */
     rb_gc_mark(c->busy_handler);
+    rb_gc_mark(c->trace_handler);
+    rb_gc_mark(c->authorizer);
+
+    rb_sqlite3_pin_array_and_contents(c->functions);
+    pin_hash_and_contents(c->collations);
+    pin_aggregators(c->aggregators);
 }
 
 static void
@@ -252,9 +305,8 @@ total_changes(VALUE self)
 static void
 tracefunc(void *data, const char *sql)
 {
-    VALUE self = (VALUE)data;
-    VALUE thing = rb_iv_get(self, "@tracefunc");
-    rb_funcall(thing, rb_intern("call"), 1, rb_str_new2(sql));
+    sqlite3RubyPtr ctx = (sqlite3RubyPtr)data;
+    rb_funcall(ctx->trace_handler, rb_intern("call"), 1, rb_str_new2(sql));
 }
 
 /* call-seq:
@@ -278,9 +330,9 @@ trace(int argc, VALUE *argv, VALUE self)
 
     if (NIL_P(block) && rb_block_given_p()) { block = rb_block_proc(); }
 
-    rb_iv_set(self, "@tracefunc", block);
+    RB_OBJ_WRITE(self, &ctx->trace_handler, block);
 
-    sqlite3_trace(ctx->db, NIL_P(block) ? NULL : tracefunc, (void *)self);
+    sqlite3_trace(ctx->db, NIL_P(block) ? NULL : tracefunc, (void *)ctx);
 
     return self;
 }
@@ -509,7 +561,7 @@ static VALUE
 define_function_with_flags(VALUE self, VALUE name, VALUE flags)
 {
     sqlite3RubyPtr ctx;
-    VALUE block;
+    VALUE block, functions;
     int status;
 
     TypedData_Get_Struct(self, sqlite3Ruby, &database_type, ctx);
@@ -530,7 +582,9 @@ define_function_with_flags(VALUE self, VALUE name, VALUE flags)
 
     CHECK(ctx->db, status);
 
-    rb_ary_push(rb_iv_get(self, "@functions"), block);
+    functions = rb_iv_get(self, "@functions");
+    rb_ary_push(functions, block);
+    RB_OBJ_WRITE(self, &ctx->functions, functions);
 
     return self;
 }
@@ -632,14 +686,13 @@ rb_sqlite3_auth(
     const char *_c,
     const char *_d)
 {
-    VALUE self   = (VALUE)ctx;
+    sqlite3RubyPtr db_ctx = (sqlite3RubyPtr)ctx;
     VALUE action = INT2NUM(_action);
     VALUE a      = _a ? rb_str_new2(_a) : Qnil;
     VALUE b      = _b ? rb_str_new2(_b) : Qnil;
     VALUE c      = _c ? rb_str_new2(_c) : Qnil;
     VALUE d      = _d ? rb_str_new2(_d) : Qnil;
-    VALUE callback = rb_iv_get(self, "@authorizer");
-    VALUE result = rb_funcall(callback, rb_intern("call"), 5, action, a, b, c, d);
+    VALUE result = rb_funcall(db_ctx->authorizer, rb_intern("call"), 5, action, a, b, c, d);
 
     if (T_FIXNUM == TYPE(result)) { return (int)NUM2INT(result); }
     if (Qtrue == result) { return SQLITE_OK; }
@@ -668,12 +721,12 @@ set_authorizer(VALUE self, VALUE authorizer)
     REQUIRE_OPEN_DB(ctx);
 
     status = sqlite3_set_authorizer(
-                 ctx->db, NIL_P(authorizer) ? NULL : rb_sqlite3_auth, (void *)self
+                 ctx->db, NIL_P(authorizer) ? NULL : rb_sqlite3_auth, (void *)ctx
              );
 
     CHECK(ctx->db, status);
 
-    rb_iv_set(self, "@authorizer", authorizer);
+    RB_OBJ_WRITE(self, &ctx->authorizer, authorizer);
 
     return self;
 }
@@ -756,6 +809,7 @@ static VALUE
 collation(VALUE self, VALUE name, VALUE comparator)
 {
     sqlite3RubyPtr ctx;
+    VALUE collations;
     TypedData_Get_Struct(self, sqlite3Ruby, &database_type, ctx);
     REQUIRE_OPEN_DB(ctx);
 
@@ -766,8 +820,10 @@ collation(VALUE self, VALUE name, VALUE comparator)
               (void *)comparator,
               NIL_P(comparator) ? NULL : rb_comparator_func));
 
-    /* Make sure our comparator doesn't get garbage collected. */
-    rb_hash_aset(rb_iv_get(self, "@collations"), name, comparator);
+    /* sqlite holds a raw pointer to the comparator, so keep it alive and unmoved. */
+    collations = rb_iv_get(self, "@collations");
+    rb_hash_aset(collations, name, comparator);
+    RB_OBJ_WRITE(self, &ctx->collations, collations);
 
     return self;
 }

ext/sqlite3/database.h

@@ -10,6 +10,11 @@
 struct _sqlite3Ruby {
     sqlite3 *db;
     VALUE busy_handler;
+    VALUE functions;
+    VALUE collations;
+    VALUE aggregators;
+    VALUE trace_handler;
+    VALUE authorizer;
     int stmt_timeout;
     struct timespec stmt_deadline;
     rb_pid_t owner;
@@ -19,6 +24,9 @@ struct _sqlite3Ruby {
 typedef struct _sqlite3Ruby sqlite3Ruby;
 typedef sqlite3Ruby *sqlite3RubyPtr;
 
+/* Pinning a collection doesn't pin what's in it, hence both. */
+void rb_sqlite3_pin_array_and_contents(VALUE ary);
+
 void init_sqlite3_database();
 void set_sqlite3_func_result(sqlite3_context *ctx, VALUE result);
 

lib/sqlite3/database.rb

@@ -171,8 +171,6 @@ def initialize file, options = {}, zvfs = nil
         end
       end
 
-      @tracefunc = nil
-      @authorizer = nil
       @progress_handler = nil
       @collations = {}
       @functions = []

test/test_collation.rb

@@ -14,6 +14,13 @@ def compare left, right
       end
     end
 
+    # Used by one test only, so a live count of 1 means just the current one.
+    class ReleasableComparator
+      def compare(left, right)
+        left <=> right
+      end
+    end
+
     def setup
       @db = SQLite3::Database.new(":memory:")
       @create = "create table ex(id int, data string)"
@@ -33,6 +40,31 @@ def test_custom_collation
       assert_equal 1, comparator.calls.length
     end
 
+    def test_collation_does_not_use_moved_comparator_after_gc_compaction
+      skip_unless_compaction_supported
+
+      # Registered on another thread so that no reference to the comparator is
+      # left on this thread's machine stack, where conservative scanning would
+      # pin it. A pinned comparator never moves and the test passes without
+      # exercising the fix.
+      Thread.new { @db.collation "foo", Comparator.new }.join
+
+      gc_verify_compaction_references
+
+      @db.execute("select data from ex order by 1 collate foo")
+      assert_equal 1, @db.collations["foo"].calls.length
+    end
+
+    # Passes on main as well. It guards the design rather than the fix: pinning
+    # every comparator makes the lifetime of @collations matter, so this rules
+    # out ever making it an append-only array.
+    def test_replacing_a_collation_releases_the_previous_comparator
+      3.times { @db.collation "foo", ReleasableComparator.new }
+      GC.start(full_mark: true, immediate_sweep: true)
+
+      assert_equal 1, ObjectSpace.each_object(ReleasableComparator).count
+    end
+
     def test_remove_collation
       comparator = Comparator.new
 

test/test_integration.rb

@@ -82,6 +82,29 @@ def test_trace
     assert_equal "select * from foo", result
   end
 
+  def test_trace_does_not_use_moved_block_after_gc_compaction
+    skip_unless_compaction_supported
+
+    result = nil
+    @db.trace { |sql| result = sql }
+
+    gc_verify_compaction_references
+
+    @db.execute "select * from foo"
+    assert_equal "select * from foo", result
+  end
+
+  def test_authorizer_does_not_use_moved_block_after_gc_compaction
+    skip_unless_compaction_supported
+
+    @db.authorizer { |type, a, b, c, d| 0 }
+
+    gc_verify_compaction_references
+
+    rows = @db.execute "select * from foo"
+    assert_equal 3, rows.length
+  end
+
   def test_authorizer_okay
     @db.authorizer { |type, a, b, c, d| 0 }
     rows = @db.execute "select * from foo"
@@ -564,6 +587,19 @@ def test_create_function
     assert_match(/>>>.*<<</, value)
   end
 
+  def test_create_function_does_not_use_moved_block_after_gc_compaction
+    skip_unless_compaction_supported
+
+    @db.create_function("munge", 1) do |func, x|
+      func.result = ">>>#{x}<<<"
+    end
+
+    gc_verify_compaction_references
+
+    value = @db.get_first_value("select munge(b) from foo where a=1")
+    assert_match(/>>>.*<<</, value)
+  end
+
   def test_bind_array_parameter
     result = @db.get_first_value("select b from foo where a=? and b=?",
       [1, "foo"])

test/test_integration_aggregate.rb

@@ -342,6 +342,18 @@ def finalize
     end
   end
 
+  class CompactingAggregator
+    def step(*args)
+      @sum ||= 0
+      args.each { |a| @sum += a.to_i }
+      GC.verify_compaction_references(expand_heap: true, toward: :empty)
+    end
+
+    def finalize
+      @sum
+    end
+  end
+
   class AccumulateAggregator2
     def step(a, b)
       @sum ||= 1
@@ -364,6 +376,26 @@ def test_define_aggregator_with_two_different_arities
     assert_equal 2145, values[1]
   end
 
+  def test_define_aggregator_does_not_use_moved_aggregator_after_gc_compaction
+    skip_unless_compaction_supported
+
+    @db.define_aggregator("accumulate", AccumulateAggregator.new)
+
+    gc_verify_compaction_references
+
+    assert_equal 33, @db.get_first_value("select accumulate(c) from foo")
+  end
+
+  def test_define_aggregator_does_not_use_moved_instances_after_gc_compaction
+    skip_unless_compaction_supported
+
+    @db.define_aggregator("accumulate", CompactingAggregator.new)
+
+    values = @db.get_first_row("select accumulate(a), accumulate(c) from foo")
+    assert_equal 6, values[0]
+    assert_equal 33, values[1]
+  end
+
   def test_step_on_statement_whose_database_was_closed_does_not_use_freed_aggregator
     @db.define_aggregator("accumulate", AccumulateAggregator.new)
     stmt = @db.prepare("select accumulate(c) from foo")