with_advisory_lock 5.1.0 → 7.0.0

data/lib/with_advisory_lock/concern.rb

@@ -19,30 +19,41 @@ module WithAdvisoryLock
       end
 
       def with_advisory_lock_result(lock_name, options = {}, &block)
-        impl = impl_class.new(connection, lock_name, options)
-        impl.with_advisory_lock_if_needed(&block)
+        with_connection do |conn|
+          conn.with_advisory_lock_if_needed(lock_name, options, &block)
+        end
       end
 
       def advisory_lock_exists?(lock_name)
-        impl = impl_class.new(connection, lock_name, 0)
-        impl.already_locked? || !impl.yield_with_lock.lock_was_acquired?
+        with_connection do |conn|
+          lock_str = "#{ENV.fetch(CoreAdvisory::LOCK_PREFIX_ENV, nil)}#{lock_name}"
+          lock_stack_item = LockStackItem.new(lock_str, false)
+
+          if conn.advisory_lock_stack.include?(lock_stack_item)
+            true
+          else
+            # For PostgreSQL, try non-blocking query first to avoid race conditions
+            if conn.respond_to?(:advisory_lock_exists_for?)
+              query_result = conn.advisory_lock_exists_for?(lock_name)
+              return query_result unless query_result.nil?
+            end
+
+            # Fall back to the original implementation
+            result = conn.with_advisory_lock_if_needed(lock_name, { timeout_seconds: 0 })
+            !result.lock_was_acquired?
+          end
+        end
       end
 
       def current_advisory_lock
-        lock_stack_key = WithAdvisoryLock::Base.lock_stack.first
-        lock_stack_key && lock_stack_key[0]
+        with_connection do |conn|
+          conn.advisory_lock_stack.first&.name
+        end
       end
 
-      private
-
-      def impl_class
-        adapter = WithAdvisoryLock::DatabaseAdapterSupport.new(connection)
-        if adapter.postgresql?
-          WithAdvisoryLock::PostgreSQL
-        elsif adapter.mysql?
-          WithAdvisoryLock::MySQL
-        else
-          WithAdvisoryLock::Flock
+      def current_advisory_locks
+        with_connection do |conn|
+          conn.advisory_lock_stack.map(&:name)
         end
       end
     end

data/lib/with_advisory_lock/core_advisory.rb

