1 /*
  2  * Copyright (c) 2026, Oracle and/or its affiliates. All rights reserved.
  3  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  4  *
  5  * This code is free software; you can redistribute it and/or modify it
  6  * under the terms of the GNU General Public License version 2 only, as
  7  * published by the Free Software Foundation.
  8  *
  9  * This code is distributed in the hope that it will be useful, but WITHOUT
 10  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 11  * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 12  * version 2 for more details (a copy is included in the LICENSE file that
 13  * accompanied this code).
 14  *
 15  * You should have received a copy of the GNU General Public License version
 16  * 2 along with this work; if not, write to the Free Software Foundation,
 17  * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 18  *
 19  * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 20  * or visit www.oracle.com if you need additional information or have any
 21  * questions.
 22  */
 23 
 24 #include "jvmti.h"
 25 #include "jvmti_common.hpp"
 26 
 27 #include <atomic>
 28 
 29 extern "C" {
 30 
 31 static std::atomic<int> events_counter(0);
 32 static jvmtiEnv* jvmti;
 33 static jclass expected_class;
 34 
 35 static bool
 36 is_test_class(jvmtiEnv* jvmti, JNIEnv* jni, jclass cls) {
 37   char* sig = nullptr;
 38   check_jvmti_error(jvmti->GetClassSignature(cls, &sig, nullptr), "GetClassSignature");
 39 
 40   LOG("Object class: %s\n", sig);
 41   jvmti->Deallocate((unsigned char *)sig);
 42   bool res = (jni->IsSameObject(cls, expected_class) == JNI_TRUE);
 43   return res;
 44 }
 45 
 46 JNIEXPORT void JNICALL
 47 Java_SampledObjectAllocValue_enableEvents(JNIEnv* jni, jclass klass, jthread thread, jclass tested_class) {
 48   if (events_counter != 0) {
 49     fatal(jni, "SampledObjectAlloc events counter should be zero");
 50   }
 51   LOG("enableEvents: events_counter: %d\n", events_counter.load());
 52 
 53   expected_class = (jclass)jni->NewGlobalRef(tested_class);
 54   jvmtiError err = jvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_SAMPLED_OBJECT_ALLOC, thread);
 55   check_jvmti_error(err, "SetEventNotificationMode for SAMPLED_OBJECT_ALLOC");
 56 }
 57 
 58 JNIEXPORT void JNICALL
 59 SampledObjectAlloc(jvmtiEnv* jvmti, JNIEnv* jni, jthread thread, jobject object, jclass klass, jlong size) {
 60   if (klass == nullptr) {
 61     fatal(jni, "klass in SampledObjectAlloc callback is not expected to be null");
 62   }
 63   if (!is_test_class(jvmti, jni, klass)) {
 64     return; // interested in tested class only
 65   }
 66   if (size == 0L) {
 67     fatal(jni, "size in SampledObjectAlloc callback is not expected to be 0");
 68   }
 69   if (object != nullptr) {
 70     fatal(jni, "object in SampledObjectAlloc callback is expected to be null for value object allocations");
 71   }
 72   events_counter++;
 73   LOG("SampledObjectAlloc: events_counter: %d\n", events_counter.load());
 74 }
 75 
 76 void JNICALL
 77 VMDeath(jvmtiEnv* jvmti, JNIEnv* jni) {
 78   jvmtiError err = jvmti->SetEventNotificationMode(JVMTI_DISABLE, JVMTI_EVENT_SAMPLED_OBJECT_ALLOC, nullptr);
 79   check_jvmti_error(err, "SetEventNotificationMode for SAMPLED_OBJECT_ALLOC");
 80 
 81   if (events_counter == 0) {
 82     fatal(jni, "SampledObjectAlloc events counter shouldn't be zero");
 83   }
 84   LOG("VMDeath: events_counter: %d\n", events_counter.load());
 85 }
 86 
 87 jint Agent_Initialize(JavaVM *jvm, char *options, void *reserved) {
 88   jvmtiCapabilities caps;
 89   jvmtiEventCallbacks callbacks;
 90   jvmtiError err;
 91   jint res;
 92 
 93   LOG("Agent_Initialize\n");
 94   res = jvm->GetEnv((void **) &jvmti, JVMTI_VERSION_9);
 95   if (res != JNI_OK || jvmti == nullptr) {
 96     LOG("Wrong result of a valid call to GetEnv!\n");
 97     return JNI_ERR;
 98   }
 99 
100   memset(&caps, 0, sizeof(caps));
101   caps.can_generate_sampled_object_alloc_events = 1;
102   if (jvmti->AddCapabilities(&caps) != JVMTI_ERROR_NONE) {
103     return JNI_ERR;
104   }
105 
106   memset(&callbacks, 0, sizeof(callbacks));
107   callbacks.VMDeath = &VMDeath;
108   callbacks.SampledObjectAlloc = &SampledObjectAlloc;
109 
110   err = jvmti->SetEventCallbacks(&callbacks, sizeof(callbacks));
111   check_jvmti_error(err, "SetEventCallbacks");
112 
113   // Interval should be small enough to triggger sampling event.
114   err = jvmti->SetHeapSamplingInterval(100);
115   check_jvmti_error(err, "SetHeapSamplingInterval");
116 
117   err = jvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_VM_DEATH, nullptr);
118   check_jvmti_error(err, "SetEventNotificationMode for VM_DEATH");
119 
120   return JNI_OK;
121 }
122 
123 JNIEXPORT jint JNICALL Agent_OnLoad(JavaVM *jvm, char *options, void *reserved) {
124   return Agent_Initialize(jvm, options, reserved);
125 }
126 
127 JNIEXPORT jint JNICALL Agent_OnAttach(JavaVM *jvm, char *options, void *reserved) {
128   return Agent_Initialize(jvm, options, reserved);
129 }
130 }