Add context level tracking

  - Add track_event_in_context, prefer a separate method,
    less PFADD operations
  - Add valid_context_list
  - Adjust valid_context_list with license plans in ee

doc/development/product_analytics/usage_ping.md

@@ -429,13 +429,22 @@ w
    - `values`: One value or array of values we count. For example: user_id, visitor_id, user_ids.
    - `event_name`: event name.
 
-1. Get event data using `Gitlab::UsageDataCounters::HLLRedisCounter.unique_events(event_names:, start_date:, end_date)`.
+1. Track event on context level using base module `Gitlab::UsageDataCounters::HLLRedisCounter.track_event_in_context(entity_id, event_name, context)`.
+
+   Arguments:
+
+   - `entity_id`: value we count. For example: user_id, visitor_id.
+   - `event_name`: event name.
+   - `context`: context value. Allowed values are `default`, `free`, `bronze`, `silver`, `gold`, `starter`, `premium`, `ultimate`
+
+1. Get event data using `Gitlab::UsageDataCounters::HLLRedisCounter.unique_events(event_names:, start_date:, end_date:, context: '')`.
 
    Arguments:
 
    - `event_names`: the list of event names.
    - `start_date`: start date of the period for which we want to get event data.
    - `end_date`: end date of the period for which we want to get event data.
+   - `context`: context of the event. Allowed values are `default`, `free`, `bronze`, `silver`, `gold`, `starter`, `premium`, `ultimate`.
 
 Recommendations:
 

ee/app/models/license.rb

@@ -9,6 +9,8 @@ class License < ApplicationRecord
   ULTIMATE_PLAN = 'ultimate'.freeze
   ALLOWED_PERCENTAGE_OF_USERS_OVERAGE = (10 / 100.0).freeze
 
