Add unique userid count for Geo proxied requests

This adds a custom header for Geo proxied events
through Workhorse, and uses a HLLRedisCounter on
the #show action for major controllers to store the
unique users count going through the proxy.

Changelog: added
EE: true

config/feature_flags/development/track_geo_proxy_events.yml

+---
+name: track_geo_proxy_events
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/76587
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/348414
+milestone: '14.6'
+type: development
+group: group::geo
+default_enabled: false

ee/app/controllers/concerns/geo_instrumentation.rb

+# frozen_string_literal: true
+
+module GeoInstrumentation
+  extend ActiveSupport::Concern
+
+  included do
+    before_action :track_geo_proxy_event, only: [:show]
+  end
+
+  def track_geo_proxy_event
+    return unless Feature.enabled?(:track_geo_proxy_events, default_enabled: :yaml) && ::Gitlab::Geo.proxied_request?(request.env)
+    return unless current_user
+
+    Gitlab::UsageDataCounters::HLLRedisCounter.track_event('g_geo_proxied_requests', values: current_user.id)
+  end
+end

ee/app/controllers/ee/groups_controller.rb

@@ -8,6 +8,8 @@ module EE
     include GroupInviteMembers
 
     prepended do
+      include GeoInstrumentation
+
       alias_method :ee_authorize_admin_group!, :authorize_admin_group!
 
       before_action :ee_authorize_admin_group!, only: [:restore]

ee/app/controllers/ee/projects/issues_controller.rb

@@ -8,6 +8,7 @@ module EE
 
       prepended do
         include DescriptionDiffActions
+        include GeoInstrumentation
 
         before_action :disable_query_limiting_ee, only: [:update]
         before_action only: [:new, :create] do

ee/app/controllers/ee/projects/merge_requests_controller.rb

@@ -7,6 +7,7 @@ module EE
 
       prepended do
         include DescriptionDiffActions
+        include GeoInstrumentation
 
         before_action only: [:show] do
           if can_run_sast_experiments_on?(@project)

ee/app/controllers/ee/projects_controller.rb

@@ -6,6 +6,8 @@ module EE
     extend ::Gitlab::Utils::Override
 
     prepended do
+      include GeoInstrumentation
+
       before_action :log_download_export_audit_event, only: [:download_export]
       before_action :log_archive_audit_event, only: [:archive]
       before_action :log_unarchive_audit_event, only: [:unarchive]

ee/app/controllers/ee/users_controller.rb

@@ -2,8 +2,13 @@
 
 module EE
   module UsersController
+    extend ActiveSupport::Concern
     extend ::Gitlab::Utils::Override
 
+    prepended do
+      include GeoInstrumentation
+    end
+
     def available_project_templates
       load_custom_project_templates
     end

ee/config/metrics/counts_28d/20211214160234_g_geo_proxied_requests_monthly.yml

+---
+key_path: redis_hll_counters.geo.g_geo_proxied_requests_monthly
+description: Unique signed in users through Geo proxied requests
+product_section: enablement
+product_stage: enablement
+product_group: group::geo
+product_category: geo_replication
+value_type: number
+status: active
+milestone: "14.6"
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/76587
+time_frame: 28d
+data_source: redis_hll
+data_category: optional
+instrumentation_class: RedisHLLMetric
+performance_indicator_type: []
+options:
+  events:
+  - g_geo_proxied_requests
+distribution:
+- ee
+tier:
+- premium
+- ultimate

ee/config/metrics/counts_7d/20211214160228_g_geo_proxied_requests_weekly.yml

+---
+key_path: redis_hll_counters.geo.g_geo_proxied_requests_weekly
+description: Unique signed in users through Geo proxied requests
+product_section: enablement
+product_stage: enablement
+product_group: group::geo
+product_category: geo_replication
+value_type: number
+status: active
+milestone: "14.6"
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/76587
+time_frame: 7d
+data_source: redis_hll
+data_category: optional
+instrumentation_class: RedisHLLMetric
+performance_indicator_type: []
+options:
+  events:
+  - g_geo_proxied_requests
+distribution:
+- ee
+tier:
+- premium
+- ultimate

ee/lib/gitlab/geo.rb

@@ -90,6 +90,10 @@ module Gitlab
       self.secondary_with_primary? && self.primary_node.url == self.current_node.url
     end
 
+    def self.proxied_request?(env)
+      env['HTTP_GITLAB_WORKHORSE_GEO_PROXY'] == '1'
+    end
+
     def self.license_allows?
       ::License.feature_available?(:geo)
     end

