activerecord-health 0.1.0

data/lib/activerecord/health.rb

@@ -0,0 +1,134 @@
+# frozen_string_literal: true
+
+require_relative "health/version"
+require_relative "health/configuration"
+require_relative "health/adapters/postgresql_adapter"
+require_relative "health/adapters/mysql_adapter"
+require_relative "health/railtie" if defined?(Rails::Railtie)
+
+module ActiveRecord
+  module Health
+    QUERY_TIMEOUT = 1
+
+    class << self
+      def configure
+        yield(configuration)
+      end
+
+      def configuration
+        @configuration ||= Configuration.new
+      end
+
+      def reset_configuration!
+        @configuration = nil
+      end
+
+      def ok?(model: ActiveRecord::Base)
+        load_pct(model: model) <= config_for(model).threshold
+      end
+
+      def load_pct(model: ActiveRecord::Base)
+        db_config_name = model.connection_db_config.name
+        cache_key = "activerecord_health:load_pct:#{db_config_name}"
+
+        read_from_cache(cache_key) { query_load_pct(model) }
+      end
+
+      def sheddable(model: ActiveRecord::Base)
+        return false unless ok?(model: model)
+        yield
+        true
+      end
+
+      def sheddable_pct(pct:, model: ActiveRecord::Base)
+        return false if load_pct(model: model) > pct
+        yield
+        true
+      end
+
+      private
+
+      def config_for(model)
+        model_class = model.is_a?(Class) ? model : model.class
+        configuration.for_model(model_class)
+      end
+
+      def read_from_cache(cache_key)
+        cached_value = configuration.cache.read(cache_key)
+        return cached_value unless cached_value.nil?
+
+        value = yield
+        configuration.cache.write(cache_key, value, expires_in: configuration.cache_ttl)
+        value
+      rescue
+        0.0
+      end
+
+      def query_load_pct(model)
+        connection = model.connection
+        adapter = adapter_for(connection)
+        config = config_for(model)
+        db_config_name = model.connection_db_config.name
+
+        active_sessions = execute_with_timeout(connection, adapter.active_session_count_query)
+        load_pct = active_sessions.to_f / config.vcpu_count
+
+        instrument(db_config_name, load_pct, active_sessions)
+
+        load_pct
+      rescue
+        1.0
+      end
+
+      def execute_with_timeout(connection, query)
+        adapter_name = connection.adapter_name.downcase
+
+        case adapter_name
+        when /postgresql/
+          execute_with_postgresql_timeout(connection, query)
+        when /mysql/
+          execute_with_mysql_timeout(connection, query)
+        else
+          connection.select_value(query)
+        end
+      end
+
+      def execute_with_postgresql_timeout(connection, query)
+        connection.transaction do
+          connection.execute("SET LOCAL statement_timeout = '#{QUERY_TIMEOUT}s'")
+          connection.select_value(query)
+        end
+      end
+
+      def execute_with_mysql_timeout(connection, query)
+        connection.transaction do
+          connection.execute("SET max_execution_time = #{QUERY_TIMEOUT * 1000}")
+          connection.select_value(query)
+        end
+      end
+
+      def adapter_for(connection)
+        case connection.adapter_name.downcase
+        when /postgresql/
+          Adapters::PostgreSQLAdapter.new
+        when /mysql/
+          version = connection.select_value("SELECT VERSION()")
+          Adapters::MySQLAdapter.new(version)
+        else
+          raise "Unsupported database adapter: #{connection.adapter_name}"
+        end
+      end
+
+      def instrument(db_config_name, load_pct, active_sessions)
+        return unless defined?(ActiveSupport::Notifications)
+
+        ActiveSupport::Notifications.instrument(
+          "health_check.activerecord_health",
+          database: db_config_name,
+          load_pct: load_pct,
+          active_sessions: active_sessions
+        )
+      end
+    end
+  end
+end

data/lib/activerecord-health.rb

@@ -0,0 +1,3 @@
+# frozen_string_literal: true
+
+require "activerecord/health"

data/test/integration/mysql_integration_test.rb

