counting_semaphore 0.1.0

data/lib/counting_semaphore/shared_semaphore.rb

@@ -0,0 +1,381 @@
+# A distributed counting semaphore that allows up to N concurrent operations across multiple processes.
+# Uses Redis for coordination and automatically handles lease expiration for crashed processes.
+# Uses Redis Lua scripts for atomic operations to prevent race conditions.
+require "digest"
+require "securerandom"
+
+module CountingSemaphore
+  class RedisSemaphore
+    # Custom exception for lease acquisition timeouts
+    class LeaseTimeout < StandardError
+      def initialize(token_count, timeout_seconds)
+        super("Failed to acquire #{token_count} tokens within #{timeout_seconds} seconds")
+      end
+    end
+
+    LEASE_EXPIRATION_SECONDS = 5
+
+    # Lua script for atomic lease acquisition
+    # Returns: [success, lease_key, current_usage]
+    # success: 1 if lease was acquired, 0 if no capacity
+    # lease_key: the key of the acquired lease (if successful)
+    # current_usage: current usage count after operation
+    GET_LEASE_SCRIPT = <<~LUA
+      local lease_key = KEYS[1]
+      local lease_set_key = KEYS[2]
+      local capacity = tonumber(ARGV[1])
+      local token_count = tonumber(ARGV[2])
+      local expiration_seconds = tonumber(ARGV[3])
+      
+      -- Get all active leases from the set and calculate current usage
+      local lease_keys = redis.call('SMEMBERS', lease_set_key)
+      local current_usage = 0
+      local valid_leases = {}
+      
+      for i, key in ipairs(lease_keys) do
+        local tokens = redis.call('GET', key)
+        if tokens then
+          local tokens_from_lease = tonumber(tokens)
+          if tokens_from_lease then
+            current_usage = current_usage + tokens_from_lease
+            table.insert(valid_leases, key)
+          else
+            -- Remove lease with invalid token count
+            redis.call('DEL', key)
+            redis.call('SREM', lease_set_key, key)
+          end
+        else
+          -- Lease key doesn't exist, remove from set
+          redis.call('SREM', lease_set_key, key)
+        end
+      end
+      
+      -- Check if we have capacity
+      local available = capacity - current_usage
+      if available >= token_count then
+        -- Set lease with TTL (value is just the token count)
+        redis.call('SETEX', lease_key, expiration_seconds, token_count)
+        -- Add lease key to the set
+        redis.call('SADD', lease_set_key, lease_key)
+        -- Set TTL on the set (4x the lease TTL to ensure cleanup)
+        redis.call('EXPIRE', lease_set_key, expiration_seconds * 4)
+        
+        return {1, lease_key, current_usage + token_count}
+      else
+        return {0, '', current_usage}
+      end
+    LUA
+
+    # Lua script for getting current usage
+    # Returns: current_usage (integer)
+    GET_USAGE_SCRIPT = <<~LUA
+      local lease_set_key = KEYS[1]
+      local expiration_seconds = tonumber(ARGV[1])
+      
+      -- Get all active leases from the set and calculate current usage
+      local lease_keys = redis.call('SMEMBERS', lease_set_key)
+      local current_usage = 0
+      local has_valid_leases = false
+      
+      for i, lease_key in ipairs(lease_keys) do
+        local tokens = redis.call('GET', lease_key)
+        if tokens then
+          local tokens_from_lease = tonumber(tokens)
+          if tokens_from_lease then
+            current_usage = current_usage + tokens_from_lease
+            has_valid_leases = true
+          else
+            -- Remove lease with invalid token count
+            redis.call('DEL', lease_key)
+            redis.call('SREM', lease_set_key, lease_key)
+          end
+        else
+          -- Lease key doesn't exist, remove from set
+          redis.call('SREM', lease_set_key, lease_key)
+        end
+      end
+      
+      -- Refresh TTL on the set if there are valid leases (4x the lease TTL)
+      if has_valid_leases then
+        redis.call('EXPIRE', lease_set_key, expiration_seconds * 4)
+      end
+      
+      return current_usage
+    LUA
+
+    # Lua script for atomic lease release and signal
+    # Returns: 1 (success)
+    RELEASE_LEASE_SCRIPT = <<~LUA
+      local lease_key = KEYS[1]
+      local queue_key = KEYS[2]
+      local lease_set_key = KEYS[3]
+      local token_count = tonumber(ARGV[1])
+      local max_signals = tonumber(ARGV[2])
+      
+      -- Remove the lease
+      redis.call('DEL', lease_key)
+      -- Remove from the lease set
+      redis.call('SREM', lease_set_key, lease_key)
+      
+      -- Signal waiting clients about the released tokens
+      redis.call('LPUSH', queue_key, 'tokens:' .. token_count)
+      
+      -- Trim queue to prevent indefinite growth (atomic)
+      redis.call('LTRIM', queue_key, 0, max_signals - 1)
+      
+      return 1
+    LUA
+
+    # Precomputed script SHAs
+    GET_LEASE_SCRIPT_SHA = Digest::SHA1.hexdigest(GET_LEASE_SCRIPT)
+    GET_USAGE_SCRIPT_SHA = Digest::SHA1.hexdigest(GET_USAGE_SCRIPT)
+    RELEASE_LEASE_SCRIPT_SHA = Digest::SHA1.hexdigest(RELEASE_LEASE_SCRIPT)
+
+    # @return [Integer]
+    attr_reader :capacity
+
+    # Initialize the semaphore with a maximum capacity and required namespace.
+    #
+    # @param capacity [Integer] Maximum number of concurrent operations allowed
+    # @param namespace [String] Required namespace for Redis keys
+    # @param redis [Redis, ConnectionPool] Optional Redis client or connection pool (defaults to new Redis instance)
+    # @param logger [Logger] the logger
+    # @raise [ArgumentError] if capacity is not positive
+    def initialize(capacity, namespace, redis: nil, logger: CountingSemaphore::NullLogger, lease_expiration_seconds: LEASE_EXPIRATION_SECONDS)
+      raise ArgumentError, "Capacity must be positive, got #{capacity}" unless capacity > 0
+
+      # Require Redis only when SharedSemaphore is used
+      require "redis" unless defined?(Redis)
+
+      @capacity = capacity
+      @redis_connection_pool = wrap_redis_client_with_pool(redis || Redis.new)
+      @namespace = namespace
+      @lease_expiration_seconds = lease_expiration_seconds
+      @logger = logger
+
+      # Scripts are precomputed and will be loaded on-demand if needed
+    end
+
+    # Null pool for bare Redis connections that don't need connection pooling
+    class NullPool
+      def initialize(redis_connection)
+        @redis_connection = redis_connection
+      end
+
+      def with(&block)
+        block.call(@redis_connection)
+      end
+    end
+
+    # Acquire a lease for the specified number of tokens and execute the block.
+    # Blocks until sufficient resources are available.
+    #
+    # @param token_count [Integer] Number of tokens to acquire
+    # @param timeout_seconds [Integer] Maximum time to wait for lease acquisition (default: 30 seconds)
+    # @yield The block to execute while holding the lease
+    # @return The result of the block
+    # @raise [ArgumentError] if token_count is negative or exceeds the semaphore capacity
+    # @raise [LeaseTimeout] if lease cannot be acquired within timeout
+    def with_lease(token_count, timeout_seconds: 30)
+      raise ArgumentError, "Token count must be non-negative, got #{token_count}" if token_count < 0
+      if token_count > @capacity
+        raise ArgumentError, "Cannot lease #{token_count} slots as we only allow #{@capacity}"
+      end
+
+      # Handle zero tokens case - no Redis coordination needed
+      return yield if token_count.zero?
+
+      lease_key = acquire_lease(token_count, timeout_seconds: timeout_seconds)
+      begin
+        @logger.debug "🚦Leased #{token_count} tokens with lease #{lease_key}"
+        yield
+      ensure
+        release_lease(lease_key, token_count)
+      end
+    end
+
+    # Get current usage and active leases for debugging
+    def debug_info
+      usage = get_current_usage
+      lease_set_key = "#{@namespace}:lease_set"
+      lease_keys = with_redis { |redis| redis.smembers(lease_set_key) }
+      active_leases = []
+
+      lease_keys.each do |lease_key|
+        tokens = with_redis { |redis| redis.get(lease_key) }
+        next unless tokens
+
+        active_leases << {
+          key: lease_key,
+          tokens: tokens.to_i
+        }
+      end
+
+      {
+        usage: usage,
+        capacity: @capacity,
+        available: @capacity - usage,
+        active_leases: active_leases
+      }
+    end
+
+    private
+
+    # Wraps a Redis client to support both ConnectionPool and bare Redis connections
+    # @param redis [Redis, ConnectionPool] The Redis client or connection pool
+    # @return [Object] A wrapper that supports the `with` method
+    def wrap_redis_client_with_pool(redis)
+      # If it's already a ConnectionPool, return it as-is
+      return redis if redis.respond_to?(:with)
+
+      # For bare Redis connections, wrap in a NullPool
+      NullPool.new(redis)
+    end
+
+    # Executes a block with a Redis connection from the pool
+    # @yield [redis] The Redis connection
+    # @return The result of the block
+    def with_redis(&block)
+      @redis_connection_pool.with(&block)
+    end
+
+    # Executes a Redis script with automatic fallback to EVAL on NOSCRIPT error
+    # @param script_type [Symbol] The type of script (:get_lease, :release_lease, :get_usage)
+    # @param keys [Array] Keys for the script
+    # @param argv [Array] Arguments for the script
+    # @return The result of the script execution
+    def execute_script(script_type, keys: [], argv: [])
+      script_sha, script_body = case script_type
+      when :get_lease then [GET_LEASE_SCRIPT_SHA, GET_LEASE_SCRIPT]
+      when :release_lease then [RELEASE_LEASE_SCRIPT_SHA, RELEASE_LEASE_SCRIPT]
+      when :get_usage then [GET_USAGE_SCRIPT_SHA, GET_USAGE_SCRIPT]
+      else raise ArgumentError, "Unknown script type: #{script_type}"
+      end
+
+      with_redis do |redis|
+        redis.evalsha(script_sha, keys: keys, argv: argv)
+      end
+    rescue Redis::CommandError => e
+      if e.message.include?("NOSCRIPT")
+        @logger.debug "🚦Script not found, using EVAL: #{e.message}"
+        # Fall back to EVAL with the script body
+        with_redis do |redis|
+          redis.eval(script_body, keys: keys, argv: argv)
+        end
+      else
+        raise
+      end
+    end
+
+    def acquire_lease(token_count, timeout_seconds: 30)
+      start_time = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+
+      loop do
+        # Check if we've exceeded the timeout using monotonic time
+        elapsed_time = Process.clock_gettime(Process::CLOCK_MONOTONIC) - start_time
+        if elapsed_time >= timeout_seconds
+          raise LeaseTimeout.new(token_count, timeout_seconds)
+        end
+
+        # Try optimistic acquisition first
+        lease_key = attempt_lease_acquisition(token_count)
+        return lease_key if lease_key
+
+        # If failed, wait for signals with timeout
+        lease_key = wait_for_tokens(token_count, timeout_seconds - elapsed_time)
+        return lease_key if lease_key
+      end
+    end
+
+    def wait_for_tokens(token_count, remaining_timeout)
+      # Ensure minimum timeout to prevent infinite blocking
+      # BLPOP with timeout 0 blocks forever, so we need at least a small positive timeout
+      minimum_timeout = 0.1
+      if remaining_timeout <= minimum_timeout
+        @logger.debug "🚦Remaining timeout (#{remaining_timeout}s) too small, not waiting"
+        return nil
+      end
+
+      # Block with timeout (longer than lease expiration to handle stale leases)
+      # But don't exceed the remaining timeout
+      timeout = [@lease_expiration_seconds + 2, remaining_timeout].min
+      @logger.debug "🚦Unable to lease #{token_count}, waiting for signals (timeout: #{timeout}s)"
+
+      with_redis { |redis| redis.blpop("#{@namespace}:waiting_queue", timeout: timeout.to_f) }
+
+      # Try to acquire after any signal or timeout
+      lease_key = attempt_lease_acquisition(token_count)
+      if lease_key
+        return lease_key
+      end
+
+      # If still can't acquire, return nil to continue the loop in acquire_lease
+      @logger.debug "🚦Still unable to lease #{token_count} after signal/timeout, continuing to wait"
+      nil
+    end
+
+    def attempt_lease_acquisition(token_count)
+      lease_id = generate_lease_id
+      lease_key = "#{@namespace}:leases:#{lease_id}"
+      lease_set_key = "#{@namespace}:lease_set"
+
+      # Use Lua script for atomic lease acquisition
+      result = execute_script(
+        :get_lease,
+        keys: [lease_key, lease_set_key],
+        argv: [
+          @capacity.to_s,
+          token_count.to_s,
+          @lease_expiration_seconds.to_s
+        ]
+      )
+
+      success, full_lease_key, current_usage = result
+
+      if success == 1
+        # Extract just the lease ID from the full key for return value
+        lease_id = full_lease_key.split(":").last
+        @logger.debug "🚦Acquired lease #{lease_id}, current usage: #{current_usage}/#{@capacity}"
+        lease_id
+      else
+        @logger.debug "🚦No capacity available, current usage: #{current_usage}/#{@capacity}"
+        nil
+      end
+    end
+
+    def release_lease(lease_key, token_count)
+      return if lease_key.nil?
+
+      full_lease_key = "#{@namespace}:leases:#{lease_key}"
+      queue_key = "#{@namespace}:waiting_queue"
+      lease_set_key = "#{@namespace}:lease_set"
+
+      # Use Lua script for atomic lease release and signal
+      execute_script(
+        :release_lease,
+        keys: [full_lease_key, queue_key, lease_set_key],
+        argv: [
+          token_count.to_s,
+          (@capacity * 2).to_s
+        ]
+      )
+
+      @logger.debug "🚦Returned #{token_count} leased tokens (lease: #{lease_key}) and signaled waiting clients"
+    end
+
+    def get_current_usage
+      lease_set_key = "#{@namespace}:lease_set"
+
+      # Use the dedicated usage script that calculates current usage
+      execute_script(
+        :get_usage,
+        keys: [lease_set_key],
+        argv: [@lease_expiration_seconds.to_s]
+      )
+    end
+
+    def generate_lease_id
+      SecureRandom.uuid
+    end
+  end
+end