@@ -0,0 +1,110 @@
+# frozen_string_literal: true
+
+require 'zlib'
+
+module WithAdvisoryLock
+  module CoreAdvisory
+    extend ActiveSupport::Concern
+
+    LOCK_PREFIX_ENV = 'WITH_ADVISORY_LOCK_PREFIX'
+
+    # Thread-local lock stack management
+    def advisory_lock_stack
+      Thread.current[:with_advisory_lock_stack] ||= []
+    end
+
+    def with_advisory_lock_if_needed(lock_name, options = {}, &block)
+      options = { timeout_seconds: options } unless options.respond_to?(:fetch)
+      options.assert_valid_keys :timeout_seconds, :shared, :transaction, :disable_query_cache
+
+      # Validate transaction-level locks are used within a transaction
+      if options.fetch(:transaction, false) && !transaction_open?
+        raise ArgumentError, 'transaction-level advisory locks require an active transaction'
+      end
+
+      lock_str = "#{ENV.fetch(LOCK_PREFIX_ENV, nil)}#{lock_name}"
+      lock_stack_item = LockStackItem.new(lock_str, options.fetch(:shared, false))
+
+      if advisory_lock_stack.include?(lock_stack_item)
+        # Already have this exact lock (same name and type), just yield
+        return Result.new(lock_was_acquired: true, result: yield)
+      end
+
+      # Check if we have a lock with the same name but different type (for upgrade/downgrade)
+      same_name_different_type = advisory_lock_stack.any? do |item|
+        item.name == lock_str && item.shared != options.fetch(:shared, false)
+      end
+      if same_name_different_type && options.fetch(:transaction, false)
+        # PostgreSQL doesn't support upgrading/downgrading transaction-level locks
+        return Result.new(lock_was_acquired: false)
+      end
+
+      disable_query_cache = options.fetch(:disable_query_cache, false)
+
+      if disable_query_cache
+        uncached do
+          advisory_lock_and_yield(lock_name, lock_str, lock_stack_item, options, &block)
+        end
+      else
+        advisory_lock_and_yield(lock_name, lock_str, lock_stack_item, options, &block)
+      end
+    end
+
+    private
+
+    def advisory_lock_and_yield(lock_name, lock_str, lock_stack_item, options, &)
+      timeout_seconds = options.fetch(:timeout_seconds, nil)
+      shared = options.fetch(:shared, false)
+      transaction = options.fetch(:transaction, false)
+
+      lock_keys = lock_keys_for(lock_name)
+
+      # MySQL supports database-level timeout in GET_LOCK, skip Ruby-level polling
+      if supports_database_timeout? || timeout_seconds&.zero?
+        yield_with_lock(lock_keys, lock_name, lock_str, lock_stack_item, shared, transaction, timeout_seconds, &)
+      else
+        yield_with_lock_and_timeout(lock_keys, lock_name, lock_str, lock_stack_item, shared, transaction,
+                                    timeout_seconds, &)
+      end
+    end
+
+    def yield_with_lock_and_timeout(lock_keys, lock_name, lock_str, lock_stack_item, shared, transaction,
+                                    timeout_seconds, &)
+      give_up_at = timeout_seconds ? Time.now + timeout_seconds : nil
+      while give_up_at.nil? || Time.now < give_up_at
+        r = yield_with_lock(lock_keys, lock_name, lock_str, lock_stack_item, shared, transaction, 0, &)
+        return r if r.lock_was_acquired?
+
+        # Randomizing sleep time may help reduce contention.
+        sleep(rand(0.05..0.15))
+      end
+      Result.new(lock_was_acquired: false)
+    end
+
+    def yield_with_lock(lock_keys, lock_name, _lock_str, lock_stack_item, shared, transaction, timeout_seconds = nil)
+      if try_advisory_lock(lock_keys, lock_name: lock_name, shared: shared, transaction: transaction,
+                                      timeout_seconds: timeout_seconds)
+        begin
+          advisory_lock_stack.push(lock_stack_item)
+          result = block_given? ? yield : nil
+          Result.new(lock_was_acquired: true, result: result)
+        ensure
+          advisory_lock_stack.pop
+          release_advisory_lock(lock_keys, lock_name: lock_name, shared: shared, transaction: transaction)
+        end
+      else
+        Result.new(lock_was_acquired: false)
+      end
+    end
+
+    def stable_hashcode(input)
+      if input.is_a? Numeric
+        input.to_i
+      else
+        # Ruby MRI's String#hash is randomly seeded as of Ruby 1.9 so
+        # make sure we use a deterministic hash.
+        Zlib.crc32(input.to_s, 0)
+      end
+    end
+  end
+end

data/lib/with_advisory_lock/jruby_adapter.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module WithAdvisoryLock
+  module JRubyAdapter
+    # JRuby compatibility - ensure adapters are patched after they're loaded
+    def self.install!
+      ActiveSupport.on_load :active_record do
+        ActiveRecord::Base.singleton_class.prepend(Module.new do
+          def connection
+            super.tap do |conn|
+              case conn
+              when ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
+                unless conn.class.include?(WithAdvisoryLock::CoreAdvisory)
+                  conn.class.prepend WithAdvisoryLock::CoreAdvisory
+                  conn.class.prepend WithAdvisoryLock::PostgreSQLAdvisory
+                end
+              when ActiveRecord::ConnectionAdapters::Mysql2Adapter, ActiveRecord::ConnectionAdapters::TrilogyAdapter
+                unless conn.class.include?(WithAdvisoryLock::CoreAdvisory)
+                  conn.class.prepend WithAdvisoryLock::CoreAdvisory
+                  conn.class.prepend WithAdvisoryLock::MySQLAdvisory
+                end
+              end
+            end
+          end
+        end)
+      end
+    end
+  end
+end

data/lib/with_advisory_lock/lock_stack_item.rb

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+module WithAdvisoryLock
+  # Lock stack item to track acquired locks
+  LockStackItem = Data.define(:name, :shared)
+end

