Context ID functionality

src/api/one/profiler/AsyncProfiler.java

@@ -208,8 +208,27 @@ private void filterThread(Thread thread, boolean enable) {
         }
     }
 
+    /**
+     * Passing context identifier to a profiler. This ID is thread-local and is dumped in
+     * the JFR output only. 0 is a reserved value for "no-context". The context functionality
+     * is available for 64bit Java only.
+     *
+     * @param contextId Context identifier that should be stored for current thread
+     */
+    public void setContextId(long contextId) {
+        setContextId0(contextId);
+    }
+
+    /**
+     * Clears context identifier for current thread.
+     */
+    public void clearContextId() {
+        setContextId0(0);
+    }
+
     private native void start0(String event, long interval, boolean reset) throws IllegalStateException;
     private native void stop0() throws IllegalStateException;
     private native String execute0(String command) throws IllegalArgumentException, IllegalStateException, IOException;
     private native void filterThread0(Thread thread, boolean enable);
+    private native void setContextId0(long contextId);
 }

src/flightRecorder.cpp

@@ -1075,6 +1075,7 @@ class Recording {
         buf->putVar32(tid);
         buf->putVar32(call_trace_id);
         buf->putVar32(event->_thread_state);
+        buf->putVar64(Profiler::instance()->getContextId());
         buf->put8(start, buf->offset() - start);
     }
 
@@ -1087,6 +1088,7 @@ class Recording {
         buf->putVar32(event->_class_id);
         buf->putVar64(event->_instance_size);
         buf->putVar64(event->_total_size);
+        buf->putVar64(Profiler::instance()->getContextId());
         buf->put8(start, buf->offset() - start);
     }
 
@@ -1098,6 +1100,7 @@ class Recording {
         buf->putVar32(call_trace_id);
         buf->putVar32(event->_class_id);
         buf->putVar64(event->_total_size);
+        buf->putVar64(Profiler::instance()->getContextId());
         buf->put8(start, buf->offset() - start);
     }
 
@@ -1111,6 +1114,7 @@ class Recording {
         buf->putVar32(event->_class_id);
         buf->put8(0);
         buf->putVar64(event->_address);
+        buf->putVar64(Profiler::instance()->getContextId());
         buf->put8(start, buf->offset() - start);
     }
 

src/javaApi.cpp

@@ -67,6 +67,15 @@ Java_one_profiler_AsyncProfiler_stop0(JNIEnv* env, jobject unused) {
     }
 }
 
