Extract agent token usage tracking logic into service

app/models/clusters/agent_token.rb

@@ -10,9 +10,6 @@ module Clusters
 
     self.table_name = 'cluster_agent_tokens'
 
-    # The `UPDATE_USED_COLUMN_EVERY` defines how often the token DB entry can be updated
-    UPDATE_USED_COLUMN_EVERY = (40.minutes..55.minutes).freeze
-
     belongs_to :agent, class_name: 'Clusters::Agent', optional: false
     belongs_to :created_by_user, class_name: 'User', optional: true
 
@@ -28,41 +25,5 @@ module Clusters
       active: 0,
       revoked: 1
     }
-
-    def track_usage
-      track_values = { last_used_at: Time.current.utc }
-
-      cache_attributes(track_values)
-
-      if can_update_track_values?
-        log_activity_event!(track_values[:last_used_at]) unless agent.connected?
-
-        # Use update_column so updated_at is skipped
-        update_columns(track_values)
-      end
-    end
-
-    private
-
-    def can_update_track_values?
-      # Use a random threshold to prevent beating DB updates.
-      last_used_at_max_age = Random.rand(UPDATE_USED_COLUMN_EVERY)
-
-      real_last_used_at = read_attribute(:last_used_at)
-
-      # Handle too many updates from high token traffic
-      real_last_used_at.nil? ||
-        (Time.current - real_last_used_at) >= last_used_at_max_age
-    end
-
-    def log_activity_event!(recorded_at)
-      Clusters::Agents::CreateActivityEventService.new( # rubocop: disable CodeReuse/ServiceClass
-        agent,
-        kind: :agent_connected,
-        level: :info,
-        recorded_at: recorded_at,
-        agent_token: self
-      ).execute
-    end
   end
 end

app/services/clusters/agent_tokens/track_usage_service.rb

+# frozen_string_literal: true
+
+module Clusters
+  module AgentTokens
+    class TrackUsageService
+      # The `UPDATE_USED_COLUMN_EVERY` defines how often the token DB entry can be updated
+      UPDATE_USED_COLUMN_EVERY = (40.minutes..55.minutes).freeze
+
+      delegate :agent, to: :token
+
+      def initialize(token)
+        @token = token
+      end
+
+      def execute
+        track_values = { last_used_at: Time.current.utc }
+
+        token.cache_attributes(track_values)
+
+        if can_update_track_values?
+          log_activity_event!(track_values[:last_used_at]) unless agent.connected?
+
+          # Use update_column so updated_at is skipped
+          token.update_columns(track_values)
+        end
+      end
+
+      private
+
+      attr_reader :token
+
+      def can_update_track_values?
+        # Use a random threshold to prevent beating DB updates.
+        last_used_at_max_age = Random.rand(UPDATE_USED_COLUMN_EVERY)
+
+        real_last_used_at = token.read_attribute(:last_used_at)
+
+        # Handle too many updates from high token traffic
+        real_last_used_at.nil? ||
+          (Time.current - real_last_used_at) >= last_used_at_max_age
+      end
+
+      def log_activity_event!(recorded_at)
+        Clusters::Agents::CreateActivityEventService.new(
+          agent,
+          kind: :agent_connected,
+          level: :info,
+          recorded_at: recorded_at,
+          agent_token: token
+        ).execute
+      end
+    end
+  end
+end

lib/api/internal/kubernetes.rb

@@ -53,7 +53,7 @@ module API
         def check_agent_token
           unauthorized! unless agent_token
 
-          agent_token.track_usage
+          Clusters::AgentTokens::TrackUsageService.new(agent_token).execute
         end
       end
 

spec/models/clusters/agent_token_spec.rb

@@ -49,83 +49,4 @@ RSpec.describe Clusters::AgentToken do
       expect(agent_token.token.length).to be >= 50
     end
   end
