schematichq 1.5.2 → 1.5.3

data/lib/schematic/credits/leases/lease_store.rb

@@ -0,0 +1,240 @@
+# frozen_string_literal: true
+
+module Schematic
+  module Credits
+    module Leases
+      def self.lease_key(company_id, credit_type_id)
+        "#{company_id}:#{credit_type_id}"
+      end
+
+      # Per-process lease store, keyed by "<company_id>:<credit_type_id>".
+      #
+      # Holds the lease id, the server-authoritative granted total, the expiry,
+      # and the SDK's local view of local_remaining_credits: the portion of the
+      # lease not yet carved out by an open reservation. A per-key mutex keeps
+      # reserve, refund, extend, and replace atomic per slot.
+      #
+      # The store never talks to the API. LeaseManager drives acquire, extend,
+      # and release over the wire and uses these methods to mirror remote state.
+      #
+      # For cross-process deployments use RedisLeaseStore instead: both answer
+      # the same calls with the same semantics.
+      class LeaseStore
+        def initialize(clock: DEFAULT_CLOCK)
+          @clock = clock
+          @leases = {}
+          @locks = {}
+          # Guards the two hashes themselves. Always taken inside a key lock,
+          # never the other way round, so the pair cannot deadlock.
+          @table_mutex = Mutex.new
+        end
+
+        # Snapshot of the current entry, or nil when no lease occupies the slot.
+        # A copy, so a caller cannot mutate store state by holding the result.
+        def get(company_id, credit_type_id)
+          key = Leases.lease_key(company_id, credit_type_id)
+          with_lock_if_present(key) { read(key)&.dup }
+        end
+
+        # Install a fresh lease for the slot, but only if no live lease already
+        # occupies it.
+        #
+        # A live lease is left untouched even when its lease_id differs (a
+        # sibling process won the race), so its already-debited
+        # local_remaining_credits wins. An expired row carrying the SAME
+        # lease_id is not rewritten either: rewriting would reset
+        # local_remaining_credits to the full grant and erase debits whose
+        # reservations are still open, and the idempotent server hands a racing
+        # acquire that same active lease back. Such a row is reconciled like an
+        # extend instead: granted to the incoming total, expiry only forward,
+        # balance untouched.
+        #
+        # Returns true when a fresh row was written, false when an existing
+        # lease was kept or reconciled.
+        def replace(entry)
+          key = Leases.lease_key(entry.company_id, entry.credit_type_id)
+          with_lock(key) do
+            existing = read(key)
+            if existing && !existing.expired?(@clock.call)
+              false
+            elsif existing && existing.lease_id == entry.lease_id
+              reconcile(existing, entry.granted_amount, entry.expires_at)
+              false
+            else
+              write(key, LeaseEntry.new(
+                lease_id: entry.lease_id,
+                company_id: entry.company_id,
+                credit_type_id: entry.credit_type_id,
+                granted_amount: entry.granted_amount,
+                expires_at: entry.expires_at
+              ))
+              true
+            end
+          end
+        end
+
+        # Reconcile the slot to the server-authoritative granted total after a
+        # remote extend, crediting the difference to local_remaining_credits.
+        #
+        # The delta is computed here against the CURRENT stored total, never by
+        # the caller from a pre-wire-call read: two writers extending the same
+        # lease concurrently would each apply a delta against the same stale
+        # read and mint phantom credits. Reconciling to the absolute total
+        # converges in any order, and the expiry only ever moves forward so an
+        # out-of-order apply cannot shorten a lease a sibling just extended.
+        #
+        # When pin_lease_id is given the extend applies only if the slot still
+        # holds that lease: the server extended lease A, so its credits must not
+        # land on a successor B that replaced A after it expired mid-extend.
+        def extend(company_id, credit_type_id, granted_amount, new_expires_at = nil, pin_lease_id = nil)
+          key = Leases.lease_key(company_id, credit_type_id)
+          with_lock_if_present(key) do
+            entry = read(key)
+            next if entry.nil?
+            next if pin_lease_id && entry.lease_id != pin_lease_id
+
+            reconcile(entry, granted_amount, new_expires_at)
+          end
+          nil
+        end
+
+        # Drop the slot entry, after a remote release. An explicit drop is the
+        # only thing that removes one: a lease that merely expired stays here,
+        # readable, until it is dropped or replaced, which is what the spec
+        # requires, so every path re-guards on expiry rather than trusting
+        # presence. The slot's mutex goes with the entry, so a long-lived
+        # process does not accumulate one per slot it has ever leased.
+        def drop(company_id, credit_type_id)
+          key = Leases.lease_key(company_id, credit_type_id)
+          with_lock(key) do
+            @table_mutex.synchronize do
+              @leases.delete(key)
+              @locks.delete(key)
+            end
+          end
+          nil
+        end
+
+        # Atomically check and debit `credits` from the lease's remaining
+        # balance. Returns a ReserveResult carrying the post-debit balance and
+        # the lease the debit landed on, or nil when there is no live lease, the
+        # balance is short, or `credits` is not a finite non-negative number.
+        #
+        # The debit is not keyed by lease id: it charges whichever lease holds
+        # the slot at that moment, which need not be the one the caller's
+        # acquire returned. The caller must therefore pin its reservation to the
+        # returned lease_id, since the settle refund, the sweep refund, and the
+        # track event's lease_id all have to name the lease that was charged.
+        def try_reserve(company_id, credit_type_id, credits)
+          # NaN passes every comparison below, and a NaN balance would approve
+          # every later reserve, so reject it before it reaches the arithmetic.
+          return nil unless Leases.valid_quantity?(credits)
+
+          key = Leases.lease_key(company_id, credit_type_id)
+          with_lock_if_present(key) do
+            entry = read(key)
+            next nil if entry.nil?
+            # An expired lease is released server-side and its grant refunded to
+            # the company balance, so its local balance is stale.
+            next nil if entry.expired?(@clock.call)
+            next nil if entry.local_remaining_credits < credits
+
+            entry.local_remaining_credits -= credits
+            # The lease id is read under the same lock as the debit: a read
+            # afterwards could name a lease that replaced this one in between.
+            ReserveResult.new(entry.local_remaining_credits, entry.lease_id)
+          end
+        end
+
+        # Refund credits to the slot's lease balance, clamped at granted_amount.
+        #
+        # When pin_lease_id is given the refund applies only if the slot still
+        # holds that lease: a hold carved out of expired lease A must never
+        # inflate a successor B, because A's unspent remainder was already
+        # returned to the company balance server-side when A expired.
+        def refund(company_id, credit_type_id, credits, pin_lease_id = nil)
+          return nil if credits.nil? || credits <= 0
+
+          key = Leases.lease_key(company_id, credit_type_id)
+          with_lock_if_present(key) do
+            entry = read(key)
+            next if entry.nil?
+            next if pin_lease_id && entry.lease_id != pin_lease_id
+
+            entry.local_remaining_credits = [entry.local_remaining_credits + credits, entry.granted_amount].min
+          end
+          nil
+        end
+
+        # Snapshot of every entry. Only the per-process store implements this:
+        # close releases the leases this process exclusively holds, and a shared
+        # backend must never enumerate, since sibling processes may still be
+        # drawing on those leases.
+        def list
+          keys = @table_mutex.synchronize { @leases.keys }
+          keys.filter_map { |key| with_lock(key) { read(key)&.dup } }
+        end
+
+        private
+
+        # Granted to the incoming total, expiry only forward, balance untouched
+        # except for the credited difference. Shared by replace's same-lease
+        # path and extend, which reconcile identically.
+        def reconcile(entry, granted_amount, new_expires_at)
+          add = granted_amount.to_f - entry.granted_amount
+          if add.positive?
+            entry.granted_amount = granted_amount.to_f
+            entry.local_remaining_credits += add
+          end
+          return unless new_expires_at && new_expires_at.to_f > entry.expires_at.to_f
+
+          entry.expires_at = new_expires_at
+        end
+
+        def read(key)
+          @table_mutex.synchronize { @leases[key] }
+        end
+
+        def write(key, entry)
+          @table_mutex.synchronize { @leases[key] = entry }
+        end
+
+        # A read of a slot nothing has leased must not leave a mutex behind, or
+        # pruning would only half work: a process checking flags for companies
+        # that never lease would still collect one per slot it asked about. With
+        # no lock registered and no lease stored there is nothing to serialize
+        # against, so answer without taking one. A writer landing in that window
+        # is the same race as reading a moment earlier.
+        def with_lock_if_present(key, &)
+          present = @table_mutex.synchronize { @locks.key?(key) || @leases.key?(key) }
+          return nil unless present
+
+          with_lock(key, &)
+        end
+
+        # Serialize on the slot's mutex, re-checking after the acquire that it
+        # is still the registered one.
+        #
+        # Pruning is what makes the re-check necessary. A drop deletes the mutex
+        # while holding it, so a thread that was already blocked on it wakes
+        # owning an object the table no longer knows about, while a thread
+        # arriving afterwards takes a fresh mutex for the same slot. Without the
+        # check those two would run side by side on one slot. The waiter sees
+        # its mutex is no longer the registered one and retries against the
+        # current one instead. Every retry follows a drop, which happens once
+        # per lease, so the loop cannot spin.
+        def with_lock(key, &block)
+          loop do
+            lock = @table_mutex.synchronize { @locks[key] ||= Mutex.new }
+            stale = false
+            result = lock.synchronize do
+              stale = @table_mutex.synchronize { !@locks[key].equal?(lock) }
+              block.call unless stale
+            end
+            return result unless stale
+          end
+        end
+      end
+    end
+  end
+end

