Gauge primitive RubyVM::YJIT.runtime_stats, if YJIT is enabled

.github/workflows/test-yjit.yaml

@@ -0,0 +1,32 @@
+name: Test YJIT
+on: [push]
+jobs:
+  test-yjit:
+    strategy:
+      fail-fast: false
+      matrix:
+        os:
+          - ubuntu-latest
+        ruby:
+          - '3.2'
+          # ADD NEW RUBIES HERE
+    name: Test (${{ matrix.os }}, ${{ matrix.ruby }})
+    runs-on: ${{ matrix.os }}
+    env:
+      RUBYOPT: "--yjit"
+      SKIP_SIMPLECOV: 1
+      DD_INSTRUMENTATION_TELEMETRY_ENABLED: false
+    steps:
+      - uses: actions/checkout@v3
+      # bundler appears to match both prerelease and release rubies when we
+      # want the former only. relax the constraint to allow any version for
+      # head rubies
+      - if: ${{ matrix.ruby == 'head' }}
+        run: sed -i~ -e '/spec\.required_ruby_version/d' ddtrace.gemspec
+      - uses: ruby/setup-ruby@9669f3ee51dc3f4eda8447ab696b3ab19a90d14b # v1.144.0
+        with:
+          ruby-version: ${{ matrix.ruby }}
+          bundler-cache: true # runs 'bundle install' and caches installed gems automatically
+          bundler: latest # needed to fix issue with steep on Ruby 3.0/3.1
+          cache-version: v2 # bump this to invalidate cache
+      - run: bundle exec rake spec:yjit

Rakefile

@@ -94,6 +94,11 @@ namespace :spec do
     t.rspec_opts = args.to_a.join(' ')
   end
 
+  RSpec::Core::RakeTask.new(:yjit) do |t, args|
+    t.pattern = 'spec/datadog/core/runtime/metrics_spec.rb'
+    t.rspec_opts = args.to_a.join(' ')
+  end
+
   # rails_semantic_logger is the dog at the dog park that doesnt play nicely with other
   # logging gems, aka it tries to bite/monkeypatch them, so we have to put it in its own appraisal and rake task
   # in order to isolate its effects for rails logs auto injection

Steepfile

@@ -636,6 +636,9 @@ target :ddtrace do
   ignore 'lib/ddtrace/transport/traces.rb'
   ignore 'lib/ddtrace/version.rb'
 
+  # References `RubyVM::YJIT`, which does not have type information.
+  ignore 'lib/datadog/core/environment/yjit.rb'
+
   library 'pathname'
   library 'cgi'
   library 'logger', 'monitor'

docs/GettingStarted.md