data/lib/with_advisory_lock/mysql_advisory.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+require 'securerandom'
+
+module WithAdvisoryLock
+  module MySQLAdvisory
+    extend ActiveSupport::Concern
+
+    LOCK_PREFIX_ENV = 'WITH_ADVISORY_LOCK_PREFIX'
+
+    def try_advisory_lock(lock_keys, lock_name:, shared:, transaction:, timeout_seconds: nil)
+      raise ArgumentError, 'shared locks are not supported on MySQL' if shared
+      raise ArgumentError, 'transaction level locks are not supported on MySQL' if transaction
+
+      # MySQL GET_LOCK supports native timeout:
+      # - timeout_seconds = nil: wait indefinitely (-1)
+      # - timeout_seconds = 0: try once, no wait (0)
+      # - timeout_seconds > 0: wait up to timeout_seconds
+      mysql_timeout = case timeout_seconds
+                      when nil then -1
+                      when 0 then 0
+                      else timeout_seconds.to_i
+                      end
+
+      execute_successful?("GET_LOCK(#{quote(lock_keys.first)}, #{mysql_timeout})")
+    end
+
+    def release_advisory_lock(lock_keys, lock_name:, **)
+      execute_successful?("RELEASE_LOCK(#{quote(lock_keys.first)})")
+    rescue ActiveRecord::StatementInvalid => e
+      # If the connection is broken, the lock is automatically released by MySQL
+      # No need to fail the release operation
+      connection_lost = case e.cause
+                        when defined?(Mysql2::Error::ConnectionError) && Mysql2::Error::ConnectionError
+                          true
+                        when defined?(Trilogy::ConnectionError) && Trilogy::ConnectionError
+                          true
+                        else
+                          e.message =~ /Lost connection|MySQL server has gone away|Connection refused/i
+                        end
+
+      return if connection_lost
+
+      raise
+    end
+
+    def lock_keys_for(lock_name)
+      lock_str = "#{ENV.fetch(LOCK_PREFIX_ENV, nil)}#{lock_name}"
+      [lock_str]
+    end
+
+    def supports_database_timeout?
+      true
+    end
+
+    private
+
+    def execute_successful?(mysql_function)
+      select_value("SELECT #{mysql_function}") == 1
+    end
+  end
+end

data/lib/with_advisory_lock/postgresql_advisory.rb

@@ -0,0 +1,112 @@
+# frozen_string_literal: true
+
+require 'securerandom'
+
+module WithAdvisoryLock
+  module PostgreSQLAdvisory
+    extend ActiveSupport::Concern
+
+    LOCK_PREFIX_ENV = 'WITH_ADVISORY_LOCK_PREFIX'
+    LOCK_RESULT_VALUES = ['t', true].freeze
+    ERROR_MESSAGE_REGEX = / ERROR: +current transaction is aborted,/
+
+    def try_advisory_lock(lock_keys, lock_name:, shared:, transaction:, timeout_seconds: nil)
+      # timeout_seconds is accepted for compatibility but ignored - PostgreSQL doesn't support
+      # native timeouts with pg_try_advisory_lock, requiring Ruby-level polling instead
+      function = advisory_try_lock_function(transaction, shared)
+      execute_advisory(function, lock_keys, lock_name)
+    end
+
+    def release_advisory_lock(*args)
+      # Handle both signatures - ActiveRecord's built-in and ours
+      if args.length == 1 && args[0].is_a?(Integer)
+        # ActiveRecord's built-in signature: release_advisory_lock(lock_id)
+        super
+      else
+        # Our signature: release_advisory_lock(lock_keys, lock_name:, shared:, transaction:)
+        lock_keys, options = args
+        return if options[:transaction]
+
+        function = advisory_unlock_function(options[:shared])
+        execute_advisory(function, lock_keys, options[:lock_name])
+      end
+    rescue ActiveRecord::StatementInvalid => e
+      # If the connection is broken, the lock is automatically released by PostgreSQL
+      # No need to fail the release operation
+      return if e.cause.is_a?(PG::ConnectionBad) || e.message =~ /PG::ConnectionBad/
+
+      raise unless e.message =~ ERROR_MESSAGE_REGEX
+
+      begin
+        rollback_db_transaction
+        execute_advisory(function, lock_keys, options[:lock_name])
+      ensure
+        begin_db_transaction
+      end
+    end
+
+    def lock_keys_for(lock_name)
+      [
+        stable_hashcode(lock_name),
+        ENV.fetch(LOCK_PREFIX_ENV, nil)
+      ].map { |ea| ea.to_i & 0x7fffffff }
+    end
+
+    def supports_database_timeout?
+      false
+    end
+
+    # Non-blocking check for advisory lock existence to avoid race conditions
+    # This queries pg_locks directly instead of trying to acquire the lock
+    def advisory_lock_exists_for?(lock_name, shared: false)
+      lock_keys = lock_keys_for(lock_name)
+
+      query = <<~SQL.squish
+        SELECT 1 FROM pg_locks
+        WHERE locktype = 'advisory'
+          AND database = (SELECT oid FROM pg_database WHERE datname = CURRENT_DATABASE())
+          AND classid = #{lock_keys.first}
+          AND objid = #{lock_keys.last}
+          AND mode = '#{shared ? 'ShareLock' : 'ExclusiveLock'}'
+        LIMIT 1
+      SQL
+
+      select_value(query).present?
+    rescue ActiveRecord::StatementInvalid
+      # If pg_locks is not accessible, fall back to nil to indicate we should use the default method
+      nil
+    end
+
+    private
+
+    def advisory_try_lock_function(transaction_scope, shared)
+      [
+        'pg_try_advisory',
+        transaction_scope ? '_xact' : nil,
+        '_lock',
+        shared ? '_shared' : nil
+      ].compact.join
+    end
+
+    def advisory_unlock_function(shared)
+      [
+        'pg_advisory_unlock',
+        shared ? '_shared' : nil
+      ].compact.join
+    end
+
+    def execute_advisory(function, lock_keys, lock_name)
+      result = select_value(prepare_sql(function, lock_keys, lock_name))
+      LOCK_RESULT_VALUES.include?(result)
+    end
+
+    def prepare_sql(function, lock_keys, lock_name)
+      comment = lock_name.to_s.gsub(%r{(/\*)|(\*/)}, '--')
+      "SELECT #{function}(#{lock_keys.join(',')}) AS #{unique_column_name} /* #{comment} */"
+    end
+
+    def unique_column_name
+      "t#{SecureRandom.hex}"
+    end
+  end
+end

