upkeep-rails 0.2.5-aarch64-linux-gnu

data/lib/upkeep/subscriptions/active_record_subscription_persistence.rb

@@ -0,0 +1,411 @@
+# frozen_string_literal: true
+
+require "active_record"
+require "active_support/notifications"
+require "logger"
+require_relative "json_snapshot"
+require_relative "persistent_reverse_index"
+require_relative "store"
+
+module Upkeep
+  module Subscriptions
+    class ActiveRecordSubscriptionPersistence
+      PERSIST_NOTIFICATION = "persist_subscription_store.upkeep"
+      INDEX_ENTRIES_SNAPSHOT_KEY = "__upkeep_index_entries"
+      Job = Data.define(:subscription, :entries, :operation)
+
+      def initialize(subscription_record:, index_record:, shape_index_record:, index_builder:)
+        @subscription_record = subscription_record
+        @index_record = index_record
+        @shape_index_record = shape_index_record
+        @index_builder = index_builder
+        @count_mutex = Mutex.new
+        @count_cache = nil
+      end
+
+      def persist_jobs(jobs)
+        if ActiveSupport::Notifications.notifier.listening?(PERSIST_NOTIFICATION)
+          payload = {
+            store: "active_record",
+            jobs: jobs.size,
+            subscriptions: jobs.count { |job| persist_subscription?(job) },
+            index_jobs: jobs.count { |job| persist_index?(job) },
+            dependency_entries: jobs.sum { |job| persist_index?(job) ? job.entries.size : 0 },
+            pending_index_entries: jobs.sum { |job| persist_subscription?(job) ? job.entries.size : 0 },
+            operations: operation_counts(jobs)
+          }
+          ActiveSupport::Notifications.instrument(PERSIST_NOTIFICATION, payload) do
+            result = persist_jobs_without_instrumentation(jobs)
+            payload[:subscription_rows] = result.fetch(:subscription_rows)
+            payload[:index_rows] = result.fetch(:index_rows)
+            payload[:direct_index_rows] = result.fetch(:direct_index_rows)
+            payload[:shape_index_rows] = result.fetch(:shape_index_rows)
+          end
+        else
+          persist_jobs_without_instrumentation(jobs)
+        end
+      end
+
+      def touch(id, metadata:, now:)
+        subscription_record.where(id: id).find_each do |record|
+          record.update_columns(
+            metadata: record.metadata.to_h.merge(metadata),
+            updated_at: now
+          )
+        end
+      end
+
+      def prune_stale!(older_than:, limit: nil)
+        stale = subscription_record.where(subscription_record.arel_table[:updated_at].lt(older_than))
+        stale = stale.order(:updated_at).limit(limit) if limit
+        stale_ids = stale.pluck(:id)
+        return [] if stale_ids.empty?
+
+        delete(stale_ids)
+        stale_ids
+      end
+
+      def delete(ids)
+        ids = Array(ids)
+        return if ids.empty?
+
+        silence_active_record_logging do
+          ActiveRecord::Base.connection_pool.with_connection do
+            ActiveRecord::Base.transaction do
+              shape_keys = shape_keys_for_subscriptions(ids)
+              index_record.where(subscription_id: ids).delete_all
+              deleted = subscription_record.where(id: ids).delete_all
+              delete_orphaned_shape_index_rows(shape_keys)
+              decrement_count_cache(deleted)
+            end
+          end
+        end
+      end
+
+      def fetch(id)
+        record = subscription_record.find(id)
+        subscription_with_metadata(record)
+      end
+
+      def fetch_with_index_entries(id)
+        record = subscription_record.find(id)
+        [subscription_with_metadata(record), index_entries_from_snapshot(record.recorder_snapshot)]
+      end
+
+      def subscriptions
+        subscription_record.order(:created_at, :id).map { |record| subscription_with_metadata(record) }
+      end
+
+      def reset
+        index_record.delete_all
+        shape_index_record.delete_all
+        subscription_record.delete_all
+        write_count_cache(0)
+      end
+
+      def count
+        @count_mutex.synchronize do
+          @count_cache ||= subscription_record.count
+        end
+      end
+
+      private
+
+      attr_reader :subscription_record, :index_record, :shape_index_record, :index_builder
+
+      def persist_jobs_without_instrumentation(jobs)
+        subscription_jobs = jobs.select { |job| persist_subscription?(job) }
+        index_jobs = jobs.select { |job| persist_index?(job) }
+
+        result = silence_active_record_logging do
+          ActiveRecord::Base.connection_pool.with_connection do
+            ActiveRecord::Base.transaction do
+              subscription_rows = persist_subscription_records(subscription_jobs)
+              index_result = index_subscriptions(index_jobs)
+              {
+                subscription_rows: subscription_rows,
+                index_rows: index_result.fetch(:index_rows),
+                direct_index_rows: index_result.fetch(:direct_index_rows),
+                shape_index_rows: index_result.fetch(:shape_index_rows)
+              }
+            end
+          end
+        end
+        increment_count_cache(result.fetch(:subscription_rows))
+        result
+      end
+
+      def persist_subscription_records(jobs)
+        now = Time.now
+        rows = jobs.map do |job|
+          subscription = job.subscription
+          {
+            id: subscription.id,
+            subscriber_id: subscription.subscriber_id,
+            metadata: metadata_without_shape_key(subscription.metadata),
+            subscription_shape_key: shape_key_for_metadata(subscription.metadata),
+            recorder_snapshot: dump(persistent_snapshot_for(job)),
+            created_at: now,
+            updated_at: now
+          }
+        end
+
+        subscription_record.insert_all!(rows) if rows.any?
+        rows.size
+      end
+
+      def index_subscriptions(jobs)
+        return index_result if jobs.empty?
+
+        now = Time.now
+        subscription_ids = jobs.map { |job| job.subscription.id }.uniq
+        shape_jobs, direct_jobs = jobs.partition { |job| shape_indexable?(job) }
+
+        index_record.where(subscription_id: subscription_ids).delete_all
+
+        direct_rows = index_direct_subscriptions(direct_jobs, now: now)
+        shape_rows = index_shape_subscriptions(shape_jobs, now: now)
+        index_result(direct_index_rows: direct_rows, shape_index_rows: shape_rows)
+      end
+
+      def index_direct_subscriptions(jobs, now:)
+        return 0 if jobs.empty?
+
+        grouped_rows = {}
+
+        jobs.each do |job|
+          job.entries.each do |entry|
+            index_builder.lookup_keys_for_dependency(entry.dependency).each do |lookup_key|
+              lookup_attributes = typed_lookup_attributes(entry.dependency, lookup_key)
+              key = [job.subscription.id, lookup_attributes]
+              row = grouped_rows[key] ||= {
+                subscription_id: job.subscription.id,
+                lookup_key_digest: PersistentReverseIndex.digest(lookup_key),
+                owner_ids: [],
+                created_at: now,
+                updated_at: now
+              }.merge(lookup_attributes)
+              row.fetch(:owner_ids) << entry.owner_id
+            end
+          end
+        end
+
+        rows = grouped_rows.values.map do |row|
+          row.merge(owner_ids_snapshot: dump(row.delete(:owner_ids).uniq))
+        end
+
+        index_record.insert_all!(rows) if rows.any?
+        rows.size
+      end
+
+      def index_shape_subscriptions(jobs, now:)
+        return 0 if jobs.empty?
+
+        grouped_rows = {}
+        shape_keys = jobs.map { |job| shape_index_key(job) }.compact.uniq
+
+        jobs.each do |job|
+          shape_key = shape_index_key(job)
+          job.entries.each do |entry|
+            index_builder.lookup_keys_for_dependency(entry.dependency).each do |lookup_key|
+              lookup_attributes = typed_lookup_attributes(entry.dependency, lookup_key)
+              key = [shape_key, lookup_attributes]
+              row = grouped_rows[key] ||= {
+                subscription_shape_key: shape_key,
+                lookup_key_digest: PersistentReverseIndex.digest(lookup_key),
+                owner_ids: [],
+                created_at: now,
+                updated_at: now
+              }.merge(lookup_attributes)
+              row.fetch(:owner_ids) << entry.owner_id
+            end
+          end
+        end
+
+        rows = grouped_rows.values.map do |row|
+          row.merge(owner_ids_snapshot: dump(row.delete(:owner_ids).uniq))
+        end
+
+        shape_index_record.where(subscription_shape_key: shape_keys).delete_all
+        shape_index_record.insert_all!(rows) if rows.any?
+        rows.size
+      end
+
+      def typed_lookup_attributes(dependency, lookup_key)
+        lookup_type = lookup_key.fetch(0)
+        source = dependency.source.to_s
+        dependency_key = dependency.key
+
+        case lookup_type
+        when :active_record_attribute
+          _type, table, record_id, attribute = lookup_key
+          {
+            dependency_source: source,
+            lookup_table: table.to_s,
+            lookup_record_id_snapshot: dump(record_id),
+            lookup_attribute: attribute.to_s,
+            dependency_table: dependency_key.fetch(:table).to_s,
+            dependency_predicate_digest: nil,
+            dependency_metadata_snapshot: nil
+          }
+        when :active_record_attribute_any_id
+          _type, table, attribute = lookup_key
+          {
+            dependency_source: source,
+            lookup_table: table.to_s,
+            lookup_record_id_snapshot: nil,
+            lookup_attribute: attribute.to_s,
+            dependency_table: dependency_key.fetch(:table).to_s,
+            dependency_predicate_digest: nil,
+            dependency_metadata_snapshot: dependency_key[:scope] ? dump(scope: dependency_key[:scope]) : nil
+          }
+        when :active_record_collection_column
+          _type, table, attribute = lookup_key
+          {
+            dependency_source: source,
+            lookup_table: table.to_s,
+            lookup_record_id_snapshot: nil,
+            lookup_attribute: attribute.to_s,
+            dependency_table: dependency_key.fetch(:table).to_s,
+            dependency_predicate_digest: dependency_key.fetch(:predicate_digest).to_s,
+            dependency_metadata_snapshot: dump(dependency.metadata)
+          }
+        else
+          raise ArgumentError, "unsupported persistent lookup key: #{lookup_key.inspect}"
+        end
+      end
+
+      def subscription_with_metadata(record)
+        subscription = Subscription.from_h(load(record.recorder_snapshot))
+        metadata = subscription.metadata.merge(record.metadata.to_h)
+        metadata = metadata.merge(subscription_shape_key: record.subscription_shape_key) if record.subscription_shape_key
+
+        Subscription.new(
+          subscription.id,
+          subscription.subscriber_id,
+          subscription.recorder,
+          subscription.graph,
+          metadata
+        )
+      end
+
+      def persistent_snapshot_for(job)
+        subscription = job.subscription
+        subscription.to_persistent_h.merge(
+          INDEX_ENTRIES_SNAPSHOT_KEY => index_entries_snapshot(job.entries)
+        )
+      end
+
+      def index_entries_snapshot(entries)
+        entries.map do |entry|
+          {
+            subscription_id: entry.subscription_id,
+            owner_id: entry.owner_id,
+            dependency: entry.dependency.to_h
+          }
+        end
+      end
+
+      def index_entries_from_snapshot(recorder_snapshot)
+        snapshot = load(recorder_snapshot)
+        Array(snapshot.fetch(INDEX_ENTRIES_SNAPSHOT_KEY)).map do |entry_snapshot|
+          entry_snapshot = Dependencies.symbolize_keys(entry_snapshot)
+          dependency = Dependencies.from_h(entry_snapshot.fetch(:dependency))
+          ReverseIndex::Entry.new(
+            entry_snapshot.fetch(:subscription_id),
+            entry_snapshot.fetch(:owner_id),
+            dependency.cache_key,
+            dependency,
+            nil,
+            nil,
+            nil
+          )
+        end
+      end
+
+      def dump(value)
+        JsonSnapshot.dump(value)
+      end
+
+      def load(snapshot)
+        JsonSnapshot.load(snapshot)
+      end
+
+      def increment_count_cache(value)
+        @count_mutex.synchronize { @count_cache += value if @count_cache }
+      end
+
+      def decrement_count_cache(value)
+        @count_mutex.synchronize { @count_cache -= value if @count_cache }
+      end
+
+      def write_count_cache(value)
+        @count_mutex.synchronize { @count_cache = value }
+      end
+
+      def shape_keys_for_subscriptions(ids)
+        subscription_record.where(id: ids).pluck(:subscription_shape_key).compact.uniq
+      end
+
+      def delete_orphaned_shape_index_rows(shape_keys)
+        shape_keys.each do |shape_key|
+          next if subscription_record.where(subscription_shape_key: shape_key).exists?
+
+          shape_index_record.where(subscription_shape_key: shape_key).delete_all
+        end
+      end
+
+      def shape_index_key(job)
+        shape_key_for_metadata(job.subscription.metadata)
+      end
+
+      def shape_indexable?(job)
+        shape_index_key(job) && job.entries.any? && job.entries.all?(&:cohort?)
+      end
+
+      def index_result(direct_index_rows: 0, shape_index_rows: 0)
+        {
+          index_rows: direct_index_rows + shape_index_rows,
+          direct_index_rows: direct_index_rows,
+          shape_index_rows: shape_index_rows
+        }
+      end
+
+      def shape_key_for_metadata(metadata)
+        metadata_value(metadata, :subscription_shape_key)
+      end
+
+      def metadata_without_shape_key(metadata)
+        metadata.to_h.dup.tap do |attributes|
+          attributes.delete(:subscription_shape_key)
+          attributes.delete("subscription_shape_key")
+        end
+      end
+
+      def metadata_value(metadata, key)
+        metadata.to_h[key] || metadata.to_h[key.to_s]
+      end
+
+      def persist_subscription?(job)
+        job.operation == :persist_subscription
+      end
+
+      def persist_index?(job)
+        job.operation == :persist_index
+      end
+
+      def operation_counts(jobs)
+        jobs.each_with_object(Hash.new(0)) { |job, counts| counts[job.operation.to_s] += 1 }.to_h
+      end
+
+      def silence_active_record_logging
+        logger = ActiveRecord::Base.logger
+        if logger.respond_to?(:silence)
+          logger.silence(::Logger::WARN) { yield }
+        else
+          yield
+        end
+      end
+    end
+  end
+end

