improve performance using Unsafe to activate/deactivate thread filter

ddprof-lib/src/main/cpp/javaApi.cpp

@@ -406,3 +406,24 @@ Java_com_datadoghq_profiler_JVMAccess_healthCheck0(JNIEnv *env,
                                                          jobject unused) {
   return true;
 }
+
+extern "C" DLLEXPORT jlong JNICALL
+Java_com_datadoghq_profiler_ActiveBitmap_bitmapAddressFor0(JNIEnv *env,
+		                                                    jclass unused,
+						                                    jint tid) {
+  u64* bitmap = Profiler::instance()->threadFilter()->bitmapAddressFor((int)tid);
+  return (jlong)bitmap;
+}
+
+extern "C" DLLEXPORT jboolean JNICALL
+Java_com_datadoghq_profiler_ActiveBitmap_isActive0(JNIEnv *env,
+                                                   jclass unused,
+                                                   jint tid) {
+  return Profiler::instance()->threadFilter()->accept((int)tid) ? JNI_TRUE : JNI_FALSE;
+}
+
+extern "C" DLLEXPORT jlong JNICALL
+Java_com_datadoghq_profiler_ActiveBitmap_getActiveCountAddr0(JNIEnv *env,
+                                                              jclass unused) {
+  return (jlong)Profiler::instance()->threadFilter()->addressOfSize();
+}

ddprof-lib/src/main/cpp/reverse_bits.h

@@ -0,0 +1,23 @@
+//
+// Borrow the implementation from openjdk
+// https://github.com/openjdk/jdk/blob/master/src/hotspot/share/utilities/reverse_bits.hpp
+//
+
+#ifndef REVERSE_BITS_H
+#define REVERSE_BITS_H
+#include "arch_dd.h"
+#include <stdint.h>
+
+static constexpr u32 rep_5555 = static_cast<u32>(UINT64_C(0x5555555555555555));
+static constexpr u32 rep_3333 = static_cast<u32>(UINT64_C(0x3333333333333333));
+static constexpr u32 rep_0F0F = static_cast<u32>(UINT64_C(0x0F0F0F0F0F0F0F0F));
+
+inline u16 reverse16(u16 v) {
+  u32 x = static_cast<u32>(v);
+  x = ((x & rep_5555) << 1) | ((x >> 1) & rep_5555);
+  x = ((x & rep_3333) << 2) | ((x >> 2) & rep_3333);
+  x = ((x & rep_0F0F) << 4) | ((x >> 4) & rep_0F0F);
+  return __builtin_bswap16(static_cast<u16>(x));
+}
+
+#endif //REVERSE_BITS_H

ddprof-lib/src/main/cpp/threadFilter.cpp

@@ -17,8 +17,11 @@
 #include "threadFilter.h"
 #include "counters.h"
 #include "os.h"
+#include "reverse_bits.h"
+#include <cassert>
 #include <stdlib.h>
 #include <string.h>
