Merge pull request #37 from prog-supdex/feat/add-stats

feat: add ability to fetch subscription stats

Author: palkan

README.md

@@ -239,6 +239,105 @@ As in AnyCable there is no place to store subscription data in-memory, it should
     => 52ee8d65-275e-4d22-94af-313129116388
     ```
 
+## Stats
+
+You can grab Redis subscription statistics by calling
+
+```ruby
+    GraphQL::AnyCable.stats
+```
+
+It will return a total of the amount of the key with the following prefixes
+
+```
+    graphql-subscription
+    graphql-fingerprints
+    graphql-subscriptions
+    graphql-channel
+```
+
+The response will look like this
+
+```json
+  {
+    "total": {
+      "subscription":22646,
+      "fingerprints":3200,
+      "subscriptions":20101,
+      "channel": 4900
+    }
+  }
+```
+
+You can also grab the number of subscribers grouped by subscriptions
+
+```ruby
+    GraphQL::AnyCable.stats(include_subscriptions: true)
+```
+
+It will return the response that contains `subscriptions`
+
+```json
+  {
+    "total": {
+      "subscription":22646,
+      "fingerprints":3200,
+      "subscriptions":20101,
+      "channel": 4900
+    },
+    "subscriptions": {
+      "productCreated": 11323,
+      "productUpdated": 11323
+    }
+  }
+```
+
+Also, you can set another `scan_count`, if needed.
+The default value is 1_000
+
+```ruby
+    GraphQL::AnyCable.stats(scan_count: 100)
+```
+
+We can set statistics data to [Yabeda][] for tracking amount of subscriptions
+
+```ruby
+  # config/initializers/metrics.rb
+  Yabeda.configure do
+    group :graphql_anycable_statistics do
+      gauge :subscriptions_count,  comment: "Number of graphql-anycable subscriptions"
+    end
+  end
+```
+
+```ruby
+  # in your app
+  statistics = GraphQL::AnyCable.stats[:total]
+
+  statistics.each do |key , value|
+    Yabeda.graphql_anycable_statistics.subscriptions_count.set({name: key}, value)
+  end
+```
+
+Or you can use `collect`
+```ruby
+  # config/initializers/metrics.rb
+  Yabeda.configure do
+    group :graphql_anycable_statistics do
+      gauge :subscriptions_count,  comment: "Number of graphql-anycable subscriptions"
+    end
+    
+    collect do
+      statistics = GraphQL::AnyCable.stats[:total]
+
+      statistics.each do |redis_prefix, value|
+        graphql_anycable_statistics.subscriptions_count.set({name: redis_prefix}, value)
+      end
+    end
+  end
+```
+
+
 ## Testing applications which use `graphql-anycable`
 
 You can pass custom redis-server URL to AnyCable using ENV variable.
@@ -298,3 +397,4 @@ The gem is available as open source under the terms of the [MIT License](https:/
 [AnyCable]: https://github.com/anycable/anycable "Polyglot replacement for Ruby WebSocket servers with Action Cable protocol"
 [LiteCable]: https://github.com/palkan/litecable "Lightweight Action Cable implementation (Rails-free)"
 [anyway_config]: https://github.com/palkan/anyway_config "Ruby libraries and applications configuration on steroids!"
+[Yabeda]: https://github.com/yabeda-rb/yabeda "Extendable solution for easy setup of monitoring in your Ruby apps"

lib/graphql-anycable.rb

@@ -6,6 +6,7 @@
 require_relative "graphql/anycable/cleaner"
 require_relative "graphql/anycable/config"
 require_relative "graphql/anycable/railtie" if defined?(Rails)
+require_relative "graphql/anycable/stats"
 require_relative "graphql/subscriptions/anycable_subscriptions"
 
 module GraphQL
@@ -20,6 +21,10 @@ def self.use(schema, **options)
       schema.use GraphQL::Subscriptions::AnyCableSubscriptions, **options
     end
 
+    def self.stats(**options)
+      Stats.new(**options).collect
+    end
+
     module_function
 
     def redis

lib/graphql/anycable/stats.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+module GraphQL
+  module AnyCable
+    # Calculates  amount of Graphql Redis keys
+    # (graphql-subscription, graphql-fingerprints, graphql-subscriptions, graphql-channel)
+    # Also, calculate the number of subscribers grouped by subscriptions
+    class Stats
+      SCAN_COUNT_RECORDS_AMOUNT = 1_000
+
+      attr_reader :scan_count, :include_subscriptions
+
+      def initialize(scan_count: SCAN_COUNT_RECORDS_AMOUNT, include_subscriptions: false)
+        @scan_count = scan_count
+        @include_subscriptions = include_subscriptions
+      end
+
+      def collect
+        total_subscriptions_result = {total: {}}
+
+        list_prefixes_keys.each do |name, prefix|
+          total_subscriptions_result[:total][name] = count_by_scan(match: "#{prefix}*")
+        end
+
+        if include_subscriptions
+          total_subscriptions_result[:subscriptions] = group_subscription_stats
+        end
+
+        total_subscriptions_result
+      end
+
+      private
+
+      # Counting all keys, that match the pattern with iterating by count
+      def count_by_scan(match:)
+        sb_amount = 0
+        cursor = '0'
+
+        loop do
+          cursor, result = redis.scan(cursor, match: match, count: scan_count)
+          sb_amount += result.count
+
+          break if cursor == '0'
+        end
+
+        sb_amount
+      end
+
+      # Calculate subscribes, grouped by subscriptions
+      def group_subscription_stats
+        subscription_groups = {}
+
+        redis.scan_each(match: "#{list_prefixes_keys[:fingerprints]}*", count: scan_count) do |fingerprint_key|
+          subscription_name = fingerprint_key.gsub(/#{list_prefixes_keys[:fingerprints]}|:/, "")
+          subscription_groups[subscription_name] = 0
+
+          redis.zscan_each(fingerprint_key) do |data|
+            redis.sscan_each("#{list_prefixes_keys[:subscriptions]}#{data[0]}") do |subscription_key|
+              next unless redis.exists?("#{list_prefixes_keys[:subscription]}#{subscription_key}")
+
+              subscription_groups[subscription_name] += 1
+            end
+          end
+        end
+
+        subscription_groups
+      end
+
+      def list_prefixes_keys
+        {
+          subscription: redis_key(adapter::SUBSCRIPTION_PREFIX),
+          fingerprints: redis_key(adapter::FINGERPRINTS_PREFIX),
+          subscriptions: redis_key(adapter::SUBSCRIPTIONS_PREFIX),
+          channel: redis_key(adapter::CHANNEL_PREFIX)
+        }
+      end
+
+      def adapter
+        GraphQL::Subscriptions::AnyCableSubscriptions
+      end
+
+      def redis
+        GraphQL::AnyCable.redis
+      end
+
+      def config
+        GraphQL::AnyCable.config
+      end
+
+      def redis_key(prefix)
+        "#{config.redis_prefix}-#{prefix}"
+      end
+    end
+  end
+end

spec/graphql/anycable_spec.rb

@@ -260,4 +260,12 @@
       end
     end
   end
+
+  describe ".stats" do
+    it "calls Graphql::AnyCable::Stats" do
+      allow_any_instance_of(GraphQL::AnyCable::Stats).to receive(:collect)
+
+      described_class.stats
+    end
+  end
 end

spec/graphql/stats_spec.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+RSpec.describe GraphQL::AnyCable::Stats do
+  describe "#collect" do
+    let(:query) do
+      <<~GRAPHQL
+        subscription SomeSubscription {
+          productCreated { id title }
+          productUpdated { id }
+        }
+      GRAPHQL
+    end
+
+    let(:channel) do
+      socket = double("Socket", istate: AnyCable::Socket::State.new({}))
+      connection = double("Connection", anycable_socket: socket)
+      double("Channel", id: "legacy_id", params: { "channelId" => "legacy_id" }, stream_from: nil, connection: connection)
+    end
+
+    let(:subscription_id) do
+      "some-truly-random-number"
+    end
+
+    before do
+      AnycableSchema.execute(
+        query: query,
+        context: { channel: channel, subscription_id: subscription_id },
+        variables: {},
+        operation_name: "SomeSubscription",
+      )
+    end
+
+    context "when include_subscriptions is false" do
+      let(:expected_result) do
+        {total: {subscription: 1, fingerprints: 2, subscriptions: 2, channel: 1}}
+      end
+
+      it "returns total stat" do
+        expect(subject.collect).to eq(expected_result)
+      end
+    end
+
+    context "when include_subscriptions is true" do
+      subject { described_class.new(include_subscriptions: true) }
+
+      let(:expected_result) do
+        {
+          total: {subscription: 1, fingerprints: 2, subscriptions: 2, channel: 1},
+          subscriptions: {
+            "productCreated"=> 1,
+            "productUpdated"=> 1
+          }
+        }
+      end
+
+      it "returns total stat with grouped subscription stats" do
+        expect(subject.collect).to eq(expected_result)
+      end
+    end
+  end
+end