Merge "Separate guardrails from heapprofd."
diff --git a/Android.bp b/Android.bp
index ba600ad..9b85dde 100644
--- a/Android.bp
+++ b/Android.bp
@@ -77,6 +77,7 @@
     ":perfetto_src_profiling_common_interner",
     ":perfetto_src_profiling_common_interning_output",
     ":perfetto_src_profiling_common_proc_utils",
+    ":perfetto_src_profiling_common_profiler_guardrails",
     ":perfetto_src_profiling_common_unwind_support",
     ":perfetto_src_profiling_memory_daemon",
     ":perfetto_src_profiling_memory_ring_buffer",
@@ -303,6 +304,7 @@
     ":perfetto_src_profiling_common_interner",
     ":perfetto_src_profiling_common_interning_output",
     ":perfetto_src_profiling_common_proc_utils",
+    ":perfetto_src_profiling_common_profiler_guardrails",
     ":perfetto_src_profiling_common_unwind_support",
     ":perfetto_src_profiling_memory_client",
     ":perfetto_src_profiling_memory_client_api",
@@ -1640,6 +1642,7 @@
     ":perfetto_src_profiling_common_interner",
     ":perfetto_src_profiling_common_interning_output",
     ":perfetto_src_profiling_common_proc_utils",
+    ":perfetto_src_profiling_common_profiler_guardrails",
     ":perfetto_src_profiling_common_unwind_support",
     ":perfetto_src_profiling_memory_client",
     ":perfetto_src_profiling_memory_daemon",
@@ -6538,6 +6541,14 @@
   ],
 }
 
+// GN: //src/profiling/common:profiler_guardrails
+filegroup {
+  name: "perfetto_src_profiling_common_profiler_guardrails",
+  srcs: [
+    "src/profiling/common/profiler_guardrails.cc",
+  ],
+}
+
 // GN: //src/profiling/common:unittests
 filegroup {
   name: "perfetto_src_profiling_common_unittests",
@@ -8309,6 +8320,7 @@
     ":perfetto_src_profiling_common_interner",
     ":perfetto_src_profiling_common_interning_output",
     ":perfetto_src_profiling_common_proc_utils",
+    ":perfetto_src_profiling_common_profiler_guardrails",
     ":perfetto_src_profiling_common_unittests",
     ":perfetto_src_profiling_common_unwind_support",
     ":perfetto_src_profiling_deobfuscator",
diff --git a/src/profiling/common/BUILD.gn b/src/profiling/common/BUILD.gn
index e999f29..5a35838 100644
--- a/src/profiling/common/BUILD.gn
+++ b/src/profiling/common/BUILD.gn
@@ -76,6 +76,19 @@
   ]
 }
 