@@ -2780,16 +2780,18 @@ See the [Dogstatsd documentation](https://www.rubydoc.info/github/DataDog/dogsta
 
 The stats are VM specific and will include:
 
-| Name                        | Type    | Description                                              | Available on |
-| --------------------------  | ------- | -------------------------------------------------------- | ------------ |
-| `runtime.ruby.class_count`  | `gauge` | Number of classes in memory space.                       | CRuby        |
-| `runtime.ruby.gc.*`         | `gauge` | Garbage collection statistics: collected from `GC.stat`. | All runtimes |
-| `runtime.ruby.thread_count` | `gauge` | Number of threads.                                       | All runtimes |
+| Name                        | Type    | Description                                                   | Available on       |
+| --------------------------  | ------- | ------------------------------------------------------------- | ------------------ |
+| `runtime.ruby.class_count`  | `gauge` | Number of classes in memory space.                            | CRuby              |
+| `runtime.ruby.gc.*`         | `gauge` | Garbage collection statistics: collected from `GC.stat`.      | All runtimes       |
+| `runtime.ruby.yjit.*`       | `gauge` | YJIT statistics collected from `RubyVM::YJIT.runtime_stats`.  | CRuby (if enabled) |
+| `runtime.ruby.thread_count` | `gauge` | Number of threads.                                            | All runtimes       |
 | `runtime.ruby.global_constant_state`        | `gauge` | Global constant cache generation.                                                                     | CRuby ≤ 3.1                                                                                     |
 | `runtime.ruby.global_method_state`          | `gauge` | [Global method cache generation.](https://tenderlovemaking.com/2015/12/23/inline-caching-in-mri.html) | [CRuby 2.x](https://docs.ruby-lang.org/en/3.0.0/NEWS_md.html#label-Implementation+improvements) |
 | `runtime.ruby.constant_cache_invalidations` | `gauge` | Constant cache invalidations.                                                                         | CRuby ≥ 3.2                                                                                     |
 | `runtime.ruby.constant_cache_misses`        | `gauge` | Constant cache misses.                                                                                | CRuby ≥ 3.2                                                                                     |
 
+
 In addition, all metrics include the following tags:
 
 | Name         | Description                                             |

lib/datadog/core/environment/yjit.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Core
+    module Environment
+      # Reports YJIT primitive runtime statistics.
+      module YJIT
+        module_function
+
+        # Inline code size
+        def inline_code_size
+          ::RubyVM::YJIT.runtime_stats[:inline_code_size]
+        end
+
+        # Outlined code size
+        def outlined_code_size
+          ::RubyVM::YJIT.runtime_stats[:outlined_code_size]
+        end
+
+        # GCed pages
+        def freed_page_count
+          ::RubyVM::YJIT.runtime_stats[:freed_page_count]
+        end
+
+        # GCed code size
+        def freed_code_size
+          ::RubyVM::YJIT.runtime_stats[:freed_code_size]
+        end
+
+        # Live pages
+        def live_page_count
+          ::RubyVM::YJIT.runtime_stats[:live_page_count]
+        end
+
+        # Code GC count
+        def code_gc_count
+          ::RubyVM::YJIT.runtime_stats[:code_gc_count]
+        end
+
+        # Size of memory region allocated for JIT code
+        def code_region_size
+          ::RubyVM::YJIT.runtime_stats[:code_region_size]
+        end
+
+        # Total number of object shapes
+        def object_shape_count
+          ::RubyVM::YJIT.runtime_stats[:object_shape_count]
+        end
+
+        def available?
+          defined?(::RubyVM::YJIT) \
+            && ::RubyVM::YJIT.enabled? \
+            && ::RubyVM::YJIT.respond_to?(:runtime_stats)
+        end
+      end
+    end
+  end
+end

lib/datadog/core/runtime/ext.rb

@@ -21,6 +21,14 @@ module Metrics
           METRIC_GLOBAL_METHOD_STATE = 'runtime.ruby.global_method_state'
           METRIC_CONSTANT_CACHE_INVALIDATIONS = 'runtime.ruby.constant_cache_invalidations'
           METRIC_CONSTANT_CACHE_MISSES = 'runtime.ruby.constant_cache_misses'
+          METRIC_YJIT_CODE_GC_COUNT = 'runtime.ruby.yjit.code_gc_count'
+          METRIC_YJIT_CODE_REGION_SIZE = 'runtime.ruby.yjit.code_region_size'
+          METRIC_YJIT_FREED_CODE_SIZE = 'runtime.ruby.yjit.freed_code_size'
+          METRIC_YJIT_FREED_PAGE_COUNT = 'runtime.ruby.yjit.freed_page_count'
+          METRIC_YJIT_INLINE_CODE_SIZE = 'runtime.ruby.yjit.inline_code_size'
+          METRIC_YJIT_LIVE_PAGE_COUNT = 'runtime.ruby.yjit.live_page_count'
+          METRIC_YJIT_OBJECT_SHAPE_COUNT = 'runtime.ruby.yjit.object_shape_count'
+          METRIC_YJIT_OUTLINED_CODE_SIZE = 'runtime.ruby.yjit.outlined_code_size'
 
           TAG_SERVICE = 'service'
         end

lib/datadog/core/runtime/metrics.rb

@@ -5,6 +5,7 @@
 require_relative '../environment/gc'
 require_relative '../environment/thread_count'
 require_relative '../environment/vm_cache'
+require_relative '../environment/yjit'
 
 module Datadog
   module Core
@@ -78,6 +79,8 @@ def flush
               )
             end
           end
+
+          flush_yjit_stats
         end
 
         def gc_metrics
@@ -134,6 +137,46 @@ def to_metric_name(str)
         def gauge_if_not_nil(metric_name, metric_value)
           gauge(metric_name, metric_value) if metric_value
         end
+
+        def flush_yjit_stats
+          # Only on Ruby >= 3.2
+          try_flush do
+            if Core::Environment::YJIT.available?
+              gauge_if_not_nil(
+                Core::Runtime::Ext::Metrics::METRIC_YJIT_CODE_GC_COUNT,
+                Core::Environment::YJIT.code_gc_count
+              )
+              gauge_if_not_nil(
+                Core::Runtime::Ext::Metrics::METRIC_YJIT_CODE_REGION_SIZE,
+                Core::Environment::YJIT.code_region_size
+              )
+              gauge_if_not_nil(
+                Core::Runtime::Ext::Metrics::METRIC_YJIT_FREED_CODE_SIZE,
+                Core::Environment::YJIT.freed_code_size
+              )
+              gauge_if_not_nil(
+                Core::Runtime::Ext::Metrics::METRIC_YJIT_FREED_PAGE_COUNT,
+                Core::Environment::YJIT.freed_page_count
+              )
+              gauge_if_not_nil(
+                Core::Runtime::Ext::Metrics::METRIC_YJIT_INLINE_CODE_SIZE,
+                Core::Environment::YJIT.inline_code_size
+              )
+              gauge_if_not_nil(
+                Core::Runtime::Ext::Metrics::METRIC_YJIT_LIVE_PAGE_COUNT,
+                Core::Environment::YJIT.live_page_count
+              )
+              gauge_if_not_nil(
+                Core::Runtime::Ext::Metrics::METRIC_YJIT_OBJECT_SHAPE_COUNT,
+                Core::Environment::YJIT.object_shape_count
+              )
+              gauge_if_not_nil(
+                Core::Runtime::Ext::Metrics::METRIC_YJIT_OUTLINED_CODE_SIZE,
+                Core::Environment::YJIT.outlined_code_size
+              )
+            end
+          end
+        end
       end
     end
   end

sig/datadog/core/environment/yjit.rbs

@@ -0,0 +1,18 @@
+module Datadog
+  module Core
+    module Environment
+      module YJIT
+        def self?.inline_code_size: () -> untyped
+        def self?.outlined_code_size: () -> untyped
+        def self?.freed_page_count: () -> untyped
+        def self?.freed_code_size: () -> untyped
+        def self?.live_page_count: () -> untyped
+        def self?.code_gc_count: () -> untyped
+        def self?.code_region_size: () -> untyped
+        def self?.object_shape_count: () -> untyped
+
+        def self?.available?: () -> untyped
+      end
+    end
+  end
+end

spec/datadog/core/runtime/metrics_spec.rb

@@ -191,6 +191,58 @@
           end
         end
       end
+
+      context 'including YJIT stats' do
+        before do
+          skip('This feature is only supported in CRuby') unless PlatformHelpers.mri?
+          skip('Test only runs on Ruby >= 3.2') if RUBY_VERSION < '3.2.'
+        end
+
+        context 'with YJIT enabled and RubyVM::YJIT.stats_enabled? false' do
+          before do
+            unless Datadog::Core::Environment::YJIT.available?
+              skip('Test only runs with YJIT enabled and RubyVM::YJIT.stats_enabled? false')
+            end
+            allow(runtime_metrics).to receive(:gauge)
+          end
+
+          it do
+            flush
+
+            expect(runtime_metrics).to have_received(:gauge)
+              .with(Datadog::Core::Runtime::Ext::Metrics::METRIC_YJIT_CODE_GC_COUNT, kind_of(Numeric))
+              .once
+
+            expect(runtime_metrics).to have_received(:gauge)
+              .with(Datadog::Core::Runtime::Ext::Metrics::METRIC_YJIT_CODE_REGION_SIZE, kind_of(Numeric))
+              .once
+
+            expect(runtime_metrics).to have_received(:gauge)
+              .with(Datadog::Core::Runtime::Ext::Metrics::METRIC_YJIT_FREED_CODE_SIZE, kind_of(Numeric))
+              .once
+
+            expect(runtime_metrics).to have_received(:gauge)
+              .with(Datadog::Core::Runtime::Ext::Metrics::METRIC_YJIT_FREED_PAGE_COUNT, kind_of(Numeric))
+              .once
+
+            expect(runtime_metrics).to have_received(:gauge)
+              .with(Datadog::Core::Runtime::Ext::Metrics::METRIC_YJIT_INLINE_CODE_SIZE, kind_of(Numeric))
+              .once
+
+            expect(runtime_metrics).to have_received(:gauge)
+              .with(Datadog::Core::Runtime::Ext::Metrics::METRIC_YJIT_LIVE_PAGE_COUNT, kind_of(Numeric))
+              .once
+
+            expect(runtime_metrics).to have_received(:gauge)
+              .with(Datadog::Core::Runtime::Ext::Metrics::METRIC_YJIT_OBJECT_SHAPE_COUNT, kind_of(Numeric))
+              .once
+
+            expect(runtime_metrics).to have_received(:gauge)
+              .with(Datadog::Core::Runtime::Ext::Metrics::METRIC_YJIT_OUTLINED_CODE_SIZE, kind_of(Numeric))
+              .once
+          end
+        end
+      end
     end
 
     it_behaves_like 'a flush of all runtime metrics'