data/lib/upkeep/subscriptions/active_registry.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+require_relative "reverse_index"
+
+module Upkeep
+  module Subscriptions
+    class ActiveRegistry
+      def initialize(reverse_index: ReverseIndex.new)
+        @mutex = Mutex.new
+        @subscriptions = {}
+        @reverse_index = reverse_index
+      end
+
+      def register(subscription, entries: nil)
+        @mutex.synchronize do
+          @subscriptions[subscription.id] = subscription
+          if entries
+            @reverse_index.index_entries(entries, subscription: subscription)
+          else
+            @reverse_index.index(subscription)
+          end
+        end
+      end
+
+      def fetch(id)
+        @mutex.synchronize { @subscriptions[id] }
+      end
+
+      def subscriptions
+        @mutex.synchronize { @subscriptions.values }
+      end
+
+      def unregister(ids)
+        ids = Array(ids)
+        @mutex.synchronize do
+          ids.each do |id|
+            next unless @subscriptions.delete(id)
+
+            @reverse_index.delete_subscription(id)
+          end
+        end
+      end
+
+      def touch(id, metadata:)
+        @mutex.synchronize do
+          subscription = @subscriptions[id]
+          return false unless subscription
+
+          @subscriptions[id] = subscription.with(metadata: subscription.metadata.merge(metadata))
+          true
+        end
+      end
+
+      def entries_for(changes)
+        @mutex.synchronize { @reverse_index.entries_for(changes) }
+      end
+
+      def reset
+        @mutex.synchronize do
+          @subscriptions = {}
+          @reverse_index = ReverseIndex.new
+        end
+      end
+
+      def covers?(persistent_count)
+        count >= persistent_count
+      end
+
+      def count
+        @mutex.synchronize { @subscriptions.size }
+      end
+
+      def summary
+        @mutex.synchronize do
+          @reverse_index.summary.merge(subscriptions: @subscriptions.size)
+        end
+      end
+    end
+  end
+end

