Amortize data usage using TrafficStats on Android.

components/data_usage/android/traffic_stats_amortizer.cc

Index: components/data_usage/android/traffic_stats_amortizer.cc
diff --git a/components/data_usage/android/traffic_stats_amortizer.cc b/components/data_usage/android/traffic_stats_amortizer.cc
new file mode 100644
index 0000000000000000000000000000000000000000..e9ed1fd3fa9bc71d371de005b820d84aadfc710e
--- /dev/null
+++ b/components/data_usage/android/traffic_stats_amortizer.cc
@@ -0,0 +1,292 @@
+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "components/data_usage/android/traffic_stats_amortizer.h"
+
+#include <algorithm>  // For std::min.
+#include <cmath>      // For std::modf.
+
+#include "base/location.h"
+#include "base/time/default_tick_clock.h"
+#include "base/timer/timer.h"
+#include "components/data_usage/core/data_use.h"
+#include "net/android/traffic_stats.h"
+
+namespace data_usage {
+namespace android {
+
+namespace {
+
+// Convenience typedef.
+typedef std::vector<std::pair<linked_ptr<DataUse>,
+                              DataUseAmortizer::AmortizationCompleteCallback>>
+    DataUseBuffer;
+
+// The delay between receiving DataUse and querying TrafficStats byte counts for
+// amortization.
+// TODO(sclittle): Control this with a field trial parameter.
+const int64_t kDefaultTrafficStatsQueryDelayMs = 50;
+
+// The longest amount of time that an amortization run can be delayed for.
+// TODO(sclittle): Control this with a field trial parameter.
+const int64_t kDefaultMaxAmortizationDelayMs = 500;
+
+// The maximum allowed size of the DataUse buffer. If the buffer ever exceeds
+// this size, then DataUse will be amortized immediately and the buffer will be
+// flushed.
+// TODO(sclittle): Control this with a field trial parameter.
+const size_t kDefaultMaxDataUseBufferSize = 128;
+
+// Scales |bytes| by |ratio|, using |remainder| to hold the running rounding
+// error. |bytes| must be non-negative, and multiplying |bytes| by |ratio| must
+// yield a number that's representable within the bounds of a non-negative
+// int64_t.
+int64_t ScaleByteCount(int64_t bytes, double ratio, double* remainder) {
+  DCHECK_GE(bytes, 0);
+  DCHECK_GE(ratio, 0.0);
+  DCHECK_LE(ratio, static_cast<double>(INT64_MAX));
+  DCHECK_GE(*remainder, 0.0);
+  DCHECK_LT(*remainder, 1.0);
+
+  double intpart;
+  *remainder =
+      std::modf(static_cast<double>(bytes) * ratio + (*remainder), &intpart);
+
+  DCHECK_GE(intpart, 0.0);
+  DCHECK_LE(intpart, static_cast<double>(INT64_MAX));
+  DCHECK_GE(*remainder, 0.0);
+  DCHECK_LT(*remainder, 1.0);
+
+  // Due to floating point error, casting the double |intpart| to an int64_t
+  // could cause it to overflow, even though it's already been checked to be
+  // less than the double representation of INT64_MAX. If this happens, cap the
+  // scaled value at INT64_MAX.
+  uint64_t scaled_bytes = std::min(static_cast<uint64_t>(intpart),
+                                   static_cast<uint64_t>(INT64_MAX));
+  return static_cast<int64_t>(scaled_bytes);
+}
+
+// Amortizes the difference between |desired_post_amortization_total| and
+// |pre_amortization_total| into each of the DataUse objects in
+// |data_use_sequence| by scaling the byte counts determined by the
+// |get_byte_count_fn| function (e.g. tx_bytes, rx_bytes) for each DataUse
+// appropriately.
+void AmortizeByteCountSequence(DataUseBuffer* data_use_sequence,
+                               int64_t* (*get_byte_count_fn)(DataUse*),
+                               int64_t pre_amortization_total,
+                               int64_t desired_post_amortization_total) {
+  DCHECK_GE(pre_amortization_total, 0);
+  DCHECK_GE(desired_post_amortization_total, 0);
+
+  if (pre_amortization_total == 0) {
+    // TODO(sclittle): If |desired_post_amortization_total| is non-zero, this
+    // could be ignoring overhead that should be amortized in somehow. Handle
+    // this case gracefully.
+    return;
+  }
+
+  const double ratio = static_cast<double>(desired_post_amortization_total) /
+                       static_cast<double>(pre_amortization_total);
+
+  double remainder = 0.0;
+  for (auto& data_use_buffer_pair : *data_use_sequence) {
+    int64_t* byte_count = get_byte_count_fn(data_use_buffer_pair.first.get());
+    *byte_count = ScaleByteCount(*byte_count, ratio, &remainder);
+    // TODO(sclittle): Record UMA about values before vs. after amortization.
+  }
+}
+
+int64_t* GetTxBytes(DataUse* data_use) {
+  return &data_use->tx_bytes;
+}
+int64_t* GetRxBytes(DataUse* data_use) {
+  return &data_use->rx_bytes;
+}
+
+// Amortizes the difference between |desired_post_amortization_total_tx_bytes|
+// and |pre_amortization_total_tx_bytes| into each of the DataUse objects in
+// |data_use_sequence| by scaling the DataUse's |tx_bytes| appropriately. Does
+// the same with the |rx_bytes| using those respective parameters.
+void AmortizeDataUseSequence(DataUseBuffer* data_use_sequence,
+                             int64_t pre_amortization_total_tx_bytes,
+                             int64_t desired_post_amortization_total_tx_bytes,
+                             int64_t pre_amortization_total_rx_bytes,
+                             int64_t desired_post_amortization_total_rx_bytes) {
+  DCHECK(data_use_sequence);
+  DCHECK(!data_use_sequence->empty());
+
+  AmortizeByteCountSequence(data_use_sequence, &GetTxBytes,
+                            pre_amortization_total_tx_bytes,
+                            desired_post_amortization_total_tx_bytes);
+
+  AmortizeByteCountSequence(data_use_sequence, &GetRxBytes,
+                            pre_amortization_total_rx_bytes,
+                            desired_post_amortization_total_rx_bytes);
+}
+
+}  // namespace
+
+TrafficStatsAmortizer::TrafficStatsAmortizer()
+    : TrafficStatsAmortizer(
+          scoped_ptr<base::TickClock>(new base::DefaultTickClock()),
+          scoped_ptr<base::Timer>(new base::Timer(false, false)),
+          base::TimeDelta::FromMilliseconds(kDefaultTrafficStatsQueryDelayMs),
+          base::TimeDelta::FromMilliseconds(kDefaultMaxAmortizationDelayMs),
+          kDefaultMaxDataUseBufferSize) {}
+
+TrafficStatsAmortizer::~TrafficStatsAmortizer() {}
+
+void TrafficStatsAmortizer::AmortizeDataUse(
+    scoped_ptr<DataUse> data_use,
+    const AmortizationCompleteCallback& callback) {
+  DCHECK(thread_checker_.CalledOnValidThread());
+  DCHECK(!callback.is_null());
+  int64_t tx_bytes = data_use->tx_bytes, rx_bytes = data_use->rx_bytes;
+
+  // TODO(sclittle): Combine consecutive buffered DataUse objects that are
+  // identical except for byte counts and have the same callback.
+  buffered_data_use_.push_back(
+      std::make_pair(linked_ptr<DataUse>(data_use.release()), callback));
+
+  AddPreAmortizationBytes(tx_bytes, rx_bytes);
+}
+
+void TrafficStatsAmortizer::OnExtraBytes(int64_t extra_tx_bytes,
+                                         int64_t extra_rx_bytes) {
+  DCHECK(thread_checker_.CalledOnValidThread());
+  AddPreAmortizationBytes(extra_tx_bytes, extra_rx_bytes);
+}
+
+base::WeakPtr<TrafficStatsAmortizer> TrafficStatsAmortizer::GetWeakPtr() {
+  DCHECK(thread_checker_.CalledOnValidThread());
+  return weak_ptr_factory_.GetWeakPtr();
+}
+
+TrafficStatsAmortizer::TrafficStatsAmortizer(
+    scoped_ptr<base::TickClock> tick_clock,
+    scoped_ptr<base::Timer> traffic_stats_query_timer,
+    const base::TimeDelta& traffic_stats_query_delay,
+    const base::TimeDelta& max_amortization_delay,
+    size_t max_data_use_buffer_size)
+    : tick_clock_(tick_clock.Pass()),
+      traffic_stats_query_timer_(traffic_stats_query_timer.Pass()),
+      traffic_stats_query_delay_(traffic_stats_query_delay),
+      max_amortization_delay_(max_amortization_delay),
+      max_data_use_buffer_size_(max_data_use_buffer_size),
+      is_amortization_in_progress_(false),
+      are_last_amortization_traffic_stats_available_(false),
+      last_amortization_traffic_stats_tx_bytes_(-1),
+      last_amortization_traffic_stats_rx_bytes_(-1),
+      pre_amortization_tx_bytes_(0),
+      pre_amortization_rx_bytes_(0),
+      weak_ptr_factory_(this) {}
+
+bool TrafficStatsAmortizer::QueryTrafficStats(int64_t* tx_bytes,
+                                              int64_t* rx_bytes) const {
+  DCHECK(thread_checker_.CalledOnValidThread());
+  return net::android::traffic_stats::GetCurrentUidTxBytes(tx_bytes) &&
+         net::android::traffic_stats::GetCurrentUidRxBytes(rx_bytes);
+}
+
+void TrafficStatsAmortizer::AddPreAmortizationBytes(int64_t tx_bytes,
+                                                    int64_t rx_bytes) {
+  DCHECK(thread_checker_.CalledOnValidThread());
+  DCHECK_GE(tx_bytes, 0);
+  DCHECK_GE(rx_bytes, 0);
+  base::TimeTicks now_ticks = tick_clock_->NowTicks();
+
+  if (!is_amortization_in_progress_) {
+    is_amortization_in_progress_ = true;
+    current_amortization_run_start_time_ = now_ticks;
+  }
+
+  pre_amortization_tx_bytes_ += tx_bytes;
+  pre_amortization_rx_bytes_ += rx_bytes;
+
+  if (buffered_data_use_.size() > max_data_use_buffer_size_) {
+    // Enforce a maximum limit on the size of |buffered_data_use_| to avoid
+    // hogging memory. Note that this will likely cause the post-amortization
+    // byte counts calculated here to be less accurate than if the amortizer
+    // waited to perform amortization.
+    // TODO(sclittle): Record UMA about how often this occurs.
+    traffic_stats_query_timer_->Stop();
+    AmortizeNow();
+    return;
+  }
+
+  // Cap any amortization delay to |max_amortization_delay_|. Note that if
+  // |max_amortization_delay_| comes earlier, then this will likely cause the
+  // post-amortization byte counts calculated here to be less accurate than if
+  // the amortizer waited to perform amortization.
+  // TODO(sclittle): Record UMA about how often |max_amortization_delay_| comes
+  // earlier.
+  base::TimeDelta query_delay = std::min(
+      traffic_stats_query_delay_, current_amortization_run_start_time_ +
+                                      max_amortization_delay_ - now_ticks);
+
+  // Set the timer to query TrafficStats and amortize after a delay, so that
+  // it's more likely that TrafficStats will be queried when the network is
+  // idle. If the timer was already set, then this overrides the previous delay.
+  traffic_stats_query_timer_->Start(
+      FROM_HERE, query_delay,
+      base::Bind(&TrafficStatsAmortizer::AmortizeNow, GetWeakPtr()));
+}
+
+void TrafficStatsAmortizer::AmortizeNow() {
+  DCHECK(thread_checker_.CalledOnValidThread());
+
+  int64_t current_traffic_stats_tx_bytes = -1;
+  int64_t current_traffic_stats_rx_bytes = -1;
+  bool are_current_traffic_stats_available = QueryTrafficStats(
+      &current_traffic_stats_tx_bytes, &current_traffic_stats_rx_bytes);
+
+  if (are_current_traffic_stats_available &&
+      are_last_amortization_traffic_stats_available_ &&
+      !buffered_data_use_.empty()) {
+    // These TrafficStats byte counts are guaranteed to increase monotonically
+    // since device boot.
+    DCHECK_GE(current_traffic_stats_tx_bytes,
+              last_amortization_traffic_stats_tx_bytes_);
+    DCHECK_GE(current_traffic_stats_rx_bytes,
+              last_amortization_traffic_stats_rx_bytes_);
+
+    int64_t desired_post_amortization_total_tx_bytes =
+        current_traffic_stats_tx_bytes -
+        last_amortization_traffic_stats_tx_bytes_;
+    int64_t desired_post_amortization_total_rx_bytes =
+        current_traffic_stats_rx_bytes -
+        last_amortization_traffic_stats_rx_bytes_;
+
+    AmortizeDataUseSequence(&buffered_data_use_, pre_amortization_tx_bytes_,
+                            desired_post_amortization_total_tx_bytes,
+                            pre_amortization_rx_bytes_,
+                            desired_post_amortization_total_rx_bytes);
+  }
+
+  // TODO(sclittle): Record some UMA about the delay before amortizing and how
+  // big the buffer was before amortizing.
+
+  // Reset state now that the amortization run has finished.
+  is_amortization_in_progress_ = false;
+  current_amortization_run_start_time_ = base::TimeTicks();
+
+  are_last_amortization_traffic_stats_available_ =
+      are_current_traffic_stats_available;
+  last_amortization_traffic_stats_tx_bytes_ = current_traffic_stats_tx_bytes;
+  last_amortization_traffic_stats_rx_bytes_ = current_traffic_stats_rx_bytes;
+
+  pre_amortization_tx_bytes_ = 0;
+  pre_amortization_rx_bytes_ = 0;
+
+  // Pass post-amortization DataUse objects to their respective callbacks.
+  DataUseBuffer data_use_sequence;
+  data_use_sequence.swap(buffered_data_use_);
+  for (auto& data_use_buffer_pair : data_use_sequence) {
+    scoped_ptr<DataUse> data_use(data_use_buffer_pair.first.release());
+    data_use_buffer_pair.second.Run(data_use.Pass());
+  }
+}
+
+}  // namespace android
+}  // namespace data_usage