@@ -0,0 +1,104 @@
+# frozen_string_literal: true
+
+require "test_helper"
+require "mysql2"
+
+class MySQLIntegrationTest < ActiveRecord::Health::TestCase
+  def setup
+    super
+    skip "MySQL not available" unless mysql_available?
+
+    ActiveRecord::Base.establish_connection(
+      adapter: "mysql2",
+      host: mysql_host,
+      port: ENV.fetch("MYSQL_PORT", 3306).to_i,
+      username: ENV.fetch("MYSQL_USER", "root"),
+      password: ENV.fetch("MYSQL_PASSWORD", "root"),
+      database: ENV.fetch("MYSQL_DB", "activerecord_health_test")
+    )
+
+    @cache = MockCache.new
+    @baseline_sessions = count_active_sessions
+    ActiveRecord::Health.configure do |config|
+      config.vcpu_count = @baseline_sessions + 4
+      config.threshold = 0.75
+      config.cache = @cache
+    end
+  end
+
+  def teardown
+    @sleep_threads&.each(&:kill)
+    @sleep_threads&.each(&:join)
+    ActiveRecord::Base.connection_pool.disconnect!
+    super
+  end
+
+  def test_load_pct_increases_with_active_sessions
+    baseline_load = ActiveRecord::Health.load_pct(model: ActiveRecord::Base)
+    @cache.clear
+
+    spawn_sleeping_connections(2)
+    with_load = ActiveRecord::Health.load_pct(model: ActiveRecord::Base)
+
+    expected_increase = 2.0 / (@baseline_sessions + 4)
+    assert_in_delta baseline_load + expected_increase, with_load, 0.01
+  end
+
+  def test_ok_returns_true_when_below_threshold
+    @cache.clear
+
+    assert ActiveRecord::Health.ok?(model: ActiveRecord::Base)
+  end
+
+  def test_ok_returns_false_when_above_threshold
+    @cache.clear
+    spawn_sleeping_connections(4)
+
+    refute ActiveRecord::Health.ok?(model: ActiveRecord::Base)
+  end
+
+  def count_active_sessions
+    version = ActiveRecord::Base.connection.select_value("SELECT VERSION()")
+    adapter = ActiveRecord::Health::Adapters::MySQLAdapter.new(version)
+    ActiveRecord::Base.connection.select_value(adapter.active_session_count_query).to_i
+  end
+
+  private
+
+  def mysql_available?
+    Mysql2::Client.new(mysql_config).close
+    true
+  rescue Mysql2::Error => e
+    warn "MySQL connection failed: #{e.message}"
+    false
+  end
+
+  def spawn_sleeping_connections(count)
+    @sleep_threads = count.times.map do
+      Thread.new do
+        client = Mysql2::Client.new(mysql_config)
+        client.query("SELECT SLEEP(30)")
+      rescue
+      ensure
+        client&.close
+      end
+    end
+
+    sleep 0.5
+  end
+
+  def mysql_config
+    {
+      host: mysql_host,
+      port: ENV.fetch("MYSQL_PORT", 3306).to_i,
+      username: ENV.fetch("MYSQL_USER", "root"),
+      password: ENV.fetch("MYSQL_PASSWORD", "root"),
+      database: ENV.fetch("MYSQL_DB", "activerecord_health_test")
+    }
+  end
+
+  def mysql_host
+    host = ENV.fetch("MYSQL_HOST", "127.0.0.1")
+    (host == "localhost") ? "127.0.0.1" : host
+  end
+end

data/test/integration/postgresql_integration_test.rb

