Automatically close any open writable connections after a fork.

Author: flavorjones

README.md

@@ -153,14 +153,14 @@ fine to share, but may require manual locking for thread safety.
 
 [Sqlite is not fork
 safe](https://www.sqlite.org/howtocorrupt.html#_carrying_an_open_database_connection_across_a_fork_)
-and you should not carry an open database connection across a `fork()`. Using an inherited
+and instructs users to not carry an open writable database connection across a `fork()`. Using an inherited
 connection in the child may corrupt your database, leak memory, or cause other undefined behavior.
 
-Instead, whenever possible, close connections in the parent before forking.
+To help protect users of this gem from accidental corruption due to this lack of fork safety, the gem will immediately close any open writable databases in the child after a fork.
 
-If that's not possible or convenient, then immediately close any inherited connections in the child
-after forking, before opening any new connections. This will incur a small one-time memory leak per
-connection, but that's preferable to potentially corrupting your database.
+Whenever possible, close writable connections in the parent before forking. Discarding writable
+connections in the child will incur a small one-time memory leak per connection, but that's
+preferable to potentially corrupting your database.
 
 See [./adr/2024-09-fork-safety.md](./adr/2024-09-fork-safety.md) for more information and context.
 

adr/2024-09-fork-safety.md

@@ -1,5 +1,5 @@
 
-# 2024-09 Discard database connections when carried across fork()of fork safety
+# 2024-09 Automatically close database connections when carried across fork()
 
 ## Status
 
@@ -15,17 +15,23 @@ SQLite is known to not be fork-safe[^howto], so this was not entirely surprising
 Advice from upstream contributors[^advice] is, essentially: don't fork if you have open database connections. Or, if you have forked, don't call `sqlite3_close` on those connections and thereby leak some amount of memory in the child process. Neither of these options are ideal, see below.
 
 
-## Decision
+## Decisions
 
-Open database connections carried across a `fork()` will not be fully closed in the child process, to avoid the risk of corrupting the database file.
+1. Open writable database connections carried across a `fork()` will automatically be closed in the child process to mitigate the risk of corrupting the database file.
+2. These connections will be incompletely closed ("discarded") which will result in a one-time memory leak in the child process.
 
-The sqlite3-ruby gem will track the ID of the process that opened each database connection. If, when the database is closed (either explicitly with `Database#close` or implicitly via GC) the current process ID is different from the original process, then we "discard" the connection.
+First, the gem will register an "after fork" handler via `Process._fork` that will close any open writable database connections in the child process. This is a best-effort attempt to avoid corruption, but it is not guaranteed to prevent corruption in all cases. Any connections closed by this handler will also emit a warning to let users know what's happening.
+
+Second, the sqlite3-ruby gem will store the ID of the process that opened each database connection. If, when a writable database is closed (either explicitly with `Database#close` or implicitly via GC or after-fork callback) the current process ID is different from the original process, then we "discard" the connection.
 
 "Discard" here means:
 
 - The `Database` object acts "closed", including returning `true` from `#closed?`.
 - `sqlite3_close_v2` is not called on the object, because it is unsafe to do so per sqlite instructions[^howto]. As a result, some memory will be lost permanently (a one-time "memory leak").
 - Open file descriptors associated with the database are closed.
+- Any memory that can be freed safely is recovered.
+
+Note that readonly databases are being treated as "fork safe" and are not affected by any of these changes.
 
 
 ## Consequences

ext/sqlite3/database.c

@@ -16,7 +16,9 @@ static void
 close_or_discard_db(sqlite3RubyPtr ctx)
 {
     if (ctx->db) {
-        if (ctx->owner == getpid()) {
+        int isReadonly = (ctx->flags & SQLITE_OPEN_READONLY);
+
+        if (isReadonly || ctx->owner == getpid()) {
             // Ordinary close.
             sqlite3_close_v2(ctx->db);
         } else {
@@ -25,9 +27,10 @@ close_or_discard_db(sqlite3RubyPtr ctx)
             sqlite3_file *sfile;
             int status;
 
-            rb_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.");
+            rb_warning("An open writable sqlite database connection was inherited from a "
+                       "forked process and is being closed to prevent the risk of corruption. "
+                       "If possible, please close all writable sqlite database connections "
+                       "before forking.");
 
             // release as much heap memory as possible by deallocating non-essential memory
             // allocations held by the database library. Memory used to cache database pages to
@@ -124,6 +127,7 @@ rb_sqlite3_open_v2(VALUE self, VALUE file, VALUE mode, VALUE zvfs)
 {
     sqlite3RubyPtr ctx;
     int status;
+    int flags;
 
     TypedData_Get_Struct(self, sqlite3Ruby, &database_type, ctx);
 
@@ -136,14 +140,16 @@ rb_sqlite3_open_v2(VALUE self, VALUE file, VALUE mode, VALUE zvfs)
 #  endif
 #endif
 
+    flags = NUM2INT(mode);
     status = sqlite3_open_v2(
                  StringValuePtr(file),
                  &ctx->db,
-                 NUM2INT(mode),
+                 flags,
                  NIL_P(zvfs) ? NULL : StringValuePtr(zvfs)
              );
 
-    CHECK(ctx->db, status)
+    CHECK(ctx->db, status);
+    ctx->flags = flags;
 
     return self;
 }
@@ -919,6 +925,10 @@ rb_sqlite3_open16(VALUE self, VALUE file)
 
     status = sqlite3_open16(utf16_string_value_ptr(file), &ctx->db);
 
+    // these are the perm flags used implicitly by sqlite3_open16,
+    // see https://www.sqlite.org/capi3ref.html#sqlite3_open
+    ctx->flags = SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE;
+
     CHECK(ctx->db, status)
 
     return INT2NUM(status);

ext/sqlite3/database.h

@@ -9,6 +9,7 @@ struct _sqlite3Ruby {
     int stmt_timeout;
     struct timespec stmt_deadline;
     rb_pid_t owner;
+    int flags;
 };
 
 typedef struct _sqlite3Ruby sqlite3Ruby;

lib/sqlite3/database.rb

@@ -5,6 +5,7 @@
 require "sqlite3/pragmas"
 require "sqlite3/statement"
 require "sqlite3/value"
+require "sqlite3/fork_safety"
 
 module SQLite3
   # The Database class encapsulates a single connection to a SQLite3 database.
@@ -134,6 +135,8 @@ def initialize file, options = {}, zvfs = nil
       @readonly = mode & Constants::Open::READONLY != 0
       @default_transaction_mode = options[:default_transaction_mode] || :deferred
 
+      ForkSafety.track(self)
+
       if block_given?
         begin
           yield self

lib/sqlite3/fork_safety.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require "weakref"
+
+# based on Rails's active_support/fork_safety.rb
+module SQLite3
+  module ForkSafety
+    module CoreExt
+      def _fork
+        pid = super
+        if pid == 0
+          ForkSafety.discard
+        end
+        pid
+      end
+    end
+
+    @databases = []
+
+    class << self
+      def hook!
+        ::Process.singleton_class.prepend(CoreExt)
+      end
+
+      def track(database)
+        @databases << WeakRef.new(database)
+      end
+
+      def discard
+        @databases.each do |db|
+          next unless db.weakref_alive?
+
+          unless db.closed? || db.readonly?
+            db.close
+          end
+        end
+        @databases.clear
+      end
+    end
+  end
+end
+
+SQLite3::ForkSafety.hook!

test/test_database.rb

@@ -721,40 +721,37 @@ def test_transaction_returns_block_result
       result = @db.transaction { :foo }
       assert_equal :foo, result
     end
+  end
 
-    def test_discard_a_connection
+  class TestDiscardDatabase < SQLite3::TestCase
+    def test_fork_discards_an_open_readwrite_connection
       skip("interpreter doesn't support fork") unless Process.respond_to?(:fork)
       skip("valgrind doesn't handle forking") if i_am_running_in_valgrind
 
+      GC.start
       begin
+        db = SQLite3::Database.new("test.db")
         read, write = IO.pipe
 
-        db = SQLite3::Database.new("test.db")
-        db.execute("attach database 'test.db' as 'foo';") # exercise sqlite3_db_name()
+        old_stderr, $stderr = $stderr, StringIO.new
         Process.fork do
           read.close
-          $stderr = StringIO.new
 
-          result = db.close
-
-          write.write((result == db) ? "ok\n" : "fail\n")
           write.write(db.closed? ? "ok\n" : "fail\n")
           write.write($stderr.string)
 
           write.close
           exit!
         end
+        $stderr = old_stderr
         write.close
-
-        assert1, assert2, *stderr = *read.readlines
+        assertion, *stderr = *read.readlines
         read.close
 
-        assert_equal("ok", assert1.chomp, "return value was not the database")
-        assert_equal("ok", assert2.chomp, "closed? did not return true")
-
+        assert_equal("ok", assertion.chomp, "closed? did not return true")
         assert_equal(1, stderr.count, "unexpected output on stderr: #{stderr.inspect}")
         assert_match(
-          /warning: An open sqlite database connection was inherited from a forked process/,
+          /warning: An open writable sqlite database connection was inherited from a forked process/,
           stderr.first,
           "expected warning was not emitted"
         )
@@ -763,5 +760,113 @@ def test_discard_a_connection
         FileUtils.rm_f("test.db")
       end
     end
+
+    def test_fork_does_not_discard_closed_connections
+      skip("interpreter doesn't support fork") unless Process.respond_to?(:fork)
+      skip("valgrind doesn't handle forking") if i_am_running_in_valgrind
+
+      GC.start
+      begin
+        db = SQLite3::Database.new("test.db")
+        read, write = IO.pipe
+
+        db.close
+
+        old_stderr, $stderr = $stderr, StringIO.new
+        Process.fork do
+          read.close
+
+          write.write($stderr.string)
+
+          write.close
+          exit!
+        end
+        $stderr = old_stderr
+        write.close
+        stderr = read.readlines
+        read.close
+
+        assert_equal(0, stderr.count, "unexpected output on stderr: #{stderr.inspect}")
+      ensure
+        db.close
+        FileUtils.rm_f("test.db")
+      end
+    end
+
+    def test_fork_does_not_discard_readonly_connections
+      skip("interpreter doesn't support fork") unless Process.respond_to?(:fork)
+      skip("valgrind doesn't handle forking") if i_am_running_in_valgrind
+
+      GC.start
+      begin
+        SQLite3::Database.open("test.db") do |db|
+          db.execute("create table foo (bar int)")
+          db.execute("insert into foo values (1)")
+        end
+
+        db = SQLite3::Database.new("test.db", readonly: true)
+        read, write = IO.pipe
+
+        old_stderr, $stderr = $stderr, StringIO.new
+        Process.fork do
+          read.close
+
+          write.write(db.closed? ? "fail\n" : "ok\n") # should be open and readable
+          write.write((db.execute("select * from foo") == [[1]]) ? "ok\n" : "fail\n")
+          write.write($stderr.string)
+
+          write.close
+          exit!
+        end
+        $stderr = old_stderr
+        write.close
+        assertion1, assertion2, *stderr = *read.readlines
+        read.close
+
+        assert_equal("ok", assertion1.chomp, "closed? did not return false")
+        assert_equal("ok", assertion2.chomp, "could not read from database")
+        assert_equal(0, stderr.count, "unexpected output on stderr: #{stderr.inspect}")
+      ensure
+        db&.close
+        FileUtils.rm_f("test.db")
+      end
+    end
+
+    def test_close_does_not_discard_readonly_connections
+      skip("interpreter doesn't support fork") unless Process.respond_to?(:fork)
+      skip("valgrind doesn't handle forking") if i_am_running_in_valgrind
+
+      GC.start
+      begin
+        SQLite3::Database.open("test.db") do |db|
+          db.execute("create table foo (bar int)")
+          db.execute("insert into foo values (1)")
+        end
+
+        db = SQLite3::Database.new("test.db", readonly: true)
+        read, write = IO.pipe
+
+        old_stderr, $stderr = $stderr, StringIO.new
+        Process.fork do
+          read.close
+
+          db.close
+
+          write.write($stderr.string)
+
+          write.close
+          exit!
+        end
+        $stderr = old_stderr
+        write.close
+        stderr = read.readlines
+        read.close
+
+        assert_equal(0, stderr.count, "unexpected output on stderr: #{stderr.inspect}")
+      ensure
+        db&.close
+        FileUtils.rm_f("test.db")
+      end
+    end
   end
 end