wip: transactions

Author: comandeo-mongo

lib/mongo/session.rb

@@ -625,6 +625,8 @@ def start_transaction(options = nil)
       @state = STARTING_TRANSACTION_STATE
       @already_committed = false
 
+      tracer.start_transaction_span(self)
+
       # This method has no explicit return value.
       # We could return nil here but true indicates to the user that the
       # operation succeeded. This is intended for interactive use.
@@ -707,6 +709,7 @@ def commit_transaction(options=nil)
           end
         end
       ensure
+        tracer.finish_transaction_span(self)
         @state = TRANSACTION_COMMITTED_STATE
         @committing_transaction = false
       end
@@ -781,6 +784,7 @@ def abort_transaction(options = nil)
         @state = TRANSACTION_ABORTED_STATE
         raise
       ensure
+        tracer.finish_transaction_span(self)
         @aborting_transaction = false
       end
 

lib/mongo/tracing/open_telemetry/command_tracer.rb

@@ -24,7 +24,6 @@ class CommandTracer
         extend Forwardable
 
         def_delegators :@parent_tracer,
-          :cursor_context_map,
           :parent_context_for,
           :transaction_context_map,
           :transaction_map_key
@@ -36,7 +35,7 @@ def initialize(otel_tracer, parent_tracer, query_text_max_length: 0)
         end
 
         def trace_command(message, operation_context, connection)
-          parent_context = parent_context_for(operation_context, cursor_id(message))
+          parent_context = parent_context_for(operation_context)
           span = @otel_tracer.start_span(
             query_summary(message),
             attributes: span_attributes(message, connection),

lib/mongo/tracing/open_telemetry/operation_tracer.rb

@@ -24,7 +24,6 @@ class OperationTracer
         extend Forwardable
 
         def_delegators :@parent_tracer,
-          :cursor_context_map,
           :parent_context_for,
           :transaction_context_map,
           :transaction_map_key
@@ -35,7 +34,7 @@ def initialize(otel_tracer, parent_tracer)
         end
 
         def trace_operation(operation, operation_context, op_name: nil)
-          parent_context = parent_context_for(operation_context, operation.cursor_id)
+          parent_context = parent_context_for(operation_context)
           span = @otel_tracer.start_span(
             operation_span_name(operation, op_name),
             attributes: span_attributes(operation, op_name),
@@ -55,6 +54,32 @@ def trace_operation(operation, operation_context, op_name: nil)
           span&.finish
         end
 
+        def start_transaction_span(session)
+          raise ArgumentError, 'Session must have a transaction' unless session.in_transaction?
+          key = transaction_map_key(session)
+          raise ArgumentError, 'Transaction already started' if transaction_context_map[key]
+
+          span = @otel_tracer.start_span(
+            'transaction',
+            attributes: {
+              'db.system' => 'mongodb',
+              'db.operation.name' => 'transaction',
+            },
+            kind: :client
+          )
+          transaction_context_map[key] = ::OpenTelemetry::Trace.context_with_span(span)
+        end
+
+        def finish_transaction_span(session)
+          raise ArgumentError, 'Session must have a transaction' unless session.in_transaction?
+          key = transaction_map_key(session)
+          context = transaction_context_map.delete(key)
+          return unless context
+
+          span = ::OpenTelemetry::Trace.current_span(context)
+          span&.finish
+        end
+
         private
 
         def operation_name(operation, op_name = nil)
@@ -79,14 +104,7 @@ def span_attributes(operation, op_name)
         def process_cursor_context(result, cursor_id, context, span)
           return unless result.is_a?(Cursor)
 
-          if result.id.zero?
-            # If the cursor is closed, remove it from the context map.
-            cursor_context_map.delete(cursor_id)
-          elsif result.id && cursor_id.nil?
-            # New cursor created, store its context.
-            cursor_context_map[result.id] = context
-            span.set_attribute('db.mongodb.cursor_id', result.id)
-          end
+          span.set_attribute('db.mongodb.cursor_id', result.id) if result.id.positive?
         end
 
         def collection_name(operation)

lib/mongo/tracing/open_telemetry/tracer.rb

@@ -61,29 +61,23 @@ def trace_operation(operation, operation_context, op_name: nil, &block)
           @operation_tracer.trace_operation(operation, operation_context, op_name: op_name, &block)
         end
 
-        def trace_command(message, operation_context, connection, &block)
-          return yield unless enabled?
-
-          @command_tracer.trace_command(message, operation_context, connection, &block)
+        def start_transaction_span(session)
+          @operation_tracer.start_transaction_span(session)
         end
 
-        def cursor_context_map
-          @cursor_context_map ||= {}
+        def finish_transaction_span(session)
+          @operation_tracer.finish_transaction_span(session)
         end
 
-        def cursor_map_key(session, cursor_id)
-          return if cursor_id.nil? || session.nil?
+        def trace_command(message, operation_context, connection, &block)
+          return yield unless enabled?
 
-          "#{session.session_id['id'].to_uuid}-#{cursor_id}"
+          @command_tracer.trace_command(message, operation_context, connection, &block)
         end
 
-        def parent_context_for(operation_context, cursor_id)
+        def parent_context_for(operation_context)
           if (key = transaction_map_key(operation_context.session))
             transaction_context_map[key]
-          elsif (_key = cursor_map_key(operation_context.session, cursor_id))
-            # We return nil here unless we decide how to nest cursor operations.
-            nil
-            # cursor_context_map[key]
           end
         end
 
@@ -92,9 +86,10 @@ def transaction_context_map
         end
 
         def transaction_map_key(session)
+          byebug
           return if session.nil? || session.implicit? || !session.in_transaction?
 
-          "#{session.session_id['id'].to_uuid}-#{session.txn_num}"
+          session.session_id['id'].to_uuid.to_s
         end
 
         private