+source_set("profiler_guardrails") {
+  deps = [
+    ":proc_utils",
+    "../../../gn:default_deps",
+    "../../../include/perfetto/ext/tracing/core",
+    "../../base",
+  ]
+  sources = [
+    "profiler_guardrails.cc",
+    "profiler_guardrails.h",
+  ]
+}
+
 perfetto_unittest_source_set("unittests") {
   testonly = true
   deps = [
diff --git a/src/profiling/common/profiler_guardrails.cc b/src/profiling/common/profiler_guardrails.cc
new file mode 100644
index 0000000..95f1a19
--- /dev/null
+++ b/src/profiling/common/profiler_guardrails.cc
@@ -0,0 +1,41 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "src/profiling/common/profiler_guardrails.h"
+
+#include <algorithm>
+#include "perfetto/ext/base/optional.h"
+#include "perfetto/ext/base/watchdog_posix.h"
+
+namespace perfetto {
+namespace profiling {
+
+base::Optional<uint64_t> ProfilerCpuGuardrails::GetCputimeSec() {
+  if (!stat_fd_) {
+    return base::nullopt;
+  }
+  lseek(stat_fd_.get(), 0, SEEK_SET);
+  base::ProcStat stat;
+  if (!ReadProcStat(stat_fd_.get(), &stat)) {
+    PERFETTO_ELOG("Failed to read stat file to enforce guardrails.");
+    return base::nullopt;
+  }
+  return (stat.utime + stat.stime) /
+         static_cast<unsigned long>(sysconf(_SC_CLK_TCK));
+}
+
+}  // namespace profiling
+}  // namespace perfetto
diff --git a/src/profiling/common/profiler_guardrails.h b/src/profiling/common/profiler_guardrails.h
new file mode 100644
index 0000000..7145936
--- /dev/null
+++ b/src/profiling/common/profiler_guardrails.h
@@ -0,0 +1,131 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef SRC_PROFILING_COMMON_PROFILER_GUARDRAILS_H_
+#define SRC_PROFILING_COMMON_PROFILER_GUARDRAILS_H_
+
+#include <inttypes.h>
+
+#include "perfetto/ext/base/file_utils.h"
+#include "perfetto/ext/base/optional.h"
+#include "perfetto/ext/base/scoped_file.h"
+#include "perfetto/ext/tracing/core/basic_types.h"
+#include "src/profiling/common/proc_utils.h"
+
+namespace perfetto {
+namespace profiling {
+
+class ProfilerCpuGuardrails {
+ public:
+  explicit ProfilerCpuGuardrails(base::ScopedFile stat_fd)
+      : stat_fd_(std::move(stat_fd)) {}
+
+  template <typename T, typename F>
+  void CheckDataSourceCpu(T begin, T end, F guardrail_hit_callback) {
+    bool any_guardrail = false;
+    for (auto it = begin; it != end; ++it) {
+      auto& ds = it->second;
+      if (ds.GetCpuGuardrailSecs() > 0) {
+        any_guardrail = true;
+        break;
+      }
+    }
+    if (!any_guardrail)
+      return;
+
+    base::Optional<uint64_t> opt_cputime_sec = GetCputimeSec();
+    if (!opt_cputime_sec) {
+      PERFETTO_ELOG("Failed to get CPU time.");
+      return;
+    }
+
+    uint64_t cputime_sec = *opt_cputime_sec;
+
+    for (auto it = begin; it != end; ++it) {
+      auto& ds = it->second;
+      uint64_t ds_max_cpu = ds.GetCpuGuardrailSecs();
+      if (ds_max_cpu > 0) {
+        auto start_cputime_sec = ds.GetCpuStartSecs();
+        // We reject data-sources with CPU guardrails if we cannot read the
+        // initial value, which means we get a non-nullopt value here.
+        PERFETTO_CHECK(start_cputime_sec);
+        uint64_t cpu_diff = cputime_sec - *start_cputime_sec;
+        if (cputime_sec > *start_cputime_sec && cpu_diff > ds_max_cpu) {
+          PERFETTO_ELOG(
+              "Exceeded data-source CPU guardrail "
+              "(%" PRIu64 " > %" PRIu64 "). Shutting down.",
+              cpu_diff, ds_max_cpu);
+          guardrail_hit_callback(&ds);
+        }
+      }
+    }
+  }
+
+  base::Optional<uint64_t> GetCputimeSec();
+
+ private:
+  base::ScopedFile stat_fd_;
+};
+
+class ProfilerMemoryGuardrails {
+ public:
+  explicit ProfilerMemoryGuardrails(base::ScopedFile status_fd)
+      : status_fd_(std::move(status_fd)) {}
+
+  template <typename T, typename F>
+  void CheckDataSourceMemory(T begin, T end, F guardrail_hit_callback) {
+    bool any_guardrail = false;
+    for (auto it = begin; it != end; ++it) {
+      auto& ds = it->second;
+      if (ds.GetMemoryGuardrailKb() > 0) {
+        any_guardrail = true;
+        break;
+      }
+    }
+    if (!any_guardrail)
+      return;
+
+    base::Optional<uint32_t> anon_and_swap;
+    std::string status;
+    lseek(status_fd_.get(), 0, SEEK_SET);
+    if (base::ReadFileDescriptor(*status_fd_, &status))
+      anon_and_swap = GetRssAnonAndSwap(status);
+
+    if (!anon_and_swap) {
+      PERFETTO_ELOG("Failed to read memory usage.");
+      return;
+    }
+
+    for (auto it = begin; it != end; ++it) {
+      auto& ds = it->second;
+      uint32_t ds_max_mem = ds.GetMemoryGuardrailKb();
+      if (ds_max_mem > 0 && *anon_and_swap > ds_max_mem) {
+        PERFETTO_ELOG("Exceeded data-source memory guardrail (%" PRIu32
+                      " > %" PRIu32 "). Shutting down.",
+                      *anon_and_swap, ds_max_mem);
+        guardrail_hit_callback(&ds);
+      }
+    }
+  }
+
+ private:
+  base::ScopedFile status_fd_;
+};
+
+}  // namespace profiling
+}  // namespace perfetto
+
+#endif  // SRC_PROFILING_COMMON_PROFILER_GUARDRAILS_H_
diff --git a/src/profiling/memory/BUILD.gn b/src/profiling/memory/BUILD.gn
index b70cd38..6536d81 100644
--- a/src/profiling/memory/BUILD.gn
+++ b/src/profiling/memory/BUILD.gn
@@ -241,6 +241,7 @@
     "../common:interner",
     "../common:interning_output",
     "../common:proc_utils",
+    "../common:profiler_guardrails",
     "../common:unwind_support",
   ]
   public_deps = [
diff --git a/src/profiling/memory/heapprofd_producer.cc b/src/profiling/memory/heapprofd_producer.cc
index 962ad5f..51d8d57 100644
--- a/src/profiling/memory/heapprofd_producer.cc
+++ b/src/profiling/memory/heapprofd_producer.cc
@@ -25,6 +25,7 @@
 #include <sys/types.h>
 #include <unistd.h>
 
+#include "perfetto/base/logging.h"
 #include "perfetto/ext/base/file_utils.h"
 #include "perfetto/ext/base/optional.h"
 #include "perfetto/ext/base/string_utils.h"
@@ -188,14 +189,24 @@
       unwinding_workers_(MakeUnwindingWorkers(this, kUnwinderThreads)),
       socket_delegate_(this),
       weak_factory_(this) {
-  CheckDataSourceMemory();  // Kick off guardrail task.
-  stat_fd_.reset(open("/proc/self/stat", O_RDONLY | O_CLOEXEC));
-  if (!stat_fd_) {
+  auto stat_fd = base::OpenFile("/proc/self/stat", O_RDONLY | O_CLOEXEC);
+  if (!stat_fd) {
     PERFETTO_ELOG(
         "Failed to open /proc/self/stat. Cannot accept profiles "
         "with CPU guardrails.");
   } else {
-    CheckDataSourceCpu();  // Kick off guardrail task.
+    profiler_cpu_guardrails_.emplace(std::move(stat_fd));
+    CheckDataSourceCpuTask();  // Kick off guardrail task.
+  }
+
+  auto status_fd = base::OpenFile("/proc/self/status", O_RDONLY | O_CLOEXEC);
+  if (!status_fd) {
+    PERFETTO_ELOG(
+        "Failed to open /proc/self/status. Cannot accept profiles "
+        "with memory guardrails.");
+  } else {
+    profiler_memory_guardrails_.emplace(std::move(status_fd));
+    CheckDataSourceMemoryTask();  // Kick off guardrail task.
   }
 }
 
@@ -438,7 +449,8 @@
 
   base::Optional<uint64_t> start_cputime_sec;
   if (heapprofd_config.max_heapprofd_cpu_secs() > 0) {
-    start_cputime_sec = GetCputimeSec();
+    if (profiler_cpu_guardrails_)
+      start_cputime_sec = profiler_cpu_guardrails_->GetCputimeSec();
 
     if (!start_cputime_sec) {
       PERFETTO_ELOG("Failed to enforce CPU guardrail. Rejecting config.");
@@ -446,6 +458,12 @@
     }
   }
 
+  if (heapprofd_config.has_max_heapprofd_memory_kb() &&
+      !profiler_memory_guardrails_) {
+    PERFETTO_ELOG("Failed to enforce memory guardrail. Rejecting config.");
+    return;
+  }
+
   auto buffer_id = static_cast<BufferID>(ds_config.target_buffer());
   DataSource data_source(endpoint_->CreateTraceWriter(buffer_id));
   data_source.id = id;
@@ -479,20 +497,6 @@
   return false;
 }
 
-base::Optional<uint64_t> HeapprofdProducer::GetCputimeSec() {
-  if (!stat_fd_) {
-    return base::nullopt;
-  }
-  lseek(stat_fd_.get(), 0, SEEK_SET);
-  base::ProcStat stat;
-  if (!ReadProcStat(stat_fd_.get(), &stat)) {
-    PERFETTO_ELOG("Failed to read stat file to enforce guardrails.");
-    return base::nullopt;
-  }
-  return (stat.utime + stat.stime) /
-         static_cast<unsigned long>(sysconf(_SC_CLK_TCK));
-}
-
 void HeapprofdProducer::SetStartupProperties(DataSource* data_source) {
   const HeapprofdConfig& heapprofd_config = data_source->config;
   if (heapprofd_config.all())
@@ -1179,93 +1183,40 @@
   MaybeFinishDataSource(&ds);
 }
 
-void HeapprofdProducer::CheckDataSourceCpu() {
+void HeapprofdProducer::CheckDataSourceCpuTask() {
   auto weak_producer = weak_factory_.GetWeakPtr();
   task_runner_->PostDelayedTask(
       [weak_producer] {
         if (!weak_producer)
           return;
-        weak_producer->CheckDataSourceCpu();
+        weak_producer->CheckDataSourceCpuTask();
       },
       kGuardrailIntervalMs);
 
-  bool any_guardrail = std::any_of(
-      data_sources_.begin(), data_sources_.end(),
-      [](const std::pair<const DataSourceInstanceID, DataSource>& id_and_ds) {
-        const DataSource& ds = id_and_ds.second;
-        return ds.config.max_heapprofd_cpu_secs() > 0;
+  PERFETTO_DCHECK(profiler_cpu_guardrails_);
+  profiler_cpu_guardrails_->CheckDataSourceCpu(
+      data_sources_.begin(), data_sources_.end(), [this](DataSource* ds) {
+        ds->hit_guardrail = true;
+        ShutdownDataSource(ds);
       });
-
-  if (!any_guardrail)
-    return;
-
-  base::Optional<uint64_t> cputime_sec = GetCputimeSec();
-  if (!cputime_sec) {
-    PERFETTO_ELOG("Failed to get CPU time.");
-    return;
-  }
-
-  for (auto& id_and_ds : data_sources_) {
-    DataSource& ds = id_and_ds.second;
-    uint64_t ds_max_cpu = ds.config.max_heapprofd_cpu_secs();
-    if (ds_max_cpu > 0) {
-      // We reject data-sources with CPU guardrails if we cannot read the
-      // initial value.
-      PERFETTO_CHECK(ds.start_cputime_sec);
-      uint64_t cpu_diff = *cputime_sec - *ds.start_cputime_sec;
-      if (*cputime_sec > *ds.start_cputime_sec && cpu_diff > ds_max_cpu) {
-        PERFETTO_ELOG(
-            "Exceeded data-source CPU guardrail "
-            "(%" PRIu64 " > %" PRIu64 "). Shutting down.",
-            cpu_diff, ds_max_cpu);
-        ds.hit_guardrail = true;
-        ShutdownDataSource(&ds);
-      }
-    }
-  }
 }
 
-void HeapprofdProducer::CheckDataSourceMemory() {
+void HeapprofdProducer::CheckDataSourceMemoryTask() {
   auto weak_producer = weak_factory_.GetWeakPtr();
   task_runner_->PostDelayedTask(
       [weak_producer] {
         if (!weak_producer)
           return;
-        weak_producer->CheckDataSourceMemory();
+        weak_producer->CheckDataSourceMemoryTask();
       },
       kGuardrailIntervalMs);
 
-  bool any_guardrail = std::any_of(
-      data_sources_.begin(), data_sources_.end(),
-      [](const std::pair<const DataSourceInstanceID, DataSource>& id_and_ds) {
-        const DataSource& ds = id_and_ds.second;
-        return ds.config.max_heapprofd_memory_kb() > 0;
+  PERFETTO_DCHECK(profiler_memory_guardrails_);
+  profiler_memory_guardrails_->CheckDataSourceMemory(
+      data_sources_.begin(), data_sources_.end(), [this](DataSource* ds) {
+        ds->hit_guardrail = true;
+        ShutdownDataSource(ds);
       });
-
-  if (!any_guardrail)
-    return;
-
-  base::Optional<uint32_t> anon_and_swap;
-  base::Optional<std::string> status = ReadStatus(getpid());
-  if (status)
-    anon_and_swap = GetRssAnonAndSwap(*status);
-
-  if (!anon_and_swap) {
-    PERFETTO_ELOG("Failed to read heapprofd memory.");
-    return;
-  }
-
-  for (auto& id_and_ds : data_sources_) {
-    DataSource& ds = id_and_ds.second;
-    uint32_t ds_max_mem = ds.config.max_heapprofd_memory_kb();
-    if (ds_max_mem > 0 && *anon_and_swap > ds_max_mem) {
-      PERFETTO_ELOG("Exceeded data-source memory guardrail (%" PRIu32
-                    " > %" PRIu32 "). Shutting down.",
-                    *anon_and_swap, ds_max_mem);
-      ds.hit_guardrail = true;
-      ShutdownDataSource(&ds);
-    }
-  }
 }
 
 }  // namespace profiling
diff --git a/src/profiling/memory/heapprofd_producer.h b/src/profiling/memory/heapprofd_producer.h
index 6c001f6..32b95e2 100644
--- a/src/profiling/memory/heapprofd_producer.h
+++ b/src/profiling/memory/heapprofd_producer.h
@@ -37,6 +37,7 @@
 
 #include "src/profiling/common/interning_output.h"
 #include "src/profiling/common/proc_utils.h"
+#include "src/profiling/common/profiler_guardrails.h"
 #include "src/profiling/memory/bookkeeping.h"
 #include "src/profiling/memory/bookkeeping_dump.h"
 #include "src/profiling/memory/system_property.h"
@@ -226,6 +227,15 @@
       memset(&client_configuration, 0, sizeof(client_configuration));
     }
 
+    // For ProfilerCpuGuardrails.
+    uint64_t GetCpuGuardrailSecs() { return config.max_heapprofd_cpu_secs(); }
+
+    // For ProfilerCpuGuardrails.
+    base::Optional<uint64_t> GetCpuStartSecs() { return start_cputime_sec; }
+
+    // For CheckDataSourceMemory.
+    uint32_t GetMemoryGuardrailKb() { return config.max_heapprofd_memory_kb(); }
+
     DataSourceInstanceID id;
     std::unique_ptr<TraceWriter> trace_writer;
     HeapprofdConfig config;
@@ -258,10 +268,8 @@
   void ResetConnectionBackoff();
   void IncreaseConnectionBackoff();
 
-  base::Optional<uint64_t> GetCputimeSec();
-
-  void CheckDataSourceMemory();
-  void CheckDataSourceCpu();
+  void CheckDataSourceMemoryTask();
+  void CheckDataSourceCpuTask();
 
   void FinishDataSourceFlush(FlushRequestID flush_id);
   void DumpProcessesInDataSource(DataSource* ds);
@@ -331,7 +339,8 @@
   base::Optional<std::function<void()>> data_source_callback_;
 
   SocketDelegate socket_delegate_;
-  base::ScopedFile stat_fd_;
+  base::Optional<ProfilerCpuGuardrails> profiler_cpu_guardrails_;
+  base::Optional<ProfilerMemoryGuardrails> profiler_memory_guardrails_;
 
   base::WeakPtrFactory<HeapprofdProducer> weak_factory_;  // Keep last.
 };