upkeep-rails 0.1.0

data/lib/upkeep/shared_streams.rb

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+require "digest"
+
+module Upkeep
+  module SharedStreams
+    PREFIX = "upkeep:shared"
+
+    module_function
+
+    def stream_name(target:, identity_signature:, sharing_signature:)
+      digest = Digest::SHA256.hexdigest([target.kind, target.id, identity_signature, sharing_signature].inspect)[0, 32]
+      "#{PREFIX}:#{digest}"
+    end
+
+    def signature_for(recipe)
+      if recipe.instance_variable_defined?(:@upkeep_shared_stream_signature)
+        recipe.instance_variable_get(:@upkeep_shared_stream_signature)
+      else
+        Digest::SHA256.hexdigest(recipe.to_h.inspect).tap do |signature|
+          recipe.instance_variable_set(:@upkeep_shared_stream_signature, signature)
+        end
+      end
+    end
+
+    def names_for_subscription(subscription)
+      names_for_graph(subscription.graph)
+    end
+
+    def names_for_recorder(recorder)
+      names_for_graph(recorder.graph)
+    end
+
+    def names_for_graph(graph)
+      graph.frame_nodes.filter_map do |frame|
+        next unless frame.payload.fetch(:kind) == "render_site"
+
+        recipe = frame.payload[:recipe]
+        next unless recipe
+
+        identity_signature = identity_signature_for(graph, frame.id)
+        next unless identity_signature == "public"
+
+        target = target_for_frame(frame)
+        next unless target
+
+        stream_name(
+          target: target,
+          identity_signature: identity_signature,
+          sharing_signature: signature_for(recipe)
+        )
+      end.uniq.sort
+    end
+
+    def identity_signature_for(graph, frame_id)
+      identity_dependencies = graph.contained_node_ids(frame_id)
+        .flat_map { |owner_id| graph.dependencies_for(owner_id) }
+        .select(&:identity?)
+        .uniq(&:cache_key)
+      return "public" if identity_dependencies.empty?
+
+      Digest::SHA256.hexdigest(identity_dependencies.map(&:identity_key).sort_by(&:inspect).inspect)[0, 16]
+    end
+
+    def target_for_frame(frame)
+      case frame.payload.fetch(:kind)
+      when "render_site"
+        Targeting::Target.new("render_site", frame.payload.fetch(:site_id), "shared render-site frame")
+      end
+    end
+  end
+end

data/lib/upkeep/subscriptions/active_record_store.rb

