Merge "profiling: Allow to set sampling rate on command line."
diff --git a/src/profiling/memory/main.cc b/src/profiling/memory/main.cc
index 850d6c0..9b859b5 100644
--- a/src/profiling/memory/main.cc
+++ b/src/profiling/memory/main.cc
@@ -15,6 +15,7 @@
  */
 
 #include <stdlib.h>
+#include <unistd.h>
 #include <array>
 #include <memory>
 #include <vector>
@@ -31,7 +32,7 @@
 constexpr size_t kUnwinderQueueSize = 1000;
 constexpr size_t kBookkeepingQueueSize = 1000;
 constexpr size_t kUnwinderThreads = 5;
-constexpr double kSamplingRate = 1;
+constexpr double kDefaultSamplingRate = 1;
 
 // We create kUnwinderThreads unwinding threads and one bookeeping thread.
 // The bookkeeping thread is singleton in order to avoid expensive and
@@ -59,6 +60,22 @@
 //           |Bookkeeping Thread|
 //           +------------------+
 int HeapprofdMain(int argc, char** argv) {
+  // TODO(fmayer): This is temporary until heapprofd is integrated with Perfetto
+  // and receives its configuration via that.
+  double sampling_rate = kDefaultSamplingRate;
+  int opt;
+  while ((opt = getopt(argc, argv, "r:")) != -1) {
+    switch (opt) {
+      case 'r': {
+        char* end;
+        sampling_rate = strtol(optarg, &end, 10);
+        if (*end != '\0' || *optarg == '\0')
+          PERFETTO_FATAL("Invalid sampling rate: %s", optarg);
+        break;
+      }
+    }
+  }
+
   GlobalCallstackTrie callsites;
   std::unique_ptr<base::UnixSocket> sock;
 
@@ -81,15 +98,16 @@
     unwinder_queues[static_cast<size_t>(r.pid) % kUnwinderThreads].Add(
         std::move(r));
   };
-  SocketListener listener({kSamplingRate}, std::move(on_record_received),
+  SocketListener listener({sampling_rate}, std::move(on_record_received),
                           &callsites);
 
   base::UnixTaskRunner read_task_runner;
-  if (argc == 2) {
+  if (optind == argc - 1) {
     // Allow to be able to manually specify the socket to listen on
     // for testing and sideloading purposes.
-    sock = base::UnixSocket::Listen(argv[1], &listener, &read_task_runner);
-  } else if (argc == 1) {
+    sock =
+        base::UnixSocket::Listen(argv[argc - 1], &listener, &read_task_runner);
+  } else if (optind == argc) {
     // When running as a service launched by init on Android, the socket
     // is created by init and passed to the application using an environment
     // variable.
@@ -106,7 +124,8 @@
     sock = base::UnixSocket::Listen(base::ScopedFile(raw_fd), &listener,
                                     &read_task_runner);
   } else {
-    PERFETTO_FATAL("Invalid number of arguments. %s [SOCKET]", argv[0]);
+    PERFETTO_FATAL("Invalid number of arguments. %s [-r rate] [SOCKET]",
+                   argv[0]);
   }
 
   if (sock->last_error() != 0)