Fix #200 Crash related to aligned_alloc and free in context (#208)

* Fix #200 Crash related to aligned_alloc and free in context

* If posix_memalign fails, use a sentinel to avoid duplicate calls to posix_memalign

Author: yanglong1010

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

@@ -47,17 +47,21 @@ Context &Contexts::get(int tid) {
 
 Context &Contexts::empty() { return DD_EMPTY_CONTEXT; }
 
-void Contexts::initialize(int pageIndex) {
+bool Contexts::initialize(int pageIndex) {
   if (pageIndex >= _max_pages) {
     Counters::increment(CounterId::CONTEXT_BOUNDS_MISS_INITS);
     // extreme edge case: pageIndex >= _max_pages if pid_max was increased
     // during the process's runtime
-    return;
+    return false;
   }
   if (__atomic_load_n(&_pages[pageIndex], __ATOMIC_ACQUIRE) == NULL) {
     u32 capacity = DD_CONTEXT_PAGE_SIZE * sizeof(Context);
-    Context *page = (Context *)aligned_alloc(sizeof(Context), capacity);
-    // need to zero the storage because there is no aligned_calloc
+    Context *page;
+    if (posix_memalign((void**)&page, sizeof(Context), capacity)) {
+      Counters::increment(CONTEXT_ALLOC_FAILS);
+      return false;
+    }
+    // need to zero the storage
     memset(page, 0, capacity);
     if (!__sync_bool_compare_and_swap(&_pages[pageIndex], NULL, page)) {
       free(page);
@@ -66,6 +70,7 @@ void Contexts::initialize(int pageIndex) {
       Counters::increment(CONTEXT_STORAGE_PAGES);
     }
   }
+  return true;
 }
 
 void Contexts::reset() {
@@ -78,9 +83,12 @@ void Contexts::reset() {
 
 ContextPage Contexts::getPage(int tid) {
   int pageIndex = tid >> DD_CONTEXT_PAGE_SHIFT;
-  initialize(pageIndex);
-  return {.capacity = DD_CONTEXT_PAGE_SIZE * sizeof(Context),
-          .storage = _pages[pageIndex]};
+  if (initialize(pageIndex)) {
+    return {.capacity = DD_CONTEXT_PAGE_SIZE * sizeof(Context),
+      .storage = _pages[pageIndex]};
+  } else {
+    return {};
+  }
 }
 
 // The number of pages that can cover all allowed thread IDs

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

@@ -52,7 +52,7 @@ class Contexts {
 private:
   static int _max_pages;
   static Context **_pages;
-  static void initialize(int pageIndex);
+  static bool initialize(int pageIndex);
 
 public:
   // get must not allocate

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

@@ -43,6 +43,7 @@
   X(CONTEXT_BOUNDS_MISS_GETS, "context_bounds_miss_gets")                      \
   X(CONTEXT_CHECKSUM_REJECT_GETS, "context_checksum_reject_gets")              \
   X(CONTEXT_NULL_PAGE_GETS, "context_null_page_gets")                          \
+  X(CONTEXT_ALLOC_FAILS, "context_alloc_fails")                                \
   X(CALLTRACE_STORAGE_BYTES, "calltrace_storage_bytes")                        \
   X(CALLTRACE_STORAGE_TRACES, "calltrace_storage_traces")                      \
   X(LINEAR_ALLOCATOR_BYTES, "linear_allocator_bytes")                          \

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

@@ -144,6 +144,9 @@ Java_com_datadoghq_profiler_JavaProfiler_getContextPage0(JNIEnv *env,
                                                          jobject unused,
                                                          jint tid) {
   ContextPage page = Contexts::getPage((int)tid);
+  if (page.storage == 0) {
+    return NULL;
+  }
   return env->NewDirectByteBuffer((void *)page.storage, (jlong)page.capacity);
 }
 

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

@@ -66,6 +66,7 @@ static final class TSCFrequencyHolder {
 
     private ByteBuffer[] contextStorage;
     private long[] contextBaseOffsets;
+    private static final ByteBuffer SENTINEL = ByteBuffer.allocate(0);
 
     private JavaProfiler() {
     }
@@ -240,6 +241,9 @@ private void setContextJDK8(int tid, long spanId, long rootSpanId) {
             return;
         }
         long pageOffset = getPageUnsafe(tid);
+        if (pageOffset == 0) {
+            return;
+        }
         int index = (tid % PAGE_SIZE) * CONTEXT_SIZE;
         long base = pageOffset + index;
         UNSAFE.putLong(base + SPAN_OFFSET, spanId);
@@ -252,20 +256,27 @@ private void setContextByteBuffer(int tid, long spanId, long rootSpanId) {
             return;
         }
         ByteBuffer page = getPage(tid);
+        if (page == SENTINEL) {
+            return;
+        }
         int index = (tid % PAGE_SIZE) * CONTEXT_SIZE;
         page.putLong(index + SPAN_OFFSET, spanId);
         page.putLong(index + ROOT_SPAN_OFFSET, rootSpanId);
         page.putLong(index + CHECKSUM_OFFSET, spanId ^ rootSpanId);
     }
 
-
-
     private ByteBuffer getPage(int tid) {
         int pageIndex = tid / PAGE_SIZE;
         ByteBuffer page = contextStorage[pageIndex];
         if (page == null) {
             // the underlying page allocation is atomic so we don't care which view we have over it
-            contextStorage[pageIndex] = page = getContextPage0(tid).order(ByteOrder.LITTLE_ENDIAN);
+            ByteBuffer buffer = getContextPage0(tid);
+            if (buffer == null) {
+                page = SENTINEL;
+            } else {
+                page = buffer.order(ByteOrder.LITTLE_ENDIAN);
+            }
+            contextStorage[pageIndex] = page;
         }
         return page;
     }
@@ -305,6 +316,9 @@ private void setContextJDK8(int tid, int offset, int value) {
             return;
         }
         long pageOffset = getPageUnsafe(tid);
+        if (pageOffset == 0) {
+            return;
+        }
         UNSAFE.putInt(pageOffset + addressOf(tid, offset), value);
     }
 
@@ -313,6 +327,9 @@ public void setContextByteBuffer(int tid, int offset, int value) {
             return;
         }
         ByteBuffer page = getPage(tid);
+        if (page == SENTINEL) {
+            return;
+        }
         page.putInt(addressOf(tid, offset), value);
     }
 
@@ -330,6 +347,9 @@ void copyTagsJDK8(int tid, int[] snapshot) {
             return;
         }
         long pageOffset = getPageUnsafe(tid);
+        if (pageOffset == 0) {
+            return;
+        }
         long address = pageOffset + addressOf(tid, 0);
         for (int i = 0; i < snapshot.length; i++) {
             snapshot[i] = UNSAFE.getInt(address);
@@ -342,6 +362,9 @@ void copyTagsByteBuffer(int tid, int[] snapshot) {
             return;
         }
         ByteBuffer page = getPage(tid);
+        if (page == SENTINEL) {
+            return;
+        }
         int address = addressOf(tid, 0);
         for (int i = 0; i < snapshot.length; i++) {
             snapshot[i] = page.getInt(address + i * Integer.BYTES);