fix schema sampling logic

DataDog · Jun 14, 2024 · 9d5f3a9 · 9d5f3a9
1 parent 33a2dcd
commit 9d5f3a9
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 11 deletions.
diff --git a/...main/java/datadog/trace/instrumentation/protobuf_java/AbstractMessageInstrumentation.java b/...main/java/datadog/trace/instrumentation/protobuf_java/AbstractMessageInstrumentation.java
@@ -36,7 +36,7 @@ public String hierarchyMarkerType() {
   public ElementMatcher<TypeDescription> hierarchyMatcher() {
     return declaresMethod(named("writeTo"))
         .and(extendsClass(named(hierarchyMarkerType())))
-        .and(not(nameStartsWith("com.google.protobuf")));
+        .and(not(nameStartsWith("com.google.protobuf")).or(named("com.google.protobuf.DynamicMessage")));
   }
 
   @Override

diff --git a/dd-trace-core/src/main/java/datadog/trace/core/datastreams/SchemaSampler.java b/dd-trace-core/src/main/java/datadog/trace/core/datastreams/SchemaSampler.java
@@ -13,21 +13,19 @@ public SchemaSampler() {
   }
 
   public int trySample(long currentTimeMillis) {
-    weight.incrementAndGet();
     if (currentTimeMillis >= lastSampleMillis + SAMPLE_INTERVAL_MILLIS) {
       synchronized (this) {
         if (currentTimeMillis >= lastSampleMillis + SAMPLE_INTERVAL_MILLIS) {
           lastSampleMillis = currentTimeMillis;
-          int currentWeight = weight.get();
-          weight.set(0);
-          return currentWeight;
+          return weight.getAndSet(0);
         }
       }
     }
     return 0;
   }
 
   public boolean canSample(long currentTimeMillis) {
+    weight.incrementAndGet();
     return currentTimeMillis >= lastSampleMillis + SAMPLE_INTERVAL_MILLIS;
   }
 }
diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/SchemaSamplerTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/datastreams/SchemaSamplerTest.groovy
@@ -13,24 +13,18 @@ class SchemaSamplerTest  extends DDCoreSpecification {
     boolean canSample1 = sampler.canSample(currentTimeMillis)
     int weight1 = sampler.trySample(currentTimeMillis)
     boolean canSample2= sampler.canSample(currentTimeMillis + 1000)
-    int weight2 = sampler.trySample(currentTimeMillis + 1000)
     boolean canSample3 = sampler.canSample(currentTimeMillis + 2000)
-    int weight3 = sampler.trySample(currentTimeMillis + 2000)
     boolean canSample4 = sampler.canSample(currentTimeMillis + 30000)
     int weight4 = sampler.trySample(currentTimeMillis + 30000)
     boolean canSample5 = sampler.canSample(currentTimeMillis + 30001)
-    int weight5 = sampler.trySample(currentTimeMillis + 30001)
 
     then:
     canSample1
     weight1 == 1
     !canSample2
-    weight2 == 0
     !canSample3
-    weight3 == 0
     canSample4
     weight4 == 3
     !canSample5
-    weight5 == 0
   }
 }