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 <string.h>
 25 #include "jvmti.h"
 26 #include "jvmti_common.hpp"
 27 
 28 #include <atomic>
 29 
 30 extern "C" {
 31 
 32 static std::atomic<int> events_counter(0);
 33 static jvmtiEnv* jvmti;
 34 static jclass expected_class;
 35 static bool expect_events = false;
 36 
 37 static bool
 38 is_test_class(jvmtiEnv* jvmti, JNIEnv* jni, jclass cls) {
 39   char* sig = nullptr;
 40   check_jvmti_error(jvmti->GetClassSignature(cls, &sig, nullptr), "GetClassSignature");
 41 
 42   LOG("Object class: %s\n", sig);
 43   jvmti->Deallocate((unsigned char *)sig);
 44   bool res = (jni->IsSameObject(cls, expected_class) == JNI_TRUE);
 45   return res;
 46 }
 47 
 48 JNIEXPORT void JNICALL
 49 Java_SampledObjectAllocValue_enableEvents(JNIEnv* jni, jclass klass, jthread thread, jclass tested_class) {
 50   if (events_counter != 0) {
 51     fatal(jni, "SampledObjectAlloc events counter should be zero");
 52   }
 53   LOG("enableEvents: events_counter: %d\n", events_counter.load());
 54 
 55   expected_class = (jclass)jni->NewGlobalRef(tested_class);
 56   jvmtiError err = jvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_SAMPLED_OBJECT_ALLOC, thread);
 57   check_jvmti_error(err, "SetEventNotificationMode for SAMPLED_OBJECT_ALLOC");
 58 }
 59 
 60 JNIEXPORT void JNICALL
 61 SampledObjectAlloc(jvmtiEnv* jvmti, JNIEnv* jni, jthread thread, jobject object, jclass klass, jlong size) {
 62   if (klass == nullptr) {
 63     fatal(jni, "klass in SampledObjectAlloc callback is not expected to be null");
 64   }
 65   if (!is_test_class(jvmti, jni, klass)) {
 66     return; // interested in tested class only
 67   }
 68   if (size == 0L) {
 69     fatal(jni, "size in SampledObjectAlloc callback is not expected to be 0");
 70   }
 71   if (object != nullptr) {
 72     fatal(jni, "object in SampledObjectAlloc callback is expected to be null for value object allocations");
 73   }
 74   events_counter++;
 75   LOG("SampledObjectAlloc: events_counter: %d\n", events_counter.load());
 76 }
 77 
 78 void JNICALL
 79 VMDeath(jvmtiEnv* jvmti, JNIEnv* jni) {
 80   jvmtiError err = jvmti->SetEventNotificationMode(JVMTI_DISABLE, JVMTI_EVENT_SAMPLED_OBJECT_ALLOC, nullptr);
 81   check_jvmti_error(err, "SetEventNotificationMode for SAMPLED_OBJECT_ALLOC");
 82 
 83   LOG("VMDeath: expect_events: %d events_counter: %d\n", expect_events, events_counter.load());
 84 
 85   if (expect_events && events_counter == 0) {
 86     fatal(jni, "SampledObjectAlloc events counter shouldn't be zero");
 87   }
 88   if (!expect_events && events_counter != 0) {
 89     fatal(jni, "SampledObjectAlloc events counter should be zero");
 90   }
 91 }
 92 
 93 jint Agent_Initialize(JavaVM *jvm, char *options, void *reserved) {
 94   bool support_value_objects = options != nullptr && strcmp(options, "can_support_value_objects") == 0;
 95   jvmtiCapabilities caps;
 96   jvmtiEventCallbacks callbacks;
 97   jvmtiError err;
 98   jint res;
 99 
100   LOG("Agent_Initialize\n");
101   res = jvm->GetEnv((void **) &jvmti, JVMTI_VERSION_9);
102   if (res != JNI_OK || jvmti == nullptr) {
103     LOG("Wrong result of a valid call to GetEnv!\n");
104     return JNI_ERR;
105   }
106 
107   memset(&caps, 0, sizeof(caps));
108   caps.can_generate_sampled_object_alloc_events = 1;
109   if (support_value_objects) {
110     caps.can_support_value_objects = 1;
111     expect_events = true;
112   }
113 
114   if (jvmti->AddCapabilities(&caps) != JVMTI_ERROR_NONE) {
115     return JNI_ERR;
116   }
117 
118   memset(&callbacks, 0, sizeof(callbacks));
119   callbacks.VMDeath = &VMDeath;
120   callbacks.SampledObjectAlloc = &SampledObjectAlloc;
121 
122   err = jvmti->SetEventCallbacks(&callbacks, sizeof(callbacks));
123   check_jvmti_error(err, "SetEventCallbacks");
124 
125   // Interval should be small enough to trigger sampling event.
126   err = jvmti->SetHeapSamplingInterval(100);
127   check_jvmti_error(err, "SetHeapSamplingInterval");
128 
129   err = jvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_VM_DEATH, nullptr);
130   check_jvmti_error(err, "SetEventNotificationMode for VM_DEATH");
131 
132   return JNI_OK;
133 }
134 
135 JNIEXPORT jint JNICALL Agent_OnLoad(JavaVM *jvm, char *options, void *reserved) {
136   return Agent_Initialize(jvm, options, reserved);
137 }
138 
139 JNIEXPORT jint JNICALL Agent_OnAttach(JavaVM *jvm, char *options, void *reserved) {
140   return Agent_Initialize(jvm, options, reserved);
141 }
142 }