From 64451ecb7ef65dc2d0f0b33413a2e8a970739259 Mon Sep 17 00:00:00 2001
From: Mike Dalessio <mike@37signals.com>
Date: Mon, 16 Sep 2024 13:23:16 -0400
Subject: [PATCH] Database#close invokes discard if the pid has changed

This will help prevent corrupting the database.

Also provide a warning to let people know they're doing something
that's not supported by sqlite.
---
 CHANGELOG.md            |  6 ++++++
 ext/sqlite3/database.c  | 18 +++++++-----------
 lib/sqlite3/database.rb | 19 +++++++++++++++++++
 test/test_database.rb   | 21 +++++++++++++++++++++
 4 files changed, 53 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e91dbdac..c9808bfb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,12 @@
 
 ## next / unreleased
 
+### Added
+
+- New method `Database#discard` is intended to be used after forking a process with an open database connection. Forking is not supported by sqlite, and as an alternativ `discard` will "close" the connection and release what resources can be safely released. [#558] @flavorjones
+- `Database#close` will detect if it's being closed from a child process, and if so will emit a warning and invoke `discard` instead. [#558] @flavorjones
+
+
 ### Improved
 
 - Use `sqlite3_close_v2` to close databases in a deferred manner if there are unclosed prepared statements. Previously closing a database while statements were open resulted in a `BusyException`. See https://www.sqlite.org/c3ref/close.html for more context. [#557] @flavorjones
diff --git a/ext/sqlite3/database.c b/ext/sqlite3/database.c
index 3829e5f9..cd313cfc 100644
--- a/ext/sqlite3/database.c
+++ b/ext/sqlite3/database.c
@@ -23,7 +23,7 @@ static void
 deallocate(void *ctx)
 {
     sqlite3RubyPtr c = (sqlite3RubyPtr)ctx;
-    sqlite3 *db     = c->db;
+    sqlite3 *db = c->db;
 
     if (db) { sqlite3_close_v2(db); }
     xfree(c);
@@ -62,8 +62,6 @@ utf16_string_value_ptr(VALUE str)
     return RSTRING_PTR(str);
 }
 
-static VALUE sqlite3_rb_close(VALUE self);
-
 sqlite3RubyPtr
 sqlite3_database_unwrap(VALUE database)
 {
@@ -119,18 +117,16 @@ rb_sqlite3_disable_quirk_mode(VALUE self)
 #endif
 }
 
-/* call-seq: db.close
- *
- * Closes this database.
- */
 static VALUE
-sqlite3_rb_close(VALUE self)
+sqlite3_rb__close(VALUE self)
 {
     sqlite3RubyPtr ctx;
     TypedData_Get_Struct(self, sqlite3Ruby, &database_type, ctx);
 
-    CHECK(ctx->db, sqlite3_close_v2(ctx->db));
-    ctx->db = NULL;
+    if (ctx->db) {
+        CHECK(ctx->db, sqlite3_close_v2(ctx->db));
+        ctx->db = NULL;
+    }
 
     rb_iv_set(self, "-aggregators", Qnil);
 
@@ -921,7 +917,7 @@ init_sqlite3_database(void)
     rb_define_private_method(cSqlite3Database, "open_v2", rb_sqlite3_open_v2, 3);
     rb_define_private_method(cSqlite3Database, "open16", rb_sqlite3_open16, 1);
     rb_define_method(cSqlite3Database, "collation", collation, 2);
-    rb_define_method(cSqlite3Database, "close", sqlite3_rb_close, 0);
+    rb_define_private_method(cSqlite3Database, "_close", sqlite3_rb__close, 0);
     rb_define_method(cSqlite3Database, "discard", sqlite3_rb_discard, 0);
     rb_define_method(cSqlite3Database, "closed?", closed_p, 0);
     rb_define_method(cSqlite3Database, "total_changes", total_changes, 0);
diff --git a/lib/sqlite3/database.rb b/lib/sqlite3/database.rb
index c4a490ab..e3a14d73 100644
--- a/lib/sqlite3/database.rb
+++ b/lib/sqlite3/database.rb
@@ -133,6 +133,7 @@ def initialize file, options = {}, zvfs = nil
       @results_as_hash = options[:results_as_hash]
       @readonly = mode & Constants::Open::READONLY != 0
       @default_transaction_mode = options[:default_transaction_mode] || :deferred
+      @owner_pid = Process.pid
 
       if block_given?
         begin
@@ -143,6 +144,24 @@ def initialize file, options = {}, zvfs = nil
       end
     end
 
+    # Close the database and release all associated resources.
+    #
+    # ⚠ If the process that created the database forks a child process, and this method is called
+    # from the child process, then this method will _not_ free memory resources and instead will
+    # call discard. This is a memory leak, but is safer than risking database corruption.
+    #
+    # See adr/2024-09-fork-safety.md for more information on fork safety.
+    def close
+      if Process.pid != @owner_pid
+        warn "WARNING: An open sqlite database connection was inherited from a forked process and " \
+          "is being discarded. This is a memory leak. If possible, please close all sqlite " \
+          "database connections before forking."
+        discard
+      else
+        _close
+      end
+    end
+
     # call-seq: db.encoding
     #
     # Fetch the encoding set on this database
diff --git a/test/test_database.rb b/test/test_database.rb
index 38ac0c87..f95c6ebf 100644
--- a/test/test_database.rb
+++ b/test/test_database.rb
@@ -740,6 +740,27 @@ def test_discard_a_connection
       end
     end
 
+    def test_close_in_a_new_process_calls_discard_and_warns
+      db = SQLite3::Database.new("test.db")
+
+      assert_equal(Process.pid, db.instance_variable_get(:@owner_pid))
+
+      called = false
+      db.define_singleton_method(:discard) do
+        called = true
+      end
+      db.instance_variable_set(:@owner_pid, Process.pid + 1)
+
+      assert_output(nil, /WARNING: An open sqlite database connection was inherited from a forked process/) do
+        db.close
+      end
+      assert(called)
+    ensure
+      db.instance_variable_set(:@owner_pid, Process.pid)
+      db.close
+      FileUtils.rm_f("test.db")
+    end
+
     def test_discard_a_closed_connection
       db = SQLite3::Database.new("test.db")
       db.close