data/lib/with_advisory_lock/result.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module WithAdvisoryLock
+  # Result object that indicates whether a lock was acquired and the result of the block
+  Result = Data.define(:lock_was_acquired, :result) do
+    def initialize(lock_was_acquired:, result: nil)
+      super
+    end
+
+    def lock_was_acquired?
+      lock_was_acquired
+    end
+  end
+end

data/lib/with_advisory_lock/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module WithAdvisoryLock
-  VERSION = Gem::Version.new('5.1.0')
+  VERSION = Gem::Version.new('7.0.0')
 end

data/lib/with_advisory_lock.rb

@@ -1,17 +1,46 @@
+# frozen_string_literal: true
+
 require 'with_advisory_lock/version'
 require 'active_support'
-require 'zeitwerk'
-
-loader = Zeitwerk::Loader.for_gem
-loader.inflector.inflect(
-  'mysql' => 'MySQL',
-  'postgresql' => 'PostgreSQL',
-  )
-loader.setup
+require 'active_support/concern'
 
 module WithAdvisoryLock
+  autoload :Concern, 'with_advisory_lock/concern'
+  autoload :Result, 'with_advisory_lock/result'
+  autoload :LockStackItem, 'with_advisory_lock/lock_stack_item'
+
+  # Modules for adapter injection
+  autoload :CoreAdvisory, 'with_advisory_lock/core_advisory'
+  autoload :PostgreSQLAdvisory, 'with_advisory_lock/postgresql_advisory'
+  autoload :MySQLAdvisory, 'with_advisory_lock/mysql_advisory'
+
+  autoload :FailedToAcquireLock, 'with_advisory_lock/failed_to_acquire_lock'
 end
 
 ActiveSupport.on_load :active_record do
-  include WithAdvisoryLock::Concern
+  require 'active_record/connection_adapters/abstract_adapter'
+  ActiveRecord::Base.include WithAdvisoryLock::Concern
+end
+
+# JRuby compatibility handling
+if RUBY_ENGINE == 'jruby'
+  require 'with_advisory_lock/jruby_adapter'
+  WithAdvisoryLock::JRubyAdapter.install!
+  # Don't set up the standard hooks for JRuby
+else
+  # Standard adapter injection for MRI and TruffleRuby
+  ActiveSupport.on_load :active_record_postgresqladapter do
+    prepend WithAdvisoryLock::CoreAdvisory
+    prepend WithAdvisoryLock::PostgreSQLAdvisory
+  end
+
+  ActiveSupport.on_load :active_record_mysql2adapter do
+    prepend WithAdvisoryLock::CoreAdvisory
+    prepend WithAdvisoryLock::MySQLAdvisory
+  end
+
+  ActiveSupport.on_load :active_record_trilogyadapter do
+    prepend WithAdvisoryLock::CoreAdvisory
+    prepend WithAdvisoryLock::MySQLAdvisory
+  end
 end

data/test/dummy/Rakefile

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+# Add your own tasks in files placed in lib/tasks ending in .rake,
+# for example lib/tasks/capistrano.rake, and they will automatically be available to Rake.
+
+require_relative 'config/application'
+
+Rails.application.load_tasks

data/test/dummy/app/controllers/application_controller.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+class ApplicationController < ActionController::Base
+  # Prevent CSRF attacks by raising an exception.
+  # For APIs, you may want to use :null_session instead.
+  protect_from_forgery with: :exception
+end

data/test/dummy/app/models/application_record.rb

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+class ApplicationRecord < ActiveRecord::Base
+  self.abstract_class = true
+  establish_connection(:primary)
+end

data/test/dummy/app/models/label.rb

