upkeep-rails 0.1.0

data/lib/upkeep/subscriptions/active_registry.rb

@@ -0,0 +1,93 @@
+# frozen_string_literal: true
+
+require_relative "reverse_index"
+require_relative "store"
+
+module Upkeep
+  module Subscriptions
+    class ActiveRegistry
+      def initialize
+        @mutex = Mutex.new
+        @subscriptions = {}
+        @reverse_index = ReverseIndex.new
+      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 unless subscription
+
+          @subscriptions[id] = Subscription.new(
+            subscription.id,
+            subscription.subscriber_id,
+            subscription.recorder,
+            subscription.graph,
+            subscription.metadata.merge(metadata)
+          )
+        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
+
+      private
+
+      def rebuild_reverse_index!
+        @reverse_index = ReverseIndex.new
+        @subscriptions.each_value { |subscription| @reverse_index.index(subscription) }
+      end
+    end
+  end
+end

data/lib/upkeep/subscriptions/async_durable_writer.rb

@@ -0,0 +1,136 @@
+# frozen_string_literal: true
+
+module Upkeep
+  module Subscriptions
+    class AsyncDurableWriter
+      DEFAULT_BATCH_SIZE = 100
+      DEFAULT_FLUSH_INTERVAL = 1.0
+      Job = Data.define(:subscription, :entries, :operation)
+
+      def initialize(batch_size: DEFAULT_BATCH_SIZE, flush_interval: DEFAULT_FLUSH_INTERVAL, &persist_batch)
+        @batch_size = batch_size
+        @flush_interval = flush_interval
+        @persist_batch = persist_batch
+        @mutex = Mutex.new
+        @available = ConditionVariable.new
+        @drained = ConditionVariable.new
+        @queue = []
+        @pending = 0
+        @inflight_ids = Hash.new(0)
+        @closed = false
+        @flush_now = false
+        @errors = []
+        @worker = Thread.new { work_loop }
+        @worker.name = "upkeep-durable-writer" if @worker.respond_to?(:name=)
+      end
+
+      def enqueue(subscription, entries:, operation:)
+        @mutex.synchronize do
+          raise IOError, "Upkeep durable writer is closed" if @closed
+
+          @queue << Job.new(subscription, entries, operation)
+          @pending += 1
+          @available.signal
+        end
+      end
+
+      def cancel(ids)
+        ids = Array(ids)
+        return [] if ids.empty?
+
+        requested_ids = ids.to_h { |id| [id, true] }
+
+        @mutex.synchronize do
+          queued_ids = {}
+          removed = 0
+          @queue.delete_if do |job|
+            id = job.subscription.id
+            requested_ids.key?(id).tap do |matched|
+              if matched
+                queued_ids[id] = true
+                removed += 1
+              end
+            end
+          end
+          @pending -= removed
+          @drained.broadcast if @pending.zero?
+          persisted_ids = ids.reject { |id| queued_ids[id] }
+          @drained.wait(@mutex) while persisted_ids.any? { |id| @inflight_ids.fetch(id, 0).positive? }
+          persisted_ids
+        end
+      end
+
+      def drain(raise_errors: true)
+        errors = @mutex.synchronize do
+          @flush_now = true
+          @available.broadcast
+          @drained.wait(@mutex) while @pending.positive?
+          drained_errors = @errors
+          @errors = [] if raise_errors
+          drained_errors
+        end
+
+        raise errors.first if raise_errors && errors.any?
+
+        errors
+      end
+
+      def shutdown
+        drain(raise_errors: false)
+        @mutex.synchronize do
+          @closed = true
+          @available.broadcast
+        end
+        @worker.join
+      end
+
+      private
+
+      def work_loop
+        loop do
+          batch = next_batch
+          break unless batch
+
+          begin
+            @persist_batch.call(batch)
+          rescue StandardError => error
+            @mutex.synchronize { @errors << error }
+          ensure
+            @mutex.synchronize do
+              batch.each do |job|
+                id = job.subscription.id
+                @inflight_ids[id] -= 1
+                @inflight_ids.delete(id) unless @inflight_ids[id].positive?
+              end
+              @pending -= batch.size
+              @drained.broadcast if @pending.zero?
+            end
+          end
+        end
+      end
+
+      def next_batch
+        @mutex.synchronize do
+          @available.wait(@mutex) while @queue.empty? && !@closed
+          return nil if @queue.empty? && @closed
+
+          wait_for_batch_fill
+          @flush_now = false
+          @queue.shift(@batch_size).tap do |batch|
+            batch.each { |job| @inflight_ids[job.subscription.id] += 1 }
+          end
+        end
+      end
+
+      def wait_for_batch_fill
+        deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + @flush_interval
+        while @queue.size < @batch_size && !@closed && !@flush_now
+          remaining = deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          break unless remaining.positive?
+
+          @available.wait(@mutex, remaining)
+        end
+      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

