From b1372e880cdfca174a44076f95e4ec8e96647c65 Mon Sep 17 00:00:00 2001
From: Sebastian Serth <Sebastian.Serth@hpi.de>
Date: Fri, 10 Feb 2023 19:16:52 +0100
Subject: [PATCH] Enable Sentry instrumentation for WebSocket connection

---
 lib/runner/connection.rb | 71 ++++++++++++++++++++++++++++++++++++----
 1 file changed, 65 insertions(+), 6 deletions(-)

diff --git a/lib/runner/connection.rb b/lib/runner/connection.rb
index 9ab3e03e..991c7a2c 100644
--- a/lib/runner/connection.rb
+++ b/lib/runner/connection.rb
@@ -8,6 +8,8 @@ class Runner::Connection
   EVENTS = %i[start exit stdout stderr].freeze
   WEBSOCKET_MESSAGE_TYPES = %i[start stdout stderr error timeout exit].freeze
   BACKEND_OUTPUT_SCHEMA = JSONSchemer.schema(JSON.parse(File.read('lib/runner/backend-output.schema.json')))
+  SENTRY_OP_NAME = 'websocket.client'
+  SENTRY_BREADCRUMB_CATEGORY = 'net.websocket'
 
   # @!attribute start_callback
   # @!attribute exit_callback
@@ -18,7 +20,12 @@ class Runner::Connection
 
   def initialize(url, strategy, event_loop, locale = I18n.locale)
     Rails.logger.debug { "#{Time.zone.now.getutc.inspect}: Opening connection to #{url}" }
-    @socket = Faye::WebSocket::Client.new(url, [], strategy.class.websocket_header)
+
+    sentry_transaction = Sentry.get_current_scope&.get_span
+    sentry_span = sentry_transaction&.start_child(op: SENTRY_OP_NAME, start_timestamp: Sentry.utc_now.to_f)
+    http_headers = strategy.class.websocket_header.merge sentry_trace_header(sentry_span)
+
+    @socket = Faye::WebSocket::Client.new(url, [], http_headers)
     @strategy = strategy
     @status = :new
     @event_loop = event_loop
@@ -31,7 +38,10 @@ class Runner::Connection
     %i[open message error close].each do |event_type|
       @socket.on(event_type) do |event|
         # The initial locale when establishing the connection is used for all callbacks
-        I18n.with_locale(@locale) { __send__(:"on_#{event_type}", event) }
+        I18n.with_locale(@locale) do
+          clone_sentry_hub_from_sentry_span(sentry_span)
+          __send__(:"on_#{event_type}", event, sentry_span)
+        end
       end
     end
 
@@ -103,7 +113,7 @@ class Runner::Connection
   # These callbacks are executed based on events indicated by Faye WebSockets and are
   # independent of the JSON specification that is used within the WebSocket once established.
 
-  def on_message(raw_event)
+  def on_message(raw_event, _sentry_span)
     Rails.logger.debug { "#{Time.zone.now.getutc.inspect}: Receiving from #{@socket.url}: #{raw_event.data.inspect}" }
     event = decode(raw_event.data)
     return unless BACKEND_OUTPUT_SCHEMA.valid?(event)
@@ -118,21 +128,23 @@ class Runner::Connection
     end
   end
 
-  def on_open(_event)
+  def on_open(_event, _sentry_span)
     Rails.logger.debug { "#{Time.zone.now.getutc.inspect}: Established connection to #{@socket.url}" }
     @status = :established
     @start_callback.call
   end
 
-  def on_error(event)
+  def on_error(event, _sentry_span)
     # In case of an WebSocket error, the connection will be closed by Faye::WebSocket::Client automatically.
     # Thus, no further handling is required here (the user will get notified).
     @status = :error
     @error = Runner::Error::Unknown.new("The WebSocket connection to #{@socket.url} was closed with an error: #{event.message}")
   end
 
-  def on_close(_event)
+  def on_close(event, sentry_span)
     Rails.logger.debug { "#{Time.zone.now.getutc.inspect}: Closing connection to #{@socket.url} with status: #{@status}" }
+    record_sentry_breadcrumb(event)
+    end_sentry_span(sentry_span, event)
     flush_buffers
 
     # Depending on the status, we might want to destroy the runner at management.
@@ -208,4 +220,51 @@ class Runner::Connection
     # The runner management stopped the execution as the permitted execution time was exceeded.
     # We set the status here and wait for the connection to be closed (by the runner management).
   end
+
+  # The methods below are inspired by the Sentry::Net:HTTP class
+  # and adapted to the Websocket protocol running with EventMachine.
+
+  def clone_sentry_hub_from_sentry_span(sentry_span)
+    Thread.current.thread_variable_set(Sentry::THREAD_LOCAL, sentry_span.transaction.hub) if sentry_span
+  end
+
+  def sentry_trace_header(sentry_span)
+    return {} unless sentry_span
+
+    http_headers = {}
+    client = Sentry.get_current_client
+
+    trace = client.generate_sentry_trace(sentry_span)
+    http_headers[Sentry::SENTRY_TRACE_HEADER_NAME] = trace if trace
+
+    baggage = client.generate_baggage(sentry_span)
+    http_headers[Sentry::BAGGAGE_HEADER_NAME] = baggage if baggage.present?
+
+    {
+      headers: http_headers,
+    }
+  end
+
+  def end_sentry_span(sentry_span, event)
+    return unless sentry_span
+
+    sentry_span.set_description("WebSocket #{@socket.url}")
+    sentry_span.set_data(:status, event.code.to_i)
+    sentry_span.finish(end_timestamp: Sentry.utc_now.to_f)
+  end
+
+  def record_sentry_breadcrumb(event)
+    return unless Sentry.initialized? && Sentry.configuration.breadcrumbs_logger.include?(:http_logger)
+
+    crumb = Sentry::Breadcrumb.new(
+      level: :info,
+      category: SENTRY_BREADCRUMB_CATEGORY,
+      type: :info,
+      data: {
+        status: event.code.to_i,
+        url: @socket.url,
+      }
+    )
+    Sentry.add_breadcrumb(crumb)
+  end
 end