mongo 2.22.0 → 2.24.0

data/lib/mongo/tracing/open_telemetry/tracer.rb

@@ -0,0 +1,236 @@
+# frozen_string_literal: true
+
+# Copyright (C) 2025-present MongoDB Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an 'AS IS' BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+module Mongo
+  module Tracing
+    module OpenTelemetry
+      # OpenTelemetry tracer for MongoDB operations and commands.
+      # @api private
+      class Tracer
+        # @return [ OpenTelemetry::Trace::Tracer ] the OpenTelemetry tracer implementation
+        #   used to create spans for MongoDB operations and commands.
+        #
+        # @api private
+        attr_reader :otel_tracer
+
+        # Initializes a new OpenTelemetry tracer.
+        #
+        # @param enabled [ Boolean | nil ] whether OpenTelemetry is enabled or not.
+        #   Defaults to nil, which means it will check the environment variable
+        #   OTEL_RUBY_INSTRUMENTATION_MONGODB_ENABLED (values: true/1/yes). If the
+        #   environment variable is not set, OpenTelemetry will be disabled by default.
+        # @param query_text_max_length [ Integer | nil ] maximum length for captured query text.
+        #   Defaults to nil, which means it will check the environment variable
+        #   OTEL_RUBY_INSTRUMENTATION_MONGODB_QUERY_TEXT_MAX_LENGTH. If the environment variable is not set,
+        #   the query text will not be captured.
+        # @param otel_tracer [ OpenTelemetry::Trace::Tracer | nil ] the OpenTelemetry tracer
+        #   implementation to use. Defaults to nil, which means it will use the default tracer
+        #   from OpenTelemetry's tracer provider.
+        def initialize(enabled: nil, query_text_max_length: nil, otel_tracer: nil)
+          @enabled = if enabled.nil?
+                       %w[true 1 yes].include?(ENV['OTEL_RUBY_INSTRUMENTATION_MONGODB_ENABLED']&.downcase)
+                     else
+                       enabled
+                     end
+          check_opentelemetry_loaded
+          @query_text_max_length = if query_text_max_length.nil?
+                                     ENV['OTEL_RUBY_INSTRUMENTATION_MONGODB_QUERY_TEXT_MAX_LENGTH'].to_i
+                                   else
+                                     query_text_max_length
+                                   end
+          @otel_tracer = otel_tracer || initialize_tracer
+          @operation_tracer = OperationTracer.new(@otel_tracer, self)
+          @command_tracer = CommandTracer.new(@otel_tracer, self, query_text_max_length: @query_text_max_length)
+        end
+
+        # Whether OpenTelemetry is enabled or not.
+        #
+        # @return [Boolean] true if OpenTelemetry is enabled, false otherwise.
+        def enabled?
+          @enabled
+        end
+
+        # Trace a MongoDB operation.
+        #
+        # @param operation [Mongo::Operation] The MongoDB operation to trace.
+        # @param operation_context [Mongo::Operation::Context] The context of the operation.
+        # @param op_name [String, nil] An optional name for the operation.
+        # @yield The block representing the operation to be traced.
+        # @return [Object] The result of the operation.
+        def trace_operation(operation, operation_context, op_name: nil, &block)
+          return yield unless enabled?
+
+          @operation_tracer.trace_operation(operation, operation_context, op_name: op_name, &block)
+        end
+
+        # Trace a MongoDB command.
+        #
+        # @param message [Mongo::Protocol::Message] The MongoDB command message to trace.
+        # @param operation_context [Mongo::Operation::Context] The context of the operation.
+        # @param connection [Mongo::Server::Connection] The connection used to send the command
+        # @yield The block representing the command to be traced.
+        # @return [Object] The result of the command.
+        def trace_command(message, operation_context, connection, &block)
+          return yield unless enabled?
+
+          @command_tracer.trace_command(message, operation_context, connection, &block)
+        end
+
+        # Start a transaction span and activate its context.
+        #
+        # @param session [Mongo::Session] The session starting the transaction.
+        def start_transaction_span(session)
+          return unless enabled?
+
+          key = transaction_map_key(session)
+          return unless key
+
+          # Create the transaction span with minimal attributes
+          span = @otel_tracer.start_span(
+            'transaction',
+            attributes: { 'db.system.name' => 'mongodb' },
+            kind: :client
+          )
+
+          # Create a context containing this span
+          context = ::OpenTelemetry::Trace.context_with_span(span)
+
+          # Activate the context and store the token for later detachment
+          token = ::OpenTelemetry::Context.attach(context)
+
+          # Store span, token, and context for later retrieval
+          transaction_span_map[key] = span
+          transaction_token_map[key] = token
+          transaction_context_map[key] = context
+        end
+
+        # Finish a transaction span and deactivate its context.
+        #
+        # @param session [Mongo::Session] The session finishing the transaction.
+        def finish_transaction_span(session)
+          return unless enabled?
+
+          key = transaction_map_key(session)
+          return unless key
+
+          span = transaction_span_map.delete(key)
+          token = transaction_token_map.delete(key)
+          transaction_context_map.delete(key)
+
+          return unless span && token
+
+          begin
+            span.finish
+          ensure
+            ::OpenTelemetry::Context.detach(token)
+          end
+        end
+
+        # Returns the cursor context map for tracking cursor-related OpenTelemetry contexts.
+        #
+        # @return [ Hash ] map of cursor IDs to OpenTelemetry contexts.
+        def cursor_context_map
+          @cursor_context_map ||= {}
+        end
+
+        # Generates a unique key for cursor tracking in the context map.
+        #
+        # @param session [ Mongo::Session ] the session associated with the cursor.
+        # @param cursor_id [ Integer ] the cursor ID.
+        #
+        # @return [ String | nil ] unique key combining session ID and cursor ID, or nil if either is nil.
+        def cursor_map_key(session, cursor_id)
+          return if cursor_id.nil? || session.nil?
+
+          "#{session.session_id['id'].to_uuid}-#{cursor_id}"
+        end
+
+        # Determines the parent OpenTelemetry context for an operation.
+        #
+        # Returns the transaction context if the operation is part of a transaction,
+        # otherwise returns nil. Cursor-based context nesting is not currently implemented.
+        #
+        # @param operation_context [ Mongo::Operation::Context ] the operation context.
+        # @param cursor_id [ Integer ] the cursor ID, if applicable.
+        #
+        # @return [ OpenTelemetry::Context | nil ] parent context or nil.
+        def parent_context_for(operation_context, cursor_id)
+          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
+          end
+        end
+
+        # Returns the transaction context map for tracking active transaction contexts.
+        #
+        # @return [ Hash ] map of transaction keys to OpenTelemetry contexts.
+        def transaction_context_map
+          @transaction_context_map ||= {}
+        end
+
+        # Returns the transaction span map for tracking active transaction spans.
+        #
+        # @return [ Hash ] map of transaction keys to OpenTelemetry spans.
+        def transaction_span_map
+          @transaction_span_map ||= {}
+        end
+
+        # Returns the transaction token map for tracking context attachment tokens.
+        #
+        # @return [ Hash ] map of transaction keys to OpenTelemetry context tokens.
+        def transaction_token_map
+          @transaction_token_map ||= {}
+        end
+
+        # Generates a unique key for transaction tracking.
+        #
+        # Returns nil for implicit sessions or sessions not in a transaction.
+        #
+        # @param session [ Mongo::Session ] the session.
+        #
+        # @return [ String | nil ] unique key combining session ID and transaction number, or nil.
+        def transaction_map_key(session)
+          return if session.nil? || session.implicit? || !session.in_transaction?
+
+          "#{session.session_id['id'].to_uuid}-#{session.txn_num}"
+        end
+
+        private
+
+        def check_opentelemetry_loaded
+          return unless @enabled
+          return if defined?(::OpenTelemetry)
+
+          Logger.logger.warn('OpenTelemetry tracing for MongoDB is enabled, ' \
+                             'but the OpenTelemetry library is not loaded. ' \
+                             'Disabling tracing.')
+          @enabled = false
+        end
+
+        def initialize_tracer
+          return unless enabled?
+
+          ::OpenTelemetry.tracer_provider.tracer(
+            'mongo-ruby-driver',
+            Mongo::VERSION
+          )
+        end
+      end
+    end
+  end
+end

