datadog 2.1.0 → 2.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: e0297728b2d8779a69af39f3dffa86af6b3765e3f2dd19cb685340ca7ad8e2ac
-  data.tar.gz: 3c7d3506df76790fea0b2e2047d1b2c4919e9d928ccb644086194197125391e7
+  metadata.gz: 5d6610c8ef7e86c023f8a3fca884807bd7e9cf6b84fc6cbdd79b98e8a8762c2e
+  data.tar.gz: 236fafc4b8e2c809552d97c8eb025654c0ef4279685c741ad9504b2af8265b94
 SHA512:
-  metadata.gz: ec2862773c6ee93e5d698d6b586a576413689a84f7d94a621c20088759469e1ad6742e13fc1e7f9250046212cec0e284abf4fb0048e255802c19f378cf7d428c
-  data.tar.gz: 507fbfd3f8879333c11336646e1a56ca453fbaf92c84fb0bae3b0e38435f3ea1ec6b3ddc756f848c76113a4e5832f8097cd521532541bf9ca8fdff6cca1e5ff3
+  metadata.gz: e633db76f69b5d151629cde5e1b7024a6bea43343aa348ea50857d5e74d049d2b6a833be252c6c1c7aee95976805a795e1607a81d42e8c3246126abd9811777d
+  data.tar.gz: fb2bcf3803689d8e499f2266e7d063f1c5b9b7d08bb28b5d27f81bd16e7e373c01c1d5224027d43229a5ca609f66c5179aa521ef272f347d069ca73e3613bdc3

data/CHANGELOG.md

@@ -2,6 +2,64 @@
 
 ## [Unreleased]
 