data/lib/upkeep/subscriptions/base_store.rb

@@ -0,0 +1,110 @@
+# frozen_string_literal: true
+
+require "securerandom"
+require "active_support/notifications"
+
+module Upkeep
+  module Subscriptions
+    # Shared pending -> active registry choreography for the in-memory and
+    # Active Record stores. Subclasses hold the @pending_registry / @active_registry
+    # instances and supply their store-specific tails through the documented hooks.
+    class BaseStore
+      # Opportunistic trim (Solid Cache/Solid Cable style): every TRIM_EVERY
+      # registrations the store deletes at most TRIM_BATCH_LIMIT subscriptions
+      # older than the configured subscription TTL. A deterministic counter is
+      # used instead of rand so the cadence is reproducible in tests.
+      TRIM_EVERY = 100
+      TRIM_BATCH_LIMIT = 500
+      PRUNE_NOTIFICATION = "prune.upkeep"
+      DEFAULT_SUBSCRIPTION_TTL = 24 * 60 * 60
+
+      def touch(id, now: Time.now)
+        fetch(id)
+        metadata = { "last_seen_at" => now.utc.iso8601 }
+        pending_registry.touch(id, metadata: metadata)
+        active_registry.touch(id, metadata: metadata)
+        after_touch(id, metadata: metadata, now: now)
+      end
+
+      def unregister(ids)
+        ids = Array(ids)
+        before_unregister(ids)
+        pending_registry.unregister(ids)
+        active_registry.unregister(ids)
+        after_unregister(ids)
+        ids.size
+      end
+
+      def fetch(id)
+        active_registry.fetch(id) || pending_registry.fetch(id) || fetch_missing(id)
+      end
+
+      def fetch_for_composition(id)
+        fetch(id)
+      end
+
+      def explain(id)
+        fetch(id).explain
+      end
+
+      private
+
+      # Wrap a unit of work in an optional ActiveSupport notification: when a
+      # listener is attached the block runs inside #instrument with the given
+      # payload, otherwise it runs with a nil payload and no instrumentation.
+      def with_optional_notification(notification, payload)
+        if ActiveSupport::Notifications.notifier.listening?(notification)
+          ActiveSupport::Notifications.instrument(notification, payload) do
+            yield payload
+          end
+        else
+          yield nil
+        end
+      end
+
+      # Runs one bounded prune batch every TRIM_EVERY registrations. Failures
+      # never propagate into the registration path: they are logged and the
+      # stale rows stay until the next trim.
+      def trim_opportunistically
+        @trim_registration_count = (@trim_registration_count || 0) + 1
+        return unless (@trim_registration_count % TRIM_EVERY).zero?
+
+        payload = { store: store_label, limit: TRIM_BATCH_LIMIT }
+        ActiveSupport::Notifications.instrument(PRUNE_NOTIFICATION, payload) do
+          payload[:pruned] = prune_stale!(limit: TRIM_BATCH_LIMIT)
+        end
+      rescue StandardError => error
+        warn_trim_failure(error)
+      end
+
+      def warn_trim_failure(error)
+        return unless defined?(::Rails) && ::Rails.respond_to?(:logger) && ::Rails.logger
+
+        ::Rails.logger.warn(
+          "Upkeep opportunistic subscription prune failed " \
+          "(#{error.class}: #{error.message}); stale subscriptions stay until the next trim"
+        )
+      end
+
+      def stale_threshold
+        Time.now - subscription_ttl
+      end
+
+      def subscription_ttl
+        if defined?(Upkeep::Rails) && Upkeep::Rails.respond_to?(:configuration)
+          Upkeep::Rails.configuration.subscription_ttl
+        else
+          DEFAULT_SUBSCRIPTION_TTL
+        end
+      end
+
+      def before_unregister(ids); end
+
+      def after_unregister(ids); end
+
+      def next_subscription_id
+        "subscription-#{SecureRandom.uuid}"
+      end
+    end
+  end
+end