+#include <iostream>
 
 void trackPage() {
   Counters::increment(THREAD_FILTER_PAGES, 1);
@@ -85,32 +88,61 @@ void ThreadFilter::clear() {
   _size = 0;
 }
 
-bool ThreadFilter::accept(int thread_id) {
-  u64 *b = bitmap(thread_id);
-  return b != NULL && (word(b, thread_id) & (1ULL << (thread_id & 0x3f)));
+// The mapping has to be reversible: f(f(x)) == x
+int ThreadFilter::mapThreadId(int thread_id) {
+  // We want to map the thread_id inside the same bitmap
+  static_assert(BITMAP_SIZE >= (u16)0xffff, "Potential verflow");
+  u16 lower16 = (u16)(thread_id & 0xffff);
+  lower16 = reverse16(lower16);
+  int tid = (thread_id & ~0xffff) | lower16;
+  return tid; 
 }
 
-void ThreadFilter::add(int thread_id) {
-  u64 *b = bitmap(thread_id);
+// Get bitmap that contains the thread id, create one if it does not exist
+u64* ThreadFilter::getBitmapFor(int thread_id) {
+  int index = thread_id / BITMAP_CAPACITY;
+  assert(index >= 0 && index < (int)_max_bitmaps);
+  u64* b = _bitmap[index];
   if (b == NULL) {
     b = (u64 *)OS::safeAlloc(BITMAP_SIZE);
     u64 *oldb = __sync_val_compare_and_swap(
-        &_bitmap[(u32)thread_id / BITMAP_CAPACITY], NULL, b);
+        &_bitmap[index], NULL, b);
     if (oldb != NULL) {
       OS::safeFree(b, BITMAP_SIZE);
       b = oldb;
     } else {
       trackPage();
     }
   }
+  return b;
+}
+
+u64* ThreadFilter::bitmapAddressFor(int thread_id) {
+  u64* b = getBitmapFor(thread_id);
+  thread_id = mapThreadId(thread_id);
+  assert(b == bitmap(thread_id));
+  return wordAddress(b, thread_id);
+}
 
+bool ThreadFilter::accept(int thread_id) {
+  thread_id = mapThreadId(thread_id);
+  u64 *b = bitmap(thread_id);
+  return b != NULL && (word(b, thread_id) & (1ULL << (thread_id & 0x3f)));
+}
+
+void ThreadFilter::add(int thread_id) {
+  u64 *b = getBitmapFor(thread_id);
+  assert (b != NULL);
+  thread_id = mapThreadId(thread_id);
+  assert(b == bitmap(thread_id));
   u64 bit = 1ULL << (thread_id & 0x3f);
   if (!(__sync_fetch_and_or(&word(b, thread_id), bit) & bit)) {
     atomicInc(_size);
   }
 }
 
 void ThreadFilter::remove(int thread_id) {
+  thread_id = mapThreadId(thread_id);
   u64 *b = bitmap(thread_id);
   if (b == NULL) {
     return;
@@ -132,7 +164,10 @@ void ThreadFilter::collect(std::vector<int> &v) {
         // order here
         u64 word = __atomic_load_n(&b[j], __ATOMIC_ACQUIRE);
         while (word != 0) {
-          v.push_back(start_id + j * 64 + __builtin_ctzl(word));
+          int tid = start_id + j * 64 + __builtin_ctzl(word);
+          // restore thread id
+          tid = mapThreadId(tid);
+          v.push_back(tid);
           word &= (word - 1);
         }
       }

ddprof-lib/src/main/cpp/threadFilter.h

@@ -45,26 +45,35 @@ class ThreadFilter {
         __ATOMIC_ACQUIRE);
   }
 
+  static int mapThreadId(int thread_id);
+
   u64 &word(u64 *bitmap, int thread_id) {
     // todo: add thread safe APIs
     return bitmap[((u32)thread_id % BITMAP_CAPACITY) >> 6];
   }
 
+  u64* const wordAddress(u64 *bitmap, int thread_id) const {
+    return &bitmap[((u32)thread_id % BITMAP_CAPACITY) >> 6];
+  }
+
+  u64* getBitmapFor(int thread_id);
 public:
   ThreadFilter();
   ThreadFilter(ThreadFilter &threadFilter) = delete;
   ~ThreadFilter();
 
-  bool enabled() { return _enabled; }
+  bool enabled() const { return _enabled; }
 
-  int size() { return _size; }
+  int size() const { return _size; }
+  const volatile int* addressOfSize() const { return &_size; }
 
   void init(const char *filter);
   void clear();
 
   bool accept(int thread_id);
   void add(int thread_id);
   void remove(int thread_id);
+  u64* bitmapAddressFor(int thread_id);
 
   void collect(std::vector<int> &v);
 };

ddprof-lib/src/main/java/com/datadoghq/profiler/ActiveBitmap.java

@@ -0,0 +1,71 @@
+package com.datadoghq.profiler;
+
+import sun.misc.Unsafe;
+import java.lang.reflect.Field;
+
+
+class ActiveBitmap {
+    private static final Unsafe UNSAFE = JavaProfiler.UNSAFE;
+
+    // Address to size field of ThreadFilter in native
+    private static long activeCountAddr;
+
+    private static final ThreadLocal<Long> Address = new ThreadLocal<Long>() {
+        @Override protected Long initialValue() {
+            return -1L;
+        }
+    };
+
+    public static void initialize() {
+        activeCountAddr = getActiveCountAddr0();
+    }
+
+    // On native side, we reverse lower 16 bits of thread id when maps to bitmap bit.
+    // So the active bit position of the specific thread id maps to the reverse order
+    // of the second lowest byte of thread id.
+    static long getBitmask(int tid) {
+        int tmp = (tid >> 8) & 0xff ;
+        int bits = 0;
+        for (int index = 0; index <= 7; index++) {
+            if ((tmp & 0x01) == 0x01) {
+                bits |= 1 << (7 - index);
+            }
+            tmp >>>= 1;
+        }
+        return 1L << (bits & 0x3f);
+    }
+
+    static void setActive(int tid, boolean active) {
+        long addr = Address.get();
+        if (addr == -1) {
+            addr = bitmapAddressFor0(tid);
+            Address.set(addr);
+        }
+        long bitmask = getBitmask(tid);
+        long value = UNSAFE.getLong(addr);
+        long newVal = active ? (value | bitmask) : (value & ~bitmask);
+
+        while (!UNSAFE.compareAndSwapLong(null, addr, value, newVal)) {
+            value = UNSAFE.getLong(addr);
+            newVal = active ? (value | bitmask) : (value & ~bitmask);
+        }
+        int delta = active ? 1 : -1;
+        assert activeCountAddr != 0;
+        UNSAFE.getAndAddInt(null, activeCountAddr, delta);
+        if (isActive0(tid) != active) {
+            throw new RuntimeException("SetActive Failed");
+        }
+
+        assert isActive0(tid) == active;
+    }
+
+
+    // Address of bitmap word that contains the active bit of this thread Id
+    private static native long bitmapAddressFor0(int tid);
+
+    private static native long getActiveCountAddr0();
+
+    // For validation
+    private static native boolean isActive0(int tid);
+}
+

ddprof-lib/src/main/java/com/datadoghq/profiler/JavaProfiler.java

@@ -34,17 +34,17 @@
  * libjavaProfiler.so.
  */
 public final class JavaProfiler {
-    private static final Unsafe UNSAFE;
+    static final Unsafe UNSAFE;
+    static final boolean isJDK8;
     static {
         Unsafe unsafe = null;
         String version = System.getProperty("java.version");
-        if (version.startsWith("1.8")) {
-            try {
-                Field f = Unsafe.class.getDeclaredField("theUnsafe");
-                f.setAccessible(true);
-                unsafe = (Unsafe) f.get(null);
-            } catch (Exception ignore) { }
-        }
+        isJDK8 = version.startsWith("1.8");
+        try {
+            Field f = Unsafe.class.getDeclaredField("theUnsafe");
+            f.setAccessible(true);
+            unsafe = (Unsafe) f.get(null);
+        } catch (Exception ignore) { }
         UNSAFE = unsafe;
     }
 
@@ -108,6 +108,7 @@ public static synchronized JavaProfiler getInstance(String libLocation, String s
             throw new IOException("Failed to load Datadog Java profiler library", result.error);
         }
         init0();
+        ActiveBitmap.initialize();
 
         profiler.initializeContextStorage();
         instance = profiler;
@@ -128,7 +129,7 @@ private void initializeContextStorage() {
         if (this.contextStorage == null) {
             int maxPages = getMaxContextPages0();
             if (maxPages > 0) {
-                if (UNSAFE != null) {
+                if (isJDK8) {
                     contextBaseOffsets = new long[maxPages];
                     // be sure to choose an illegal address as a sentinel value
                     Arrays.fill(contextBaseOffsets, Long.MIN_VALUE);
@@ -208,15 +209,23 @@ public boolean recordTraceRoot(long rootSpanId, String endpoint, int sizeLimit)
      * 'filter' option must be enabled to use this method.
      */
     public void addThread() {
-        filterThread0(true);
+        if (UNSAFE != null) {
+            ActiveBitmap.setActive(TID.get(), true);
+        } else {
+            filterThread0(true);
+        }
     }
 
     /**
      * Remove the given thread to the set of profiled threads.
      * 'filter' option must be enabled to use this method.
      */
     public void removeThread() {
-        filterThread0(false);
+        if (UNSAFE != null) {
+            ActiveBitmap.setActive(TID.get(), false);
+        } else {
+            filterThread0(false);
+        }
     }
 
 
@@ -229,7 +238,7 @@ public void removeThread() {
      */
     public void setContext(long spanId, long rootSpanId) {
         int tid = TID.get();
-        if (UNSAFE != null) {
+        if (isJDK8) {
             setContextJDK8(tid, spanId, rootSpanId);
         } else {
             setContextByteBuffer(tid, spanId, rootSpanId);
@@ -304,7 +313,7 @@ public void clearContext() {
      */
     public void setContextValue(int offset, int value) {
         int tid = TID.get();
-        if (UNSAFE != null) {
+        if (isJDK8) {
             setContextJDK8(tid, offset, value);
         } else {
             setContextByteBuffer(tid, offset, value);
@@ -335,7 +344,7 @@ public void setContextByteBuffer(int tid, int offset, int value) {
 
     void copyTags(int[] snapshot) {
         int tid = TID.get();
-        if (UNSAFE != null) {
+        if (isJDK8) {
             copyTagsJDK8(tid, snapshot);
         } else {
             copyTagsByteBuffer(tid, snapshot);