ee/spec/controllers/concerns/geo_instrumentation_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe GeoInstrumentation do
+  let(:user) { create(:user) }
+
+  controller(ActionController::Base) do
+    include ::GeoInstrumentation
+
+    def show
+      render plain: "show action"
+    end
+  end
+
+  before do
+    routes.draw { get "show" => "anonymous#show" }
+    sign_in(user) unless user.nil?
+  end
+
+  describe '.track_geo_proxy_event' do
+    context 'when the track_geo_proxy feature flag is disabled' do
+      before do
+        stub_feature_flags(track_geo_proxy_events: false)
+      end
+
+      it 'does not track an event' do
+        expect(Gitlab::UsageDataCounters::HLLRedisCounter).not_to receive(:track_event)
+        get :show
+      end
+    end
+
+    context 'when the track_geo_proxy feature flag is enabled' do
+      context 'when the request is not proxied' do
+        before do
+          allow(::Gitlab::Geo).to receive(:proxied_request?).and_return(false)
+        end
+
+        it 'does not track an event' do
+          expect(Gitlab::UsageDataCounters::HLLRedisCounter).not_to receive(:track_event)
+          get :show
+        end
+      end
+
+      context 'when the request is proxied' do
+        before do
+          allow(::Gitlab::Geo).to receive(:proxied_request?).and_return(true)
+        end
+
+        context 'when logged in' do
+          it 'tracks a HLL event for unique geo proxied requests' do
+            expect(Gitlab::UsageDataCounters::HLLRedisCounter)
+              .to receive(:track_event).with('g_geo_proxied_requests', values: user.id)
+            get :show
+          end
+        end
+
+        context 'when not logged in' do
+          let(:user) { nil }
+
+          it 'does not track an event' do
+            expect(Gitlab::UsageDataCounters::HLLRedisCounter).not_to receive(:track_event)
+            get :show
+          end
+        end
+      end
+    end
+  end
+end

ee/spec/lib/gitlab/geo_spec.rb

@@ -192,6 +192,17 @@ RSpec.describe Gitlab::Geo, :geo, :request_store do
     end
   end
 
+  describe '.proxied_request?' do
+    it 'returns true when the header is set' do
+      expect(described_class.proxied_request?({ 'HTTP_GITLAB_WORKHORSE_GEO_PROXY' => '1' })).to be_truthy
+    end
+
+    it 'returns false when the header is not present or set o an invalid value' do
+      expect(described_class.proxied_request?({})).to be_falsey
+      expect(described_class.proxied_request?({ 'HTTP_GITLAB_WORKHORSE_GEO_PROXY' => 'invalid' })).to be_falsey
+    end
+  end
+
   describe '.enabled?' do
     it_behaves_like 'a Geo cached value', :enabled?, :node_enabled
 

lib/gitlab/usage_data_counters/known_events/common.yml

@@ -373,3 +373,9 @@
   redis_slot: network_policies
   category: network_policies
   aggregation: weekly
+# Geo group
+- name: g_geo_proxied_requests
+  category: geo
+  redis_slot: geo
+  aggregation: daily
+  feature_flag: track_geo_proxy_events

spec/lib/gitlab/usage_data_counters/hll_redis_counter_spec.rb

@@ -48,7 +48,8 @@ RSpec.describe Gitlab::UsageDataCounters::HLLRedisCounter, :clean_gitlab_redis_s
         'epic_boards_usage',
         'secure',
         'importer',
-        'network_policies'
+        'network_policies',
+        'geo'
       )
     end
   end

workhorse/internal/proxy/proxy.go

@@ -18,10 +18,17 @@ type Proxy struct {
 	Version                string
 	reverseProxy           *httputil.ReverseProxy
 	AllowResponseBuffering bool
+	customHeaders          map[string]string
 }
 