@@ -0,0 +1,292 @@
+# frozen_string_literal: true
+
+require "active_record"
+require "active_support/notifications"
+require "securerandom"
+require_relative "active_record_subscription_persistence"
+require_relative "active_registry"
+require_relative "async_durable_writer"
+require_relative "json_snapshot"
+require_relative "layered_reverse_index"
+require_relative "persistent_reverse_index"
+require_relative "reverse_index"
+require_relative "store"
+
+module Upkeep
+  module Subscriptions
+    class ActiveRecordStore
+      LOOKUP_NOTIFICATION = LayeredReverseIndex::LOOKUP_NOTIFICATION
+      REGISTER_NOTIFICATION = "register_subscription_store.upkeep"
+      ACTIVATE_NOTIFICATION = "activate_subscription_store.upkeep"
+      PERSIST_NOTIFICATION = ActiveRecordSubscriptionPersistence::PERSIST_NOTIFICATION
+      DURABILITY_MODE = "async_subscription_row_index_on_subscribe"
+      INDEX_DURABILITY = "on_subscribe"
+
+      class SubscriptionRecord < ActiveRecord::Base
+        self.table_name = "upkeep_subscriptions"
+        self.primary_key = "id"
+
+        has_many :index_entries,
+          class_name: "Upkeep::Subscriptions::ActiveRecordStore::IndexEntryRecord",
+          foreign_key: "subscription_id",
+          dependent: :delete_all
+      end
+
+      class IndexEntryRecord < ActiveRecord::Base
+        self.table_name = "upkeep_subscription_index_entries"
+
+        belongs_to :subscription,
+          class_name: "Upkeep::Subscriptions::ActiveRecordStore::SubscriptionRecord",
+          foreign_key: "subscription_id"
+      end
+
+      attr_reader :reverse_index
+
+      DeferredIndexWrite = Data.define(:subscription, :entries)
+
+      def initialize(subscription_record: SubscriptionRecord, index_record: IndexEntryRecord)
+        @subscription_record = subscription_record
+        @index_record = index_record
+        @index_builder = ReverseIndex.new
+        @pending_registry = ActiveRegistry.new
+        @active_registry = ActiveRegistry.new
+        @deferred_index_writes = {}
+        @deferred_index_mutex = Mutex.new
+        @persistence = ActiveRecordSubscriptionPersistence.new(
+          subscription_record: subscription_record,
+          index_record: index_record,
+          index_builder: index_builder
+        )
+        persistent_index = PersistentReverseIndex.new(
+          reverse_index: index_builder,
+          index_record: index_record
+        )
+        @reverse_index = LayeredReverseIndex.new(
+          active_index: active_registry,
+          persistent_index: persistent_index,
+          persistent_count: -> { persistence.count },
+          store: "active_record",
+          pending_index: pending_registry
+        )
+        @durable_writer = AsyncDurableWriter.new { |jobs| persistence.persist_jobs(jobs) }
+      end
+
+      def self.available?(connect: false)
+        return false unless ActiveRecord::Base.connected? || connect
+
+        connection = ActiveRecord::Base.connection
+        connection.data_source_exists?("upkeep_subscriptions") &&
+          connection.data_source_exists?("upkeep_subscription_index_entries")
+      rescue ActiveRecord::ConnectionNotEstablished, ActiveRecord::NoDatabaseError, ActiveRecord::StatementInvalid
+        false
+      end
+
+      def register(subscriber_id:, recorder:, metadata: {}, entries: nil)
+        if ActiveSupport::Notifications.notifier.listening?(REGISTER_NOTIFICATION)
+          payload = { store: "active_record" }
+          ActiveSupport::Notifications.instrument(REGISTER_NOTIFICATION, payload) do
+            register_subscription(subscriber_id, recorder, metadata, entries: entries, payload: payload)
+          end
+        else
+          register_subscription(subscriber_id, recorder, metadata, entries: entries)
+        end
+      end
+
+      def drain = durable_writer.drain
+
+      def shutdown
+        clear_deferred_index_writes
+        durable_writer.shutdown
+      end
+
+      def activate(id)
+        if ActiveSupport::Notifications.notifier.listening?(ACTIVATE_NOTIFICATION)
+          payload = { store: "active_record", subscription_id: id }
+          ActiveSupport::Notifications.instrument(ACTIVATE_NOTIFICATION, payload) do
+            activate_subscription(id, payload: payload)
+          end
+        else
+          activate_subscription(id)
+        end
+      end
+
+      def touch(id, now: Time.now)
+        metadata = { "last_seen_at" => now.utc.iso8601 }
+        pending_registry.touch(id, metadata: metadata)
+        active_registry.touch(id, metadata: metadata)
+        activate(id)
+        durable_writer.drain
+        persistence.touch(id, metadata: metadata, now: now)
+      end
+
+      def unregister(ids)
+        ids = Array(ids)
+        pending_registry.unregister(ids)
+        active_registry.unregister(ids)
+        delete_deferred_index_writes(ids)
+        persisted_ids = durable_writer.cancel(ids)
+        persistence.delete(persisted_ids)
+        ids.size
+      end
+
+      def prune_stale!(older_than:)
+        durable_writer.drain
+        stale_ids = persistence.prune_stale!(older_than: older_than)
+        active_registry.unregister(stale_ids)
+        stale_ids.size
+      end
+
+      def fetch(id)
+        active_registry.fetch(id) || pending_registry.fetch(id) || persistence.fetch(id)
+      end
+
+      def subscriptions
+        persistent_count = persistence.count
+        in_memory_subscriptions = (active_registry.subscriptions + pending_registry.subscriptions).to_h do |subscription|
+          [subscription.id, subscription]
+        end
+        return in_memory_subscriptions.values if in_memory_subscriptions.size >= persistent_count
+
+        seen_ids = {}
+        persisted = persistence.subscriptions.map do |subscription|
+          seen_ids[subscription.id] = true
+          in_memory_subscriptions.fetch(subscription.id, subscription)
+        end
+        persisted + in_memory_subscriptions.values.reject { |subscription| seen_ids[subscription.id] }
+      end
+
+      def reset
+        clear_deferred_index_writes
+        durable_writer.drain
+        pending_registry.reset
+        active_registry.reset
+        persistence.reset
+      end
+
+      def summary
+        persistent_count = persistence.count
+        pending_count = pending_registry.count
+        active_count = active_registry.count
+        {
+          subscriptions: [persistent_count, active_count + pending_count].max,
+          persistent_subscriptions: persistent_count,
+          pending_subscriptions: pending_count,
+          active_subscriptions: active_count,
+          deferred_index_subscriptions: deferred_index_count,
+          reverse_index: reverse_index.summary
+        }
+      end
+
+      private
+
+      attr_reader :subscription_record, :index_record, :index_builder, :pending_registry, :active_registry, :persistence, :durable_writer
+
+      def register_subscription(subscriber_id, recorder, metadata, entries: nil, payload: nil)
+        recorder.flush_pending_dependencies if recorder.respond_to?(:flush_pending_dependencies)
+        subscription = Subscription.new(
+          next_subscription_id,
+          subscriber_id,
+          recorder,
+          recorder.graph,
+          metadata
+        )
+
+        entries = unique_entries(materialize_entries(subscription, entries))
+        if payload
+          payload[:subscription_id] = subscription.id
+          payload[:dependency_entries] = entries.size
+          payload[:mode] = "pending_activation"
+          payload[:durability] = DURABILITY_MODE
+          payload[:index_durability] = INDEX_DURABILITY
+        end
+
+        pending_registry.register(subscription, entries: entries)
+        durable_writer.enqueue(subscription, entries: entries, operation: :persist_subscription)
+        defer_index_write(subscription, entries)
+        subscription
+      end
+
+      def activate_subscription(id, payload: nil)
+        subscription, entries, source = activation_index_write(id)
+        unless subscription
+          payload[:activated] = false if payload
+          payload[:miss_reason] = "no_subscription" if payload
+          return false
+        end
+
+        unless source == :active
+          active_registry.register(subscription, entries: entries)
+          pending_registry.unregister(id)
+          durable_writer.enqueue(subscription, entries: entries, operation: :persist_index)
+        end
+
+        if payload
+          payload[:activated] = true
+          payload[:activation_source] = source
+          payload[:dependency_entries] = entries.size
+          payload[:active_subscriptions] = active_registry.count
+          payload[:pending_subscriptions] = pending_registry.count
+        end
+
+        true
+      end
+
+      def defer_index_write(subscription, entries)
+        @deferred_index_mutex.synchronize do
+          @deferred_index_writes[subscription.id] = DeferredIndexWrite.new(subscription, entries)
+        end
+      end
+
+      def activation_index_write(id)
+        deferred_write = take_deferred_index_write(id)
+        if deferred_write
+          subscription = pending_registry.fetch(id) || active_registry.fetch(id) || deferred_write.subscription
+          return [subscription, deferred_write.entries, :pending]
+        end
+
+        if (subscription = active_registry.fetch(id))
+          [subscription, [], :active]
+        elsif (subscription = pending_registry.fetch(id))
+          [subscription, unique_entries(index_builder.entries_for_subscription(subscription)), :pending]
+        else
+          subscription, entries = persistence.fetch_with_index_entries(id)
+          [subscription, entries, :persistent]
+        end
+      rescue ActiveRecord::RecordNotFound
+        [nil, nil, nil]
+      end
+
+      def take_deferred_index_write(id)
+        @deferred_index_mutex.synchronize { @deferred_index_writes.delete(id) }
+      end
+
+      def delete_deferred_index_writes(ids)
+        @deferred_index_mutex.synchronize { ids.each { |id| @deferred_index_writes.delete(id) } }
+      end
+
+      def clear_deferred_index_writes
+        @deferred_index_mutex.synchronize { @deferred_index_writes.clear }
+      end
+
+      def deferred_index_count
+        @deferred_index_mutex.synchronize { @deferred_index_writes.size }
+      end
+
+      def unique_entries(entries)
+        entries.uniq { |entry| [entry.owner_id, entry.dependency_cache_key] }
+      end
+
+      def materialize_entries(subscription, entries)
+        if entries
+          index_builder.entries_for_subscription_instance(entries, subscription)
+        else
+          index_builder.entries_for_subscription(subscription)
+        end
+      end
+
+      def next_subscription_id
+        "subscription-#{SecureRandom.uuid}"
+      end
+    end
+  end
+end