@@ -0,0 +1,4 @@
+# frozen_string_literal: true
+
+class Label < ApplicationRecord
+end

data/test/dummy/app/models/mysql_label.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+class MysqlLabel < MysqlRecord
+  self.table_name = 'mysql_labels'
+end

data/test/dummy/app/models/mysql_record.rb

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+class MysqlRecord < ActiveRecord::Base
+  self.abstract_class = true
+  establish_connection :secondary
+end

data/test/dummy/app/models/mysql_tag.rb

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+class MysqlTag < MysqlRecord
+  self.table_name = 'mysql_tags'
+
+  after_save do
+    MysqlTagAudit.create(tag_name: name)
+    MysqlLabel.create(name: name)
+  end
+end

data/test/dummy/app/models/mysql_tag_audit.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+class MysqlTagAudit < MysqlRecord
+  self.table_name = 'mysql_tag_audits'
+end

data/test/dummy/app/models/tag.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+class Tag < ApplicationRecord
+  after_save do
+    TagAudit.create(tag_name: name)
+    Label.create(name: name)
+  end
+end

data/test/dummy/app/models/tag_audit.rb

@@ -0,0 +1,4 @@
+# frozen_string_literal: true
+
+class TagAudit < ApplicationRecord
+end

data/test/dummy/config/application.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+require File.expand_path('boot', __dir__)
+
+require 'rails'
+require 'active_model/railtie'
+require 'active_record/railtie'
+require 'action_controller/railtie'
+require 'action_view/railtie'
+
+Bundler.require(*Rails.groups)
+
+module TestSystemApp
+  class Application < Rails::Application
+    config.load_defaults [Rails::VERSION::MAJOR, Rails::VERSION::MINOR].join('.')
+    config.eager_load = true
+    config.serve_static_files = false
+    config.public_file_server.enabled = false
+    config.public_file_server.headers = { 'Cache-Control' => 'public, max-age=3600' }
+    config.consider_all_requests_local       = true
+    config.action_controller.perform_caching = false
+    config.action_dispatch.show_exceptions = false
+    config.action_controller.allow_forgery_protection = false
+    config.active_support.test_order = :random
+    config.active_support.deprecation = :stderr
+    config.active_record.timestamped_migrations = false
+
+    # Disable automatic database setup since we handle it manually
+    config.active_record.maintain_test_schema = false if config.respond_to?(:active_record)
+  end
+end

data/test/dummy/config/boot.rb

@@ -0,0 +1,3 @@
+# frozen_string_literal: true
+
+require 'bundler/setup'

data/test/dummy/config/database.yml

@@ -0,0 +1,13 @@
+default: &default
+  pool: 20
+
+
+test:
+  primary:
+    <<: *default
+    url: "<%= ENV['DATABASE_URL_PG'] %>"
+  secondary:
+    <<: *default
+    url: "<%= ENV['DATABASE_URL_MYSQL'] %>"
+    properties:
+      allowPublicKeyRetrieval: true

data/test/dummy/config/environment.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+# Load the Rails application.
+require File.expand_path('application', __dir__)
+
+# Initialize the Rails application.
+Rails.application.initialize!

data/test/dummy/config/routes.rb

@@ -0,0 +1,4 @@
+# frozen_string_literal: true
+
+Rails.application.routes.draw do
+end

data/test/dummy/config.ru

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+require_relative 'config/environment'
+
+run Rails.application
+Rails.application.load_server

data/test/{test_models.rb → dummy/db/schema.rb}

@@ -1,26 +1,15 @@
 # frozen_string_literal: true
 
-ActiveRecord::Schema.define(version: 0) do
+ActiveRecord::Schema.define(version: 1) do
   create_table 'tags', force: true do |t|
     t.string 'name'
   end
+
   create_table 'tag_audits', id: false, force: true do |t|
     t.string 'tag_name'
   end
+
   create_table 'labels', id: false, force: true do |t|
     t.string 'name'
   end
 end
-
-class Tag < ActiveRecord::Base
-  after_save do
-    TagAudit.create(tag_name: name)
-    Label.create(name: name)
-  end
-end
-
-class TagAudit < ActiveRecord::Base
-end
-
-class Label < ActiveRecord::Base
-end

data/test/dummy/db/secondary_schema.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+ActiveRecord::Schema.define(version: 1) do
+  create_table 'mysql_tags', force: true do |t|
+    t.string 'name'
+  end
+
+  create_table 'mysql_tag_audits', id: false, force: true do |t|
+    t.string 'tag_name'
+  end
+
+  create_table 'mysql_labels', id: false, force: true do |t|
+    t.string 'name'
+  end
+end