perfetto: Add guardrails to perfetto_cmd
Add two guardrails to perfetto_cmd:
1) Can't trace if the last trace is less than 5mins ago.
2) Can't upload more than 10mb in a 24h period.
Bug: 73052379
Bug: 73053053
Change-Id: I6742b52510ae7a4e8a2b5c41648dd94566a37b33
diff --git a/Android.bp b/Android.bp
index 9c6a192..0d2d31f 100644
--- a/Android.bp
+++ b/Android.bp
@@ -136,6 +136,7 @@
":perfetto_protos_perfetto_trace_ps_zero_gen",
":perfetto_protos_perfetto_trace_zero_gen",
":perfetto_src_ipc_wire_protocol_gen",
+ ":perfetto_src_perfetto_cmd_protos_gen",
"src/base/android_task_runner.cc",
"src/base/file_utils.cc",
"src/base/page_allocator.cc",
@@ -152,6 +153,7 @@
"src/ipc/unix_socket.cc",
"src/perfetto_cmd/main.cc",
"src/perfetto_cmd/perfetto_cmd.cc",
+ "src/perfetto_cmd/rate_limiter.cc",
"src/protozero/message.cc",
"src/protozero/message_handle.cc",
"src/protozero/proto_utils.cc",
@@ -201,6 +203,7 @@
"perfetto_protos_perfetto_trace_ps_zero_gen_headers",
"perfetto_protos_perfetto_trace_zero_gen_headers",
"perfetto_src_ipc_wire_protocol_gen_headers",
+ "perfetto_src_perfetto_cmd_protos_gen_headers",
],
defaults: [
"perfetto_defaults",
@@ -2872,6 +2875,39 @@
],
}
+// GN target: //src/perfetto_cmd:protos_gen
+genrule {
+ name: "perfetto_src_perfetto_cmd_protos_gen",
+ srcs: [
+ "src/perfetto_cmd/perfetto_cmd_state.proto",
+ ],
+ tools: [
+ "aprotoc",
+ ],
+ cmd: "mkdir -p $(genDir)/external/perfetto && $(location aprotoc) --cpp_out=$(genDir)/external/perfetto --proto_path=external/perfetto/ $(in)",
+ out: [
+ "external/perfetto/src/perfetto_cmd/perfetto_cmd_state.pb.cc",
+ ],
+}
+
+// GN target: //src/perfetto_cmd:protos_gen
+genrule {
+ name: "perfetto_src_perfetto_cmd_protos_gen_headers",
+ srcs: [
+ "src/perfetto_cmd/perfetto_cmd_state.proto",
+ ],
+ tools: [
+ "aprotoc",
+ ],
+ cmd: "mkdir -p $(genDir)/external/perfetto && $(location aprotoc) --cpp_out=$(genDir)/external/perfetto --proto_path=external/perfetto/ $(in)",
+ out: [
+ "external/perfetto/src/perfetto_cmd/perfetto_cmd_state.pb.h",
+ ],
+ export_include_dirs: [
+ ".",
+ ],
+}
+
// GN target: //src/protozero/protoc_plugin:protoc_plugin(//gn/standalone/toolchain:gcc_like_host)
cc_binary_host {
name: "perfetto_src_protozero_protoc_plugin_protoc_plugin___gn_standalone_toolchain_gcc_like_host_",
@@ -3154,6 +3190,7 @@
":perfetto_src_ftrace_reader_ftrace_reader_test_messages_zero_gen",
":perfetto_src_ipc_test_messages_gen",
":perfetto_src_ipc_wire_protocol_gen",
+ ":perfetto_src_perfetto_cmd_protos_gen",
":perfetto_src_protozero_testing_messages_lite_gen",
":perfetto_src_protozero_testing_messages_zero_gen",
"src/base/android_task_runner.cc",
@@ -3209,6 +3246,9 @@
"src/ipc/test/ipc_integrationtest.cc",
"src/ipc/unix_socket.cc",
"src/ipc/unix_socket_unittest.cc",
+ "src/perfetto_cmd/perfetto_cmd.cc",
+ "src/perfetto_cmd/rate_limiter.cc",
+ "src/perfetto_cmd/rate_limiter_unittest.cc",
"src/process_stats/file_utils.cc",
"src/process_stats/procfs_utils.cc",
"src/protozero/message.cc",
@@ -3255,6 +3295,8 @@
"src/tracing/core/trace_packet_unittest.cc",
"src/tracing/core/trace_writer_impl.cc",
"src/tracing/core/trace_writer_impl_unittest.cc",
+ "src/tracing/ipc/consumer/consumer_ipc_client_impl.cc",
+ "src/tracing/ipc/posix_shared_memory.cc",
"src/tracing/ipc/posix_shared_memory_unittest.cc",
"src/tracing/test/aligned_buffer_test.cc",
"src/tracing/test/fake_packet.cc",
@@ -3292,6 +3334,7 @@
"perfetto_src_ftrace_reader_ftrace_reader_test_messages_zero_gen_headers",
"perfetto_src_ipc_test_messages_gen_headers",
"perfetto_src_ipc_wire_protocol_gen_headers",
+ "perfetto_src_perfetto_cmd_protos_gen_headers",
"perfetto_src_protozero_testing_messages_lite_gen_headers",
"perfetto_src_protozero_testing_messages_zero_gen_headers",
],
diff --git a/BUILD.gn b/BUILD.gn
index ea02315..959ec28 100644
--- a/BUILD.gn
+++ b/BUILD.gn
@@ -60,6 +60,7 @@
"gn:gtest_main",
"src/base:unittests",
"src/ipc:unittests",
+ "src/perfetto_cmd:unittests",
"src/protozero:unittests",
"src/traced/probes/filesystem:unittests",
"src/tracing:unittests",
diff --git a/include/perfetto/base/time.h b/include/perfetto/base/time.h
index 90e8535..bd2498c 100644
--- a/include/perfetto/base/time.h
+++ b/include/perfetto/base/time.h
@@ -26,6 +26,7 @@
namespace perfetto {
namespace base {
+using TimeSeconds = std::chrono::seconds;
using TimeMillis = std::chrono::milliseconds;
using TimeNanos = std::chrono::nanoseconds;
constexpr clockid_t kWallTimeClockSource = CLOCK_MONOTONIC;
@@ -48,6 +49,10 @@
return std::chrono::duration_cast<TimeMillis>(GetWallTimeNs());
}
+inline TimeSeconds GetWallTimeS() {
+ return std::chrono::duration_cast<TimeSeconds>(GetWallTimeNs());
+}
+
inline TimeNanos GetThreadCPUTimeNs() {
return GetTimeInternalNs(CLOCK_THREAD_CPUTIME_ID);
}
diff --git a/src/ftrace_reader/BUILD.gn b/src/ftrace_reader/BUILD.gn
index 229fc3b..7e2a6a6 100644
--- a/src/ftrace_reader/BUILD.gn
+++ b/src/ftrace_reader/BUILD.gn
@@ -50,7 +50,7 @@
"../../gn:default_deps",
"../../gn:gtest_deps",
"../../protos/perfetto/trace/ftrace:lite",
- "../../src/base:test_support",
+ "../base:test_support",
]
sources = [
"cpu_reader_unittest.cc",
diff --git a/src/perfetto_cmd/BUILD.gn b/src/perfetto_cmd/BUILD.gn
index 7474760..4026e04 100644
--- a/src/perfetto_cmd/BUILD.gn
+++ b/src/perfetto_cmd/BUILD.gn
@@ -12,8 +12,12 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+import("../../gn/perfetto.gni")
+import("../../gn/proto_library.gni")
+
source_set("perfetto_cmd") {
public_deps = [
+ ":protos",
"../../include/perfetto/traced",
]
deps = [
@@ -25,5 +29,31 @@
]
sources = [
"perfetto_cmd.cc",
+ "perfetto_cmd.h",
+ "rate_limiter.cc",
+ "rate_limiter.h",
+ ]
+}
+
+proto_library("protos") {
+ generate_python = false
+ deps = []
+ sources = [
+ "perfetto_cmd_state.proto",
+ ]
+ proto_in_dir = perfetto_root_path
+ proto_out_dir = perfetto_root_path
+}
+
+source_set("unittests") {
+ testonly = true
+ public_deps = []
+ deps = [
+ ":perfetto_cmd",
+ "../../gn:default_deps",
+ "../../gn:gtest_deps",
+ ]
+ sources = [
+ "rate_limiter_unittest.cc",
]
}
diff --git a/src/perfetto_cmd/perfetto_cmd.cc b/src/perfetto_cmd/perfetto_cmd.cc
index f27b931..61de572 100644
--- a/src/perfetto_cmd/perfetto_cmd.cc
+++ b/src/perfetto_cmd/perfetto_cmd.cc
@@ -14,37 +14,35 @@
* limitations under the License.
*/
+#include "perfetto_cmd.h"
+
#include <fcntl.h>
#include <getopt.h>
+#include <stdio.h>
#include <sys/stat.h>
+#include <time.h>
#include <unistd.h>
#include <fstream>
#include <iostream>
#include <iterator>
#include <sstream>
-#include <string>
#include "perfetto/base/file_utils.h"
#include "perfetto/base/logging.h"
-#include "perfetto/base/scoped_file.h"
-#include "perfetto/base/unix_task_runner.h"
+#include "perfetto/base/time.h"
#include "perfetto/base/utils.h"
#include "perfetto/protozero/proto_utils.h"
#include "perfetto/traced/traced.h"
-#include "perfetto/tracing/core/consumer.h"
#include "perfetto/tracing/core/data_source_config.h"
#include "perfetto/tracing/core/data_source_descriptor.h"
#include "perfetto/tracing/core/trace_config.h"
#include "perfetto/tracing/core/trace_packet.h"
-#include "perfetto/tracing/ipc/consumer_ipc_client.h"
#include "perfetto/config/trace_config.pb.h"
#include "perfetto/trace/trace.pb.h"
-#if defined(PERFETTO_OS_ANDROID)
-#include "perfetto/base/android_task_runner.h"
-#endif // defined(PERFETTO_OS_ANDROID)
+#include "google/protobuf/io/zero_copy_stream_impl_lite.h"
#if defined(PERFETTO_BUILD_WITH_ANDROID)
#include <android/os/DropBoxManager.h>
@@ -58,9 +56,6 @@
namespace perfetto {
namespace {
-// Temporary directory for DropBox traces. Note that this is automatically
-// created by the system by setting setprop persist.traced.enable=1.
-constexpr char kTempDropBoxTraceDir[] = "/data/misc/perfetto-traces";
constexpr char kDefaultDropBoxTag[] = "perfetto";
std::string GetDirName(const std::string& path) {
@@ -72,51 +67,22 @@
} // namespace
+// Temporary directory for DropBox traces. Note that this is automatically
+// created by the system by setting setprop persist.traced.enable=1.
+const char* kTempDropBoxTraceDir = "/data/misc/perfetto-traces";
+
using protozero::proto_utils::WriteVarInt;
using protozero::proto_utils::MakeTagLengthDelimited;
-#if defined(PERFETTO_OS_ANDROID)
-using PlatformTaskRunner = base::AndroidTaskRunner;
-#else
-using PlatformTaskRunner = base::UnixTaskRunner;
-#endif
-
-class PerfettoCmd : public Consumer {
- public:
- int Main(int argc, char** argv);
- int PrintUsage(const char* argv0);
- void OnStopTraceTimer();
- void OnTimeout();
-
- // perfetto::Consumer implementation.
- void OnConnect() override;
- void OnDisconnect() override;
- void OnTraceData(std::vector<TracePacket>, bool has_more) override;
-
- private:
- bool OpenOutputFile();
-
- PlatformTaskRunner task_runner_;
- std::unique_ptr<perfetto::Service::ConsumerEndpoint> consumer_endpoint_;
- std::unique_ptr<TraceConfig> trace_config_;
- base::ScopedFstream trace_out_stream_;
- std::string trace_out_path_;
-
- // Only used if linkat(AT_FDCWD) isn't available.
- std::string tmp_trace_out_path_;
-
- std::string dropbox_tag_;
- bool did_process_full_trace_ = false;
-};
-
int PerfettoCmd::PrintUsage(const char* argv0) {
PERFETTO_ELOG(R"(
Usage: %s
- --background -b : Exits immediately and continues tracing in background
- --config -c : /path/to/trace/config/file or - for stdin
- --out -o : /path/to/out/trace/file
- --dropbox -d TAG : Upload trace into DropBox using tag TAG (default: %s)
- --help -h
+ --background -b : Exits immediately and continues tracing in background
+ --config -c : /path/to/trace/config/file or - for stdin
+ --out -o : /path/to/out/trace/file
+ --dropbox -d TAG : Upload trace into DropBox using tag TAG (default: %s)
+ --no-guardrails -n : Ignore guardrails triggered when using --dropbox (for testing).
+ --help -h
)",
argv0, kDefaultDropBoxTag);
return 1;
@@ -130,14 +96,16 @@
{"out", required_argument, 0, 'o'},
{"background", no_argument, 0, 'b'},
{"dropbox", optional_argument, 0, 'd'},
+ {"no-guardrails", optional_argument, 0, 'n'},
{nullptr, 0, nullptr, 0}};
int option_index = 0;
std::string trace_config_raw;
bool background = false;
+ bool ignore_guardrails = false;
for (;;) {
int option =
- getopt_long(argc, argv, "c:o:bd::", long_options, &option_index);
+ getopt_long(argc, argv, "c:o:bd::n", long_options, &option_index);
if (option == -1)
break; // EOF.
@@ -150,7 +118,7 @@
// TODO(primiano): temporary for testing only.
perfetto::protos::TraceConfig test_config;
test_config.add_buffers()->set_size_kb(4096);
- test_config.set_duration_ms(10000);
+ test_config.set_duration_ms(2000);
auto* ds_config = test_config.add_data_sources()->mutable_config();
ds_config->set_name("com.google.perfetto.ftrace");
ds_config->mutable_ftrace_config()->add_ftrace_events("sched_switch");
@@ -186,6 +154,12 @@
background = true;
continue;
}
+
+ if (option == 'n') {
+ ignore_guardrails = true;
+ continue;
+ }
+
return PrintUsage(argv[0]);
}
@@ -224,11 +198,23 @@
PERFETTO_DLOG("Continuing in background");
}
+ RateLimiter limiter;
+ RateLimiter::Args args{};
+ args.is_dropbox = !dropbox_tag_.empty();
+ args.current_time = base::GetWallTimeS();
+ args.ignore_guardrails = ignore_guardrails;
+ if (!limiter.ShouldTrace(args))
+ return 1;
+
consumer_endpoint_ = ConsumerIPCClient::Connect(PERFETTO_CONSUMER_SOCK_NAME,
this, &task_runner_);
task_runner_.Run();
- return did_process_full_trace_ ? 0 : 1;
-} // namespace perfetto
+
+ return limiter.OnTraceDone(args, did_process_full_trace_,
+ bytes_uploaded_to_dropbox_)
+ ? 0
+ : 1;
+}
void PerfettoCmd::OnConnect() {
PERFETTO_LOG(
@@ -284,7 +270,13 @@
fflush(*trace_out_stream_);
long bytes_written = ftell(*trace_out_stream_);
- if (!dropbox_tag_.empty()) {
+ if (dropbox_tag_.empty()) {
+ trace_out_stream_.reset();
+ PERFETTO_CHECK(
+ rename(tmp_trace_out_path_.c_str(), trace_out_path_.c_str()) == 0);
+ PERFETTO_ILOG("Wrote %ld bytes into %s", bytes_written,
+ trace_out_path_.c_str());
+ } else {
#if defined(PERFETTO_BUILD_WITH_ANDROID)
android::sp<android::os::DropBoxManager> dropbox =
new android::os::DropBoxManager();
@@ -304,15 +296,11 @@
PERFETTO_ELOG("DropBox upload failed: %s", status.toString8().c_str());
return;
}
+ // TODO(hjd): Account for compression.
+ bytes_uploaded_to_dropbox_ = bytes_written;
PERFETTO_ILOG("Uploaded %ld bytes into DropBox with tag %s", bytes_written,
dropbox_tag_.c_str());
#endif // defined(PERFETTO_BUILD_WITH_ANDROID)
- } else {
- trace_out_stream_.reset();
- PERFETTO_CHECK(
- rename(tmp_trace_out_path_.c_str(), trace_out_path_.c_str()) == 0);
- PERFETTO_ILOG("Wrote %ld bytes into %s", bytes_written,
- trace_out_path_.c_str());
}
did_process_full_trace_ = true;
}
diff --git a/src/perfetto_cmd/perfetto_cmd.h b/src/perfetto_cmd/perfetto_cmd.h
new file mode 100644
index 0000000..ced889b
--- /dev/null
+++ b/src/perfetto_cmd/perfetto_cmd.h
@@ -0,0 +1,87 @@
+/*
+ * Copyright (C) 2018 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_PERFETTO_CMD_PERFETTO_CMD_H_
+#define SRC_PERFETTO_CMD_PERFETTO_CMD_H_
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <time.h>
+
+#include "perfetto/base/scoped_file.h"
+#include "perfetto/base/unix_task_runner.h"
+#include "perfetto/tracing/core/consumer.h"
+#include "perfetto/tracing/ipc/consumer_ipc_client.h"
+#include "rate_limiter.h"
+
+#include "src/perfetto_cmd/perfetto_cmd_state.pb.h"
+
+#if defined(PERFETTO_OS_ANDROID)
+#include "perfetto/base/android_task_runner.h"
+#endif // defined(PERFETTO_OS_ANDROID)
+
+#if defined(PERFETTO_BUILD_WITH_ANDROID)
+#include <android/os/DropBoxManager.h>
+#include <utils/Looper.h>
+#include <utils/StrongPointer.h>
+#endif // defined(PERFETTO_BUILD_WITH_ANDROID)
+
+namespace perfetto {
+
+// Temporary directory for DropBox traces. Note that this is automatically
+// created by the system by setting setprop persist.traced.enable=1.
+extern const char* kTempDropBoxTraceDir;
+
+#if defined(PERFETTO_OS_ANDROID)
+using PlatformTaskRunner = base::AndroidTaskRunner;
+#else
+using PlatformTaskRunner = base::UnixTaskRunner;
+#endif
+
+class PerfettoCmd : public Consumer {
+ public:
+ int Main(int argc, char** argv);
+ int PrintUsage(const char* argv0);
+ void OnStopTraceTimer();
+ void OnTimeout();
+
+ // perfetto::Consumer implementation.
+ void OnConnect() override;
+ void OnDisconnect() override;
+ void OnTraceData(std::vector<TracePacket>, bool has_more) override;
+
+ private:
+ bool OpenOutputFile();
+
+ PlatformTaskRunner task_runner_;
+ std::unique_ptr<perfetto::Service::ConsumerEndpoint> consumer_endpoint_;
+ std::unique_ptr<TraceConfig> trace_config_;
+ base::ScopedFstream trace_out_stream_;
+ std::string trace_out_path_;
+
+ // Only used if linkat(AT_FDCWD) isn't available.
+ std::string tmp_trace_out_path_;
+
+ std::string dropbox_tag_;
+ bool did_process_full_trace_ = false;
+ size_t bytes_uploaded_to_dropbox_ = 0;
+};
+
+} // namespace perfetto
+
+#endif // SRC_PERFETTO_CMD_PERFETTO_CMD_H_
diff --git a/src/perfetto_cmd/perfetto_cmd_state.proto b/src/perfetto_cmd/perfetto_cmd_state.proto
new file mode 100644
index 0000000..52b56fb
--- /dev/null
+++ b/src/perfetto_cmd/perfetto_cmd_state.proto
@@ -0,0 +1,28 @@
+/*
+ * Copyright (C) 2018 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.
+ */
+
+syntax = "proto2";
+option optimize_for = LITE_RUNTIME;
+
+package perfetto;
+
+message PerfettoCmdState {
+ // These are fixed size so that the file is always a constant length.
+ optional fixed64 first_trace_timestamp = 1;
+ optional fixed64 last_trace_timestamp = 2;
+ // Total of bytes uploaded since first_trace_timestamp.
+ optional fixed64 total_bytes_uploaded = 3;
+}
diff --git a/src/perfetto_cmd/rate_limiter.cc b/src/perfetto_cmd/rate_limiter.cc
new file mode 100644
index 0000000..a418d0f
--- /dev/null
+++ b/src/perfetto_cmd/rate_limiter.cc
@@ -0,0 +1,184 @@
+/*
+ * Copyright (C) 2018 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 "rate_limiter.h"
+
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+#include "perfetto/base/logging.h"
+#include "perfetto/base/scoped_file.h"
+#include "perfetto/base/utils.h"
+#include "src/perfetto_cmd/perfetto_cmd.h"
+
+namespace perfetto {
+namespace {
+
+// 5 mins between traces.
+const uint64_t kCooldownInSeconds = 60 * 5;
+
+// Every 24 hours we reset how much we've uploaded.
+const uint64_t kMaxUploadResetPeriodInSeconds = 60 * 60 * 24;
+
+// Maximum of 10mb every 24h.
+const uint64_t kMaxUploadInBytes = 1024 * 1024 * 10;
+
+} // namespace
+
+RateLimiter::RateLimiter() = default;
+RateLimiter::~RateLimiter() = default;
+
+bool RateLimiter::ShouldTrace(const Args& args) {
+ uint64_t now_in_s = args.current_time.count();
+
+ // Not uploading?
+ // -> We can just trace.
+ if (!args.is_dropbox)
+ return true;
+
+ // The state file is gone.
+ // Maybe we're tracing for the first time or maybe something went wrong the
+ // last time we tried to save the state. Either way reinitialize the state
+ // file.
+ if (!StateFileExists()) {
+ // We can't write the empty state file?
+ // -> Give up.
+ if (!ClearState()) {
+ PERFETTO_ELOG("Guardrail: failed to initialize guardrail state.");
+ return false;
+ }
+ }
+
+ bool loaded_state = LoadState(&state_);
+
+ // Failed to load the state?
+ // Current time is before either saved times?
+ // Last saved trace time is before first saved trace time?
+ // -> Try to save a clean state but don't trace.
+ if (!loaded_state || now_in_s < state_.first_trace_timestamp() ||
+ now_in_s < state_.last_trace_timestamp() ||
+ state_.last_trace_timestamp() < state_.first_trace_timestamp()) {
+ ClearState();
+ PERFETTO_ELOG("Guardrail: state invalid, clearing it.");
+ if (!args.ignore_guardrails)
+ return false;
+ }
+
+ // If we've uploaded in the last 5mins we shouldn't trace now.
+ if ((now_in_s - state_.last_trace_timestamp()) < kCooldownInSeconds) {
+ PERFETTO_ELOG("Guardrail: Uploaded to DropBox in the last 5mins.");
+ if (!args.ignore_guardrails)
+ return false;
+ }
+
+ // First trace was more than 24h ago? Reset state.
+ if ((now_in_s - state_.first_trace_timestamp()) >
+ kMaxUploadResetPeriodInSeconds) {
+ state_.set_first_trace_timestamp(0);
+ state_.set_last_trace_timestamp(0);
+ state_.set_total_bytes_uploaded(0);
+ return true;
+ }
+
+ // If we've uploaded more than 10mb in the last 24 hours we shouldn't trace
+ // now.
+ if (state_.total_bytes_uploaded() > kMaxUploadInBytes) {
+ PERFETTO_ELOG("Guardrail: Uploaded >10mb DropBox in the last 24h.");
+ if (!args.ignore_guardrails)
+ return false;
+ }
+
+ return true;
+}
+
+bool RateLimiter::OnTraceDone(const Args& args, bool success, size_t bytes) {
+ uint64_t now_in_s = args.current_time.count();
+
+ // Failed to upload? Don't update the state.
+ if (!success)
+ return false;
+
+ if (!args.is_dropbox)
+ return true;
+
+ // If the first trace timestamp is 0 (either because this is the
+ // first time or because it was reset for being more than 24h ago).
+ // -> We update it to the time of this trace.
+ if (state_.first_trace_timestamp() == 0)
+ state_.set_first_trace_timestamp(now_in_s);
+ // Always updated the last trace timestamp.
+ state_.set_last_trace_timestamp(now_in_s);
+ // Add the amount we uploaded to the running total.
+ state_.set_total_bytes_uploaded(state_.total_bytes_uploaded() + bytes);
+
+ if (!SaveState(state_)) {
+ PERFETTO_ELOG("Failed to save state.");
+ return false;
+ }
+
+ return true;
+}
+
+std::string RateLimiter::GetStateFilePath() const {
+ return std::string(kTempDropBoxTraceDir) + "/.guardraildata";
+}
+
+bool RateLimiter::StateFileExists() {
+ struct stat out;
+ return stat(GetStateFilePath().c_str(), &out) != -1;
+}
+
+bool RateLimiter::ClearState() {
+ PerfettoCmdState zero{};
+ zero.set_total_bytes_uploaded(0);
+ zero.set_last_trace_timestamp(0);
+ zero.set_first_trace_timestamp(0);
+ bool success = SaveState(zero);
+ if (!success && StateFileExists())
+ remove(GetStateFilePath().c_str());
+ return success;
+}
+
+bool RateLimiter::LoadState(PerfettoCmdState* state) {
+ base::ScopedFile in_fd;
+ in_fd.reset(open(GetStateFilePath().c_str(), O_RDONLY));
+
+ if (!in_fd)
+ return false;
+ char buf[1024];
+ ssize_t bytes = PERFETTO_EINTR(read(in_fd.get(), &buf, sizeof(buf)));
+ if (bytes <= 0)
+ return false;
+ return state->ParseFromArray(&buf, bytes);
+}
+
+bool RateLimiter::SaveState(const PerfettoCmdState& state) {
+ base::ScopedFile out_fd;
+ out_fd.reset(
+ open(GetStateFilePath().c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0600));
+ if (!out_fd)
+ return false;
+ char buf[1024];
+ PERFETTO_CHECK(static_cast<size_t>(state.ByteSize()) < sizeof(buf));
+ size_t size = state.ByteSize();
+ if (!state.SerializeToArray(&buf, size))
+ return false;
+ ssize_t written = PERFETTO_EINTR(write(out_fd.get(), &buf, size));
+ return written >= 0 && static_cast<size_t>(written) == size;
+}
+
+} // namespace perfetto
diff --git a/src/perfetto_cmd/rate_limiter.h b/src/perfetto_cmd/rate_limiter.h
new file mode 100644
index 0000000..c163357
--- /dev/null
+++ b/src/perfetto_cmd/rate_limiter.h
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2018 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_PERFETTO_CMD_RATE_LIMITER_H_
+#define SRC_PERFETTO_CMD_RATE_LIMITER_H_
+
+#include "perfetto/base/time.h"
+#include "src/perfetto_cmd/perfetto_cmd_state.pb.h"
+
+namespace perfetto {
+
+class RateLimiter {
+ public:
+ struct Args {
+ bool is_dropbox = false;
+ bool ignore_guardrails = false;
+ base::TimeSeconds current_time = base::TimeSeconds(0);
+ };
+
+ RateLimiter();
+ virtual ~RateLimiter();
+
+ bool ShouldTrace(const Args& args);
+ bool OnTraceDone(const Args& args, bool success, size_t bytes);
+
+ bool ClearState();
+
+ // virtual for testing.
+ virtual bool LoadState(PerfettoCmdState* state);
+
+ // virtual for testing.
+ virtual bool SaveState(const PerfettoCmdState& state);
+
+ bool StateFileExists();
+ virtual std::string GetStateFilePath() const;
+
+ private:
+ PerfettoCmdState state_{};
+};
+
+} // namespace perfetto
+
+#endif // SRC_PERFETTO_CMD_RATE_LIMITER_H_
diff --git a/src/perfetto_cmd/rate_limiter_unittest.cc b/src/perfetto_cmd/rate_limiter_unittest.cc
new file mode 100644
index 0000000..13d35ab
--- /dev/null
+++ b/src/perfetto_cmd/rate_limiter_unittest.cc
@@ -0,0 +1,338 @@
+/*
+ * Copyright (C) 2018 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 "rate_limiter.h"
+
+#include <stdio.h>
+
+#include "perfetto/base/scoped_file.h"
+#include "perfetto/base/temp_file.h"
+#include "perfetto/base/utils.h"
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+using testing::_;
+using testing::NiceMock;
+using testing::StrictMock;
+using testing::Invoke;
+using testing::Return;
+
+namespace perfetto {
+
+namespace {
+
+class MockRateLimiter : public RateLimiter {
+ public:
+ MockRateLimiter() : dir_(base::TempDir::Create()) {
+ ON_CALL(*this, LoadState(_))
+ .WillByDefault(Invoke(this, &MockRateLimiter::LoadStateConcrete));
+ ON_CALL(*this, SaveState(_))
+ .WillByDefault(Invoke(this, &MockRateLimiter::SaveStateConcrete));
+ }
+
+ virtual std::string GetStateFilePath() const {
+ return std::string(dir_.path()) + "/.guardraildata";
+ }
+
+ virtual ~MockRateLimiter() override {
+ if (StateFileExists())
+ remove(GetStateFilePath().c_str());
+ }
+
+ bool LoadStateConcrete(PerfettoCmdState* state) {
+ return RateLimiter::LoadState(state);
+ }
+
+ bool SaveStateConcrete(const PerfettoCmdState& state) {
+ return RateLimiter::SaveState(state);
+ }
+
+ MOCK_METHOD1(LoadState, bool(PerfettoCmdState*));
+ MOCK_METHOD1(SaveState, bool(const PerfettoCmdState&));
+
+ private:
+ base::TempDir dir_;
+};
+
+void WriteGarbageToFile(const std::string& path) {
+ base::ScopedFile fd;
+ fd.reset(open(path.c_str(), O_WRONLY | O_CREAT, 0600));
+ constexpr char data[] = "Some random bytes.";
+ if (write(fd.get(), data, sizeof(data)) != sizeof(data))
+ ADD_FAILURE() << "Could not write garbage";
+}
+
+TEST(RateLimiterTest, RoundTripState) {
+ NiceMock<MockRateLimiter> limiter;
+
+ PerfettoCmdState input{};
+ PerfettoCmdState output{};
+
+ input.set_total_bytes_uploaded(42);
+ ASSERT_TRUE(limiter.SaveState(input));
+ ASSERT_TRUE(limiter.LoadState(&output));
+ ASSERT_EQ(output.total_bytes_uploaded(), 42u);
+}
+
+TEST(RateLimiterTest, LoadFromEmpty) {
+ NiceMock<MockRateLimiter> limiter;
+
+ PerfettoCmdState input{};
+ input.set_total_bytes_uploaded(0);
+ input.set_last_trace_timestamp(0);
+ input.set_first_trace_timestamp(0);
+ PerfettoCmdState output{};
+
+ ASSERT_TRUE(limiter.SaveState(input));
+ ASSERT_TRUE(limiter.LoadState(&output));
+ ASSERT_EQ(output.total_bytes_uploaded(), 0u);
+}
+
+TEST(RateLimiterTest, LoadFromNoFileFails) {
+ NiceMock<MockRateLimiter> limiter;
+ PerfettoCmdState output{};
+ ASSERT_FALSE(limiter.LoadState(&output));
+ ASSERT_EQ(output.total_bytes_uploaded(), 0u);
+}
+
+TEST(RateLimiterTest, LoadFromGarbageFails) {
+ NiceMock<MockRateLimiter> limiter;
+
+ WriteGarbageToFile(limiter.GetStateFilePath().c_str());
+
+ PerfettoCmdState output{};
+ ASSERT_FALSE(limiter.LoadState(&output));
+ ASSERT_EQ(output.total_bytes_uploaded(), 0u);
+}
+
+TEST(RateLimiterTest, NotDropBox) {
+ StrictMock<MockRateLimiter> limiter;
+
+ ASSERT_TRUE(limiter.ShouldTrace({}));
+ ASSERT_TRUE(limiter.OnTraceDone({}, true, 10000));
+ ASSERT_FALSE(limiter.StateFileExists());
+}
+
+TEST(RateLimiterTest, NotDropBox_FailedToTrace) {
+ StrictMock<MockRateLimiter> limiter;
+
+ ASSERT_FALSE(limiter.OnTraceDone({}, false, 0));
+ ASSERT_FALSE(limiter.StateFileExists());
+}
+
+TEST(RateLimiterTest, DropBox_IgnoreGuardrails) {
+ StrictMock<MockRateLimiter> limiter;
+ RateLimiter::Args args;
+
+ args.is_dropbox = true;
+ args.ignore_guardrails = true;
+ args.current_time = base::TimeSeconds(41);
+
+ EXPECT_CALL(limiter, SaveState(_));
+ EXPECT_CALL(limiter, LoadState(_));
+ ASSERT_TRUE(limiter.ShouldTrace(args));
+
+ EXPECT_CALL(limiter, SaveState(_));
+ ASSERT_TRUE(limiter.OnTraceDone(args, true, 42u));
+
+ PerfettoCmdState output{};
+ ASSERT_TRUE(limiter.LoadStateConcrete(&output));
+ ASSERT_EQ(output.first_trace_timestamp(), 41u);
+ ASSERT_EQ(output.last_trace_timestamp(), 41u);
+ ASSERT_EQ(output.total_bytes_uploaded(), 42u);
+}
+
+TEST(RateLimiterTest, DropBox_EmptyState) {
+ StrictMock<MockRateLimiter> limiter;
+ RateLimiter::Args args;
+
+ args.is_dropbox = true;
+ args.current_time = base::TimeSeconds(10000);
+
+ EXPECT_CALL(limiter, SaveState(_));
+ EXPECT_CALL(limiter, LoadState(_));
+ ASSERT_TRUE(limiter.ShouldTrace(args));
+
+ EXPECT_CALL(limiter, SaveState(_));
+ ASSERT_TRUE(limiter.OnTraceDone(args, true, 1024 * 1024));
+
+ PerfettoCmdState output{};
+ ASSERT_TRUE(limiter.LoadStateConcrete(&output));
+ EXPECT_EQ(output.total_bytes_uploaded(), 1024u * 1024u);
+ EXPECT_EQ(output.first_trace_timestamp(), 10000u);
+ EXPECT_EQ(output.last_trace_timestamp(), 10000u);
+}
+
+TEST(RateLimiterTest, DropBox_NormalUpload) {
+ StrictMock<MockRateLimiter> limiter;
+ RateLimiter::Args args;
+
+ PerfettoCmdState input{};
+ input.set_first_trace_timestamp(10000);
+ input.set_last_trace_timestamp(10000 + 60 * 10);
+ input.set_total_bytes_uploaded(1024 * 1024 * 2);
+ ASSERT_TRUE(limiter.SaveStateConcrete(input));
+
+ args.is_dropbox = true;
+ args.current_time = base::TimeSeconds(input.last_trace_timestamp() + 60 * 10);
+
+ EXPECT_CALL(limiter, LoadState(_));
+ ASSERT_TRUE(limiter.ShouldTrace(args));
+
+ EXPECT_CALL(limiter, SaveState(_));
+ ASSERT_TRUE(limiter.OnTraceDone(args, true, 1024 * 1024));
+
+ PerfettoCmdState output{};
+ ASSERT_TRUE(limiter.LoadStateConcrete(&output));
+ EXPECT_EQ(output.total_bytes_uploaded(), 1024u * 1024u * 3);
+ EXPECT_EQ(output.first_trace_timestamp(), input.first_trace_timestamp());
+ EXPECT_EQ(output.last_trace_timestamp(),
+ static_cast<uint64_t>(args.current_time.count()));
+}
+
+TEST(RateLimiterTest, DropBox_FailedToLoadState) {
+ StrictMock<MockRateLimiter> limiter;
+ RateLimiter::Args args;
+
+ args.is_dropbox = true;
+
+ WriteGarbageToFile(limiter.GetStateFilePath().c_str());
+
+ EXPECT_CALL(limiter, LoadState(_));
+ EXPECT_CALL(limiter, SaveState(_));
+ ASSERT_FALSE(limiter.ShouldTrace(args));
+
+ PerfettoCmdState output{};
+ ASSERT_TRUE(limiter.LoadStateConcrete(&output));
+ EXPECT_EQ(output.total_bytes_uploaded(), 0u);
+ EXPECT_EQ(output.first_trace_timestamp(), 0u);
+ EXPECT_EQ(output.last_trace_timestamp(), 0u);
+}
+
+TEST(RateLimiterTest, DropBox_NoTimeTravel) {
+ StrictMock<MockRateLimiter> limiter;
+ RateLimiter::Args args;
+
+ PerfettoCmdState input{};
+ input.set_first_trace_timestamp(100);
+ input.set_last_trace_timestamp(100);
+ ASSERT_TRUE(limiter.SaveStateConcrete(input));
+
+ args.is_dropbox = true;
+ args.current_time = base::TimeSeconds(99);
+
+ EXPECT_CALL(limiter, LoadState(_));
+ EXPECT_CALL(limiter, SaveState(_));
+ ASSERT_FALSE(limiter.ShouldTrace(args));
+
+ PerfettoCmdState output{};
+ ASSERT_TRUE(limiter.LoadStateConcrete(&output));
+ EXPECT_EQ(output.total_bytes_uploaded(), 0u);
+ EXPECT_EQ(output.first_trace_timestamp(), 0u);
+ EXPECT_EQ(output.last_trace_timestamp(), 0u);
+}
+
+TEST(RateLimiterTest, DropBox_TooSoon) {
+ StrictMock<MockRateLimiter> limiter;
+ RateLimiter::Args args;
+
+ PerfettoCmdState input{};
+ input.set_first_trace_timestamp(10000);
+ input.set_last_trace_timestamp(10000);
+ ASSERT_TRUE(limiter.SaveStateConcrete(input));
+
+ args.is_dropbox = true;
+ args.current_time = base::TimeSeconds(10000 + 60 * 4);
+
+ EXPECT_CALL(limiter, LoadState(_));
+ ASSERT_FALSE(limiter.ShouldTrace(args));
+}
+
+TEST(RateLimiterTest, DropBox_TooMuch) {
+ StrictMock<MockRateLimiter> limiter;
+ RateLimiter::Args args;
+
+ PerfettoCmdState input{};
+ input.set_total_bytes_uploaded(10 * 1024 * 1024 + 1);
+ ASSERT_TRUE(limiter.SaveStateConcrete(input));
+
+ args.is_dropbox = true;
+ args.current_time = base::TimeSeconds(60 * 60);
+
+ EXPECT_CALL(limiter, LoadState(_));
+ ASSERT_FALSE(limiter.ShouldTrace(args));
+}
+
+TEST(RateLimiterTest, DropBox_TooMuchWasUploaded) {
+ StrictMock<MockRateLimiter> limiter;
+ RateLimiter::Args args;
+
+ PerfettoCmdState input{};
+ input.set_first_trace_timestamp(1);
+ input.set_last_trace_timestamp(1);
+ input.set_total_bytes_uploaded(10 * 1024 * 1024 + 1);
+ ASSERT_TRUE(limiter.SaveStateConcrete(input));
+
+ args.is_dropbox = true;
+ args.current_time = base::TimeSeconds(60 * 60 * 24 + 2);
+
+ EXPECT_CALL(limiter, LoadState(_));
+ ASSERT_TRUE(limiter.ShouldTrace(args));
+
+ EXPECT_CALL(limiter, SaveState(_));
+ ASSERT_TRUE(limiter.OnTraceDone(args, true, 1024 * 1024));
+
+ PerfettoCmdState output{};
+ ASSERT_TRUE(limiter.LoadStateConcrete(&output));
+ EXPECT_EQ(output.total_bytes_uploaded(), 1024u * 1024u);
+ EXPECT_EQ(output.first_trace_timestamp(),
+ static_cast<uint64_t>(args.current_time.count()));
+ EXPECT_EQ(output.last_trace_timestamp(),
+ static_cast<uint64_t>(args.current_time.count()));
+}
+
+TEST(RateLimiterTest, DropBox_FailedToUpload) {
+ StrictMock<MockRateLimiter> limiter;
+ RateLimiter::Args args;
+
+ args.is_dropbox = true;
+ args.current_time = base::TimeSeconds(10000);
+
+ EXPECT_CALL(limiter, SaveState(_));
+ EXPECT_CALL(limiter, LoadState(_));
+ ASSERT_TRUE(limiter.ShouldTrace(args));
+ ASSERT_FALSE(limiter.OnTraceDone(args, false, 1024 * 1024));
+}
+
+TEST(RateLimiterTest, DropBox_FailedToSave) {
+ StrictMock<MockRateLimiter> limiter;
+ RateLimiter::Args args;
+
+ args.is_dropbox = true;
+ args.current_time = base::TimeSeconds(10000);
+
+ EXPECT_CALL(limiter, SaveState(_));
+ EXPECT_CALL(limiter, LoadState(_));
+ ASSERT_TRUE(limiter.ShouldTrace(args));
+
+ EXPECT_CALL(limiter, SaveState(_)).WillOnce(Return(false));
+ ASSERT_FALSE(limiter.OnTraceDone(args, true, 1024 * 1024));
+}
+
+} // namespace
+
+} // namespace perfetto