@@ -0,0 +1,104 @@
+# frozen_string_literal: true
+
+require "test_helper"
+require "pg"
+
+class PostgreSQLIntegrationTest < ActiveRecord::Health::TestCase
+  def setup
+    super
+    skip "PostgreSQL not available" unless postgresql_available?
+
+    ActiveRecord::Base.establish_connection(
+      adapter: "postgresql",
+      host: ENV.fetch("POSTGRES_HOST", "localhost"),
+      port: ENV.fetch("POSTGRES_PORT", 5432),
+      username: ENV.fetch("POSTGRES_USER", "postgres"),
+      password: ENV.fetch("POSTGRES_PASSWORD", "postgres"),
+      database: ENV.fetch("POSTGRES_DB", "activerecord_health_test")
+    )
+
+    @cache = MockCache.new
+    @baseline_sessions = count_active_sessions
+    ActiveRecord::Health.configure do |config|
+      config.vcpu_count = @baseline_sessions + 4
+      config.threshold = 0.75
+      config.cache = @cache
+    end
+  end
+
+  def teardown
+    @sleep_threads&.each(&:kill)
+    @sleep_threads&.each(&:join)
+    ActiveRecord::Base.connection_pool.disconnect!
+    super
+  end
+
+  def test_load_pct_increases_with_active_sessions
+    baseline_load = ActiveRecord::Health.load_pct(model: ActiveRecord::Base)
+    @cache.clear
+
+    spawn_sleeping_connections(2)
+    with_load = ActiveRecord::Health.load_pct(model: ActiveRecord::Base)
+
+    expected_increase = 2.0 / (@baseline_sessions + 4)
+    assert_in_delta baseline_load + expected_increase, with_load, 0.01
+  end
+
+  def test_ok_returns_true_when_below_threshold
+    @cache.clear
+
+    assert ActiveRecord::Health.ok?(model: ActiveRecord::Base)
+  end
+
+  def test_ok_returns_false_when_above_threshold
+    @cache.clear
+    spawn_sleeping_connections(4)
+
+    refute ActiveRecord::Health.ok?(model: ActiveRecord::Base)
+  end
+
+  def count_active_sessions
+    ActiveRecord::Base.connection.select_value(<<~SQL).to_i
+      SELECT count(*)
+      FROM pg_stat_activity
+      WHERE state = 'active'
+        AND backend_type = 'client backend'
+        AND pid != pg_backend_pid()
+    SQL
+  end
+
+  private
+
+  def postgresql_available?
+    PG.connect(
+      host: ENV.fetch("POSTGRES_HOST", "localhost"),
+      port: ENV.fetch("POSTGRES_PORT", 5432),
+      user: ENV.fetch("POSTGRES_USER", "postgres"),
+      password: ENV.fetch("POSTGRES_PASSWORD", "postgres"),
+      dbname: ENV.fetch("POSTGRES_DB", "activerecord_health_test")
+    ).close
+    true
+  rescue PG::ConnectionBad
+    false
+  end
+
+  def spawn_sleeping_connections(count)
+    @sleep_threads = count.times.map do
+      Thread.new do
+        conn = PG.connect(
+          host: ENV.fetch("POSTGRES_HOST", "localhost"),
+          port: ENV.fetch("POSTGRES_PORT", 5432),
+          user: ENV.fetch("POSTGRES_USER", "postgres"),
+          password: ENV.fetch("POSTGRES_PASSWORD", "postgres"),
+          dbname: ENV.fetch("POSTGRES_DB", "activerecord_health_test")
+        )
+        conn.exec("SELECT pg_sleep(30)")
+      rescue
+      ensure
+        conn&.close
+      end
+    end
+
+    sleep 0.5
+  end
+end

data/test/test_helper.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+$LOAD_PATH.unshift File.expand_path("../lib", __dir__)
+
+require "active_record"
+require "active_support"
+require "active_support/core_ext/string"
+require "active_support/cache"
+require "activerecord/health"
+
+require "minitest/autorun"
+require "minitest/reporters"
+
+Minitest::Reporters.use! Minitest::Reporters::DefaultReporter.new
+
+if ENV["FAIL_ON_SKIP"]
+  module Minitest
+    class << self
+      alias_method :original_run, :run
+
+      def run(args = [])
+        result = original_run(args)
+        reporter = Minitest::Reporters.reporters.first
+        if reporter && reporter.skips > 0
+          warn "\nCI failed: #{reporter.skips} tests were skipped"
+          exit 1
+        end
+        result
+      end
+    end
+  end
+end
+
+class ActiveRecord::Health::TestCase < Minitest::Test
+  def setup
+    ActiveRecord::Health.reset_configuration!
+  end
+
+  def teardown
+    ActiveRecord::Health.reset_configuration!
+  end
+end
+
+class MockCache
+  def initialize
+    @store = {}
+  end
+
+  def read(key)
+    @store[key]
+  end
+
+  def write(key, value, options = {})
+    @store[key] = value
+  end
+
+  def delete(key)
+    @store.delete(key)
+  end
+
+  def clear
+    @store.clear
+  end
+end
+
+class FailingCache
+  def read(key)
+    raise "Cache connection failed"
+  end
+
+  def write(key, value, options = {})
+    raise "Cache connection failed"
+  end
+end

data/test/unit/adapters/mysql_adapter_test.rb

