upkeep-rails 0.1.6

data/lib/upkeep/subscriptions/reverse_index.rb

@@ -0,0 +1,298 @@
+# frozen_string_literal: true
+
+module Upkeep
+  module Subscriptions
+    class ReverseIndex
+      COHORT_IDENTITY_SOURCES = %w[Current.user cookie current_attribute session warden_user].freeze
+
+      Entry = Data.define(:subscription_id, :owner_id, :dependency_cache_key, :dependency, :subscriber_ids, :cohort_key) do
+        def represented_subscriber_ids
+          Array(subscriber_ids)
+        end
+
+        def cohort?
+          !!cohort_key
+        end
+      end
+
+      def initialize
+        @entries_by_lookup_key = Hash.new { |hash, key| hash[key] = [] }
+        @entry_keys_by_lookup_key = Hash.new { |hash, key| hash[key] = {} }
+        @lookup_keys_by_subscription_id = Hash.new { |hash, key| hash[key] = {} }
+        @cohort_entries_by_index_key = {}
+        @cohort_members_by_index_key = Hash.new { |hash, key| hash[key] = {} }
+        @cohort_index_keys_by_subscription_id = Hash.new { |hash, key| hash[key] = {} }
+      end
+
+      def index(subscription)
+        index_entries(entries_for_subscription(subscription))
+      end
+
+      def index_entries(entries, subscription: nil)
+        entries = entries_for_subscription_instance(entries, subscription) if subscription
+
+        entries.each do |entry|
+          lookup_keys_for_dependency(entry.dependency).each do |lookup_key|
+            if entry.cohort?
+              index_cohort_entry(lookup_key, entry)
+            else
+              index_direct_entry(lookup_key, entry)
+            end
+          end
+        end
+      end
+
+      def delete_subscription(subscription_id)
+        lookup_keys = @lookup_keys_by_subscription_id.delete(subscription_id)&.keys || []
+        lookup_keys.each do |lookup_key|
+          entries = @entries_by_lookup_key.fetch(lookup_key, nil)
+          next unless entries
+
+          entries.reject! { |entry| entry.subscription_id == subscription_id }
+          @entry_keys_by_lookup_key[lookup_key].delete_if { |entry_key, _present| entry_key.fetch(0) == subscription_id }
+
+          next unless entries.empty?
+
+          @entries_by_lookup_key.delete(lookup_key)
+          @entry_keys_by_lookup_key.delete(lookup_key)
+        end
+
+        cohort_index_keys = @cohort_index_keys_by_subscription_id.delete(subscription_id)&.keys || []
+        cohort_index_keys.each { |index_key| delete_cohort_subscription(index_key, subscription_id) }
+      end
+
+      def entries_for(changes)
+        raw_entries = changes
+          .flat_map { |change| lookup_keys_for_change(change) }
+          .flat_map { |lookup_key| @entries_by_lookup_key.fetch(lookup_key, []) }
+          .uniq { |entry| [entry.subscription_id, entry.owner_id, entry.dependency_cache_key] }
+
+        collapse_cohort_entries(raw_entries)
+      end
+
+      def summary
+        entries = @entries_by_lookup_key.values.flatten
+        cohort_entries = entries.select(&:cohort?)
+        {
+          lookup_keys: @entries_by_lookup_key.size,
+          entries: entries.size,
+          cohort_entries: cohort_entries.size,
+          represented_subscribers: entries.flat_map(&:represented_subscriber_ids).uniq.size
+        }
+      end
+
+      def entries_for_subscription(subscription)
+        subscription.recorder.flush_pending_dependencies if subscription.recorder.respond_to?(:flush_pending_dependencies)
+        cohort_key = cohort_key_for(subscription)
+        subscription.graph.dependency_nodes.flat_map do |node|
+          subscription.graph.dependency_owner_ids(node.id).map do |owner_id|
+            Entry.new(
+              subscription.id,
+              owner_id,
+              node.payload.cache_key,
+              node.payload,
+              [subscription.subscriber_id],
+              cohort_key
+            )
+          end
+        end
+      end
+
+      def entries_for_subscription_instance(entries, subscription)
+        entries.map { |entry| entry_for_subscription(entry, subscription) }
+      end
+
+      def lookup_keys_for_dependency(dependency)
+        case dependency.source
+        when :active_record_attribute
+          active_record_attribute_lookup_keys(dependency.key)
+        when :active_record_collection, :active_record_query
+          active_record_collection_lookup_keys(dependency)
+        else
+          []
+        end
+      end
+
+      def lookup_keys_for_change(change)
+        table = change.fetch(:table)
+        attributes = change.fetch(:changed_attributes, [])
+
+        keys = attributes.map do |attribute|
+          if change[:id]
+            [:active_record_attribute, table, change.fetch(:id), attribute]
+          else
+            [:active_record_attribute_any_id, table, attribute]
+          end
+        end
+
+        keys.concat(attributes.map { |attribute| [:active_record_collection_column, table, attribute] })
+        keys.uniq
+      end
+
+      private
+
+      def index_direct_entry(lookup_key, entry)
+        entry_key = [entry.subscription_id, entry.owner_id, entry.dependency_cache_key]
+        entry_keys = @entry_keys_by_lookup_key[lookup_key]
+        return if entry_keys.key?(entry_key)
+
+        @entries_by_lookup_key[lookup_key] << entry
+        entry_keys[entry_key] = true
+        @lookup_keys_by_subscription_id[entry.subscription_id][lookup_key] = true
+      end
+
+      def index_cohort_entry(lookup_key, entry)
+        index_key = cohort_index_key(lookup_key, entry)
+        members = @cohort_members_by_index_key[index_key]
+        subscriber_ids = (members[entry.subscription_id] || []) | entry.represented_subscriber_ids
+        return if members[entry.subscription_id] == subscriber_ids
+
+        existing_entry = @cohort_entries_by_index_key[index_key]
+        members[entry.subscription_id] = subscriber_ids
+        replacement_entry = existing_entry ? append_cohort_members(existing_entry, subscriber_ids) : cohort_entry_from(entry, members)
+
+        if existing_entry
+          replace_entry(lookup_key, existing_entry, replacement_entry)
+        else
+          @entries_by_lookup_key[lookup_key] << replacement_entry
+        end
+
+        @cohort_entries_by_index_key[index_key] = replacement_entry
+        @cohort_index_keys_by_subscription_id[entry.subscription_id][index_key] = true
+      end
+
+      def delete_cohort_subscription(index_key, subscription_id)
+        lookup_key = index_key.fetch(0)
+        members = @cohort_members_by_index_key.fetch(index_key, nil)
+        return unless members&.key?(subscription_id)
+
+        existing_entry = @cohort_entries_by_index_key.fetch(index_key)
+        members.delete(subscription_id)
+
+        if members.empty?
+          remove_entry(lookup_key, existing_entry)
+          @cohort_members_by_index_key.delete(index_key)
+          @cohort_entries_by_index_key.delete(index_key)
+        else
+          replacement_entry = cohort_entry_from(existing_entry, members)
+          replace_entry(lookup_key, existing_entry, replacement_entry)
+          @cohort_entries_by_index_key[index_key] = replacement_entry
+        end
+      end
+
+      def cohort_index_key(lookup_key, entry)
+        [lookup_key, entry.cohort_key, entry.owner_id, entry.dependency_cache_key]
+      end
+
+      def cohort_entry_from(entry, members)
+        Entry.new(
+          members.keys.sort_by(&:to_s).first,
+          entry.owner_id,
+          entry.dependency_cache_key,
+          entry.dependency,
+          members.values.flatten.uniq.sort_by(&:to_s),
+          entry.cohort_key
+        )
+      end
+
+      def append_cohort_members(entry, subscriber_ids)
+        Entry.new(
+          entry.subscription_id,
+          entry.owner_id,
+          entry.dependency_cache_key,
+          entry.dependency,
+          entry.represented_subscriber_ids | subscriber_ids,
+          entry.cohort_key
+        )
+      end
+
+      def replace_entry(lookup_key, existing_entry, replacement_entry)
+        entries = @entries_by_lookup_key.fetch(lookup_key)
+        index = entries.index(existing_entry)
+        entries[index] = replacement_entry if index
+      end
+
+      def remove_entry(lookup_key, entry)
+        entries = @entries_by_lookup_key.fetch(lookup_key, nil)
+        return unless entries
+
+        entries.delete(entry)
+
+        return unless entries.empty?
+
+        @entries_by_lookup_key.delete(lookup_key)
+        @entry_keys_by_lookup_key.delete(lookup_key)
+      end
+
+      def entry_for_subscription(entry, subscription)
+        Entry.new(
+          entry.subscription_id || subscription.id,
+          entry.owner_id,
+          entry.dependency_cache_key,
+          entry.dependency,
+          entry.subscriber_ids || [subscription.subscriber_id],
+          entry.cohort_key || cohort_key_for(subscription)
+        )
+      end
+
+      def collapse_cohort_entries(entries)
+        direct_entries = []
+        cohort_entries = Hash.new { |hash, key| hash[key] = [] }
+
+        entries.each do |entry|
+          if entry.cohort?
+            cohort_entries[[entry.cohort_key, entry.owner_id, entry.dependency_cache_key]] << entry
+          else
+            direct_entries << entry
+          end
+        end
+
+        direct_entries + cohort_entries.values.map { |group| collapse_cohort_entry_group(group) }
+      end
+
+      def collapse_cohort_entry_group(entries)
+        representative = entries.first
+        Entry.new(
+          representative.subscription_id,
+          representative.owner_id,
+          representative.dependency_cache_key,
+          representative.dependency,
+          entries.flat_map(&:represented_subscriber_ids).uniq.sort_by(&:to_s),
+          representative.cohort_key
+        )
+      end
+
+      def cohort_key_for(subscription)
+        return unless identity_free_subscription?(subscription)
+
+        metadata_value(subscription, :subscription_shape_key)
+      end
+
+      def identity_free_subscription?(subscription)
+        return false if subscription.graph.dependency_nodes.any? { |node| cohort_identity_dependency?(node.payload) }
+
+        true
+      end
+
+      def metadata_value(subscription, key)
+        subscription.metadata[key] || subscription.metadata[key.to_s]
+      end
+
+      def cohort_identity_dependency?(dependency)
+        Dependencies.partitioning_identity?(dependency) && COHORT_IDENTITY_SOURCES.include?(dependency.source.to_s)
+      end
+
+      def active_record_collection_lookup_keys(dependency)
+        dependency.collection_lookup_columns.map { |table, column| [:active_record_collection_column, table, column] }
+      end
+
+      def active_record_attribute_lookup_keys(key)
+        if key.fetch(:id)
+          [[:active_record_attribute, key.fetch(:table), key.fetch(:id), key.fetch(:attribute)]]
+        else
+          [[:active_record_attribute_any_id, key.fetch(:table), key.fetch(:attribute)]]
+        end
+      end
+    end
+  end
+end

data/lib/upkeep/subscriptions/shape.rb

@@ -0,0 +1,116 @@
+# frozen_string_literal: true
+
+require "active_support/notifications"
+require_relative "../shared_streams"
+require_relative "reverse_index"
+
+module Upkeep
+  module Subscriptions
+    class ShapeCache
+      NOTIFICATION = "subscription_shape.upkeep"
+
+      Shape = Data.define(:key, :entries, :shared_stream_names)
+      Result = Data.define(:key, :entries, :shared_stream_names, :cache_state, :cacheable, :reason, :timings)
+      TemplateSubscription = Data.define(:id, :subscriber_id, :recorder, :graph, :metadata)
+
+      def initialize(index_builder: ReverseIndex.new)
+        @index_builder = index_builder
+        @shapes = {}
+        @mutex = Mutex.new
+      end
+
+      def resolve(recorder:, decision:, signature: nil)
+        if ActiveSupport::Notifications.notifier.listening?(NOTIFICATION)
+          payload = {}
+          ActiveSupport::Notifications.instrument(NOTIFICATION, payload) do
+            result = resolve_without_instrumentation(recorder: recorder, decision: decision, signature: signature)
+            payload.merge!(
+              key: result.key,
+              cache_state: result.cache_state,
+              cacheable: result.cacheable,
+              reason: result.reason,
+              entries: result.entries.size,
+              shared_stream_names: result.shared_stream_names.size
+            )
+            payload.merge!(result.timings)
+            result
+          end
+        else
+          resolve_without_instrumentation(recorder: recorder, decision: decision, signature: signature)
+        end
+      end
+
+      def reset
+        @mutex.synchronize { @shapes = {} }
+      end
+
+      private
+
+      attr_reader :index_builder
+
+      def resolve_without_instrumentation(recorder:, decision:, signature:)
+        timings = {}
+        unless cacheable?(recorder, decision)
+          shape = measure(timings, :compile_ms) { compile_shape(recorder, key: nil, timings: timings) }
+          return Result.new(nil, shape.entries, shape.shared_stream_names, "uncached", false, cache_bypass_reason(recorder, decision), timings)
+        end
+
+        key = measure(timings, :key_ms) { shape_key_for(recorder, signature: signature) }
+        @mutex.synchronize do
+          if (shape = @shapes[key])
+            Result.new(key, shape.entries, shape.shared_stream_names, "hit", true, nil, timings)
+          else
+            shape = measure(timings, :compile_ms) { compile_shape(recorder, key: key, timings: timings) }
+            @shapes[key] = shape
+            Result.new(key, shape.entries, shape.shared_stream_names, "miss", true, nil, timings)
+          end
+        end
+      end
+
+      def cacheable?(recorder, decision)
+        decision&.anonymous && recorder.reactive?
+      end
+
+      def cache_bypass_reason(recorder, decision)
+        return "identified" unless decision&.anonymous
+        return "refused_boundary" unless recorder.reactive?
+
+        "uncacheable"
+      end
+
+      def compile_shape(recorder, key:, timings:)
+        subscription = TemplateSubscription.new(nil, nil, recorder, recorder.graph, { subscription_shape_key: key }.compact)
+        entries = measure(timings, :index_template_ms) do
+          index_builder.entries_for_subscription(subscription)
+        end
+          .map { |entry| template_entry(entry) }
+          .uniq { |entry| [entry.owner_id, entry.dependency_cache_key] }
+          .freeze
+        shared_stream_names = measure(timings, :shared_stream_names_ms) { SharedStreams.names_for_recorder(recorder) }.freeze
+        Shape.new(key, entries, shared_stream_names)
+      end
+
+      def template_entry(entry)
+        ReverseIndex::Entry.new(
+          nil,
+          entry.owner_id,
+          entry.dependency_cache_key,
+          entry.dependency,
+          nil,
+          entry.cohort_key
+        )
+      end
+
+      def shape_key_for(recorder, signature:)
+        recorder.subscription_shape(request_signature: signature).signature
+      end
+
+      def measure(timings, key)
+        started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        yield
+      ensure
+        timings[key] = ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at) * 1000.0).round(3)
+      end
+    end
+  end
+end

