Commit 38e934e9 authored by Tiger Watson's avatar Tiger Watson Committed by Markus Koller

Extract agent token usage tracking logic into service

parent e1464d93
......@@ -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
# 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
......@@ -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
......
......@@ -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
......@@ -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}" })
......
# 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
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment