Merge pull request #107 from DataDog/rgs/async-sample-mutex-interleav…

…ed-test fix async-sample mutex
DataDog · Jun 12, 2024 · fe9f51f · fe9f51f
2 parents 171a540 + e54df97
commit fe9f51f
Show file tree

Hide file tree

Showing 7 changed files with 83 additions and 40 deletions.
diff --git a/ddprof-lib/src/main/cpp/asyncSampleMutex.h b/ddprof-lib/src/main/cpp/asyncSampleMutex.h
@@ -0,0 +1,38 @@
+#ifndef ASYNCSAMPLEMUTEX_H
+#define ASYNCSAMPLEMUTEX_H
+
+#include "threadLocalData.h"
+
+// controls access to AGCT
+class AsyncSampleMutex {
+private:
+    ThreadLocalData* _threadLocalData;
+    bool _acquired;
+
+    bool try_acquire() {
+        if (_threadLocalData != nullptr && !_threadLocalData->is_unwinding_Java()) {
+            _threadLocalData->set_unwinding_Java(true);
+            return true;
+        }
+        return false;
+    }
+
+public:
+    AsyncSampleMutex(ThreadLocalData* threadLocalData) : _threadLocalData(threadLocalData) {
+        _acquired = try_acquire();
+    }
+
+    AsyncSampleMutex(AsyncSampleMutex& other) = delete;
+
+    ~AsyncSampleMutex() {
+        if (_acquired) {
+            _threadLocalData->set_unwinding_Java(false);
+        }
+    }
+
+    bool acquired() {
+        return _acquired;
+    }
+};
+
+#endif //ASYNCSAMPLEMUTEX_H
diff --git a/ddprof-lib/src/main/cpp/profiler.cpp b/ddprof-lib/src/main/cpp/profiler.cpp
@@ -45,6 +45,7 @@
 #include "vmStructs.h"
 #include "context.h"
 #include "counters.h"
+#include "asyncSampleMutex.h"
 
 
 // The instance is not deleted on purpose, since profiler structures