data/lib/mongo/{error/server_api_not_supported.rb → tracing/open_telemetry.rb}

@@ -1,27 +1,27 @@
 # frozen_string_literal: true
-# rubocop:todo all
 
-# Copyright (C) 2021 MongoDB Inc.
+# Copyright (C) 2025-present MongoDB Inc.
 #
-# Licensed under the Apache License, Version 2.0 (the "License");
+# Licensed under the Apache License, Version 2.0 (the 'License');
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
+# distributed under the License is distributed on an 'AS IS' BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
 module Mongo
-  class Error
-
-    # Exception raised when a Client has :server_api configured and an
-    # operation is executed against a pre-3.6 MongoDB server using a legacy
-    # wire protocol message that does not permit sending API parameters.
-    class ServerApiNotSupported < Error
+  module Tracing
+    # This module contains OpenTelemetry tracing functionality for MongoDB operations.
+    module OpenTelemetry
     end
   end
 end
+
+require 'mongo/tracing/open_telemetry/command_tracer'
+require 'mongo/tracing/open_telemetry/operation_tracer'
+require 'mongo/tracing/open_telemetry/tracer'

data/lib/mongo/tracing.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+# Copyright (C) 2025-present MongoDB Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an 'AS IS' BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+module Mongo
+  # Provides OpenTelemetry tracing capabilities for MongoDB operations.
+  module Tracing
+    # Creates a new OpenTelemetry tracer for instrumenting MongoDB operations.
+    #
+    # @param enabled [Boolean, nil] Whether tracing is enabled. Default to nil, which
+    #   means it will check the environment variable OTEL_RUBY_INSTRUMENTATION_MONGODB_ENABLED.
+    #   See +Mongo::Tracing::OpenTelemetry::Tracer+ for details.
+    # @param query_text_max_length [Integer, nil] Maximum length for captured query text. Default to nil,
+    #   which means it will check the environment variable OTEL_RUBY_INSTRUMENTATION_MONGODB_QUERY_TEXT_MAX_LENGTH.
+    #   See +Mongo::Tracing::OpenTelemetry::Tracer+ for details.
+    # @param otel_tracer [OpenTelemetry::Trace::Tracer, nil] Custom OpenTelemetry tracer instance.
+    #
+    # @return [Mongo::Tracing::OpenTelemetry::Tracer] Configured tracer instance.
+    def create_tracer(enabled: nil, query_text_max_length: nil, otel_tracer: nil)
+      OpenTelemetry::Tracer.new(
+        enabled: enabled,
+        query_text_max_length: query_text_max_length,
+        otel_tracer: otel_tracer
+      )
+    end
+    module_function :create_tracer
+  end
+end
+
+require 'mongo/tracing/open_telemetry'