data/lib/upkeep/subscriptions/layered_reverse_index.rb

@@ -0,0 +1,122 @@
+# frozen_string_literal: true
+
+require "active_support/notifications"
+
+module Upkeep
+  module Subscriptions
+    class LayeredReverseIndex
+      LOOKUP_NOTIFICATION = "lookup_subscription_index.upkeep"
+
+      def initialize(active_index:, persistent_index:, persistent_count:, store:, pending_index: nil)
+        @active_index = active_index
+        @persistent_index = persistent_index
+        @persistent_count = persistent_count
+        @store = store
+        @pending_index = pending_index
+      end
+
+      def entries_for(changes)
+        if ActiveSupport::Notifications.notifier.listening?(LOOKUP_NOTIFICATION)
+          payload = { changes: Array(changes).size, store: store }
+          ActiveSupport::Notifications.instrument(LOOKUP_NOTIFICATION, payload) do
+            entries_for_with_payload(changes, payload)
+          end
+        else
+          entries_for_without_payload(changes)
+        end
+      end
+
+      def entries_for_without_payload(changes)
+        active_entries = active_index.entries_for(changes)
+        return persistent_index.entries_for(changes) if active_index.count.zero?
+        return active_entries if active_index.covers?(persistent_subscription_count)
+
+        merge_entries(active_entries, persistent_index.entries_for(changes))
+      end
+
+      def entries_for_with_payload(changes, payload)
+        active_entries = active_index.entries_for(changes)
+        active_count = active_index.count
+        pending_entries = pending_entries_for(changes)
+        pending_count = pending_count_for_payload
+        payload[:active_entries] = active_entries.size
+        payload[:active_subscriptions] = active_count
+        payload[:pending_entries] = pending_entries.size
+        payload[:pending_subscriptions] = pending_count
+
+        if active_count.zero?
+          persistent_entries = persistent_index.entries_for(changes)
+          payload[:mode] = persistent_entries.empty? && pending_entries.any? ? "pending_activation" : "persistent"
+          payload[:persistent_entries] = persistent_entries.size
+          apply_miss_reason(payload, active_entries: active_entries, persistent_entries: persistent_entries, pending_entries: pending_entries)
+          return persistent_entries
+        end
+
+        if active_index.covers?(persistent_subscription_count)
+          payload[:mode] = "active"
+          payload[:persistent_entries] = 0
+          apply_miss_reason(payload, active_entries: active_entries, persistent_entries: [], pending_entries: pending_entries)
+          return active_entries
+        end
+
+        persistent_entries = persistent_index.entries_for(changes)
+        payload[:mode] = "active_and_persistent"
+        payload[:persistent_entries] = persistent_entries.size
+        apply_miss_reason(payload, active_entries: active_entries, persistent_entries: persistent_entries, pending_entries: pending_entries)
+        merge_entries(active_entries, persistent_entries)
+      end
+
+      def summary
+        persistent = persistent_index.summary
+        active = active_index.summary
+        pending = pending_index&.summary || { lookup_keys: 0, entries: 0, subscriptions: 0 }
+        mode = active_index.covers?(persistent_subscription_count) ? :active : :active_and_persistent
+        totals = if mode == :active
+          active
+        else
+          {
+            lookup_keys: active.fetch(:lookup_keys) + persistent.fetch(:lookup_keys),
+            entries: active.fetch(:entries) + persistent.fetch(:entries)
+          }
+        end
+
+        {
+          lookup_keys: totals.fetch(:lookup_keys),
+          entries: totals.fetch(:entries),
+          mode: mode,
+          active: active,
+          pending: pending,
+          persistent: persistent
+        }
+      end
+
+      private
+
+      attr_reader :active_index, :persistent_index, :persistent_count, :store, :pending_index
+
+      def persistent_subscription_count
+        persistent_count.call
+      end
+
+      def merge_entries(active_entries, persistent_entries)
+        (active_entries + persistent_entries).uniq do |entry|
+          [entry.subscription_id, entry.owner_id, entry.dependency_cache_key]
+        end
+      end
+
+      def pending_entries_for(changes)
+        pending_index ? pending_index.entries_for(changes) : []
+      end
+
+      def pending_count_for_payload
+        pending_index&.count || 0
+      end
+
+      def apply_miss_reason(payload, active_entries:, persistent_entries:, pending_entries:)
+        return if active_entries.any? || persistent_entries.any?
+
+        payload[:miss_reason] = pending_entries.any? ? "not_activated_yet" : "no_matching_subscriber"
+      end
+    end
+  end
+end