+extern "C" DLLEXPORT void JNICALL
+Java_one_profiler_AsyncProfiler_setContextId0(JNIEnv* env, jobject unused, jlong contextId) {
+    Error error = Profiler::instance()->setContextId(contextId);
+
+    if (error) {
+        throwNew(env, "java/lang/IllegalStateException", error.message());
+    }
+}
+
 extern "C" DLLEXPORT jstring JNICALL
 Java_one_profiler_AsyncProfiler_execute0(JNIEnv* env, jobject unused, jstring command) {
     Arguments args;
@@ -134,11 +143,12 @@ Java_one_profiler_AsyncProfiler_filterThread0(JNIEnv* env, jobject unused, jthre
 #define F(name, sig)  {(char*)#name, (char*)sig, (void*)Java_one_profiler_AsyncProfiler_##name}
 
 static const JNINativeMethod profiler_natives[] = {
-    F(start0,        "(Ljava/lang/String;JZ)V"),
-    F(stop0,         "()V"),
-    F(execute0,      "(Ljava/lang/String;)Ljava/lang/String;"),
-    F(getSamples,    "()J"),
-    F(filterThread0, "(Ljava/lang/Thread;Z)V"),
+    F(start0,          "(Ljava/lang/String;JZ)V"),
+    F(stop0,           "()V"),
+    F(execute0,        "(Ljava/lang/String;)Ljava/lang/String;"),
+    F(getSamples,      "()J"),
+    F(filterThread0,   "(Ljava/lang/Thread;Z)V"),
+    F(setContextId0,   "(J)V"),
 };
 
 static const JNINativeMethod* execute0 = &profiler_natives[2];

src/jfrMetadata.cpp

@@ -90,7 +90,8 @@ JfrMetadata::JfrMetadata() : Element("root") {
                 << field("startTime", T_LONG, "Start Time", F_TIME_TICKS)
                 << field("sampledThread", T_THREAD, "Thread", F_CPOOL)
                 << field("stackTrace", T_STACK_TRACE, "Stack Trace", F_CPOOL)
-                << field("state", T_THREAD_STATE, "Thread State", F_CPOOL))
+                << field("state", T_THREAD_STATE, "Thread State", F_CPOOL)
+                << field("contextId", T_LONG, "Context ID"))
 
             << (type("jdk.ObjectAllocationInNewTLAB", T_ALLOC_IN_NEW_TLAB, "Allocation in new TLAB")
                 << category("Java Application")
@@ -99,15 +100,17 @@ JfrMetadata::JfrMetadata() : Element("root") {
                 << field("stackTrace", T_STACK_TRACE, "Stack Trace", F_CPOOL)
                 << field("objectClass", T_CLASS, "Object Class", F_CPOOL)
                 << field("allocationSize", T_LONG, "Allocation Size", F_BYTES)
-                << field("tlabSize", T_LONG, "TLAB Size", F_BYTES))
+                << field("tlabSize", T_LONG, "TLAB Size", F_BYTES)
+                << field("contextId", T_LONG, "Context ID"))
 
             << (type("jdk.ObjectAllocationOutsideTLAB", T_ALLOC_OUTSIDE_TLAB, "Allocation outside TLAB")
                 << category("Java Application")
                 << field("startTime", T_LONG, "Start Time", F_TIME_TICKS)
                 << field("eventThread", T_THREAD, "Event Thread", F_CPOOL)
                 << field("stackTrace", T_STACK_TRACE, "Stack Trace", F_CPOOL)
                 << field("objectClass", T_CLASS, "Object Class", F_CPOOL)
-                << field("allocationSize", T_LONG, "Allocation Size", F_BYTES))
+                << field("allocationSize", T_LONG, "Allocation Size", F_BYTES)
+                << field("contextId", T_LONG, "Context ID"))
 
             << (type("jdk.JavaMonitorEnter", T_MONITOR_ENTER, "Java Monitor Blocked")
                 << category("Java Application")
@@ -117,7 +120,8 @@ JfrMetadata::JfrMetadata() : Element("root") {
                 << field("stackTrace", T_STACK_TRACE, "Stack Trace", F_CPOOL)
                 << field("monitorClass", T_CLASS, "Monitor Class", F_CPOOL)
                 << field("previousOwner", T_THREAD, "Previous Monitor Owner", F_CPOOL)
-                << field("address", T_LONG, "Monitor Address", F_ADDRESS))
+                << field("address", T_LONG, "Monitor Address", F_ADDRESS)
+                << field("contextId", T_LONG, "Context ID"))
 
             << (type("jdk.ThreadPark", T_THREAD_PARK, "Java Thread Park")
                 << category("Java Application")

src/profiler.cpp

@@ -22,6 +22,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <pthread.h>
 #include <sys/param.h>
 #include "profiler.h"
 #include "perfEvents.h"
@@ -45,6 +46,7 @@
 #include "symbols.h"
 #include "vmStructs.h"
 
+void __attribute__((constructor)) createContextId();
 
 // The instance is not deleted on purpose, since profiler structures
 // can be still accessed concurrently during VM termination
@@ -64,6 +66,11 @@ static J9WallClock j9_wall_clock;
 static ITimer itimer;
 static Instrument instrument;
 
+static pthread_key_t local_context_id_key;
+
+void createContextId() {
+    pthread_key_create(&local_context_id_key, NULL);
+}
 
 // Stack recovery techniques used to workaround AsyncGetCallTrace flaws.
 // Can be disabled with 'safemode' option.
@@ -1008,6 +1015,19 @@ Error Profiler::start(Arguments& args, bool reset) {
     return error;
 }
 
+Error Profiler::setContextId(u64 contextId) {
+    if (sizeof(uintptr_t) < sizeof(u64)) {
+        return Error("Setting contextId is supported on 64bit Java only");
+    }
+    pthread_setspecific(local_context_id_key, (void *) contextId);
+    return Error::OK;
+}
+
+u64 Profiler::getContextId() {
+    void* value = pthread_getspecific(local_context_id_key);
+    return (u64) value;
+}
+
 Error Profiler::stop() {
     MutexLocker ml(_state_lock);
     if (_state != RUNNING) {

src/profiler.h

@@ -72,6 +72,7 @@ class Profiler {
     // TODO: single map?
     std::map<int, std::string> _thread_names;
     std::map<int, jlong> _thread_ids;
+
     Dictionary _class_map;
     Dictionary _symbol_map;
     ThreadFilter _thread_filter;
@@ -190,6 +191,8 @@ class Profiler {
     void shutdown(Arguments& args);
     Error check(Arguments& args);
     Error start(Arguments& args, bool reset);
+    Error setContextId(u64 contextId);
+    u64 getContextId();
     Error stop();
     Error flushJfr();
     Error dump(std::ostream& out, Arguments& args);