Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 32 additions & 30 deletions app/services/heartbeat_import_service.rb
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,15 @@ def self.import_from_file(file_content, user, on_progress: nil, progress_interva
heartbeat_batch.clear
end

handler = HeartbeatSaxHandler.new do |hb|
handler = HeartbeatStreamHandler.new do |hb|
total_count += 1
on_progress&.call(total_count) if progress_interval.positive? && (total_count % progress_interval).zero?

heartbeat_batch << hb
flush.call if heartbeat_batch.size >= BATCH_SIZE
end

Oj.saj_parse(handler, file_content)
Oj.sc_parse(handler, file_content)
on_progress&.call(total_count)

raise StandardError, "Expected a heartbeat export JSON file." if total_count.zero?
Expand All @@ -40,54 +40,56 @@ def self.import_from_file(file_content, user, on_progress: nil, progress_interva
skipped_count: total_count - imported_count, errors: errors + [ e.message ] }
end

class HeartbeatSaxHandler < Oj::Saj
# Retain only the current heartbeat, not the surrounding dump or day arrays.
# Unlike SAJ, ScHandler preserves embedded NULs for ingestion to sanitize.
class HeartbeatStreamHandler < Oj::ScHandler
def initialize(&block)
@block = block
@depth = 0
@current_heartbeat = nil
@heartbeat_array_depths = []
@field_array_stack = []
end

def hash_start(key)
@current_heartbeat = {} if inside_heartbeat_array? && @depth == @heartbeat_array_depths.last + 1
def hash_start
@key = nil
@depth += 1
{} if @heartbeat_array_depths.any?
end

def hash_end(key)
def hash_end
@depth -= 1
if inside_heartbeat_array? && @depth == @heartbeat_array_depths.last + 1 && @current_heartbeat
@block.call(@current_heartbeat)
@current_heartbeat = nil
end
end

def array_start(key)
@heartbeat_array_depths << @depth if key == "heartbeats"
if @current_heartbeat && key.present?
@current_heartbeat[key] = []
@field_array_stack << key
def hash_key(key)
@key = key
end

def hash_set(hash, key, value)
hash[key] = value if hash
end

def array_start
container = if @key == "heartbeats" && @heartbeat_array_depths.empty?
@heartbeat_array_depths << @depth
:heartbeats
elsif @heartbeat_array_depths.any?
[]
end
@key = nil
@depth += 1
container
end

def array_end(key)
def array_end
@depth -= 1
@heartbeat_array_depths.pop if key == "heartbeats" && @heartbeat_array_depths.last == @depth
@field_array_stack.pop if @field_array_stack.last == key
@heartbeat_array_depths.pop if @heartbeat_array_depths.last == @depth
end

def add_value(value, key)
return unless @current_heartbeat
if key
@current_heartbeat[key] = value
elsif @field_array_stack.any?
@current_heartbeat[@field_array_stack.last] << value
def array_append(array, value)
if array == :heartbeats
@block.call(value)
elsif array
array << value
end
end

private

def inside_heartbeat_array? = @heartbeat_array_depths.any?
end
end
5 changes: 3 additions & 2 deletions app/services/heartbeat_ingest.rb
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ def initialize(user:, mode:, heartbeats:, request_context: {}, user_agents_by_id
@mode = mode
@heartbeats = heartbeats
@request_context = request_context.with_indifferent_access
@user_agents_by_id = user_agents_by_id
@user_agents_by_id = user_agents_by_id.transform_keys { |id| strip_null_bytes(id.to_s) }
@schedule_rollup_refresh = schedule_rollup_refresh
end

Expand Down Expand Up @@ -237,7 +237,8 @@ def ingest_import

def normalize_imported_heartbeat(heartbeat, placeholder_state: { contexts: {}, last_project: nil })
hb = heartbeat.respond_to?(:with_indifferent_access) ? heartbeat.with_indifferent_access : heartbeat.to_h.with_indifferent_access
user_agent_info = (@user_agents_by_id[hb[:user_agent_id].to_s] || {}).with_indifferent_access
hb = strip_null_bytes(hb)
user_agent_info = strip_null_bytes((@user_agents_by_id[hb[:user_agent_id].to_s] || {}).with_indifferent_access)
Comment thread
skyfallwastaken marked this conversation as resolved.
resolved_user_agent = hb[:user_agent].presence || user_agent_info[:value].presence || hb[:user_agent_id].presence
parsed_user_agent = parse_user_agent(resolved_user_agent, category: hb[:category])
derived_ai_editor = parsed_user_agent[:editor].presence if parsed_user_agent[:ai_model].present?
Expand Down
23 changes: 23 additions & 0 deletions test/services/heartbeat_import_service_test.rb
Original file line number Diff line number Diff line change
@@ -1,6 +1,29 @@
require "test_helper"

class HeartbeatImportServiceTest < ActiveSupport::TestCase
test "sanitizes null bytes without losing valid batch rows and deduplicates replays" do
user = create(:user)
rows = [
{ entity: "first.rb", project: "api", time: 1_700_000_000.0, type: "file" },
{ entity: "sec\0ond.rb", project: "api", branch: "ma\0in", dependencies: [ "ra\0ils" ], time: 1_700_000_060.0, type: "file" }
]

result = HeartbeatImportService.import_from_file(StringIO.new({ heartbeats: rows }.to_json), user)
assert result[:success], result[:error]
assert_equal 2, result[:imported_count]
heartbeat = user.heartbeats.find_by!(entity: "second.rb")
assert_equal "main", heartbeat.branch
assert_equal [ "rails" ], heartbeat.dependencies

sanitized = [ rows.first, rows.second.merge(entity: "second.rb", branch: "main", dependencies: [ "rails" ]) ]
[ rows, sanitized ].each do |replay|
result = HeartbeatImportService.import_from_file({ heartbeats: replay }.to_json, user)
assert result[:success], result[:error]
assert_equal 0, result[:imported_count]
end
assert_equal 2, user.heartbeats.count
end

test "deduplicates imported heartbeats by fields hash" do
user = create(:user)
file_content = {
Expand Down
24 changes: 24 additions & 0 deletions test/services/heartbeat_ingest_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,30 @@ class HeartbeatIngestTest < ActiveSupport::TestCase
ActiveJob::Base.queue_adapter = @original_queue_adapter
end

test "import resolves sanitised user agent identifiers and deduplicates replays" do
user = create(:user)
payload = { time: 1_700_000_000.0, entity: "main.rb", type: "file", user_agent_id: "agent\0-1" }
metadata = { value: "custom\0-client", editor: "ze\0d", os: "lin\0ux" }
agents = { "agent\0-1" => metadata }

result = HeartbeatIngest.call(user:, mode: :import, heartbeats: [ payload ], user_agents_by_id: agents)
assert_equal 1, result.persisted_count
heartbeat = user.heartbeats.sole
assert_equal "custom-client", heartbeat.user_agent
assert_equal "zed", heartbeat.editor
assert_equal "linux", heartbeat.operating_system
assert_equal [ "agent\0-1" ], agents.keys

[ agents, { "agent-1" => metadata } ].each do |replay_agents|
assert_no_difference("user.heartbeats.count") do
replay = HeartbeatIngest.call(user:, mode: :import,
heartbeats: [ payload.merge(user_agent_id: "agent-1") ], user_agents_by_id: replay_agents)
assert_equal 1, replay.duplicate_count
assert_equal 0, replay.failed_count
end
end
end

test "direct and imported heartbeats share authoritative rules and deduplicate corrected replays" do
expected = {
".env" => "Dotenv", ".gitignore" => "Ignore List", ".rspec" => "Option List",
Expand Down
Loading