data/lib/upkeep/subscriptions/store.rb

@@ -0,0 +1,171 @@
+# frozen_string_literal: true
+
+require "securerandom"
+require "time"
+require_relative "active_registry"
+
+module Upkeep
+  module Subscriptions
+    class NotFound < KeyError; end
+
+    Subscription = Data.define(:id, :subscriber_id, :recorder, :graph, :metadata) do
+      def identity_signature(frame_id)
+        recorder.identity_signature(frame_id)
+      end
+
+      def replay_recipe(frame_id)
+        graph.node(frame_id).payload[:recipe]
+      end
+
+      def to_h
+        {
+          id: id,
+          subscriber_id: subscriber_id,
+          recorder: recorder.to_h,
+          metadata: metadata
+        }
+      end
+
+      def to_persistent_h
+        {
+          id: id,
+          subscriber_id: subscriber_id,
+          recorder: recorder.to_persistent_h,
+          metadata: metadata
+        }
+      end
+
+      def self.from_h(snapshot)
+        snapshot = Dependencies.symbolize_keys(snapshot)
+        recorder = Runtime::Recorder.from_h(snapshot.fetch(:recorder))
+
+        new(
+          snapshot.fetch(:id),
+          snapshot.fetch(:subscriber_id),
+          recorder,
+          recorder.graph,
+          snapshot.fetch(:metadata)
+        )
+      end
+    end
+
+    class Store
+      attr_reader :reverse_index
+
+      def initialize(reverse_index: ReverseIndex.new)
+        @active_registry = ActiveRegistry.new(reverse_index: reverse_index)
+        @pending_registry = ActiveRegistry.new
+        @pending_index_entries = {}
+        @reverse_index = active_registry
+        @next_id = 0
+      end
+
+      def register(subscriber_id:, recorder:, metadata: {}, entries: nil)
+        recorder.flush_pending_dependencies if recorder.respond_to?(:flush_pending_dependencies)
+        subscription = Subscription.new(
+          next_subscription_id,
+          subscriber_id,
+          recorder,
+          recorder.graph,
+          metadata
+        )
+
+        pending_registry.register(subscription, entries: entries)
+        @pending_index_entries[subscription.id] = entries if entries
+        subscription
+      end
+
+      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)
+        true
+      end
+
+      def prune_stale!(older_than:)
+        stale_ids = subscriptions.filter_map do |subscription|
+          id = subscription.id
+          id if last_seen_at(subscription) && last_seen_at(subscription) < older_than
+        end
+
+        unregister(stale_ids)
+        stale_ids.size
+      end
+
+      def unregister(ids)
+        ids = Array(ids)
+        ids.each { |id| @pending_index_entries.delete(id) }
+        pending_registry.unregister(ids)
+        active_registry.unregister(ids)
+        ids.size
+      end
+
+      def activate(id)
+        return true if active_registry.fetch(id)
+
+        subscription = pending_registry.fetch(id)
+        return false unless subscription
+
+        entries = @pending_index_entries.delete(id)
+        active_registry.register(subscription, entries: entries)
+        pending_registry.unregister(id)
+        true
+      end
+
+      def drain
+        true
+      end
+
+      def shutdown
+        true
+      end
+
+      def fetch(id)
+        active_registry.fetch(id) || pending_registry.fetch(id) || raise(NotFound, id)
+      end
+
+      def subscriptions
+        active_registry.subscriptions + pending_registry.subscriptions
+      end
+
+      def reset
+        @active_registry = ActiveRegistry.new
+        @pending_registry = ActiveRegistry.new
+        @pending_index_entries = {}
+        @reverse_index = active_registry
+        @next_id = 0
+      end
+
+      def summary
+        active = active_registry.summary
+        pending = pending_registry.summary
+        {
+          subscriptions: subscriptions.size,
+          pending_subscriptions: pending_registry.count,
+          active_subscriptions: active_registry.count,
+          deferred_index_subscriptions: 0,
+          reverse_index: active.merge(
+            mode: :active,
+            active: active,
+            pending: pending,
+            persistent: { lookup_keys: 0, entries: 0 }
+          )
+        }
+      end
+
+      private
+
+      attr_reader :pending_registry, :active_registry
+
+      def last_seen_at(subscription)
+        value = subscription.metadata["last_seen_at"] || subscription.metadata[:last_seen_at]
+        Time.parse(value.to_s) if value
+      end
+
+      def next_subscription_id
+        "subscription-#{SecureRandom.uuid}"
+      end
+    end
+  end
+end

data/lib/upkeep/subscriptions.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+require_relative "subscriptions/reverse_index"
+require_relative "subscriptions/shape"
+require_relative "subscriptions/registrar"
+require_relative "subscriptions/store"
+require_relative "subscriptions/active_record_store"