From 43be174888684ef3404a43d8434015193c656cce Mon Sep 17 00:00:00 2001
From: Jeff Sharkey <jsharkey@android.com>
Date: Thu, 21 Apr 2011 18:45:43 -0700
Subject: [PATCH] Add Socket tagging for granular data accounting.

Introduces public API to apply "tags" to track data traffic originating
from the current thread.  (Under the hood, the tags are maintained and
applied in BlockGuard.)  Also adds tag/untag methods for developers who
maintain their own Socket pools.

Change-Id: Ic2dd3155559a93a7b613c7853748d4c44fb3a39e
---
 api/current.txt                         |  4 ++
 core/java/android/net/TrafficStats.java | 64 +++++++++++++++++++++++--
 2 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/api/current.txt b/api/current.txt
index 91dbd4bad2c9d..13fc169b9b1cb 100644
--- a/api/current.txt
+++ b/api/current.txt
@@ -10951,6 +10951,7 @@ package android.net {
 
   public class TrafficStats {
     ctor public TrafficStats();
+    method public static void clearThreadStatsTag();
     method public static long getMobileRxBytes();
     method public static long getMobileRxPackets();
     method public static long getMobileTxBytes();
@@ -10971,6 +10972,9 @@ package android.net {
     method public static long getUidUdpRxPackets(int);
     method public static long getUidUdpTxBytes(int);
     method public static long getUidUdpTxPackets(int);
+    method public static void setThreadStatsTag(java.lang.String);
+    method public static void tagSocket(java.net.Socket) throws java.net.SocketException;
+    method public static void untagSocket(java.net.Socket) throws java.net.SocketException;
     field public static final int UNSUPPORTED = -1; // 0xffffffff
   }
 
diff --git a/core/java/android/net/TrafficStats.java b/core/java/android/net/TrafficStats.java
index eca06c50dd82f..7ee7a81f40304 100644
--- a/core/java/android/net/TrafficStats.java
+++ b/core/java/android/net/TrafficStats.java
@@ -16,11 +16,10 @@
 
 package android.net;
 
-import android.util.Log;
+import dalvik.system.BlockGuard;
 
-import java.io.File;
-import java.io.RandomAccessFile;
-import java.io.IOException;
+import java.net.Socket;
+import java.net.SocketException;
 
 /**
  * Class that provides network traffic statistics.  These statistics include
@@ -36,6 +35,63 @@ public class TrafficStats {
      */
     public final static int UNSUPPORTED = -1;
 
+    /**
+     * Set active tag to use when accounting {@link Socket} traffic originating
+     * from the current thread. Only one active tag per thread is supported.
+     * <p>
+     * Changes only take effect during subsequent calls to
+     * {@link #tagSocket(Socket)}.
+     */
+    public static void setThreadStatsTag(String tag) {
+        BlockGuard.setThreadSocketStatsTag(tag);
+    }
+
+    public static void clearThreadStatsTag() {
+        BlockGuard.setThreadSocketStatsTag(null);
+    }
+
+    /**
+     * Set specific UID to use when accounting {@link Socket} traffic
+     * originating from the current thread. Designed for use when performing an
+     * operation on behalf of another application.
+     * <p>
+     * Changes only take effect during subsequent calls to
+     * {@link #tagSocket(Socket)}.
+     * <p>
+     * To take effect, caller must hold
+     * {@link android.Manifest.permission#UPDATE_DEVICE_STATS} permission.
+     *
+     * {@hide}
+     */
+    public static void setThreadStatsUid(int uid) {
+        BlockGuard.setThreadSocketStatsUid(uid);
+    }
+
+    /** {@hide} */
+    public static void clearThreadStatsUid() {
+        BlockGuard.setThreadSocketStatsUid(-1);
+    }
+
+    /**
+     * Tag the given {@link Socket} with any statistics parameters active for
+     * the current thread. Subsequent calls always replace any existing
+     * parameters. When finished, call {@link #untagSocket(Socket)} to remove
+     * statistics parameters.
+     *
+     * @see #setThreadStatsTag(String)
+     * @see #setThreadStatsUid(int)
+     */
+    public static void tagSocket(Socket socket) throws SocketException {
+        BlockGuard.tagSocketFd(socket.getFileDescriptor$());
+    }
+
+    /**
+     * Remove any statistics parameters from the given {@link Socket}.
+     */
+    public static void untagSocket(Socket socket) throws SocketException {
+        BlockGuard.untagSocketFd(socket.getFileDescriptor$());
+    }
+
     /**
      * Get the total number of packets transmitted through the mobile interface.
      *