1 files changed, 20 insertions, 8 deletions
diff --git a/test/core/memory_usage/callback_client.cc b/test/core/memory_usage/callback_client.cc
index 59d3ec96dc..723b57f30f 100644
--- a/test/core/memory_usage/callback_client.cc
+++ b/test/core/memory_usage/callback_client.cc
@@ -27,7 +27,8 @@
 
 #include "absl/flags/flag.h"
 #include "absl/flags/parse.h"
-#include "absl/strings/string_view.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
 
 #include <grpc/impl/channel_arg_names.h>
 #include <grpc/support/log.h>
@@ -44,7 +45,7 @@
 
 ABSL_FLAG(std::string, target, "", "Target host:port");
 ABSL_FLAG(bool, secure, false, "Use SSL Credentials");
-ABSL_FLAG(int, server_pid, 99999, "Server's pid");
+ABSL_FLAG(int, server_pid, 0, "Server's pid");
 ABSL_FLAG(int, size, 50, "Number of channels");
 
 std::shared_ptr<grpc::Channel> CreateChannelForTest(int index) {
@@ -150,7 +151,9 @@ int main(int argc, char** argv) {
   }
 
   // Getting peak memory usage
-  long peak_server_memory = GetMemUsage(absl::GetFlag(FLAGS_server_pid));
+  long peak_server_memory = absl::GetFlag(FLAGS_server_pid) > 0
+                                ? GetMemUsage(absl::GetFlag(FLAGS_server_pid))
+                                : 0;
   long peak_client_memory = GetMemUsage();
 
   // Checking that all channels are still open
@@ -161,14 +164,23 @@ int main(int argc, char** argv) {
                                              std::chrono::milliseconds(1)));
   }
 
+  std::string prefix;
+  if (absl::StartsWith(absl::GetFlag(FLAGS_target), "xds:")) prefix = "xds ";
+  if (absl::GetFlag(FLAGS_server_pid) == 0) {
+    absl::StrAppend(&prefix, "multi_address ");
+  }
   printf("---------Client channel stats--------\n");
-  printf("client channel memory usage: %f bytes per channel\n",
+  printf("%sclient channel memory usage: %f bytes per channel\n",
+         prefix.c_str(),
          static_cast<double>(peak_client_memory - before_client_memory) / size *
              1024);
-  printf("---------Server channel stats--------\n");
-  printf("server channel memory usage: %f bytes per channel\n",
-         static_cast<double>(peak_server_memory - before_server_memory) / size *
-             1024);
+  if (absl::GetFlag(FLAGS_server_pid) > 0) {
+    printf("---------Server channel stats--------\n");
+    printf("%sserver channel memory usage: %f bytes per channel\n",
+           prefix.c_str(),
+           static_cast<double>(peak_server_memory - before_server_memory) /
+               size * 1024);
+  }
   gpr_log(GPR_INFO, "Client Done");
   return 0;
 }