@@ -0,0 +1,94 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+class MySQLAdapterTest < ActiveRecord::Health::TestCase
+  def test_active_session_count_query_for_mysql_8_0_22_plus
+    adapter = ActiveRecord::Health::Adapters::MySQLAdapter.new("8.0.22")
+
+    expected_query = <<~SQL.squish
+      SELECT COUNT(*)
+      FROM performance_schema.processlist
+      WHERE COMMAND != 'Sleep'
+        AND ID != CONNECTION_ID()
+        AND USER NOT IN ('event_scheduler', 'system user')
+    SQL
+
+    assert_equal expected_query, adapter.active_session_count_query
+  end
+
+  def test_active_session_count_query_for_mysql_8_0_21
+    adapter = ActiveRecord::Health::Adapters::MySQLAdapter.new("8.0.21")
+
+    expected_query = <<~SQL.squish
+      SELECT COUNT(*)
+      FROM information_schema.processlist
+      WHERE Command != 'Sleep'
+        AND ID != CONNECTION_ID()
+        AND User NOT IN ('event_scheduler', 'system user')
+        AND Command NOT IN ('Binlog Dump', 'Binlog Dump GTID')
+    SQL
+
+    assert_equal expected_query, adapter.active_session_count_query
+  end
+
+  def test_active_session_count_query_for_mysql_5
+    adapter = ActiveRecord::Health::Adapters::MySQLAdapter.new("5.7.35")
+
+    expected_query = <<~SQL.squish
+      SELECT COUNT(*)
+      FROM information_schema.processlist
+      WHERE Command != 'Sleep'
+        AND ID != CONNECTION_ID()
+        AND User NOT IN ('event_scheduler', 'system user')
+        AND Command NOT IN ('Binlog Dump', 'Binlog Dump GTID')
+    SQL
+
+    assert_equal expected_query, adapter.active_session_count_query
+  end
+
+  def test_active_session_count_query_for_mariadb
+    adapter = ActiveRecord::Health::Adapters::MySQLAdapter.new("10.5.12-MariaDB")
+
+    expected_query = <<~SQL.squish
+      SELECT COUNT(*)
+      FROM information_schema.processlist
+      WHERE Command != 'Sleep'
+        AND ID != CONNECTION_ID()
+        AND User NOT IN ('event_scheduler', 'system user')
+        AND Command NOT IN ('Binlog Dump', 'Binlog Dump GTID')
+    SQL
+
+    assert_equal expected_query, adapter.active_session_count_query
+  end
+
+  def test_adapter_name
+    adapter = ActiveRecord::Health::Adapters::MySQLAdapter.new("8.0.22")
+
+    assert_equal :mysql, adapter.name
+  end
+
+  def test_uses_performance_schema_returns_true_for_mysql_8_0_22_plus
+    adapter = ActiveRecord::Health::Adapters::MySQLAdapter.new("8.0.22")
+
+    assert adapter.uses_performance_schema?
+  end
+
+  def test_uses_performance_schema_returns_true_for_mysql_8_1
+    adapter = ActiveRecord::Health::Adapters::MySQLAdapter.new("8.1.0")
+
+    assert adapter.uses_performance_schema?
+  end
+
+  def test_uses_performance_schema_returns_false_for_mysql_8_0_21
+    adapter = ActiveRecord::Health::Adapters::MySQLAdapter.new("8.0.21")
+
+    refute adapter.uses_performance_schema?
+  end
+
+  def test_uses_performance_schema_returns_false_for_mariadb
+    adapter = ActiveRecord::Health::Adapters::MySQLAdapter.new("10.5.12-MariaDB")
+
+    refute adapter.uses_performance_schema?
+  end
+end

data/test/unit/adapters/postgresql_adapter_test.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+class PostgreSQLAdapterTest < ActiveRecord::Health::TestCase
+  def test_active_session_count_query
+    adapter = ActiveRecord::Health::Adapters::PostgreSQLAdapter.new
+
+    expected_query = <<~SQL.squish
+      SELECT count(*)
+      FROM pg_stat_activity
+      WHERE state = 'active'
+        AND backend_type = 'client backend'
+        AND pid != pg_backend_pid()
+    SQL
+
+    assert_equal expected_query, adapter.active_session_count_query
+  end
+
+  def test_adapter_name
+    adapter = ActiveRecord::Health::Adapters::PostgreSQLAdapter.new
+
+    assert_equal :postgresql, adapter.name
+  end
+end

data/test/unit/configuration_test.rb

