Include tab IDs when reporting data use accounting.

chrome/browser/data_usage/tab_id_annotator.cc

Index: chrome/browser/data_usage/tab_id_annotator.cc
diff --git a/chrome/browser/data_usage/tab_id_annotator.cc b/chrome/browser/data_usage/tab_id_annotator.cc
new file mode 100644
index 0000000000000000000000000000000000000000..4a2243fb45aef2fb63257212a74be67f2f74b5d2
--- /dev/null
+++ b/chrome/browser/data_usage/tab_id_annotator.cc
@@ -0,0 +1,94 @@
+// 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 "chrome/browser/data_usage/tab_id_annotator.h"
+
+#include <stdint.h>
+
+#include "base/bind.h"
+#include "base/callback.h"
+#include "base/location.h"
+#include "base/memory/ref_counted.h"
+#include "base/single_thread_task_runner.h"
+#include "chrome/browser/data_usage/tab_id_provider.h"
+#include "chrome/browser/sessions/session_tab_helper.h"
+#include "components/data_usage/core/data_use.h"
+#include "content/public/browser/browser_thread.h"
+#include "content/public/browser/render_frame_host.h"
+#include "content/public/browser/resource_request_info.h"
+#include "content/public/browser/web_contents.h"
+#include "net/url_request/url_request.h"
+
+using content::BrowserThread;
+using data_usage::DataUse;
+
+namespace chrome_browser_data_usage {
+
+namespace {
+
+// Attempts to get the associated tab ID for a given render frame. Returns -1 if
+// no associated tab was found.
+int32_t GetTabIdForRenderFrame(int render_process_id, int render_frame_id) {
+  DCHECK_CURRENTLY_ON(BrowserThread::UI);
+  // TODO(sclittle): For prerendering tabs, investigate if it's possible to find
+  // the original tab that initiated the prerender.
+  return SessionTabHelper::IdForTab(content::WebContents::FromRenderFrameHost(
+      content::RenderFrameHost::FromID(render_process_id, render_frame_id)));
+}
+
+// Annotates |data_use| with the given |tab_id|, then passes it to |callback|.
+// This is done in a separate function instead of as a method on TabIdAnnotator
+// so that an in-progress annotation can complete even if the TabIdAnnotator is
+// destroyed. This doesn't make much of a difference for production code, but
+// makes it easier to test the TabIdAnnotator.
+void AnnotateDataUse(
+    scoped_ptr<DataUse> data_use,
+    const data_usage::DataUseAnnotator::DataUseConsumerCallback& callback,
+    int32_t tab_id) {
+  DCHECK(data_use);
+  data_use->tab_id = tab_id;
+  callback.Run(data_use.Pass());
+}
+
+}  // namespace
+
+TabIdAnnotator::TabIdAnnotator() {}
+TabIdAnnotator::~TabIdAnnotator() {}
+
+void TabIdAnnotator::Annotate(net::URLRequest* request,
+                              scoped_ptr<DataUse> data_use,
+                              const DataUseConsumerCallback& callback) {
+  DCHECK(thread_checker_.CalledOnValidThread());
+  DCHECK(data_use);
+
+  TabIdProvider* existing_tab_id_provider = reinterpret_cast<TabIdProvider*>(
+      request->GetUserData(TabIdProvider::kUserDataKey));
+  if (existing_tab_id_provider) {
+    existing_tab_id_provider->ProvideTabId(
+        base::Bind(&AnnotateDataUse, base::Passed(&data_use), callback));
+    return;
+  }
+
+  int render_process_id = -1, render_frame_id = -1;
+  if (!content::ResourceRequestInfo::GetRenderFrameForRequest(
+          request, &render_process_id, &render_frame_id)) {
+    // Run the callback immediately with a tab ID of -1 if the request has no
+    // render frame.
+    AnnotateDataUse(data_use.Pass(), callback, -1 /* tab_id */);
+    return;
+  }
+
+  scoped_refptr<base::SingleThreadTaskRunner> ui_thread_task_runner =
+      BrowserThread::GetMessageLoopProxyForThread(BrowserThread::UI);
+  scoped_ptr<TabIdProvider> tab_id_provider(new TabIdProvider(
+      ui_thread_task_runner.get(), FROM_HERE,
+      base::Bind(&GetTabIdForRenderFrame, render_process_id, render_frame_id)));
+  tab_id_provider->ProvideTabId(
+      base::Bind(&AnnotateDataUse, base::Passed(&data_use), callback));
+
+  // |request| takes ownership of |tab_id_provider|.
+  request->SetUserData(TabIdProvider::kUserDataKey, tab_id_provider.release());
+}
+
+}  // namespace chrome_browser_data_usage