Expose GlobalRequestID in NavigationHandle and ResourceRequestInfo.

chrome/browser/data_use_measurement/chrome_data_use_ascriber.cc

Index: chrome/browser/data_use_measurement/chrome_data_use_ascriber.cc
diff --git a/chrome/browser/data_use_measurement/chrome_data_use_ascriber.cc b/chrome/browser/data_use_measurement/chrome_data_use_ascriber.cc
index 1a745b2501a511c0222bf81e89164910f8bf6da1..ded5cff9fa6b14b7d33f9a15eda62e6909376d4a 100644
--- a/chrome/browser/data_use_measurement/chrome_data_use_ascriber.cc
+++ b/chrome/browser/data_use_measurement/chrome_data_use_ascriber.cc
@@ -4,18 +4,37 @@
 
 #include "chrome/browser/data_use_measurement/chrome_data_use_ascriber.h"
 
+#include "base/memory/ptr_util.h"
+#include "components/data_use_measurement/core/data_use_recorder.h"
+#include "components/data_use_measurement/core/data_use_user_data.h"
 #include "content/public/browser/browser_thread.h"
 #include "content/public/browser/navigation_handle.h"
 #include "content/public/browser/render_frame_host.h"
+#include "content/public/browser/resource_request_info.h"
+#include "net/url_request/url_request.h"
 
 namespace data_use_measurement {
 
+// static
+const void* ChromeDataUseAscriber::DataUseRecorderEntryAsUserData::
+    kUserDataKey = static_cast<void*>(
+        &ChromeDataUseAscriber::DataUseRecorderEntryAsUserData::kUserDataKey);
+
+ChromeDataUseAscriber::DataUseRecorderEntryAsUserData::
+    DataUseRecorderEntryAsUserData(DataUseRecorderEntry entry)
+    : entry_(entry) {}
+
+ChromeDataUseAscriber::DataUseRecorderEntryAsUserData::
+    ~DataUseRecorderEntryAsUserData() {}
+
 ChromeDataUseAscriber::ChromeDataUseAscriber() {
   DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
 }
 
 ChromeDataUseAscriber::~ChromeDataUseAscriber() {
   DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
+
+  DCHECK_EQ(0u, data_use_recorders_.size());
 }
 
 DataUseRecorder* ChromeDataUseAscriber::GetDataUseRecorder(
@@ -24,11 +43,95 @@ DataUseRecorder* ChromeDataUseAscriber::GetDataUseRecorder(
   return nullptr;
 }
 
+void ChromeDataUseAscriber::OnBeforeUrlRequest(net::URLRequest* request) {
+  DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
+
+  if (!request->url().SchemeIsHTTPOrHTTPS())
+    return;
+
+  DataUseUserData* service = static_cast<DataUseUserData*>(
+      request->GetUserData(DataUseUserData::kUserDataKey));
+  if (service)
+    return;
+
+  const content::ResourceRequestInfo* request_info =
+      content::ResourceRequestInfo::ForRequest(request);
+  content::ResourceType resource_type = request_info
+                                            ? request_info->GetResourceType()
+                                            : content::RESOURCE_TYPE_LAST_TYPE;
+
+  if (resource_type == content::RESOURCE_TYPE_MAIN_FRAME) {
+    int render_process_id = -1;
+    int render_frame_id = -1;
+    bool has_valid_render_frame_id_ =
+        content::ResourceRequestInfo::GetRenderFrameForRequest(
+            request, &render_process_id, &render_frame_id);
+    DCHECK(has_valid_render_frame_id_);
+    // Browser tests may not set up DataUseWebContentsObservers in which case
+    // this class never sees navigation and frame events so DataUseRecorders
+    // will never be destroyed. To avoid this, we ignore requests whose
+    // render frames don't have a record. However, this can also be caused by
+    // URLRequests racing the frame create events.
+    // TODO(kundaji): Add UMA.
+    if (render_frame_data_use_map_.find(
+            RenderFrameHostID(render_process_id, render_frame_id)) ==
+        render_frame_data_use_map_.end()) {
+      return;
+    }
+
+    auto user_data = static_cast<DataUseRecorderEntryAsUserData*>(

[RyanSturm, 2016/10/27 22:41:21]

nit: Can you either change auto to DataUseRecorderEntryAsUserData* or change the
DataUseUserData* above to auto?

[Not at Google. Contact bengr, 2016/10/28 18:43:33]

Done.

+        request->GetUserData(DataUseRecorderEntryAsUserData::kUserDataKey));
+    if (user_data)
+      return;
+
+    DataUseRecorderEntry entry = data_use_recorders_.insert(
+        data_use_recorders_.end(), base::MakeUnique<DataUseRecorder>());
+    request->SetUserData(DataUseRecorderEntryAsUserData::kUserDataKey,
+                         new DataUseRecorderEntryAsUserData(entry));
+
+    pending_navigation_data_use_map_.insert(
+        std::make_pair(request_info->GetGlobalRequestID(), entry));
+  }
+}
+
+void ChromeDataUseAscriber::OnUrlRequestCompleted(net::URLRequest* request,
+                                                  bool started) {
+  const content::ResourceRequestInfo* request_info =
+      content::ResourceRequestInfo::ForRequest(request);
+  content::ResourceType resource_type = request_info
+                                            ? request_info->GetResourceType()
+                                            : content::RESOURCE_TYPE_LAST_TYPE;
+
+  if (resource_type == content::RESOURCE_TYPE_MAIN_FRAME) {
+    // If request was not successful, then NavigationHandle in
+    // DidFinishMainFrameNavigation will not have GlobalRequestID. So we erase
+    // the DataUseRecorderEntry here.
+    if (!request->status().is_success()) {
+      auto navigation_iter = pending_navigation_data_use_map_.find(

[Raj, 2016/10/27 20:33:56]

nit: This seems to be common code. Maybe wrap it in a function in anonymous
namespace. L110-L118 and L166-L173

[Not at Google. Contact bengr, 2016/10/28 18:43:33]

Done.

+          request_info->GetGlobalRequestID());
+      // Pending navigation entry will not be found if finish navigation
+      // raced this request completed event.
+      if (navigation_iter != pending_navigation_data_use_map_.end()) {
+        auto entry = navigation_iter->second;
+        pending_navigation_data_use_map_.erase(navigation_iter);
+        data_use_recorders_.erase(entry);
+      }
+    }
+  }
+}
+
 void ChromeDataUseAscriber::RenderFrameCreated(int render_process_id,
                                                int render_frame_id,
                                                int parent_render_process_id,
                                                int parent_render_frame_id) {
   DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
+
+  // TODO(kundaji): Point child render frames to the same DataUseRecorder as
+  // parent render frame.
+  DataUseRecorderEntry entry = data_use_recorders_.insert(
+      data_use_recorders_.end(), base::MakeUnique<DataUseRecorder>());
+  render_frame_data_use_map_.insert(std::make_pair(
+      RenderFrameHostID(render_process_id, render_frame_id), entry));
 }
 
 void ChromeDataUseAscriber::RenderFrameDeleted(int render_process_id,
@@ -36,6 +139,12 @@ void ChromeDataUseAscriber::RenderFrameDeleted(int render_process_id,
                                                int parent_render_process_id,
                                                int parent_render_frame_id) {
   DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
+  RenderFrameHostID key(render_process_id, render_frame_id);
+  auto frame_iter = render_frame_data_use_map_.find(key);
+  DCHECK(frame_iter != render_frame_data_use_map_.end());
+  DataUseRecorderEntry entry = frame_iter->second;
+  render_frame_data_use_map_.erase(frame_iter);
+  data_use_recorders_.erase(entry);
 }
 
 void ChromeDataUseAscriber::DidStartMainFrameNavigation(
@@ -48,11 +157,20 @@ void ChromeDataUseAscriber::DidStartMainFrameNavigation(
 
 void ChromeDataUseAscriber::DidFinishMainFrameNavigation(
     GURL gurl,
+    content::GlobalRequestID global_request_id,
     int render_process_id,
     int render_frame_id,
     bool is_same_page_navigation,
     void* navigation_handle) {
   DCHECK_CURRENTLY_ON(content::BrowserThread::IO);
+
+  auto navigation_iter =
+      pending_navigation_data_use_map_.find(global_request_id);
+  if (navigation_iter != pending_navigation_data_use_map_.end()) {
+    auto entry = navigation_iter->second;
+    pending_navigation_data_use_map_.erase(navigation_iter);
+    data_use_recorders_.erase(entry);
+  }
 }
 
 void ChromeDataUseAscriber::DidRedirectMainFrameNavigation(