-func NewProxy(myURL *url.URL, version string, roundTripper http.RoundTripper) *Proxy {
-	p := Proxy{Version: version, AllowResponseBuffering: true}
+func WithCustomHeaders(customHeaders map[string]string) func(*Proxy) {
+	return func(proxy *Proxy) {
+		proxy.customHeaders = customHeaders
+	}
+}
+
+func NewProxy(myURL *url.URL, version string, roundTripper http.RoundTripper, options ...func(*Proxy)) *Proxy {
+	p := Proxy{Version: version, AllowResponseBuffering: true, customHeaders: make(map[string]string)}
 
 	if myURL == nil {
 		myURL = defaultTarget
@@ -31,6 +38,11 @@ func NewProxy(myURL *url.URL, version string, roundTripper http.RoundTripper) *P
 	u.Path = ""
 	p.reverseProxy = httputil.NewSingleHostReverseProxy(&u)
 	p.reverseProxy.Transport = roundTripper
+
+	for _, option := range options {
+		option(&p)
+	}
+
 	return &p
 }
 
@@ -43,6 +55,10 @@ func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	req.Header.Set("Gitlab-Workhorse", p.Version)
 	req.Header.Set("Gitlab-Workhorse-Proxy-Start", fmt.Sprintf("%d", time.Now().UnixNano()))
 
+	for k, v := range p.customHeaders {
+		req.Header.Set(k, v)
+	}
+
 	if p.AllowResponseBuffering {
 		helper.AllowResponseBuffering(w)
 	}

workhorse/internal/upstream/upstream.go

@@ -37,6 +37,7 @@ var (
 		upload.RewrittenFieldsHeader,
 	}
 	geoProxyApiPollingInterval = 10 * time.Second
+	geoProxyWorkhorseHeaders   = map[string]string{"Gitlab-Workhorse-Geo-Proxy": "1"}
 )
 
 type upstream struct {
@@ -237,7 +238,12 @@ func (u *upstream) updateGeoProxyFields(geoProxyURL *url.URL) {
 	}
 
 	geoProxyRoundTripper := roundtripper.NewBackendRoundTripper(u.geoProxyBackend, "", u.ProxyHeadersTimeout, u.DevelopmentMode)
-	geoProxyUpstream := proxypkg.NewProxy(u.geoProxyBackend, u.Version, geoProxyRoundTripper)
+	geoProxyUpstream := proxypkg.NewProxy(
+		u.geoProxyBackend,
+		u.Version,
+		geoProxyRoundTripper,
+		proxypkg.WithCustomHeaders(geoProxyWorkhorseHeaders),
+	)
 	u.geoProxyCableRoute = u.wsRoute(`^/-/cable\z`, geoProxyUpstream)
 	u.geoProxyRoute = u.route("", "", geoProxyUpstream)
 }

workhorse/internal/upstream/upstream_test.go

@@ -209,6 +209,23 @@ func TestGeoProxyFeatureEnablingAndDisabling(t *testing.T) {
 	runTestCases(t, ws, testCasesProxied)
 }
 
+func TestGeoProxySetsCustomHeader(t *testing.T) {
+	remoteServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		require.Equal(t, "1", r.Header.Get("Gitlab-Workhorse-Geo-Proxy"), "custom proxy header")
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer remoteServer.Close()
+
+	geoProxyEndpointResponseBody := fmt.Sprintf(`{"geo_proxy_url":"%v"}`, remoteServer.URL)
+	railsServer, deferredClose := startRailsServer("Local Rails server", &geoProxyEndpointResponseBody)
+	defer deferredClose()
+
+	ws, wsDeferredClose, _ := startWorkhorseServer(railsServer.URL, true)
+	defer wsDeferredClose()
+
+	http.Get(ws.URL)
+}
+
 func runTestCases(t *testing.T, ws *httptest.Server, testCases []testCase) {
 	t.Helper()
 	for _, tc := range testCases {

workhorse/proxy_test.go

@@ -22,12 +22,12 @@ import (
 
 const testVersion = "123"
 
-func newProxy(url string, rt http.RoundTripper) *proxy.Proxy {
+func newProxy(url string, rt http.RoundTripper, opts ...func(*proxy.Proxy)) *proxy.Proxy {
 	parsedURL := helper.URLMustParse(url)
 	if rt == nil {
 		rt = roundtripper.NewTestBackendRoundTripper(parsedURL)
 	}
-	return proxy.NewProxy(parsedURL, testVersion, rt)
+	return proxy.NewProxy(parsedURL, testVersion, rt, opts...)
 }
 
 func TestProxyRequest(t *testing.T) {
@@ -64,6 +64,24 @@ func TestProxyRequest(t *testing.T) {
 	require.Equal(t, "test", w.Header().Get("Custom-Response-Header"), "custom response header")
 }
 
+func TestProxyWithCustomHeaders(t *testing.T) {
+	ts := testhelper.TestServerWithHandler(regexp.MustCompile(`/url/path\z`), func(w http.ResponseWriter, r *http.Request) {
+		require.Equal(t, "value", r.Header.Get("Custom-Header"), "custom proxy header")
+		require.Equal(t, testVersion, r.Header.Get("Gitlab-Workhorse"), "version header")
+
+		_, err := w.Write([]byte(`ok`))
+		require.NoError(t, err, "write ok response")
+	})
+
+	httpRequest, err := http.NewRequest("POST", ts.URL+"/url/path", nil)
+	require.NoError(t, err)
+
+	w := httptest.NewRecorder()
+	testProxy := newProxy(ts.URL, nil, proxy.WithCustomHeaders(map[string]string{"Custom-Header": "value"}))
+	testProxy.ServeHTTP(w, httpRequest)
+	testhelper.RequireResponseBody(t, w, "ok")
+}
+
 func TestProxyError(t *testing.T) {
 	httpRequest, err := http.NewRequest("POST", "/url/path", bytes.NewBufferString("REQUEST"))
 	require.NoError(t, err)