data/lib/upkeep/subscriptions/persistent_reverse_index.rb

@@ -0,0 +1,144 @@
+# frozen_string_literal: true
+
+require "digest"
+require "json"
+require_relative "../dependencies"
+require_relative "json_snapshot"
+require_relative "reverse_index"
+
+module Upkeep
+  module Subscriptions
+    class PersistentReverseIndex
+      LOOKUP_COLUMNS = [
+        :subscription_id,
+        :lookup_key_digest,
+        :dependency_source,
+        :lookup_table,
+        :lookup_record_id_snapshot,
+        :lookup_attribute,
+        :dependency_table,
+        :dependency_predicate_digest,
+        :dependency_metadata_snapshot,
+        :owner_ids_snapshot
+      ].freeze
+
+      def initialize(reverse_index:, index_record:)
+        @reverse_index = reverse_index
+        @index_record = index_record
+      end
+
+      def entries_for(changes)
+        persistent_entries_for(changes)
+      end
+
+      def summary
+        {
+          lookup_keys: index_record.distinct.count(:lookup_key_digest),
+          entries: index_record.count
+        }
+      end
+
+      def self.digest(value)
+        Digest::SHA256.hexdigest(JSON.generate(canonical_lookup_value(value)))
+      end
+
+      def self.canonical_lookup_value(value)
+        case value
+        when Array
+          value.map { |item| canonical_lookup_value(item) }
+        when Hash
+          value.keys.sort_by(&:to_s).map do |key|
+            [canonical_lookup_value(key), canonical_lookup_value(value.fetch(key))]
+          end
+        when Symbol
+          ["symbol", value.to_s]
+        when String
+          ["string", value.encode(Encoding::UTF_8)]
+        else
+          value
+        end
+      end
+
+      private
+
+      attr_reader :reverse_index, :index_record
+
+      def persistent_entries_for(changes)
+        lookup_keys = Array(changes).flat_map { |change| reverse_index.lookup_keys_for_change(change) }.uniq
+        lookup_keys_by_digest = Hash.new { |hash, digest| hash[digest] = [] }
+        lookup_keys.each do |lookup_key|
+          lookup_keys_by_digest[self.class.digest(lookup_key)] << lookup_key
+        end
+        lookup_key_digests = lookup_keys_by_digest.keys
+
+        index_record
+          .where(lookup_key_digest: lookup_key_digests)
+            .pluck(*LOOKUP_COLUMNS)
+            .flat_map { |row| entries_for_row(row, lookup_keys_by_digest) }
+            .uniq { |entry| [entry.subscription_id, entry.owner_id, entry.dependency_cache_key] }
+      end
+
+      def entries_for_row(row, lookup_keys_by_digest)
+        attributes = LOOKUP_COLUMNS.zip(row).to_h
+        lookup_keys = lookup_keys_by_digest.fetch(attributes.fetch(:lookup_key_digest)) { return [] }
+        return [] unless lookup_keys.any? { |lookup_key| lookup_key_matches_row?(lookup_key, attributes) }
+
+        dependency = dependency_for_row(attributes)
+        dependency_cache_key = dependency.cache_key
+        JsonSnapshot.load(attributes.fetch(:owner_ids_snapshot)).map do |owner_id|
+          ReverseIndex::Entry.new(
+            attributes.fetch(:subscription_id),
+            owner_id,
+            dependency_cache_key,
+            dependency,
+            nil,
+            nil
+          )
+        end
+      end
+
+      def lookup_key_matches_row?(lookup_key, attributes)
+        case lookup_key.fetch(0)
+        when :active_record_attribute
+          lookup_key.fetch(1).to_s == attributes.fetch(:lookup_table).to_s &&
+            lookup_key.fetch(2) == JsonSnapshot.load(attributes.fetch(:lookup_record_id_snapshot)) &&
+            lookup_key.fetch(3).to_s == attributes.fetch(:lookup_attribute).to_s
+        when :active_record_attribute_any_id
+          lookup_key.fetch(1).to_s == attributes.fetch(:lookup_table).to_s &&
+            attributes.fetch(:lookup_record_id_snapshot).nil? &&
+            lookup_key.fetch(2).to_s == attributes.fetch(:lookup_attribute).to_s
+        when :active_record_collection_column
+          lookup_key.fetch(1).to_s == attributes.fetch(:lookup_table).to_s &&
+            attributes.fetch(:lookup_record_id_snapshot).nil? &&
+            lookup_key.fetch(2).to_s == attributes.fetch(:lookup_attribute).to_s
+        else
+          false
+        end
+      end
+
+      def dependency_for_row(attributes)
+        source = attributes.fetch(:dependency_source).to_sym
+        case source
+        when :active_record_attribute
+          Dependencies::ActiveRecordAttribute.new(
+            table: attributes.fetch(:dependency_table),
+            id: attributes[:lookup_record_id_snapshot] && JsonSnapshot.load(attributes.fetch(:lookup_record_id_snapshot)),
+            attribute: attributes.fetch(:lookup_attribute)
+          )
+        when :active_record_collection, :active_record_query
+          metadata = JsonSnapshot.load(attributes.fetch(:dependency_metadata_snapshot))
+          dependency_class = source == :active_record_query ? Dependencies::ActiveRecordQuery : Dependencies::ActiveRecordCollection
+          dependency_class.new(
+            primary_table: attributes.fetch(:dependency_table),
+            table_columns: metadata.fetch(:table_columns),
+            coverage: metadata.fetch(:coverage),
+            sql: metadata.fetch(:sql),
+            predicates: metadata.fetch(:predicates)
+          )
+        else
+          raise ArgumentError, "unsupported persistent dependency source: #{source.inspect}"
+        end
+      end
+    end
+  end
+end

data/lib/upkeep/subscriptions/registrar.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+require_relative "shape"
+
+module Upkeep
+  module Subscriptions
+    class Registrar
+      Registration = Data.define(:identity, :decision, :subscription, :shape)
+
+      def initialize(store:, shape_cache: ShapeCache.new)
+        @store = store
+        @shape_cache = shape_cache
+      end
+
+      def register(identity:, decision:, recorder:, metadata: {}, signature: nil)
+        shape = shape_cache.resolve(recorder: recorder, decision: decision, signature: signature)
+        subscription = store.register(
+          subscriber_id: identity.subscriber_id,
+          recorder: recorder,
+          metadata: metadata.merge(
+            shared_stream_names: shape.shared_stream_names,
+            subscription_shape_key: shape.key,
+            subscription_shape_cache: shape.cache_state
+          ).compact,
+          entries: shape.entries
+        )
+
+        Registration.new(identity, decision, subscription, shape)
+      end
+
+      private
+
+      attr_reader :store, :shape_cache
+    end
+  end
+end