Update java agent

Makefile

@@ -216,7 +216,7 @@ $(TARGET_AGENT): $(SOURCES) $(HEADERS)
 
 $(TARGET_NOTICES): $(JAVA_AGENT_PATH)/NOTICES
 	mkdir -p $(dir $@)
-	cp -f $< $@
+	cp -fp $< $@
 
 $(GENFILES_PATH)/%profiler.pb.h $(GENFILES_PATH)/%profiler.pb.cc : third_party/googleapis/%profiler.proto
 	mkdir -p $(dir $@)

src/entry.cc

@@ -60,6 +60,7 @@ static void JNICALL OnThreadEnd(jvmtiEnv *jvmti_env, JNIEnv *jni_env,
   IMPLICITLY_USE(jvmti_env);
   IMPLICITLY_USE(jni_env);
   IMPLICITLY_USE(thread);
+  google::javaprofiler::Accessors::SetCurrentJniEnv(nullptr);
   threads->UnregisterCurrent();
 }
 

src/throttler_api.cc

@@ -32,7 +32,6 @@
 #include "grpcpp/create_channel.h"
 #include "grpcpp/security/credentials.h"
 #include "grpcpp/support/channel_arguments.h"
-#include "third_party/javaprofiler/heap_sampler.h"
 
 // API curated profiling configuration.
 DEFINE_string(cprof_api_address, "cloudprofiler.googleapis.com",
@@ -48,11 +47,10 @@ DEFINE_bool(cprof_use_insecure_creds_for_testing, false,
 
 namespace cloud {
 namespace profiler {
+namespace {
 
 namespace api = google::devtools::cloudprofiler::v2;
 
-namespace {
-
 // Initial value for backoffs where the duration is not server-guided.
 const int64_t kBackoffNanos = 60 * kNanosPerSecond;  // 1 minute
 // Backoff envelope exponential growth factor.
@@ -99,7 +97,7 @@ grpc_ssl_roots_override_result OverrideSSLRoots(char** pem_root_certs) {
 
 // Creates the profiler gRPC API stub. Returns nullptr on error.
 std::unique_ptr<api::grpc::ProfilerService::StubInterface>
-NewProfilerServiceStub(const std::string& addr) {
+NewProfilerServiceStub(const std::string& addr, const std::string& language) {
   std::shared_ptr<grpc::ChannelCredentials> creds;
   if (FLAGS_cprof_use_insecure_creds_for_testing) {
     creds = grpc::InsecureChannelCredentials();
@@ -114,7 +112,8 @@ NewProfilerServiceStub(const std::string& addr) {
 
   grpc::ChannelArguments channel_arguments;
   channel_arguments.SetUserAgentPrefix(
-      "gcloud-java-profiler/" + std::string(CLOUD_PROFILER_AGENT_VERSION));
+      "gcloud-" + language + "-profiler/" +
+      std::string(CLOUD_PROFILER_AGENT_VERSION));
 
   std::shared_ptr<grpc::ChannelInterface> ch =
       grpc::CreateCustomChannel(addr, creds, channel_arguments);
@@ -162,7 +161,7 @@ bool AbortedBackoffDuration(const grpc::ClientContext& ctx,
 // Initializes deployment information from environment properties and label
 // string in "name1=val1,name2=val2,..." format. Returns false on error.
 bool InitializeDeployment(CloudEnv* env, const std::string& labels,
-                          api::Deployment* d) {
+                          const std::string& language, api::Deployment* d) {
   std::string project_id = env->ProjectID();
   if (project_id.empty()) {
     LOG(ERROR) << "Project ID is unknown";
@@ -199,7 +198,7 @@ bool InitializeDeployment(CloudEnv* env, const std::string& labels,
     label_kvs[kZoneNameLabel] = zone_name;
   }
 
-  label_kvs[kLanguageLabel] = "java";
+  label_kvs[kLanguageLabel] = language;
   for (const auto& kv : label_kvs) {
     (*d->mutable_labels())[kv.first] = kv.second;
   }
@@ -249,71 +248,41 @@ bool IsValidServiceName(std::string s) {
   return true;
 }
 
-APIThrottler::APIThrottler(JNIEnv* jni)
-    : APIThrottler(DefaultCloudEnv(), DefaultClock(), nullptr, jni) {}
+APIThrottler::APIThrottler(
+    const std::vector<google::devtools::cloudprofiler::v2::ProfileType>& types,
+    const std::string& language, const std::string& language_version)
+    : APIThrottler(types, language, language_version, DefaultCloudEnv(),
+                   DefaultClock(), nullptr) {}
 
 APIThrottler::APIThrottler(
+    const std::vector<google::devtools::cloudprofiler::v2::ProfileType>& types,
+    const std::string& language, const std::string& language_version,
     CloudEnv* env, Clock* clock,
     std::unique_ptr<google::devtools::cloudprofiler::v2::grpc::ProfilerService::
                         StubInterface>
-        stub,
-    JNIEnv* jni)
-    : env_(env),
+        stub)
+    : types_(types),
+      language_(language),
+      language_version_(language_version),
+      env_(env),
       clock_(clock),
       stub_(std::move(stub)),
-      types_({api::CPU, api::WALL}),
-      java_version_(JavaVersion(jni)),
       creation_backoff_envelope_ns_(kBackoffNanos),
       closed_(false) {
   grpc_init();
   gpr_set_log_function(GRPCLog);
 
-  LOG(INFO) << "Java version: " << java_version_;
-
   // Create a random number generator.
   gen_ = std::default_random_engine(clock_->Now().tv_nsec / 1000);
   dist_ = std::uniform_int_distribution<int64_t>(0, kRandomRange);
 
   if (!stub_) {  // Set in tests
     LOG(INFO) << "Will use profiler service " << FLAGS_cprof_api_address
               << " to create and upload profiles";
-    stub_ = NewProfilerServiceStub(FLAGS_cprof_api_address);
-  }
-
-  if (google::javaprofiler::HeapMonitor::Enabled()) {
-    LOG(INFO) << "Heap allocation sampling supported for this JDK";
-    types_.push_back(api::HEAP);
+    stub_ = NewProfilerServiceStub(FLAGS_cprof_api_address, language_);
   }
 }
 
-std::string APIThrottler::JavaVersion(JNIEnv* jni) {
-  const std::string kUnknownVersion = "unknown_version";
-
-  jclass system_class = jni->FindClass("java/lang/System");
-  if (system_class == nullptr) {
-    return kUnknownVersion;
-  }
-  jmethodID get_property_method = jni->GetStaticMethodID(
-      system_class, "getProperty", "(Ljava/lang/String;)Ljava/lang/String;");
-  if (get_property_method == nullptr) {
-    return kUnknownVersion;
-  }
-  jstring jstr = reinterpret_cast<jstring>(jni->CallStaticObjectMethod(
-      system_class, get_property_method, jni->NewStringUTF("java.version")));
-  if (jstr == nullptr) {
-    return kUnknownVersion;
-  }
-  // Copy the returned value and release the memory allocated by JNI.
-  const char* s = jni->GetStringUTFChars(jstr, nullptr);
-  std::string ret = std::string(s);
-  jni->ReleaseStringUTFChars(jstr, s);
-  return ret;
-}
-
-void APIThrottler::SetProfileTypes(const std::vector<api::ProfileType>& types) {
-  types_ = types;
-}
-
 bool APIThrottler::WaitNext() {
   if (stub_ == nullptr) {
     LOG(ERROR) << "Profiler API is not initialized, stop profiling";
@@ -324,7 +293,7 @@ bool APIThrottler::WaitNext() {
   for (const auto& type : types_) {
     req.add_profile_type(type);
   }
-  if (!InitializeDeployment(env_, FLAGS_cprof_deployment_labels,
+  if (!InitializeDeployment(env_, FLAGS_cprof_deployment_labels, language_,
                             req.mutable_deployment())) {
     LOG(ERROR) << "Failed to initialize deployment, stop profiling";
     return false;
@@ -442,7 +411,7 @@ void APIThrottler::ResetClientContext() {
   ctx_.reset(new grpc::ClientContext());  // NOLINT
   ctx_->AddMetadata("x-goog-api-client",
                     "gccl/" + std::string(CLOUD_PROFILER_AGENT_VERSION) +
-                        " gl-java/" + java_version_);
+                        "gl-" + language_ + "/" + language_version_);
 
   if (closed_) {
     ctx_->TryCancel();

src/throttler_api.h

@@ -39,19 +39,20 @@ namespace profiler {
 // Throttler implementation using the Cloud Profiler API.
 class APIThrottler : public Throttler {
  public:
-  explicit APIThrottler(JNIEnv* jni_env);
+  // Language should be one of "cpp" or "java", and language_version can be
+  // omitted if unknown.
+  explicit APIThrottler(
+      const std::vector<google::devtools::cloudprofiler::v2::ProfileType>&
+          types,
+      const std::string& language, const std::string& language_version = "");
   // Testing-only constructor.
-  APIThrottler(CloudEnv* env, Clock* clock,
+  APIThrottler(const std::vector<
+                   google::devtools::cloudprofiler::v2::ProfileType>& types,
+               const std::string& language, const std::string& language_version,
+               CloudEnv* env, Clock* clock,
                std::unique_ptr<google::devtools::cloudprofiler::v2::grpc::
                                    ProfilerService::StubInterface>
-                   stub,
-               JNIEnv* jni_env);
-
-  // Set the list of supported profile types. The list is used in the profile
-  // creation call to the server to specify the supported types.
-  void SetProfileTypes(
-      const std::vector<google::devtools::cloudprofiler::v2::ProfileType>&
-          types);
+                   stub);
 
   bool WaitNext() override;
   std::string ProfileType() override;
@@ -69,18 +70,16 @@ class APIThrottler : public Throttler {
   // Resets the client gRPC context for the next call.
   void ResetClientContext();
 
-  // Determines and returns Java version.
-  static std::string JavaVersion(JNIEnv* jni_env);
-
  private:
+  const std::vector<google::devtools::cloudprofiler::v2::ProfileType> types_;
+  const std::string language_;
+  const std::string language_version_;
   CloudEnv* env_;
   Clock* clock_;
   std::unique_ptr<
       google::devtools::cloudprofiler::v2::grpc::ProfilerService::StubInterface>
       stub_;
   google::devtools::cloudprofiler::v2::Profile profile_;
-  std::vector<google::devtools::cloudprofiler::v2::ProfileType> types_;
-  const std::string java_version_;
 
   // Profile creation error handling.
   int64_t creation_backoff_envelope_ns_;

src/worker.cc

@@ -18,6 +18,7 @@
 #include "src/profiler.h"
 #include "src/throttler_api.h"
 #include "src/throttler_timed.h"
+#include "google/devtools/cloudprofiler/v2/profiler.grpc.pb.h"
 #include "third_party/javaprofiler/heap_sampler.h"
 
 DEFINE_bool(cprof_enabled, true,
@@ -32,6 +33,35 @@ DEFINE_int32(cprof_wall_sampling_period_msec, 100,
 
 namespace cloud {
 namespace profiler {
+namespace {
+
+namespace api = google::devtools::cloudprofiler::v2;
+
+std::string JavaVersion(JNIEnv *jni) {
+  const std::string kUnknownVersion = "unknown_version";
+
+  jclass system_class = jni->FindClass("java/lang/System");
+  if (system_class == nullptr) {
+    return kUnknownVersion;
+  }
+  jmethodID get_property_method = jni->GetStaticMethodID(
+      system_class, "getProperty", "(Ljava/lang/String;)Ljava/lang/String;");
+  if (get_property_method == nullptr) {
+    return kUnknownVersion;
+  }
+  jstring jstr = reinterpret_cast<jstring>(jni->CallStaticObjectMethod(
+      system_class, get_property_method, jni->NewStringUTF("java.version")));
+  if (jstr == nullptr) {
+    return kUnknownVersion;
+  }
+  // Copy the returned value and release the memory allocated by JNI.
+  const char *s = jni->GetStringUTFChars(jstr, nullptr);
+  std::string ret = std::string(s);
+  jni->ReleaseStringUTFChars(jstr, s);
+  return ret;
+}
+
+}  // namespace
 
 std::atomic<bool> Worker::enabled_;
 
@@ -44,10 +74,20 @@ void Worker::Start(JNIEnv *jni) {
     return;
   }
 
+  std::string java_version = JavaVersion(jni);
+  LOG(INFO) << "Java version: " << java_version;
+  std::vector<google::devtools::cloudprofiler::v2::ProfileType> types = {
+      api::CPU, api::WALL};
+  if (google::javaprofiler::HeapMonitor::Enabled()) {
+    LOG(INFO) << "Heap allocation sampling supported for this JDK";
+    types.push_back(api::HEAP);
+  }
+
   // Initialize the throttler here rather in the constructor, since the
   // constructor is invoked too early, before the heap profiler is initialized.
   throttler_ = FLAGS_cprof_profile_filename.empty()
-                   ? std::unique_ptr<Throttler>(new APIThrottler(jni))
+                   ? std::unique_ptr<Throttler>(
+                         new APIThrottler(types, "java", java_version))
                    : std::unique_ptr<Throttler>(
                          new TimedThrottler(FLAGS_cprof_profile_filename));
 

third_party/javaprofiler/heap_sampler.cc

@@ -145,10 +145,12 @@ std::unique_ptr<perftools::profiles::Profile> HeapEventStorage::ConvertToProto(
     auto *frames = object->Frames();
     *call_target = {nullptr, static_cast<int>(frames->size()), frames->data()};
     *trace_target = {call_target, object->Size()};
+
+    // Add the size as a label to help post-processing filtering.
+    trace_target->trace_and_labels.AddLabel("bytes", object->Size(), "bytes");
   }
 
-  builder->AddTracesAppendingMetricAsLabel(stack_trace_data.get(),
-                                           objects_size);
+  builder->AddTraces(stack_trace_data.get(), objects_size);
 
   return builder->CreateProto();
 }