From fba7ee0fa9b456e2e35f0bd19fb0f20e01476cd0 Mon Sep 17 00:00:00 2001
From: Alejandro Estringana Ruiz <alejandro.estringanaruiz@datadoghq.com>
Date: Fri, 8 Nov 2024 15:19:52 +0100
Subject: [PATCH] Set sampling priority as soon as changed on appsec

---
 appsec/src/extension/tags.c | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/appsec/src/extension/tags.c b/appsec/src/extension/tags.c
index a82c144cc2d..f042c1831af 100644
--- a/appsec/src/extension/tags.c
+++ b/appsec/src/extension/tags.c
@@ -358,12 +358,6 @@ void dd_tags_add_tags(
     // tag _dd.runtime_family
     _set_runtime_family(span);
 
-    if (_force_keep) {
-        dd_trace_set_priority_sampling_on_span_zobj(
-            span, PRIORITY_SAMPLING_USER_KEEP, DD_MECHANISM_MANUAL);
-        mlog(dd_log_debug, "Updated sampling priority to user_keep");
-    }
-
     if (zend_llist_count(&_appsec_json_frags) == 0) {
         if (!server) {
             return;
@@ -406,7 +400,19 @@ void dd_tags_add_tags(
 
 void dd_tags_add_blocked() { _blocked = true; }
 
-void dd_tags_set_sampling_priority() { _force_keep = true; }
+void dd_tags_set_sampling_priority()
+{
+    if (_force_keep) {
+        return;
+    }
+
+    zend_object *nullable span = dd_req_lifecycle_get_cur_span();
+    dd_trace_set_priority_sampling_on_span_zobj(
+        span, PRIORITY_SAMPLING_USER_KEEP, DD_MECHANISM_MANUAL);
+    mlog(dd_log_debug, "Updated sampling priority to user_keep");
+
+    _force_keep = true;
+}
 
 static void _zend_string_release_indirect(void *s)
 {