data/lib/upkeep/subscriptions/json_snapshot.rb

@@ -0,0 +1,98 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Upkeep
+  module Subscriptions
+    module JsonSnapshot
+      VERSION = 2
+      VERSION_KEY = "__upkeep_snapshot_version"
+      VALUE_KEY = "value"
+      SYMBOL_VALUE_KEY = "$sym"
+      SYMBOL_KEY_PREFIX = "$sym:"
+      STRING_KEY_PREFIX = "$str:"
+      JSON_KEY_PREFIX = "$json:"
+      RESERVED_STRING_KEYS = [SYMBOL_VALUE_KEY].freeze
+      RESERVED_STRING_KEY_PREFIXES = [SYMBOL_KEY_PREFIX, STRING_KEY_PREFIX, JSON_KEY_PREFIX].freeze
+
+      module_function
+
+      def dump(value)
+        {
+          VERSION_KEY => VERSION,
+          VALUE_KEY => encode(value)
+        }
+      end
+
+      def load(snapshot)
+        snapshot = JSON.parse(snapshot) if snapshot.is_a?(String)
+        version = snapshot.fetch(VERSION_KEY)
+        unless version.to_i == VERSION
+          raise ArgumentError, "unsupported Upkeep JSON snapshot version: #{version.inspect}"
+        end
+
+        decode(snapshot.fetch(VALUE_KEY))
+      end
+
+      def encode(value)
+        case value
+        when Symbol
+          { SYMBOL_VALUE_KEY => value.to_s }
+        when Hash
+          value.each_with_object({}) { |(key, nested_value), encoded| encoded[encode_key(key)] = encode(nested_value) }
+        when Array
+          value.map { |nested_value| encode(nested_value) }
+        when nil, true, false, Numeric, String
+          value
+        else
+          raise TypeError, "cannot persist #{value.class.name} in an Upkeep JSON snapshot"
+        end
+      end
+
+      def decode(value)
+        case value
+        when Hash
+          return value.fetch(SYMBOL_VALUE_KEY).to_sym if value.size == 1 && value.key?(SYMBOL_VALUE_KEY)
+
+          value.each_with_object({}) do |(key, nested_value), decoded|
+            decoded[decode_key(key)] = decode(nested_value)
+          end
+        when Array
+          value.map { |nested_value| decode(nested_value) }
+        else
+          value
+        end
+      end
+
+      def encode_key(key)
+        case key
+        when Symbol
+          "#{SYMBOL_KEY_PREFIX}#{key}"
+        when String
+          reserved_string_key?(key) ? "#{STRING_KEY_PREFIX}#{key}" : key
+        when nil, true, false, Numeric
+          "#{JSON_KEY_PREFIX}#{JSON.generate(encode(key))}"
+        else
+          raise TypeError, "cannot persist #{key.class.name} as an Upkeep JSON snapshot key"
+        end
+      end
+
+      def decode_key(key)
+        if key.start_with?(SYMBOL_KEY_PREFIX)
+          key.delete_prefix(SYMBOL_KEY_PREFIX).to_sym
+        elsif key.start_with?(STRING_KEY_PREFIX)
+          key.delete_prefix(STRING_KEY_PREFIX)
+        elsif key.start_with?(JSON_KEY_PREFIX)
+          decode(JSON.parse(key.delete_prefix(JSON_KEY_PREFIX)))
+        else
+          key
+        end
+      end
+
+      def reserved_string_key?(key)
+        RESERVED_STRING_KEYS.include?(key) || RESERVED_STRING_KEY_PREFIXES.any? { |prefix| key.start_with?(prefix) }
+      end
+
+    end
+  end
+end