-
-  describe '#track_usage', :clean_gitlab_redis_cache do
-    let_it_be(:agent) { create(:cluster_agent) }
-
-    let(:agent_token) { create(:cluster_agent_token, agent: agent) }
-
-    subject { agent_token.track_usage }
-
-    context 'when last_used_at was updated recently' do
-      before do
-        agent_token.update!(last_used_at: 10.minutes.ago)
-      end
-
-      it 'updates cache but not database' do
-        expect { subject }.not_to change { agent_token.reload.read_attribute(:last_used_at) }
-
-        expect_redis_update
-      end
-    end
-
-    context 'when last_used_at was not updated recently' do
-      it 'updates cache and database' do
-        does_db_update
-        expect_redis_update
-      end
-
-      context 'with invalid token' do
-        before do
-          agent_token.description = SecureRandom.hex(2000)
-        end
-
-        it 'still updates caches and database' do
-          expect(agent_token).to be_invalid
-
-          does_db_update
-          expect_redis_update
-        end
-      end
-
-      context 'agent is not connected' do
-        before do
-          allow(agent).to receive(:connected?).and_return(false)
-        end
-
-        it 'creates an activity event' do
-          expect { subject }.to change { agent.activity_events.count }
-
-          event = agent.activity_events.last
-          expect(event).to have_attributes(
-            kind: 'agent_connected',
-            level: 'info',
-            recorded_at: agent_token.reload.read_attribute(:last_used_at),
-            agent_token: agent_token
-          )
-        end
-      end
-
-      context 'agent is connected' do
-        before do
-          allow(agent).to receive(:connected?).and_return(true)
-        end
-
-        it 'does not create an activity event' do
-          expect { subject }.not_to change { agent.activity_events.count }
-        end
-      end
-    end
-
-    def expect_redis_update
-      Gitlab::Redis::Cache.with do |redis|
-        redis_key = "cache:#{described_class.name}:#{agent_token.id}:attributes"
-        expect(redis.get(redis_key)).to be_present
-      end
-    end
-
-    def does_db_update
-      expect { subject }.to change { agent_token.reload.read_attribute(:last_used_at) }
-    end
-  end
 end

spec/requests/api/internal/kubernetes_spec.rb

@@ -53,7 +53,9 @@ RSpec.describe API::Internal::Kubernetes do
 
   shared_examples 'agent token tracking' do
     it 'tracks token usage' do
-      expect { response }.to change { agent_token.reload.read_attribute(:last_used_at) }
+      expect do
+        send_request(headers: { 'Authorization' => "Bearer #{agent_token.token}" })
+      end.to change { agent_token.reload.read_attribute(:last_used_at) }
     end
   end
 
@@ -149,7 +151,7 @@ RSpec.describe API::Internal::Kubernetes do
       let(:agent) { agent_token.agent }
       let(:project) { agent.project }
 
-      shared_examples 'agent token tracking'
+      include_examples 'agent token tracking'
 
       it 'returns expected data', :aggregate_failures do
         send_request(headers: { 'Authorization' => "Bearer #{agent_token.token}" })

spec/services/clusters/agent_tokens/track_usage_service_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Clusters::AgentTokens::TrackUsageService do
+  let_it_be(:agent) { create(:cluster_agent) }
+
+  describe '#execute', :clean_gitlab_redis_cache do
+    let(:agent_token) { create(:cluster_agent_token, agent: agent) }
+
+    subject { described_class.new(agent_token).execute }
+
+    context 'when last_used_at was updated recently' do
+      before do
+        agent_token.update!(last_used_at: 10.minutes.ago)
+      end
+
+      it 'updates cache but not database' do
+        expect { subject }.not_to change { agent_token.reload.read_attribute(:last_used_at) }
+
+        expect_redis_update
+      end
+    end
+
+    context 'when last_used_at was not updated recently' do
+      it 'updates cache and database' do
+        does_db_update
+        expect_redis_update
+      end
+
+      context 'with invalid token' do
+        before do
+          agent_token.description = SecureRandom.hex(2000)
+        end
+
+        it 'still updates caches and database' do
+          expect(agent_token).to be_invalid
+
+          does_db_update
+          expect_redis_update
+        end
+      end
+
+      context 'agent is not connected' do
+        before do
+          allow(agent).to receive(:connected?).and_return(false)
+        end
+
+        it 'creates an activity event' do
+          expect { subject }.to change { agent.activity_events.count }
+
+          event = agent.activity_events.last
+          expect(event).to have_attributes(
+            kind: 'agent_connected',
+            level: 'info',
+            recorded_at: agent_token.reload.read_attribute(:last_used_at),
+            agent_token: agent_token
+          )
+        end
+      end
+
+      context 'agent is connected' do
+        before do
+          allow(agent).to receive(:connected?).and_return(true)
+        end
+
+        it 'does not create an activity event' do
+          expect { subject }.not_to change { agent.activity_events.count }
+        end
+      end
+    end
+
+    def expect_redis_update
+      Gitlab::Redis::Cache.with do |redis|
+        redis_key = "cache:#{agent_token.class}:#{agent_token.id}:attributes"
+        expect(redis.get(redis_key)).to be_present
+      end
+    end
+
+    def does_db_update
+      expect { subject }.to change { agent_token.reload.read_attribute(:last_used_at) }
+    end
+  end
+end