Support moving pending requests from one DataUseRecorder to another

components/data_use_measurement/core/data_use_recorder.cc

Index: components/data_use_measurement/core/data_use_recorder.cc
diff --git a/components/data_use_measurement/core/data_use_recorder.cc b/components/data_use_measurement/core/data_use_recorder.cc
index 0967f358b3a0621747526b651e22b1969a43d072..c4230c240305189552f98c4019a61ee98f8ddde8 100644
--- a/components/data_use_measurement/core/data_use_recorder.cc
+++ b/components/data_use_measurement/core/data_use_recorder.cc
@@ -14,17 +14,38 @@ DataUseRecorder::DataUseRecorder(DataUse::TrafficType traffic_type)
 DataUseRecorder::~DataUseRecorder() {}
 
 bool DataUseRecorder::IsDataUseComplete() {
-  return pending_url_requests_.empty() && pending_data_sources_.empty();
+  return pending_url_requests_.empty();
 }
 
 void DataUseRecorder::AddPendingURLRequest(net::URLRequest* request) {
-  pending_url_requests_.insert(request);
+  pending_url_requests_.emplace(std::piecewise_construct,
+                                std::forward_as_tuple(request),
+                                std::forward_as_tuple());
 }
 
 void DataUseRecorder::OnUrlRequestDestroyed(net::URLRequest* request) {
   pending_url_requests_.erase(request);
 }
 
+void DataUseRecorder::MovePendingURLRequest(DataUseRecorder* other,
+                                            net::URLRequest* request) {
+  auto request_it = pending_url_requests_.find(request);
+  DCHECK(request_it != pending_url_requests_.end());
+  DCHECK(other->pending_url_requests_.find(request) ==
+         other->pending_url_requests_.end());
+
+  // Increment the bytes of the request in |other|, and decrement the bytes in
+  // |this|.
+  // TODO(rajendrant): Check if the moving the bytes in |data_use_| needs to be
+  // propogated to observers, which could store per-request user data.
+  other->AddPendingURLRequest(request);
+  other->UpdateNetworkByteCounts(request, request_it->second.bytes_received,
+                                 request_it->second.bytes_sent);
+  data_use_.IncrementTotalBytes(-request_it->second.bytes_received,
+                                -request_it->second.bytes_sent);
+  pending_url_requests_.erase(request_it);
+}
+
 void DataUseRecorder::RemoveAllPendingURLRequests() {
   pending_url_requests_.clear();
 }
@@ -33,16 +54,21 @@ void DataUseRecorder::OnBeforeUrlRequest(net::URLRequest* request) {}
 
 void DataUseRecorder::OnNetworkBytesReceived(net::URLRequest* request,
                                              int64_t bytes_received) {
-  data_use_.IncrementTotalBytes(bytes_received, 0);
+  UpdateNetworkByteCounts(request, bytes_received, 0);
 }
 
 void DataUseRecorder::OnNetworkBytesSent(net::URLRequest* request,
                                          int64_t bytes_sent) {
-  data_use_.IncrementTotalBytes(0, bytes_sent);
+  UpdateNetworkByteCounts(request, 0, bytes_sent);
 }
 
-void DataUseRecorder::MergeFrom(DataUseRecorder* other) {
-  data_use_.MergeFrom(other->data_use());
+void DataUseRecorder::UpdateNetworkByteCounts(net::URLRequest* request,
+                                              int64_t bytes_received,
+                                              int64_t bytes_sent) {
+  data_use_.IncrementTotalBytes(bytes_received, bytes_sent);
+  auto request_it = pending_url_requests_.find(request);
+  request_it->second.bytes_received += bytes_received;
+  request_it->second.bytes_sent += bytes_sent;
 }
 
 }  // namespace data_use_measurement