data/lib/schematic/credits/leases/redis_lease_store.rb

@@ -0,0 +1,358 @@
+# frozen_string_literal: true
+
+require "digest"
+
+module Schematic
+  module Credits
+    module Leases
+      # Minimal interface describing the Redis client methods the lease and
+      # reservation stores use. Compatible with the "redis" gem's client, which
+      # is deliberately not a dependency: the caller injects a client it already
+      # has, the same way RedisCacheProvider takes one.
+      module RedisLeaseClientInterface
+        def evalsha(sha, keys:, argv:)
+          raise NotImplementedError
+        end
+
+        def eval(script, keys:, argv:)
+          raise NotImplementedError
+        end
+
+        def hgetall(key)
+          raise NotImplementedError
+        end
+
+        # Called as hset(key, field1, value1, field2, value2, ...).
+        def hset(key, *pairs)
+          raise NotImplementedError
+        end
+
+        def hdel(key, field)
+          raise NotImplementedError
+        end
+
+        def del(*keys)
+          raise NotImplementedError
+        end
+
+        def pexpireat(key, millis)
+          raise NotImplementedError
+        end
+
+        def zadd(key, score, member)
+          raise NotImplementedError
+        end
+
+        def zrem(key, member)
+          raise NotImplementedError
+        end
+
+        def zrangebyscore(key, min, max, limit: nil)
+          raise NotImplementedError
+        end
+
+        def zcard(key)
+          raise NotImplementedError
+        end
+      end
+
+      # A Lua script addressed by its SHA. EVALSHA first so the script body is
+      # not resent on every call; a Redis that has never seen it (a restart, or
+      # a node this process has not talked to) answers NOSCRIPT and the full
+      # body goes out once to load it.
+      class Script
+        attr_reader :source, :sha
+
+        def initialize(source)
+          @source = source
+          @sha = Digest::SHA1.hexdigest(source)
+        end
+
+        def call(client, keys:, argv:)
+          client.evalsha(@sha, keys: keys, argv: argv)
+        rescue StandardError => e
+          raise unless e.message.to_s.include?("NOSCRIPT")
+
+          client.eval(@source, keys: keys, argv: argv)
+        end
+      end
+
+      # Redis-backed lease store. One hash per (company_id, credit_type_id)
+      # slot, mutated by single-key Lua scripts so the store stays correct on
+      # standalone and clustered Redis alike.
+      #
+      # The scripts below are byte-identical to the ones the Node, Go, and
+      # Python SDKs ship, and the key layout matches, so a mixed-language fleet
+      # shares one lease per slot. Do not re-derive them.
+      class RedisLeaseStore
+        DEFAULT_KEY_PREFIX = "schematic:"
+        LEASE_KEY_NAMESPACE = "credit-lease:"
+        # How long after the declared expiry the Redis row is kept before
+        # auto-eviction. Gives the sweeper a window to refund expired
+        # reservations before the underlying lease state disappears.
+        LEASE_TTL_GRACE_MS = 60_000
+
+        # Every script below touches exactly ONE key (the lease hash), which
+        # keeps them safe under Redis Cluster, where a multi-key script spanning
+        # slots raises CROSSSLOT. Only the lease hash needs atomic mutation;
+        # cross-key bookkeeping uses ordinary single-key commands.
+        #
+        # Expiry is decided against the Redis server's clock (redis.call('TIME')),
+        # not the calling process's: with many processes sharing one lease, local
+        # clock skew would let them disagree on whether the lease is live.
+
+        # Atomic replace. Writes the lease hash only when the slot is empty or
+        # the existing lease has expired. Returns 1 on write, 0 when a LIVE
+        # lease already occupies the slot, even one with a different leaseId
+        # installed by a sibling that raced this acquire. An expired row with
+        # the SAME leaseId is reconciled like an extend instead of rewritten,
+        # since rewriting would reset the balance and erase debits whose
+        # reservations are still open.
+        REPLACE_SCRIPT = Script.new(<<~LUA)
+
+          redis.replicate_commands()
+          local t = redis.call('TIME')
+          local now = (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000)
+
+          local existing_id = redis.call('HGET', KEYS[1], 'leaseId')
+          local existing_expiry = tonumber(redis.call('HGET', KEYS[1], 'expiresAt') or '0')
+          local new_id = ARGV[1]
+          local new_granted = ARGV[2]
+          local new_expiry = tonumber(ARGV[3])
+          local grace = tonumber(ARGV[4])
+
+          if existing_id and existing_expiry > now then
+              return 0
+          end
+
+          if existing_id == new_id then
+              local granted = tonumber(redis.call('HGET', KEYS[1], 'grantedAmount') or '0')
+              local add = tonumber(new_granted) - granted
+              if add > 0 then
+                  local remaining = tonumber(redis.call('HGET', KEYS[1], 'localRemainingCredits') or '0')
+                  redis.call('HSET', KEYS[1],
+                      'grantedAmount', new_granted,
+                      'localRemainingCredits', tostring(remaining + add))
+              end
+              if new_expiry > existing_expiry then
+                  redis.call('HSET', KEYS[1], 'expiresAt', ARGV[3])
+                  redis.call('PEXPIREAT', KEYS[1], new_expiry + grace)
+              end
+              return 0
+          end
+
+          redis.call('DEL', KEYS[1])
+          redis.call('HSET', KEYS[1],
+              'leaseId', new_id,
+              'companyId', ARGV[5],
+              'creditTypeId', ARGV[6],
+              'grantedAmount', new_granted,
+              'localRemainingCredits', new_granted,
+              'expiresAt', ARGV[3])
+          redis.call('PEXPIREAT', KEYS[1], new_expiry + grace)
+          return 1
+        LUA
+
+        # Atomic check and decrement on localRemainingCredits. Returns
+        # [post-debit balance, charged leaseId] on success, with the balance as
+        # a string because a Lua number reply truncates to integer and would
+        # corrupt fractional credit costs; a nil reply when there is no lease,
+        # the lease has expired, or the remaining balance is short. Returning
+        # the lease id read inside the same script is what lets the caller pin
+        # its reservation to the lease the debit actually landed on.
+        TRY_RESERVE_SCRIPT = Script.new(<<~LUA)
+
+          redis.replicate_commands()
+          local t = redis.call('TIME')
+          local now = (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000)
+
+          local raw = redis.call('HGET', KEYS[1], 'localRemainingCredits')
+          if not raw then return false end
+          local lease_id = redis.call('HGET', KEYS[1], 'leaseId')
+          if not lease_id then return false end
+          local expiry = tonumber(redis.call('HGET', KEYS[1], 'expiresAt') or '0')
+          if expiry <= now then return false end
+          local remaining = tonumber(raw)
+          local requested = tonumber(ARGV[1])
+          if remaining < requested then return false end
+          local new_remaining = remaining - requested
+          redis.call('HSET', KEYS[1], 'localRemainingCredits', tostring(new_remaining))
+          return { tostring(new_remaining), lease_id }
+        LUA
+
+        # Refund credits, clamped at grantedAmount. ARGV[2], when non-empty,
+        # pins the refund to a leaseId: if the slot now holds a different lease,
+        # the refund is dropped, because the expired lease's unspent remainder
+        # was already returned to the company balance server-side.
+        REFUND_SCRIPT = Script.new(<<~LUA)
+
+          local raw_remaining = redis.call('HGET', KEYS[1], 'localRemainingCredits')
+          if not raw_remaining then return 0 end
+          local required_lease = ARGV[2]
+          if required_lease and required_lease ~= '' then
+              local current_lease = redis.call('HGET', KEYS[1], 'leaseId')
+              if current_lease ~= required_lease then return 0 end
+          end
+          local remaining = tonumber(raw_remaining)
+          local granted = tonumber(redis.call('HGET', KEYS[1], 'grantedAmount') or '0')
+          local refund = tonumber(ARGV[1])
+          local new_balance = remaining + refund
+          if new_balance > granted then new_balance = granted end
+          redis.call('HSET', KEYS[1], 'localRemainingCredits', tostring(new_balance))
+          return 1
+        LUA
+
+        # Reconcile the lease to the server-authoritative grantedAmount total,
+        # crediting the difference to localRemainingCredits. The delta is
+        # computed HERE, atomically against the hash's current total, so two
+        # processes extending the same lease concurrently converge instead of
+        # minting phantom credits. Expiry only moves forward. ARGV[4] pins the
+        # extend to a leaseId, mirroring the pin on the refund.
+        EXTEND_SCRIPT = Script.new(<<~LUA)
+
+          local raw_granted = redis.call('HGET', KEYS[1], 'grantedAmount')
+          if not raw_granted then return 0 end
+          local required_lease = ARGV[4]
+          if required_lease and required_lease ~= '' then
+              local current_lease = redis.call('HGET', KEYS[1], 'leaseId')
+              if current_lease ~= required_lease then return 0 end
+          end
+          local granted = tonumber(raw_granted)
+          local target = tonumber(ARGV[1])
+          local add = target - granted
+          if add > 0 then
+              local remaining = tonumber(redis.call('HGET', KEYS[1], 'localRemainingCredits') or '0')
+              redis.call('HSET', KEYS[1],
+                  'grantedAmount', tostring(target),
+                  'localRemainingCredits', tostring(remaining + add))
+          end
+          local new_expiry = tonumber(ARGV[2])
+          local grace = tonumber(ARGV[3])
+          local current_expiry = tonumber(redis.call('HGET', KEYS[1], 'expiresAt') or '0')
+          if new_expiry > current_expiry then
+              redis.call('HSET', KEYS[1], 'expiresAt', ARGV[2])
+              redis.call('PEXPIREAT', KEYS[1], new_expiry + grace)
+          end
+          return 1
+        LUA
+
+        def initialize(client:, key_prefix: nil, default_lease_duration_ms: DEFAULT_LEASE_DURATION_MS,
+                       clock: DEFAULT_CLOCK)
+          @client = client
+          @key_prefix = key_prefix || DEFAULT_KEY_PREFIX
+          # Defensive fallback for a direct caller that extends without naming a
+          # new expiry. The lease manager always passes one.
+          @default_lease_duration_ms = default_lease_duration_ms
+          @clock = clock
+        end
+
+        # Public so the reservation store can target the same lease hash.
+        def hash_key(company_id, credit_type_id)
+          "#{@key_prefix}#{LEASE_KEY_NAMESPACE}#{Leases.lease_key(company_id, credit_type_id)}"
+        end
+
+        def get(company_id, credit_type_id)
+          raw = @client.hgetall(hash_key(company_id, credit_type_id))
+          return nil if raw.nil? || raw["leaseId"].nil?
+
+          decode_entry(raw)
+        end
+
+        # rubocop:disable Naming/PredicateMethod
+        # replace is the store method name every Schematic SDK shares; its
+        # boolean says whether a fresh row was written or an existing lease kept.
+        def replace(entry)
+          # No process clock here: the script reads now from the Redis server
+          # via TIME, so every process agrees on expiry.
+          result = REPLACE_SCRIPT.call(
+            @client,
+            keys: [hash_key(entry.company_id, entry.credit_type_id)],
+            argv: [
+              entry.lease_id,
+              num(entry.granted_amount),
+              millis(entry.expires_at).to_s,
+              LEASE_TTL_GRACE_MS.to_s,
+              entry.company_id,
+              entry.credit_type_id
+            ]
+          )
+          result.to_i == 1
+        end
+        # rubocop:enable Naming/PredicateMethod
+
+        def extend(company_id, credit_type_id, granted_amount, new_expires_at = nil, pin_lease_id = nil)
+          expiry = new_expires_at ? millis(new_expires_at) : millis(@clock.call) + @default_lease_duration_ms
+          # granted_amount is the server-authoritative TOTAL; the script computes
+          # the credit delta atomically against the stored total. An empty string
+          # disables the lease pin, since Lua has no nil ARGV.
+          EXTEND_SCRIPT.call(
+            @client,
+            keys: [hash_key(company_id, credit_type_id)],
+            argv: [num(granted_amount), expiry.to_s, LEASE_TTL_GRACE_MS.to_s, pin_lease_id.to_s]
+          )
+          nil
+        end
+
+        def drop(company_id, credit_type_id)
+          # A plain single-key delete: there is no secondary index to keep in sync.
+          @client.del(hash_key(company_id, credit_type_id))
+          nil
+        end
+
+        def try_reserve(company_id, credit_type_id, credits)
+          # Reject non-finite or negative debits before they reach the script:
+          # NaN.to_s parses back to a Lua nan, slips through the comparison, and
+          # would poison the SHARED balance for every process.
+          return nil unless Leases.valid_quantity?(credits)
+
+          result = TRY_RESERVE_SCRIPT.call(
+            @client,
+            keys: [hash_key(company_id, credit_type_id)],
+            # Only the requested amount: now comes from the Redis server clock.
+            argv: [num(credits)]
+          )
+          return nil if result.nil? || result == false
+
+          ReserveResult.new(result[0].to_f, result[1].to_s)
+        end
+
+        def refund(company_id, credit_type_id, credits, pin_lease_id = nil)
+          return nil if credits.nil? || credits <= 0
+
+          REFUND_SCRIPT.call(
+            @client,
+            keys: [hash_key(company_id, credit_type_id)],
+            # An empty string disables the lease pin, since Lua has no nil ARGV.
+            argv: [num(credits), pin_lease_id.to_s]
+          )
+          nil
+        end
+
+        private
+
+        # Lua's tonumber reads a decimal string, and Ruby's to_s does not always
+        # produce one: a Rational quantity renders as "1/2", which the script
+        # reads as nil and treats as a zero. Integers are already safe and stay
+        # exact, so only the rest is forced through Float.
+        def num(value)
+          value.is_a?(Integer) ? value.to_s : value.to_f.to_s
+        end
+
+        def millis(time)
+          (time.to_f * 1000).round
+        end
+
+        def decode_entry(raw)
+          LeaseEntry.new(
+            lease_id: raw["leaseId"],
+            company_id: raw["companyId"],
+            credit_type_id: raw["creditTypeId"],
+            granted_amount: (raw["grantedAmount"] || 0).to_f,
+            local_remaining_credits: (raw["localRemainingCredits"] || 0).to_f,
+            expires_at: Time.at((raw["expiresAt"] || 0).to_f / 1000.0)
+          )
+        end
+      end
+    end
+  end
+end