@@ -0,0 +1,140 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+class ConfigurationTest < ActiveRecord::Health::TestCase
+  def test_configure_sets_vcpu_count
+    ActiveRecord::Health.configure do |config|
+      config.vcpu_count = 16
+      config.cache = MockCache.new
+    end
+
+    assert_equal 16, ActiveRecord::Health.configuration.vcpu_count
+  end
+
+  def test_configure_sets_threshold_with_default
+    ActiveRecord::Health.configure do |config|
+      config.vcpu_count = 16
+      config.cache = MockCache.new
+    end
+
+    assert_equal 0.75, ActiveRecord::Health.configuration.threshold
+  end
+
+  def test_configure_sets_custom_threshold
+    ActiveRecord::Health.configure do |config|
+      config.vcpu_count = 16
+      config.threshold = 0.5
+      config.cache = MockCache.new
+    end
+
+    assert_equal 0.5, ActiveRecord::Health.configuration.threshold
+  end
+
+  def test_configure_sets_cache
+    cache = MockCache.new
+    ActiveRecord::Health.configure do |config|
+      config.vcpu_count = 16
+      config.cache = cache
+    end
+
+    assert_same cache, ActiveRecord::Health.configuration.cache
+  end
+
+  def test_configure_sets_cache_ttl_with_default
+    ActiveRecord::Health.configure do |config|
+      config.vcpu_count = 16
+      config.cache = MockCache.new
+    end
+
+    assert_equal 60, ActiveRecord::Health.configuration.cache_ttl
+  end
+
+  def test_configure_sets_custom_cache_ttl
+    ActiveRecord::Health.configure do |config|
+      config.vcpu_count = 16
+      config.cache = MockCache.new
+      config.cache_ttl = 120
+    end
+
+    assert_equal 120, ActiveRecord::Health.configuration.cache_ttl
+  end
+
+  def test_raises_without_vcpu_count
+    ActiveRecord::Health.configure do |config|
+      config.cache = MockCache.new
+    end
+
+    error = assert_raises(ActiveRecord::Health::ConfigurationError) do
+      ActiveRecord::Health.configuration.validate!
+    end
+    assert_match(/vcpu_count/, error.message)
+  end
+
+  def test_raises_without_cache
+    ActiveRecord::Health.configure do |config|
+      config.vcpu_count = 16
+    end
+
+    error = assert_raises(ActiveRecord::Health::ConfigurationError) do
+      ActiveRecord::Health.configuration.validate!
+    end
+    assert_match(/cache/, error.message)
+  end
+
+  def test_for_model_configures_specific_database
+    ActiveRecord::Health.configure do |config|
+      config.vcpu_count = 16
+      config.cache = MockCache.new
+
+      config.for_model(AnimalsRecord) do |db|
+        db.vcpu_count = 8
+        db.threshold = 0.5
+      end
+    end
+
+    assert_equal 16, ActiveRecord::Health.configuration.vcpu_count
+    assert_equal 8, ActiveRecord::Health.configuration.for_model(AnimalsRecord).vcpu_count
+    assert_equal 0.5, ActiveRecord::Health.configuration.for_model(AnimalsRecord).threshold
+  end
+
+  def test_for_model_inherits_cache_from_default
+    cache = MockCache.new
+    ActiveRecord::Health.configure do |config|
+      config.vcpu_count = 16
+      config.cache = cache
+
+      config.for_model(AnimalsRecord) do |db|
+        db.vcpu_count = 8
+      end
+    end
+
+    assert_same cache, ActiveRecord::Health.configuration.for_model(AnimalsRecord).cache
+  end
+
+  def test_for_model_inherits_threshold_from_default_when_not_specified
+    ActiveRecord::Health.configure do |config|
+      config.vcpu_count = 16
+      config.cache = MockCache.new
+
+      config.for_model(AnimalsRecord) do |db|
+        db.vcpu_count = 8
+      end
+    end
+
+    assert_equal 0.75, ActiveRecord::Health.configuration.for_model(AnimalsRecord).threshold
+  end
+
+  def test_max_healthy_sessions_calculated_correctly
+    ActiveRecord::Health.configure do |config|
+      config.vcpu_count = 16
+      config.threshold = 0.75
+      config.cache = MockCache.new
+    end
+
+    assert_equal 12, ActiveRecord::Health.configuration.max_healthy_sessions
+  end
+end
+
+class AnimalsRecord
+end