+  EE_ALL_PLANS = [STARTER_PLAN, PREMIUM_PLAN, ULTIMATE_PLAN].freeze
+
   EES_FEATURES = %i[
     audit_events
     blocked_issues
@@ -257,6 +259,10 @@ class License < ApplicationRecord
       end
     end
 
+    def all_plans
+      EE_ALL_PLANS
+    end
+
     delegate :block_changes?, :feature_available?, to: :current, allow_nil: true
 
     def reset_current

ee/lib/ee/gitlab/usage_data_counters/hll_redis_counter.rb

+# frozen_string_literal: true
+
+module EE
+  module Gitlab
+    module UsageDataCounters
+      module HLLRedisCounter
+        extend ActiveSupport::Concern
+        class_methods do
+          extend ::Gitlab::Utils::Override
+
+          override :valid_context_list
+          def valid_context_list
+            super + License.all_plans
+          end
+        end
+      end
+    end
+  end
+end

ee/spec/lib/ee/gitlab/usage_data_counters/hll_redis_counter_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::UsageDataCounters::HLLRedisCounter, :clean_gitlab_redis_shared_state do
+  using RSpec::Parameterized::TableSyntax
+
+  let(:entity1) { 'dfb9d2d2-f56c-4c77-8aeb-6cddc4a1f857' }
+  let(:entity2) { '1dd9afb2-a3ee-4de1-8ae3-a405579c8584' }
+  let(:entity3) { '34rfjuuy-ce56-sa35-ds34-dfer567dfrf2' }
+
+  let(:default_context) { 'default' }
+  let(:ultimate_context) { 'ultimate' }
+  let(:gold_context) { 'gold' }
+  let(:invalid_context) { 'invalid' }
+
+  let(:context_event) { 'context_event' }
+  let(:other_context_event) { 'other_context_event' }
+
+  let(:known_events) do
+    [
+      { name: context_event, category: 'other', expiry: 6, aggregation: 'weekly' },
+      { name: other_context_event, category: 'other', expiry: 6, aggregation: 'weekly' }
+    ].map(&:with_indifferent_access)
+  end
+
+  around do |example|
+    # We need to freeze to a reference time
+    # because visits are grouped by the week number in the year
+    # Without freezing the time, the test may behave inconsistently
+    # depending on which day of the week test is run.
+    # Monday 6th of June
+    reference_time = Time.utc(2020, 6, 1)
+    travel_to(reference_time) { example.run }
+  end
+
+  before do
+    allow(described_class).to receive(:known_events).and_return(known_events)
+  end
+
+  describe '.track_event_in_context' do
+    context 'with valid context' do
+      where(:entity, :event_name, :context) do
+        entity1 | context_event | default_context
+        entity1 | context_event | ultimate_context
+        entity1 | context_event | gold_context
+      end
+
+      with_them do
+        it 'increments context event counter' do
+          expect(Gitlab::Redis::HLL).to receive(:add) do |kwargs|
+            expect(kwargs[:key]).to match(/^#{context}\_.*/)
+          end
+
+          described_class.track_event_in_context(entity, event_name, context)
+        end
+      end
+    end
+
+    context 'when sending empty context' do
+      it 'is not incrementing the counter' do
+        expect(Gitlab::Redis::HLL).not_to receive(:add)
+
+        described_class.track_event_in_context(entity1, context_event, '')
+      end
+    end
+  end
+
+  describe '.unique_events' do
+    before do
+      described_class.track_event_in_context([entity1, entity3], context_event, default_context, 2.days.ago)
+      described_class.track_event_in_context(entity3, context_event, ultimate_context, 2.days.ago)
+      described_class.track_event_in_context(entity3, context_event, gold_context, 2.days.ago)
+      described_class.track_event_in_context(entity3, context_event, invalid_context, 2.days.ago)
+      described_class.track_event_in_context([entity1, entity2], context_event, '', 2.weeks.ago)
+    end
+
+    context 'with correct arguments' do
+      subject(:unique_events) { described_class.unique_events(event_names: event_names, start_date: 4.weeks.ago, end_date: Date.current, context: context) }
+
+      where(:event_names, :context, :value) do
+        context_event | default_context  | 2
+        context_event | ultimate_context | 1
+        context_event | gold_context     | 1
+        context_event | ''               | 0
+      end
+
+      with_them do
+        it { is_expected.to eq value }
+      end
+    end
+
+    context 'with empty invalid context' do
+      it 'raise error' do
+        expect { described_class.unique_events(event_names: context_event, start_date: 4.weeks.ago, end_date: Date.current, context: invalid_context) }.to raise_error(Gitlab::UsageDataCounters::HLLRedisCounter::InvalidContext)
+      end
+    end
+  end
+end

lib/gitlab/usage_data_counters/hll_redis_counter.rb

@@ -14,6 +14,7 @@ module Gitlab
       SlotMismatch = Class.new(EventError)
       CategoryMismatch = Class.new(EventError)
       UnknownAggregationOperator = Class.new(EventError)
+      InvalidContext = Class.new(EventError)
 
       KNOWN_EVENTS_PATH = File.expand_path('known_events/*.yml', __dir__)
       ALLOWED_AGGREGATIONS = %i(daily weekly).freeze
@@ -42,21 +43,23 @@ module Gitlab
       class << self
         include Gitlab::Utils::UsageData
 
-        def track_event(entity_id, event_name, time = Time.zone.now)
-          return unless Gitlab::CurrentSettings.usage_ping_enabled?
-
-          event = event_for(event_name)
+        def track_event(value, event_name, time = Time.zone.now)
+          track(value, event_name, time: time)
+        end
 
-          raise UnknownEvent, "Unknown event #{event_name}" unless event.present?
+        def track_event_in_context(value, event_name, context, time = Time.zone.now)
+          return if context.blank?
+          return if context.present? && !context.in?(valid_context_list)
 
-          Gitlab::Redis::HLL.add(key: redis_key(event, time), value: entity_id, expiry: expiry(event))
+          track(value, event_name, context: context, time: time)
         end
 
-        def unique_events(event_names:, start_date:, end_date:)
-          count_unique_events(event_names: event_names, start_date: start_date, end_date: end_date) do |events|
+        def unique_events(event_names:, start_date:, end_date:, context: '')
+          count_unique_events(event_names: event_names, start_date: start_date, end_date: end_date, context: context) do |events|
             raise SlotMismatch, events unless events_in_same_slot?(events)
             raise CategoryMismatch, events unless events_in_same_category?(events)
             raise AggregationMismatch, events unless events_same_aggregation?(events)
+            raise InvalidContext if context.present? && !context.in?(valid_context_list)
           end
         end
 
@@ -114,6 +117,20 @@ module Gitlab
 
         private
 
+        def track(value, event_name, context: '', time: Time.zone.now)
+          return unless Gitlab::CurrentSettings.usage_ping_enabled?
+
+          event = event_for(event_name)
+          raise UnknownEvent, "Unknown event #{event_name}" unless event.present?
+
+          Gitlab::Redis::HLL.add(key: redis_key(event, time, context), value: value, expiry: expiry(event))
+        end
+
+        # The aray of valid context on which we allow tracking
+        def valid_context_list
+          Plan.all_plans
+        end
+
         def calculate_count_for_aggregation(aggregation, start_date:, end_date:)
           case aggregation[:operator]
           when UNION_OF_AGGREGATED_METRICS
@@ -193,14 +210,14 @@ module Gitlab
           end
         end
 
-        def count_unique_events(event_names:, start_date:, end_date:)
+        def count_unique_events(event_names:, start_date:, end_date:, context: '')
           events = events_for(Array(event_names).map(&:to_s))
 
           yield events if block_given?
 
           aggregation = events.first[:aggregation]
 
-          keys = keys_for_aggregation(aggregation, events: events, start_date: start_date, end_date: end_date)
+          keys = keys_for_aggregation(aggregation, events: events, start_date: start_date, end_date: end_date, context: context)
           redis_usage_data { Gitlab::Redis::HLL.count(keys: keys) }
         end
 
@@ -213,11 +230,11 @@ module Gitlab
           events_in_same_slot?(events) && events_in_same_category?(events) && events_same_aggregation?(events)
         end
 
-        def keys_for_aggregation(aggregation, events:, start_date:, end_date:)
+        def keys_for_aggregation(aggregation, events:, start_date:, end_date:, context: '')
           if aggregation.to_sym == :daily
-            daily_redis_keys(events: events, start_date: start_date, end_date: end_date)
+            daily_redis_keys(events: events, start_date: start_date, end_date: end_date, context: context)
           else
-            weekly_redis_keys(events: events, start_date: start_date, end_date: end_date)
+            weekly_redis_keys(events: events, start_date: start_date, end_date: end_date, context: context)
           end
         end
 
@@ -272,17 +289,26 @@ module Gitlab
         end
 
         # Compose the key in order to store events daily or weekly
-        def redis_key(event, time)
+        def redis_key(event, time, context = '')
           raise UnknownEvent.new("Unknown event #{event[:name]}") unless known_events_names.include?(event[:name].to_s)
           raise UnknownAggregation.new("Use :daily or :weekly aggregation") unless ALLOWED_AGGREGATIONS.include?(event[:aggregation].to_sym)
 
+          key = apply_slot(event)
+          key = apply_time_aggregation(key, time, event)
+          key = "#{context}_#{key}" if context.present?
+          key
+        end
+
+        def apply_slot(event)
           slot = redis_slot(event)
-          key = if slot.present?
-                  event[:name].to_s.gsub(slot, "{#{slot}}")
-                else
-                  "{#{event[:name]}}"
-                end
+          if slot.present?
+            event[:name].to_s.gsub(slot, "{#{slot}}")
+          else
+            "{#{event[:name]}}"
+          end
+        end
 
+        def apply_time_aggregation(key, time, event)
           if event[:aggregation].to_sym == :daily
             year_day = time.strftime('%G-%j')
             "#{year_day}-#{key}"
@@ -292,21 +318,29 @@ module Gitlab
           end
         end
 
-        def daily_redis_keys(events:, start_date:, end_date:)
+        def daily_redis_keys(events:, start_date:, end_date:, context: '')
           (start_date.to_date..end_date.to_date).map do |date|
-            events.map { |event| redis_key(event, date) }
+            events.map { |event| redis_key(event, date, context) }
           end.flatten
         end
 
-        def weekly_redis_keys(events:, start_date:, end_date:)
+        def validate_aggregation_operator!(operator)
+          return true if ALLOWED_METRICS_AGGREGATIONS.include?(operator)
+
+          raise UnknownAggregationOperator.new("Events should be aggregated with one of operators #{ALLOWED_METRICS_AGGREGATIONS}")
+        end
+
+        def weekly_redis_keys(events:, start_date:, end_date:, context: '')
           weeks = end_date.to_date.cweek - start_date.to_date.cweek
           weeks = 1 if weeks == 0
 
           (0..(weeks - 1)).map do |week_increment|
-            events.map { |event| redis_key(event, start_date + week_increment * 7.days) }
+            events.map { |event| redis_key(event, start_date + week_increment * 7.days, context) }
           end.flatten
         end
       end
     end
   end
 end
+
+Gitlab::UsageDataCounters::HLLRedisCounter.prepend_if_ee('EE::Gitlab::UsageDataCounters::HLLRedisCounter')

spec/lib/gitlab/usage_data_counters/hll_redis_counter_spec.rb

@@ -8,6 +8,9 @@ RSpec.describe Gitlab::UsageDataCounters::HLLRedisCounter, :clean_gitlab_redis_s
   let(:entity3) { '34rfjuuy-ce56-sa35-ds34-dfer567dfrf2' }
   let(:entity4) { '8b9a2671-2abf-4bec-a682-22f6a8f7bf31' }
 
+  let(:default_context) { 'default' }
+  let(:invalid_context) { 'invalid' }
+
   around do |example|
     # We need to freeze to a reference time
     # because visits are grouped by the week number in the year
@@ -55,11 +58,13 @@ RSpec.describe Gitlab::UsageDataCounters::HLLRedisCounter, :clean_gitlab_redis_s
     let(:no_slot) { 'no_slot' }
     let(:different_aggregation) { 'different_aggregation' }
     let(:custom_daily_event) { 'g_analytics_custom' }
+    let(:context_event) { 'context_event' }
 
     let(:global_category) { 'global' }
     let(:compliance_category) { 'compliance' }
     let(:productivity_category) { 'productivity' }
     let(:analytics_category) { 'analytics' }
+    let(:other_category) { 'other' }
 
     let(:known_events) do
       [
@@ -68,7 +73,8 @@ RSpec.describe Gitlab::UsageDataCounters::HLLRedisCounter, :clean_gitlab_redis_s
         { name: category_productivity_event, redis_slot: "analytics", category: productivity_category, aggregation: "weekly" },
         { name: compliance_slot_event, redis_slot: "compliance", category: compliance_category, aggregation: "weekly" },
         { name: no_slot, category: global_category, aggregation: "daily" },
-        { name: different_aggregation, category: global_category, aggregation: "monthly" }
+        { name: different_aggregation, category: global_category, aggregation: "monthly" },
+        { name: context_event, category: other_category, expiry: 6, aggregation: 'weekly' }
       ].map(&:with_indifferent_access)
     end
 
@@ -170,6 +176,34 @@ RSpec.describe Gitlab::UsageDataCounters::HLLRedisCounter, :clean_gitlab_redis_s
       end
     end
 
+    describe '.track_event_in_context' do
+      context 'with valid contex' do
+        it 'increments conext event counte' do
+          expect(Gitlab::Redis::HLL).to receive(:add) do |kwargs|
+            expect(kwargs[:key]).to match(/^#{default_context}\_.*/)
+          end
+
+          described_class.track_event_in_context(entity1, context_event, default_context)
+        end
+      end
+
+      context 'with empty context' do
+        it 'does not increment a counter' do
+          expect(Gitlab::Redis::HLL).not_to receive(:add)
+
+          described_class.track_event_in_context(entity1, context_event, '')
+        end
+      end
+
+      context 'when sending invalid context' do
+        it 'does not increment a counter' do
+          expect(Gitlab::Redis::HLL).not_to receive(:add)
+
+          described_class.track_event_in_context(entity1, context_event, invalid_context)
+        end
+      end
+    end
+
     describe '.unique_events' do
       before do
         # events in current week, should not be counted as week is not complete
@@ -250,6 +284,48 @@ RSpec.describe Gitlab::UsageDataCounters::HLLRedisCounter, :clean_gitlab_redis_s
     end
   end
 
+  describe 'context level tracking' do
+    using RSpec::Parameterized::TableSyntax
+
+    let(:known_events) do
+      [
+        { name: 'event_name_1', redis_slot: 'event', category: 'category1', aggregation: "weekly" },
+        { name: 'event_name_2', redis_slot: 'event', category: 'category1', aggregation: "weekly" },
+        { name: 'event_name_3', redis_slot: 'event', category: 'category1', aggregation: "weekly" }
+      ].map(&:with_indifferent_access)
+    end
+
+    before do
+      allow(described_class).to receive(:known_events).and_return(known_events)
+      allow(described_class).to receive(:categories).and_return(%w(category1 category2))
+
+      described_class.track_event_in_context([entity1, entity3], 'event_name_1', default_context, 2.days.ago)
+      described_class.track_event_in_context(entity3, 'event_name_1', default_context, 2.days.ago)
+      described_class.track_event_in_context(entity3, 'event_name_1', invalid_context, 2.days.ago)
+      described_class.track_event_in_context([entity1, entity2], 'event_name_2', '', 2.weeks.ago)
+    end
+
+    subject(:unique_events) { described_class.unique_events(event_names: event_names, start_date: 4.weeks.ago, end_date: Date.current, context: context) }
+
+    context 'with correct arguments' do
+      where(:event_names, :context, :value) do
+        ['event_name_1'] | 'default' | 2
+        ['event_name_1'] | ''        | 0
+        ['event_name_2'] | ''        | 0
+      end
+
+      with_them do
+        it { is_expected.to eq value }
+      end
+    end
+
+    context 'with empty invalid context' do
+      it 'raise error' do
+        expect { described_class.unique_events(event_names: 'event_name_1', start_date: 4.weeks.ago, end_date: Date.current, context: invalid_context) }.to raise_error(Gitlab::UsageDataCounters::HLLRedisCounter::InvalidContext)
+      end
+    end
+  end
+
   describe 'unique_events_data' do
     let(:known_events) do
       [