data/lib/counting_semaphore/version.rb

@@ -0,0 +1,3 @@
+module CountingSemaphore
+  VERSION = "0.1.0"
+end

data/lib/counting_semaphore.rb

@@ -0,0 +1,19 @@
+require_relative "counting_semaphore/version"
+
+module CountingSemaphore
+  # Custom exception for lease acquisition timeouts
+  class LeaseTimeout < StandardError
+    attr_reader :semaphore, :token_count, :timeout_seconds
+
+    def initialize(token_count, timeout_seconds, semaphore = nil)
+      @token_count = token_count
+      @timeout_seconds = timeout_seconds
+      @semaphore = semaphore
+      super("Failed to acquire #{token_count} tokens within #{timeout_seconds} seconds")
+    end
+  end
+
+  autoload :LocalSemaphore, "counting_semaphore/local_semaphore"
+  autoload :RedisSemaphore, "counting_semaphore/redis_semaphore"
+  autoload :NullLogger, "counting_semaphore/null_logger"
+end

data/test/counting_semaphore/local_semaphore_test.rb

@@ -0,0 +1,304 @@
+require "test_helper"
+require "counting_semaphore"
+
+class LocalSemaphoreTest < Minitest::Test
+  def test_initializes_with_correct_capacity
+    semaphore = CountingSemaphore::LocalSemaphore.new(5)
+    assert_equal 5, semaphore.instance_variable_get(:@capacity)
+  end
+
+  def test_capacity_attribute_returns_the_initialized_capacity
+    semaphore = CountingSemaphore::LocalSemaphore.new(10)
+    assert_equal 10, semaphore.capacity
+  end
+
+  def test_capacity_attribute_is_immutable
+    semaphore = CountingSemaphore::LocalSemaphore.new(7)
+    assert_equal 7, semaphore.capacity
+
+    # Verify that capacity cannot be modified directly
+    assert_raises(NoMethodError) do
+      semaphore.capacity = 5
+    end
+  end
+
+  def test_raises_error_for_negative_capacity
+    assert_raises(ArgumentError, "Capacity must be positive, got -3") do
+      CountingSemaphore::LocalSemaphore.new(-3)
+    end
+  end
+
+  def test_raises_error_for_zero_capacity
+    assert_raises(ArgumentError, "Capacity must be positive, got 0") do
+      CountingSemaphore::LocalSemaphore.new(0)
+    end
+  end
+
+  def test_allows_operations_within_capacity
+    semaphore = CountingSemaphore::LocalSemaphore.new(2)
+    result = nil
+
+    semaphore.with_lease(1) do
+      result = "success"
+    end
+
+    assert_equal "success", result
+  end
+
+  def test_blocks_when_capacity_exceeded
+    semaphore = CountingSemaphore::LocalSemaphore.new(1)
+    start_time = Time.now
+    completed = false
+
+    # Start a thread that will hold the semaphore
+    thread1 = Thread.new do
+      semaphore.with_lease(1) do
+        sleep(0.1) # Hold the semaphore briefly
+        "thread1"
+      end
+    end
+
+    # Start another thread that should block
+    thread2 = Thread.new do
+      semaphore.with_lease(1) do
+        completed = true
+        "thread2"
+      end
+    end
+
+    thread1.join
+    thread2.join
+
+    # The second thread should have completed after the first released
+    assert completed
+    assert (Time.now - start_time) >= 0.1
+  end
+
+  def test_raises_error_when_requesting_more_tokens_than_capacity
+    semaphore = CountingSemaphore::LocalSemaphore.new(2)
+
+    assert_raises(ArgumentError) do
+      semaphore.with_lease(3) do
+        "should not reach here"
+      end
+    end
+  end
+
+  def test_raises_error_for_negative_token_count
+    semaphore = CountingSemaphore::LocalSemaphore.new(2)
+
+    assert_raises(ArgumentError, "Token count must be non-negative, got -1") do
+      semaphore.with_lease(-1) do
+        "should not reach here"
+      end
+    end
+  end
+
+  def test_allows_zero_token_count
+    semaphore = CountingSemaphore::LocalSemaphore.new(2)
+    result = nil
+
+    semaphore.with_lease(0) do
+      result = "success"
+    end
+
+    assert_equal "success", result
+  end
+
+  def test_properly_releases_tokens_after_block_completion
+    semaphore = CountingSemaphore::LocalSemaphore.new(1)
+    results = []
+
+    # First operation should succeed immediately
+    semaphore.with_lease(1) do
+      results << "first"
+    end
+
+    # Second operation should also succeed immediately (tokens were released)
+    semaphore.with_lease(1) do
+      results << "second"
+    end
+
+    assert_equal ["first", "second"], results
+  end
+
+  def test_handles_exceptions_and_still_releases_tokens
+    semaphore = CountingSemaphore::LocalSemaphore.new(1)
+    results = []
+
+    # First operation succeeds
+    semaphore.with_lease(1) do
+      results << "first"
+    end
+
+    # Second operation should fail but still release tokens
+    assert_raises(RuntimeError) do
+      semaphore.with_lease(1) do
+        results << "second"
+        raise "test error"
+      end
+    end
+
+    # Third operation should succeed (tokens were released despite exception)
+    semaphore.with_lease(1) do
+      results << "third"
+    end
+
+    assert_equal ["first", "second", "third"], results
+  end
+
+  def test_supports_multiple_tokens_per_lease
+    semaphore = CountingSemaphore::LocalSemaphore.new(3)
+    results = []
+
+    # Should be able to lease 2 tokens at once
+    semaphore.with_lease(2) do
+      results << "two_tokens"
+    end
+
+    # Should be able to lease 1 more token
+    semaphore.with_lease(1) do
+      results << "one_token"
+    end
+
+    assert_equal ["two_tokens", "one_token"], results
+  end
+
+  def test_with_lease_accepts_timeout_parameter
+    semaphore = CountingSemaphore::LocalSemaphore.new(1)
+
+    # Test that timeout parameter is accepted
+    result = semaphore.with_lease(1, timeout_seconds: 5) do
+      "success"
+    end
+
+    assert_equal "success", result
+  end
+
+  def test_with_lease_raises_timeout_error_when_timeout_is_exceeded
+    semaphore = CountingSemaphore::LocalSemaphore.new(1)
+
+    # Fill the semaphore
+    semaphore.with_lease(1) do
+      # Try to acquire another token with a very short timeout
+      assert_raises(CountingSemaphore::LeaseTimeout) do
+        semaphore.with_lease(1, timeout_seconds: 0.1) do
+          "should not reach here"
+        end
+      end
+    end
+  end
+
+  def test_lease_timeout_includes_semaphore_reference
+    semaphore = CountingSemaphore::LocalSemaphore.new(1)
+    exception = nil
+
+    # Fill the semaphore
+    semaphore.with_lease(1) do
+      # Try to acquire another token with a very short timeout
+
+      semaphore.with_lease(1, timeout_seconds: 0.1) do
+        "should not reach here"
+      end
+    rescue CountingSemaphore::LeaseTimeout => e
+      exception = e
+    end
+
+    refute_nil exception
+    assert_equal semaphore, exception.semaphore
+    assert_equal 1, exception.token_count
+    assert_equal 0.1, exception.timeout_seconds
+  end
+
+  def test_with_lease_uses_default_token_count_of_1
+    semaphore = CountingSemaphore::LocalSemaphore.new(2)
+    result = nil
+
+    # Should work with default token count (1)
+    semaphore.with_lease do
+      result = "success"
+    end
+
+    assert_equal "success", result
+  end
+
+  def test_with_lease_default_blocks_when_capacity_exceeded
+    semaphore = CountingSemaphore::LocalSemaphore.new(1)
+    start_time = Time.now
+    completed = false
+
+    # Start a thread that will hold the semaphore
+    thread1 = Thread.new do
+      semaphore.with_lease do  # Uses default token count of 1
+        sleep(0.1) # Hold the semaphore briefly
+        "thread1"
+      end
+    end
+
+    # Start another thread that should block
+    thread2 = Thread.new do
+      semaphore.with_lease do  # Uses default token count of 1
+        completed = true
+        "thread2"
+      end
+    end
+
+    thread1.join
+    thread2.join
+
+    # The second thread should have completed after the first released
+    assert completed
+    assert (Time.now - start_time) >= 0.1
+  end
+
+  def test_currently_leased_returns_zero_initially
+    semaphore = CountingSemaphore::LocalSemaphore.new(5)
+    assert_equal 0, semaphore.currently_leased
+  end
+
+  def test_currently_leased_increases_during_lease
+    semaphore = CountingSemaphore::LocalSemaphore.new(5)
+    usage_during_lease = nil
+
+    semaphore.with_lease(2) do
+      usage_during_lease = semaphore.currently_leased
+    end
+
+    assert_equal 2, usage_during_lease
+    assert_equal 0, semaphore.currently_leased
+  end
+
+  def test_currently_leased_returns_to_zero_after_lease
+    semaphore = CountingSemaphore::LocalSemaphore.new(3)
+
+    semaphore.with_lease(2) do
+      assert_equal 2, semaphore.currently_leased
+    end
+
+    assert_equal 0, semaphore.currently_leased
+  end
+
+  def test_currently_leased_with_multiple_concurrent_leases
+    semaphore = CountingSemaphore::LocalSemaphore.new(5)
+    usage_values = []
+    mutex = Mutex.new
+
+    # Start multiple threads that will hold leases
+    threads = []
+    3.times do |i|
+      threads << Thread.new do
+        semaphore.with_lease(1) do
+          mutex.synchronize { usage_values << semaphore.currently_leased }
+          sleep(0.1) # Hold the lease briefly
+        end
+      end
+    end
+
+    threads.each(&:join)
+
+    # Should have seen usage values of 1, 2, and 3 (or some combination)
+    assert usage_values.any? { |usage| usage >= 1 }
+    assert usage_values.any? { |usage| usage <= 3 }
+    assert_equal 0, semaphore.currently_leased
+  end
+end