@@ -683,7 +684,7 @@ void Profiler::recordSample(void* ucontext, u64 counter, int tid, jint event_typ
             int java_frames = 0;
             {
                 // Async events
-                AsyncSampleMutex mutex;
+                AsyncSampleMutex mutex(ProfiledThread::current());
                 if (mutex.acquired()) {
                     java_frames = getJavaTraceAsync(ucontext, frames + num_frames, _max_stack_depth, &java_ctx,
                                                     &truncated);

diff --git a/ddprof-lib/src/main/cpp/profiler.h b/ddprof-lib/src/main/cpp/profiler.h
@@ -59,34 +59,6 @@ struct CallTraceBuffer {
     ASGCT_CallFrame _asgct_frames[1];
 };
 
-// controls access to AGCT
-class AsyncSampleMutex {
-private:
-    bool _acquired;
-    bool try_set(bool flag) {
-        ProfiledThread* current = ProfiledThread::current();
-        if (current != NULL) {
-            bool was_set = current->is_unwinding_Java();
-            current->set_unwinding_Java(flag);
-            return !was_set;
-        }
-        return false;
-    }
-public:
-    AsyncSampleMutex() {
-        _acquired = try_set(true);
-    }
-    AsyncSampleMutex(AsyncSampleMutex& other) = delete;
-    ~AsyncSampleMutex() {
-        if (_acquired) {
-            try_set(false);
-        }
-    }
-    bool acquired() {
-        return _acquired;
-    }
-};
-
 
 class FrameName;
 class NMethod;

diff --git a/ddprof-lib/src/main/cpp/thread.h b/ddprof-lib/src/main/cpp/thread.h
@@ -8,8 +8,9 @@
 #include <vector>
 #include "os.h"
 #include <jvmti.h>
+#include "threadLocalData.h"
 
-class ProfiledThread {
+class ProfiledThread : public ThreadLocalData {
   private:
     static pthread_key_t _tls_key;
     static int _buffer_size;
@@ -32,18 +33,17 @@ class ProfiledThread {
     u32 _call_trace_id;
     u32 _recording_epoch;
     u64 _span_id;
-    bool _unwinding_java;
 
     ProfiledThread(int buffer_pos, int tid) :
+        ThreadLocalData(),
         _buffer_pos(buffer_pos),
         _tid(tid),
         _cpu_epoch(0),
         _wall_epoch(0),
         _pc(0),
         _call_trace_id(0),
         _recording_epoch(0),
-        _span_id(0),
-        _unwinding_java(false){};
+        _span_id(0) {};
 
     void releaseFromBuffer();
   public:
@@ -61,12 +61,6 @@ class ProfiledThread {
 
     static ProfiledThread* current();
     static int currentTid();
-    bool is_unwinding_Java() {
-        return _unwinding_java;
-    }
-    void set_unwinding_Java(bool flag) {
-        _unwinding_java = flag;
-    }
 
     inline int tid() {
         return _tid;

diff --git a/ddprof-lib/src/main/cpp/threadLocalData.h b/ddprof-lib/src/main/cpp/threadLocalData.h
@@ -0,0 +1,18 @@
+#ifndef THREADLOCALDATA_H
+#define THREADLOCALDATA_H
+
+class ThreadLocalData {
+protected:
+    bool _unwinding_Java;
+public:
+    ThreadLocalData() : _unwinding_Java(false) {}
+    virtual bool is_unwinding_Java() final {
+        return _unwinding_Java;
+    }
+
+    virtual void set_unwinding_Java(bool unwinding_Java) final {
+        _unwinding_Java = unwinding_Java;
+    }
+};
+
+#endif //THREADLOCALDATA_H
diff --git a/ddprof-lib/src/test/cpp/ddprof_ut.cpp b/ddprof-lib/src/test/cpp/ddprof_ut.cpp
@@ -1,12 +1,14 @@
     #include <gtest/gtest.h>
 
+    #include "asyncSampleMutex.h"
     #include "buffers.h"
     #include "context.h"
     #include "counters.h"
     #include "mutex.h"
     #include "os.h"
     #include "threadFilter.h"
     #include "threadInfo.h"
+    #include "threadLocalData.h"
     #include <vector>
 
     ssize_t callback(char* ptr, int len) {
@@ -172,6 +174,24 @@
         ASSERT_EQ(0, info.size());
     }
 
+    TEST(AsyncSampleMutex, testAsyncSampleMutexInterleaving) {
+        ThreadLocalData data;
+        EXPECT_FALSE(data.is_unwinding_Java());
+        {
+            AsyncSampleMutex first(&data);
+            EXPECT_TRUE(first.acquired());
+            EXPECT_TRUE(data.is_unwinding_Java());
+            {
+                AsyncSampleMutex second(&data);
+                EXPECT_FALSE(second.acquired());
+                EXPECT_TRUE(first.acquired());
+                EXPECT_TRUE(data.is_unwinding_Java());
+            }
+            EXPECT_TRUE(first.acquired());
+        }
+        EXPECT_FALSE(data.is_unwinding_Java());
+    }
+
     int main(int argc, char **argv) {
       ::testing::InitGoogleTest(&argc, argv);
       return RUN_ALL_TESTS();

diff --git a/ddprof-lib/test.sh b/ddprof-lib/test.sh
@@ -22,7 +22,7 @@ function build_and_test() {
   cd ${HERE}/src/test
   cmake -DCMAKE_BUILD_TYPE=${BUILD_TYPE} -Wno-dev -S . -B ${TARGET}
   cmake --build ${TARGET}
-  cd ${TARGET} && ctest
+  cd ${TARGET} && ctest --output-on-failure
 }
 
 if [ -z "${BUILD_TYPE}" ]; then