+## [2.3.0] - 2024-08-22
+
+### Added
+
+* Core: Support agentless telemetry ([#3779][])
+* Tracing: Add support for span events ([#3776][])
+* Tracing: Add tags to enable inferred service dependencies for databases ([#3789][])
+* Tracing: Emit log message and instructions for incompatible Lograge setup ([#3812][], [#3839][])
+* Tracing: Add `append_comment` option to append SQL comment propagation for `mysql`, `pg` and `trilogy` ([#3809][])
+* AppSec: Add threat detection and protection for `graphql` ([#3769][], [#3814][])
+
+### Changed
+
+* Core: Enable crashtracking by default ([#3826][])
+* Profiling: Reduce allocation overhead ([#3805][], [#3797][])
+* Profiling: Speed up stack sampling ([#3837][])
+* Profiling: Upgrade to libdatadog 11 ([#3799][])
+* Profiling: Disable allocation counting feature by default ([#3798][])
+* Profiling: Reduce the maximum biased result for allocation samples ([#3793][])
+* Tracing: Reduce noisy integration logs ([#3785][])
+
+### Fixed
+
+* Tracing: Fix `require` issue for `graphql` ([#3813][])
+* AppSec: Fix an error when parsing http headers with integer value ([#3790][])
+* AppSec: Fix an error when tracking login failure without `user_id` ([#3841][])
+* Fix a syntax error for Ruby < 2.4 during single step instrumentation ([#3795][])
+
+## [2.2.0] - 2024-07-11
+
+### Added
+
+* Tracing: Add `Rails` Runner instrumentation ([#2509][])
+* Tracing: Introduce a new, reworked `GraphQL` tracer to comply with span attributes specification ([#3672][])
+* Tracing: Enhance `ActiveSupport::Cache` instrumentation with `ActiveSupport::Notifications` subscription ([#3772][])
+* Profiling: Track unscaled allocation counts in allocation profiler ([#3770][])
+
+### Changed
+
+* Core: Send Telemetry events in batches ([#3749][])
+* Tracing: Populate spans from `ActiveSupport::Notifications` as early as possible ([#3725][])
+* Profiling: Upgrade to `libdatadog` 10 ([#3753][])
+* Profiling: Optimize `CodeProvenance#record_loaded_files` to avoid allocations ([#3757][])
+
+### Fixed
+
+* Core: Fix Telemetry events blocking main thread ([#3718][])
+* Core: Fix deadlock from Telemetry threads ([#3743][])
+* Tracing: Fix empty log correlation when tracing is disabled ([#3731][])
+* Tracing: Fix HTTP propagation when missing parent span id ([#3730][])
+* Tracing: Ensure `_dd.p.tid` tag with fixed size ([#3729][])
+* OTel: Fix ids encoding/decoding for propagation ([#3709][])
+* Profiling: Workaround Ruby `Dir` returning incorrect results ([#3720][])
+* Profiling: Fix `Phusion Passenger` detection when missing from `Gemfile`/`gems.rb` ([#3750][])
+* Profiling: Fix `rpath` for linking to libdatadog when loading extension ([#3706][])
+* Profiling: Fix incorrect code provenance due to broken JSON monkey patch ([#3695][])
+* Profiling: Fix aggregation of actionview template classes ([#3759][], [#3774][])
+
 ## [2.1.0] - 2024-06-10
 
 ### Added
@@ -2904,7 +2962,9 @@ Release notes: https://github.com/DataDog/dd-trace-rb/releases/tag/v0.3.1
 Git diff: https://github.com/DataDog/dd-trace-rb/compare/v0.3.0...v0.3.1
 
 
-[Unreleased]: https://github.com/DataDog/dd-trace-rb/compare/v2.1.0...master
+[Unreleased]: https://github.com/DataDog/dd-trace-rb/compare/v2.3.0...master
+[2.3.0]: https://github.com/DataDog/dd-trace-rb/compare/v2.2.0...v2.3.0
+[2.2.0]: https://github.com/DataDog/dd-trace-rb/compare/v2.1.0...v2.2.0
 [2.1.0]: https://github.com/DataDog/dd-trace-rb/compare/v2.0.0...v2.1.0
 [2.0.0]: https://github.com/DataDog/dd-trace-rb/compare/v2.0.0.rc1...v2.0.0
 [2.0.0.rc1]: https://github.com/DataDog/dd-trace-rb/compare/v2.0.0.beta2...v2.0.0.rc1
@@ -3901,6 +3961,7 @@ Git diff: https://github.com/DataDog/dd-trace-rb/compare/v0.3.0...v0.3.1
 [#2497]: https://github.com/DataDog/dd-trace-rb/issues/2497
 [#2501]: https://github.com/DataDog/dd-trace-rb/issues/2501
 [#2504]: https://github.com/DataDog/dd-trace-rb/issues/2504
+[#2509]: https://github.com/DataDog/dd-trace-rb/issues/2509
 [#2512]: https://github.com/DataDog/dd-trace-rb/issues/2512
 [#2513]: https://github.com/DataDog/dd-trace-rb/issues/2513
 [#2522]: https://github.com/DataDog/dd-trace-rb/issues/2522
@@ -4283,6 +4344,45 @@ Git diff: https://github.com/DataDog/dd-trace-rb/compare/v0.3.0...v0.3.1
 [#3651]: https://github.com/DataDog/dd-trace-rb/issues/3651
 [#3657]: https://github.com/DataDog/dd-trace-rb/issues/3657
 [#3664]: https://github.com/DataDog/dd-trace-rb/issues/3664
+[#3672]: https://github.com/DataDog/dd-trace-rb/issues/3672
+[#3695]: https://github.com/DataDog/dd-trace-rb/issues/3695
+[#3706]: https://github.com/DataDog/dd-trace-rb/issues/3706
+[#3709]: https://github.com/DataDog/dd-trace-rb/issues/3709
+[#3718]: https://github.com/DataDog/dd-trace-rb/issues/3718
+[#3720]: https://github.com/DataDog/dd-trace-rb/issues/3720
+[#3725]: https://github.com/DataDog/dd-trace-rb/issues/3725
+[#3729]: https://github.com/DataDog/dd-trace-rb/issues/3729
+[#3730]: https://github.com/DataDog/dd-trace-rb/issues/3730
+[#3731]: https://github.com/DataDog/dd-trace-rb/issues/3731
+[#3743]: https://github.com/DataDog/dd-trace-rb/issues/3743
+[#3749]: https://github.com/DataDog/dd-trace-rb/issues/3749
+[#3750]: https://github.com/DataDog/dd-trace-rb/issues/3750
+[#3753]: https://github.com/DataDog/dd-trace-rb/issues/3753
+[#3757]: https://github.com/DataDog/dd-trace-rb/issues/3757
+[#3759]: https://github.com/DataDog/dd-trace-rb/issues/3759
+[#3769]: https://github.com/DataDog/dd-trace-rb/issues/3769
+[#3770]: https://github.com/DataDog/dd-trace-rb/issues/3770
+[#3772]: https://github.com/DataDog/dd-trace-rb/issues/3772
+[#3774]: https://github.com/DataDog/dd-trace-rb/issues/3774
+[#3776]: https://github.com/DataDog/dd-trace-rb/issues/3776
+[#3779]: https://github.com/DataDog/dd-trace-rb/issues/3779
+[#3785]: https://github.com/DataDog/dd-trace-rb/issues/3785
+[#3789]: https://github.com/DataDog/dd-trace-rb/issues/3789
+[#3790]: https://github.com/DataDog/dd-trace-rb/issues/3790
+[#3793]: https://github.com/DataDog/dd-trace-rb/issues/3793
+[#3795]: https://github.com/DataDog/dd-trace-rb/issues/3795
+[#3797]: https://github.com/DataDog/dd-trace-rb/issues/3797
+[#3798]: https://github.com/DataDog/dd-trace-rb/issues/3798
+[#3799]: https://github.com/DataDog/dd-trace-rb/issues/3799
+[#3805]: https://github.com/DataDog/dd-trace-rb/issues/3805
+[#3809]: https://github.com/DataDog/dd-trace-rb/issues/3809
+[#3812]: https://github.com/DataDog/dd-trace-rb/issues/3812
+[#3813]: https://github.com/DataDog/dd-trace-rb/issues/3813
+[#3814]: https://github.com/DataDog/dd-trace-rb/issues/3814
+[#3826]: https://github.com/DataDog/dd-trace-rb/issues/3826
+[#3837]: https://github.com/DataDog/dd-trace-rb/issues/3837
+[#3839]: https://github.com/DataDog/dd-trace-rb/issues/3839
+[#3841]: https://github.com/DataDog/dd-trace-rb/issues/3841
 [@AdrianLC]: https://github.com/AdrianLC
 [@Azure7111]: https://github.com/Azure7111
 [@BabyGroot]: https://github.com/BabyGroot

data/ext/datadog_profiling_loader/extconf.rb

@@ -1,22 +1,22 @@
 # rubocop:disable Style/StderrPuts
 # rubocop:disable Style/GlobalVars
 
-if RUBY_ENGINE != 'ruby' || Gem.win_platform?
+if RUBY_ENGINE != "ruby" || Gem.win_platform?
   $stderr.puts(
-    'WARN: Skipping build of Datadog profiling loader. See Datadog profiling native extension note for details.'
+    "WARN: Skipping build of Datadog profiling loader. See Datadog profiling native extension note for details."
   )
 
-  File.write('Makefile', 'all install clean: # dummy makefile that does nothing')
+  File.write("Makefile", "all install clean: # dummy makefile that does nothing")
   exit
 end
 
-require 'mkmf'
+require "mkmf"
 
 # mkmf on modern Rubies actually has an append_cflags that does something similar
 # (see https://github.com/ruby/ruby/pull/5760), but as usual we need a bit more boilerplate to deal with legacy Rubies
 def add_compiler_flag(flag)
   if try_cflags(flag)
-    $CFLAGS << ' ' << flag
+    $CFLAGS << " " << flag
   else
     $stderr.puts("WARNING: '#{flag}' not accepted by compiler, skipping it")
   end
@@ -24,26 +24,26 @@ end
 
 # Because we can't control what compiler versions our customers use, shipping with -Werror by default is a no-go.
 # But we can enable it in CI, so that we quickly spot any new warnings that just got introduced.
-add_compiler_flag '-Werror' if ENV['DATADOG_GEM_CI'] == 'true'
+add_compiler_flag "-Werror" if ENV["DATADOG_GEM_CI"] == "true"
 
 # Older gcc releases may not default to C99 and we need to ask for this. This is also used:
 # * by upstream Ruby -- search for gnu99 in the codebase
 # * by msgpack, another datadog gem dependency
 #   (https://github.com/msgpack/msgpack-ruby/blob/18ce08f6d612fe973843c366ac9a0b74c4e50599/ext/msgpack/extconf.rb#L8)
-add_compiler_flag '-std=gnu99'
+add_compiler_flag "-std=gnu99"
 
 # Gets really noisy when we include the MJIT header, let's omit it (TODO: Use #pragma GCC diagnostic instead?)
-add_compiler_flag '-Wno-unused-function'
+add_compiler_flag "-Wno-unused-function"
 
 # Allow defining variables at any point in a function
-add_compiler_flag '-Wno-declaration-after-statement'
+add_compiler_flag "-Wno-declaration-after-statement"
 
 # If we forget to include a Ruby header, the function call may still appear to work, but then
 # cause a segfault later. Let's ensure that never happens.
-add_compiler_flag '-Werror-implicit-function-declaration'
+add_compiler_flag "-Werror-implicit-function-declaration"
 
 # Warn on unused parameters to functions. Use `DDTRACE_UNUSED` to mark things as known-to-not-be-used.
-add_compiler_flag '-Wunused-parameter'
+add_compiler_flag "-Wunused-parameter"
 
 # The native extension is not intended to expose any symbols/functions for other native libraries to use;
 # the sole exception being `Init_datadog_profiling_loader` which needs to be visible for Ruby to call it when
@@ -51,14 +51,14 @@ add_compiler_flag '-Wunused-parameter'
 #
 # By setting this compiler flag, we tell it to assume that everything is private unless explicitly stated.
 # For more details see https://gcc.gnu.org/wiki/Visibility
-add_compiler_flag '-fvisibility=hidden'
+add_compiler_flag "-fvisibility=hidden"
 
 # Avoid legacy C definitions
-add_compiler_flag '-Wold-style-definition'
+add_compiler_flag "-Wold-style-definition"
 
 # Enable all other compiler warnings
-add_compiler_flag '-Wall'
-add_compiler_flag '-Wextra'
+add_compiler_flag "-Wall"
+add_compiler_flag "-Wextra"
 
 # Tag the native extension library with the Ruby version and Ruby platform.
 # This makes it easier for development (avoids "oops I forgot to rebuild when I switched my Ruby") and ensures that

data/ext/datadog_profiling_native_extension/clock_id.h

@@ -2,6 +2,7 @@
 
 #include <stdbool.h>
 #include <time.h>
+#include <ruby.h>
 
 // Contains the operating-system specific identifier needed to fetch CPU-time, and a flag to indicate if we failed to fetch it
 typedef struct thread_cpu_time_id {

data/ext/datadog_profiling_native_extension/clock_id_from_pthread.c

@@ -7,11 +7,10 @@
 #include <pthread.h>
 #include <time.h>
 #include <errno.h>
-#include <ruby.h>
 
+#include "clock_id.h"
 #include "helpers.h"
 #include "private_vm_api_access.h"
-#include "clock_id.h"
 #include "time_helpers.h"
 
 // Validate that our home-cooked pthread_id_for() matches pthread_self() for the current thread

data/ext/datadog_profiling_native_extension/clock_id_noop.c

@@ -4,10 +4,9 @@
 // is not available.
 #ifndef HAVE_PTHREAD_GETCPUCLOCKID
 
-#include <ruby.h>
-
 #include "clock_id.h"
 #include "helpers.h"
+#include "datadog_ruby_common.h"
 
 void self_test_clock_id(void) { } // Nothing to check
 

data/ext/datadog_profiling_native_extension/collectors_cpu_and_wall_time_worker.c

@@ -20,7 +20,9 @@
 #define ERR_CLOCK_FAIL "failed to get clock time"
 
 // Maximum allowed value for an allocation weight. Attempts to use higher values will result in clamping.
-unsigned int MAX_ALLOC_WEIGHT = 65535;
+// See https://docs.google.com/document/d/1lWLB714wlLBBq6T4xZyAc4a5wtWhSmr4-hgiPKeErlA/edit#heading=h.ugp0zxcj5iqh
+// (Datadog-only link) for research backing the choice of this value.
+unsigned int MAX_ALLOC_WEIGHT = 10000;
 
 // Used to trigger the execution of Collectors::ThreadState, which implements all of the sampling logic
 // itself; this class only implements the "when to do it" part.
@@ -96,6 +98,7 @@ struct cpu_and_wall_time_worker_state {
   bool no_signals_workaround_enabled;
   bool dynamic_sampling_rate_enabled;
   bool allocation_profiling_enabled;
+  bool allocation_counting_enabled;
   bool skip_idle_samples_for_testing;
   VALUE self_instance;
   VALUE thread_context_collector_instance;
@@ -104,7 +107,6 @@ struct cpu_and_wall_time_worker_state {
   dynamic_sampling_rate_state cpu_dynamic_sampling_rate;
   discrete_dynamic_sampler allocation_sampler;
   VALUE gc_tracepoint; // Used to get gc start/finish information
-  VALUE object_allocation_tracepoint; // Used to get allocation counts and allocation profiling
 
   // These are mutable and used to signal things between the worker thread and other threads
 
@@ -117,7 +119,7 @@ struct cpu_and_wall_time_worker_state {
 
   // Others
 
-  // Used to detect/avoid nested sampling, e.g. when the object_allocation_tracepoint gets triggered by a memory allocation
+  // Used to detect/avoid nested sampling, e.g. when on_newobj_event gets triggered by a memory allocation
   // that happens during another sample.
   bool during_sample;
 
@@ -181,6 +183,7 @@ static VALUE _native_initialize(
   VALUE dynamic_sampling_rate_enabled,
   VALUE dynamic_sampling_rate_overhead_target_percentage,
   VALUE allocation_profiling_enabled,
+  VALUE allocation_counting_enabled,
   VALUE skip_idle_samples_for_testing
 );
 static void cpu_and_wall_time_worker_typed_data_mark(void *state_ptr);
@@ -216,12 +219,28 @@ static void grab_gvl_and_sample(void);
 static void reset_stats_not_thread_safe(struct cpu_and_wall_time_worker_state *state);
 static void sleep_for(uint64_t time_ns);
 static VALUE _native_allocation_count(DDTRACE_UNUSED VALUE self);
-static void on_newobj_event(VALUE tracepoint_data, DDTRACE_UNUSED void *unused);
+static void on_newobj_event(DDTRACE_UNUSED VALUE unused1, DDTRACE_UNUSED void *unused2);
 static void disable_tracepoints(struct cpu_and_wall_time_worker_state *state);
 static VALUE _native_with_blocked_sigprof(DDTRACE_UNUSED VALUE self);
 static VALUE rescued_sample_allocation(VALUE tracepoint_data);
 static void delayed_error(struct cpu_and_wall_time_worker_state *state, const char *error);
 static VALUE _native_delayed_error(DDTRACE_UNUSED VALUE self, VALUE instance, VALUE error_msg);
+static VALUE _native_hold_signals(DDTRACE_UNUSED VALUE self);
+static VALUE _native_resume_signals(DDTRACE_UNUSED VALUE self);
+
+// We're using `on_newobj_event` function with `rb_add_event_hook2`, which requires in its public signature a function
+// with signature `rb_event_hook_func_t` which doesn't match `on_newobj_event`.
+//
+// But in practice, because we pass the `RUBY_EVENT_HOOK_FLAG_RAW_ARG` flag to  `rb_add_event_hook2`, it casts the
+// expected signature into a `rb_event_hook_raw_arg_func_t`:
+// > typedef void (*rb_event_hook_raw_arg_func_t)(VALUE data, const rb_trace_arg_t *arg); (from vm_trace.c)
+// which does match `on_newobj_event`.
+//
+// So TL;DR we're just doing this here to avoid the warning and explain why the apparent mismatch in function signatures.
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wcast-function-type"
+  static const rb_event_hook_func_t on_newobj_event_as_hook = (rb_event_hook_func_t) on_newobj_event;
+#pragma GCC diagnostic pop
 
 // Note on sampler global state safety:
 //
@@ -276,7 +295,7 @@ void collectors_cpu_and_wall_time_worker_init(VALUE profiling_module) {
   // https://bugs.ruby-lang.org/issues/18007 for a discussion around this.
   rb_define_alloc_func(collectors_cpu_and_wall_time_worker_class, _native_new);
 
-  rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_initialize", _native_initialize, 9);
+  rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_initialize", _native_initialize, 10);
   rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_sampling_loop", _native_sampling_loop, 1);
   rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_stop", _native_stop, 2);
   rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_reset_after_fork", _native_reset_after_fork, 1);
@@ -285,7 +304,9 @@ void collectors_cpu_and_wall_time_worker_init(VALUE profiling_module) {
   rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_allocation_count", _native_allocation_count, 0);
   rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_is_running?", _native_is_running, 1);
   rb_define_singleton_method(testing_module, "_native_current_sigprof_signal_handler", _native_current_sigprof_signal_handler, 0);
-  // TODO: Remove `_native_is_running` from `testing_module` once `prof-correctness` has been updated to not need it
+  rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_hold_signals", _native_hold_signals, 0);
+  rb_define_singleton_method(collectors_cpu_and_wall_time_worker_class, "_native_resume_signals", _native_resume_signals, 0);
+  // TODO: Remove `_native_is_running` from `testing_module` (should be in class) once `prof-correctness` has been updated to not need it
   rb_define_singleton_method(testing_module, "_native_is_running?", _native_is_running, 1);
   rb_define_singleton_method(testing_module, "_native_install_testing_signal_handler", _native_install_testing_signal_handler, 0);
   rb_define_singleton_method(testing_module, "_native_remove_testing_signal_handler", _native_remove_testing_signal_handler, 0);
@@ -312,6 +333,8 @@ static const rb_data_type_t cpu_and_wall_time_worker_typed_data = {
 };
 
 static VALUE _native_new(VALUE klass) {
+  long now = monotonic_wall_time_now_ns(RAISE_ON_FAILURE);
+
   struct cpu_and_wall_time_worker_state *state = ruby_xcalloc(1, sizeof(struct cpu_and_wall_time_worker_state));
 
   // Note: Any exceptions raised from this note until the TypedData_Wrap_Struct call will lead to the state memory
@@ -321,13 +344,13 @@ static VALUE _native_new(VALUE klass) {
   state->no_signals_workaround_enabled = false;
   state->dynamic_sampling_rate_enabled = true;
   state->allocation_profiling_enabled = false;
+  state->allocation_counting_enabled = false;
   state->skip_idle_samples_for_testing = false;
   state->thread_context_collector_instance = Qnil;
   state->idle_sampling_helper_instance = Qnil;
   state->owner_thread = Qnil;
   dynamic_sampling_rate_init(&state->cpu_dynamic_sampling_rate);
   state->gc_tracepoint = Qnil;
-  state->object_allocation_tracepoint = Qnil;
 
   atomic_init(&state->should_run, false);
   state->failure_exception = Qnil;
@@ -336,15 +359,12 @@ static VALUE _native_new(VALUE klass) {
   state->during_sample = false;
 
   reset_stats_not_thread_safe(state);
-
-  long now = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE);
-  if (now == 0) {
-    ruby_xfree(state);
-    rb_raise(rb_eRuntimeError, ERR_CLOCK_FAIL);
-  }
-
   discrete_dynamic_sampler_init(&state->allocation_sampler, "allocation", now);
 
+  // Note: As of this writing, no new Ruby objects get created and stored in the state. If that ever changes, remember
+  // to keep them on the stack and mark them with RB_GC_GUARD -- otherwise it's possible for a GC to run and
+  // since the instance representing the state does not yet exist, such objects will not get marked.
+
   return state->self_instance = TypedData_Wrap_Struct(klass, &cpu_and_wall_time_worker_typed_data, state);
 }
 
@@ -358,6 +378,7 @@ static VALUE _native_initialize(
   VALUE dynamic_sampling_rate_enabled,
   VALUE dynamic_sampling_rate_overhead_target_percentage,
   VALUE allocation_profiling_enabled,
+  VALUE allocation_counting_enabled,
   VALUE skip_idle_samples_for_testing
 ) {
   ENFORCE_BOOLEAN(gc_profiling_enabled);
@@ -365,6 +386,7 @@ static VALUE _native_initialize(
   ENFORCE_BOOLEAN(dynamic_sampling_rate_enabled);
   ENFORCE_TYPE(dynamic_sampling_rate_overhead_target_percentage, T_FLOAT);
   ENFORCE_BOOLEAN(allocation_profiling_enabled);
+  ENFORCE_BOOLEAN(allocation_counting_enabled);
   ENFORCE_BOOLEAN(skip_idle_samples_for_testing)
 
   struct cpu_and_wall_time_worker_state *state;
@@ -374,6 +396,7 @@ static VALUE _native_initialize(
   state->no_signals_workaround_enabled = (no_signals_workaround_enabled == Qtrue);
   state->dynamic_sampling_rate_enabled = (dynamic_sampling_rate_enabled == Qtrue);
   state->allocation_profiling_enabled = (allocation_profiling_enabled == Qtrue);
+  state->allocation_counting_enabled = (allocation_counting_enabled == Qtrue);
   state->skip_idle_samples_for_testing = (skip_idle_samples_for_testing == Qtrue);
 
   double total_overhead_target_percentage = NUM2DBL(dynamic_sampling_rate_overhead_target_percentage);
@@ -390,7 +413,6 @@ static VALUE _native_initialize(
   state->thread_context_collector_instance = enforce_thread_context_collector_instance(thread_context_collector_instance);
   state->idle_sampling_helper_instance = idle_sampling_helper_instance;
   state->gc_tracepoint = rb_tracepoint_new(Qnil, RUBY_INTERNAL_EVENT_GC_ENTER | RUBY_INTERNAL_EVENT_GC_EXIT, on_gc_event, NULL /* unused */);
-  state->object_allocation_tracepoint = rb_tracepoint_new(Qnil, RUBY_INTERNAL_EVENT_NEWOBJ, on_newobj_event, NULL /* unused */);
 
   return Qtrue;
 }
@@ -405,7 +427,6 @@ static void cpu_and_wall_time_worker_typed_data_mark(void *state_ptr) {
   rb_gc_mark(state->failure_exception);
   rb_gc_mark(state->stop_thread);
   rb_gc_mark(state->gc_tracepoint);
-  rb_gc_mark(state->object_allocation_tracepoint);
 }
 
 // Called in a background thread created in CpuAndWallTimeWorker#start
@@ -751,7 +772,14 @@ static VALUE release_gvl_and_run_sampling_trigger_loop(VALUE instance) {
   // because they may raise exceptions.
   install_sigprof_signal_handler(handle_sampling_signal, "handle_sampling_signal");
   if (state->gc_profiling_enabled) rb_tracepoint_enable(state->gc_tracepoint);
-  if (state->allocation_profiling_enabled) rb_tracepoint_enable(state->object_allocation_tracepoint);
+  if (state->allocation_profiling_enabled) {
+    rb_add_event_hook2(
+      on_newobj_event_as_hook,
+      RUBY_INTERNAL_EVENT_NEWOBJ,
+      state->self_instance,
+      RUBY_EVENT_HOOK_FLAG_SAFE | RUBY_EVENT_HOOK_FLAG_RAW_ARG)
+    ;
+  }
 
   // Flag the profiler as running before we release the GVL, in case anyone's waiting to know about it
   rb_funcall(instance, rb_intern("signal_running"), 0);
@@ -1032,46 +1060,87 @@ static void sleep_for(uint64_t time_ns) {
 }
 
 static VALUE _native_allocation_count(DDTRACE_UNUSED VALUE self) {
-  bool are_allocations_being_tracked = active_sampler_instance_state != NULL && active_sampler_instance_state->allocation_profiling_enabled;
+  struct cpu_and_wall_time_worker_state *state = active_sampler_instance_state;
+
+  bool are_allocations_being_tracked = state != NULL && state->allocation_profiling_enabled && state->allocation_counting_enabled;
 
   return are_allocations_being_tracked ? ULL2NUM(allocation_count) : Qnil;
 }
 
-// Implements memory-related profiling events. This function is called by Ruby via the `object_allocation_tracepoint`
-// when the RUBY_INTERNAL_EVENT_NEWOBJ event is triggered.
-static void on_newobj_event(VALUE tracepoint_data, DDTRACE_UNUSED void *unused) {
-  // Update thread-local allocation count
-  if (RB_UNLIKELY(allocation_count == UINT64_MAX)) {
-    allocation_count = 0;
-  } else {
-    allocation_count++;
-  }
+#define HANDLE_CLOCK_FAILURE(call) ({ \
+    long _result = (call); \
+    if (_result == 0) { \
+        delayed_error(state, ERR_CLOCK_FAIL); \
+        return; \
+    } \
+    _result; \
+})
 
+// Implements memory-related profiling events. This function is called by Ruby via the `rb_add_event_hook2`
+// when the RUBY_INTERNAL_EVENT_NEWOBJ event is triggered.
+//
+// When allocation sampling is enabled, this function gets called for almost all* objects allocated by the Ruby VM.
+// (*In some weird cases the VM may skip this tracepoint.)
+//
+// At a high level, there's two paths through this function:
+// 1. should_sample == false -> return
+// 2. should_sample == true -> sample
+//
+// On big applications, path 1. is the hottest, since we don't sample every object. So it's quite important for it to
+// be as fast as possible.
+//
+// NOTE: You may be wondering why we don't use any of the arguments to this function. It turns out it's possible to just
+// call `rb_tracearg_from_tracepoint(anything)` anywhere during this function or its callees to get the data, so that's
+// why it's not being passed as an argument.
+static void on_newobj_event(DDTRACE_UNUSED VALUE unused1, DDTRACE_UNUSED void *unused2) {
   struct cpu_and_wall_time_worker_state *state = active_sampler_instance_state; // Read from global variable, see "sampler global state safety" note above
 
   // This should not happen in a normal situation because the tracepoint is always enabled after the instance is set
   // and disabled before it is cleared, but just in case...
   if (state == NULL) return;
 
-  // In a few cases, we may actually be allocating an object as part of profiler sampling. We don't want to recursively
+  if (RB_UNLIKELY(state->allocation_counting_enabled)) {
+    // Update thread-local allocation count
+    if (RB_UNLIKELY(allocation_count == UINT64_MAX)) {
+      allocation_count = 0;
+    } else {
+      allocation_count++;
+    }
+  }
+
+  // In rare cases, we may actually be allocating an object as part of profiler sampling. We don't want to recursively
   // sample, so we just return early
   if (state->during_sample) {
     state->stats.allocations_during_sample++;
     return;
   }
 
-  if (state->dynamic_sampling_rate_enabled) {
-    long now = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE);
-    if (now == 0) {
-      delayed_error(state, ERR_CLOCK_FAIL);
-      return;
-    }
-    if (!discrete_dynamic_sampler_should_sample(&state->allocation_sampler, now)) {
-      state->stats.allocation_skipped++;
-      return;
+  // Hot path: Dynamic sampling rate is usually enabled and the sampling decision is usually false
+  if (RB_LIKELY(state->dynamic_sampling_rate_enabled && !discrete_dynamic_sampler_should_sample(&state->allocation_sampler))) {
+    state->stats.allocation_skipped++;
+
+    coarse_instant now = monotonic_coarse_wall_time_now_ns();
+    HANDLE_CLOCK_FAILURE(now.timestamp_ns);
+
+    bool needs_readjust = discrete_dynamic_sampler_skipped_sample(&state->allocation_sampler, now);
+    if (RB_UNLIKELY(needs_readjust)) {
+      // We rarely readjust, so this is a cold path
+      // Also, while above we used the cheaper monotonic_coarse, for this call we want the regular monotonic call,
+      // which is why we end up getting time "again".
+      discrete_dynamic_sampler_readjust(
+        &state->allocation_sampler, HANDLE_CLOCK_FAILURE(monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE))
+      );
     }
+
+    return;
   }
 
+  // From here on, we've decided to go ahead with the sample, which is way less common than skipping it
+
+  discrete_dynamic_sampler_before_sample(
+    &state->allocation_sampler, HANDLE_CLOCK_FAILURE(monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE))
+  );
+
   // @ivoanjo: Strictly speaking, this is not needed because Ruby should not call the same tracepoint while a previous
   // invocation is still pending, (e.g. it wouldn't call `on_newobj_event` while it's already running), but I decided
   // to keep this here for consistency -- every call to the thread context (other than the special gc calls which are
@@ -1079,7 +1148,7 @@ static void on_newobj_event(VALUE tracepoint_data, DDTRACE_UNUSED void *unused)
   state->during_sample = true;
 
   // Rescue against any exceptions that happen during sampling
-  safely_call(rescued_sample_allocation, tracepoint_data, state->self_instance);
+  safely_call(rescued_sample_allocation, Qnil, state->self_instance);
 
   if (state->dynamic_sampling_rate_enabled) {
     long now = monotonic_wall_time_now_ns(DO_NOT_RAISE_ON_FAILURE);
@@ -1104,9 +1173,7 @@ static void disable_tracepoints(struct cpu_and_wall_time_worker_state *state) {
   if (state->gc_tracepoint != Qnil) {
     rb_tracepoint_disable(state->gc_tracepoint);
   }
-  if (state->object_allocation_tracepoint != Qnil) {
-    rb_tracepoint_disable(state->object_allocation_tracepoint);
-  }
+  rb_remove_event_hook_with_data(on_newobj_event_as_hook, state->self_instance);
 }
 
 static VALUE _native_with_blocked_sigprof(DDTRACE_UNUSED VALUE self) {
@@ -1122,13 +1189,14 @@ static VALUE _native_with_blocked_sigprof(DDTRACE_UNUSED VALUE self) {
   }
 }
 
-static VALUE rescued_sample_allocation(VALUE tracepoint_data) {
+static VALUE rescued_sample_allocation(DDTRACE_UNUSED VALUE unused) {
   struct cpu_and_wall_time_worker_state *state = active_sampler_instance_state; // Read from global variable, see "sampler global state safety" note above
 
   // This should not happen in a normal situation because on_newobj_event already checked for this, but just in case...
   if (state == NULL) return Qnil;
 
-  rb_trace_arg_t *data = rb_tracearg_from_tracepoint(tracepoint_data);
+  // If we're getting called from inside a tracepoint/event hook, Ruby exposes the data using this function.
+  rb_trace_arg_t *data = rb_tracearg_from_tracepoint(Qnil);
   VALUE new_object = rb_tracearg_object(data);
 
   unsigned long allocations_since_last_sample = state->dynamic_sampling_rate_enabled ?
@@ -1136,9 +1204,15 @@ static VALUE rescued_sample_allocation(VALUE tracepoint_data) {
     discrete_dynamic_sampler_events_since_last_sample(&state->allocation_sampler) :
     // if we aren't, then we're sampling every event
     1;
-  // TODO: Signal in the profile that clamping happened?
+
+  // To control bias from sampling, we clamp the maximum weight attributed to a single allocation sample. This avoids
+  // assigning a very large number to a sample, if for instance the dynamic sampling mechanism chose a really big interval.
   unsigned int weight = allocations_since_last_sample > MAX_ALLOC_WEIGHT ? MAX_ALLOC_WEIGHT : (unsigned int) allocations_since_last_sample;
   thread_context_collector_sample_allocation(state->thread_context_collector_instance, weight, new_object);
+  // ...but we still represent the skipped samples in the profile, thus the data will account for all allocations.
+  if (weight < allocations_since_last_sample) {
+    thread_context_collector_sample_skipped_allocation_samples(state->thread_context_collector_instance, allocations_since_last_sample - weight);
+  }
 
   // Return a dummy VALUE because we're called from rb_rescue2 which requires it
   return Qnil;
@@ -1159,3 +1233,17 @@ static VALUE _native_delayed_error(DDTRACE_UNUSED VALUE self, VALUE instance, VA
 
   return Qnil;
 }
+
+// Masks SIGPROF interruptions for the current thread. Please don't use this -- you may end up with incomplete
+// profiling data.
+static VALUE _native_hold_signals(DDTRACE_UNUSED VALUE self) {
+  block_sigprof_signal_handler_from_running_in_current_thread();
+  return Qtrue;
+}
+
+// Unmasks SIGPROF interruptions for the current thread. If there's a pending sample, it'll be triggered inside this
+// method.
+static VALUE _native_resume_signals(DDTRACE_UNUSED VALUE self) {
+  unblock_sigprof_signal_handler_from_running_in_current_thread();
+  return Qtrue;
+}