data/lib/upkeep/subscriptions/active_record_subscription_persistence.rb

@@ -0,0 +1,291 @@
+# frozen_string_literal: true
+
+require "active_record"
+require "active_support/notifications"
+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"
+
+      def initialize(subscription_record:, index_record:, index_builder:)
+        @subscription_record = subscription_record
+        @index_record = 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)
+          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:)
+        stale_ids = subscription_record.where(subscription_record.arel_table[:updated_at].lt(older_than)).pluck(:id)
+        return [] if stale_ids.empty?
+
+        delete(stale_ids)
+        stale_ids
+      end
+
+      def delete(ids)
+        ids = Array(ids)
+        return if ids.empty?
+
+        ActiveRecord::Base.connection_pool.with_connection do
+          ActiveRecord::Base.transaction do
+            index_record.where(subscription_id: ids).delete_all
+            deleted = subscription_record.where(id: ids).delete_all
+            decrement_count_cache(deleted)
+          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
+        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, :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 = ActiveRecord::Base.connection_pool.with_connection do
+          ActiveRecord::Base.transaction do
+            {
+              subscription_rows: persist_subscription_records(subscription_jobs),
+              index_rows: index_subscriptions(index_jobs)
+            }
+          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: 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 0 if jobs.empty?
+
+        now = Time.now
+        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.where(subscription_id: jobs.map { |job| job.subscription.id }.uniq).delete_all
+        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: 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))
+        Subscription.new(
+          subscription.id,
+          subscription.subscriber_id,
+          subscription.recorder,
+          subscription.graph,
+          subscription.metadata.merge(record.metadata.to_h)
+        )
+      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
+          )
+        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 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
+    end
+  end
+end