Add tagging variants for DatagramSocket.

Oops, we missed these along the way.

Bug: 25799174
Change-Id: I4ad368f9faee4b1996d605534dce4c2b23dbe200

4 files changed, 28 insertions, 0 deletions

diff --git a/api/current.txt b/api/current.txt
index 16737a947c3a..db97250c2ef2 100644
--- a/api/current.txt
+++ b/api/current.txt
@@ -22931,7 +22931,9 @@ package android.net {
     method public static void incrementOperationCount(int);
     method public static void incrementOperationCount(int, int);
     method public static void setThreadStatsTag(int);
+    method public static void tagDatagramSocket(java.net.DatagramSocket) throws java.net.SocketException;
     method public static void tagSocket(java.net.Socket) throws java.net.SocketException;
+    method public static void untagDatagramSocket(java.net.DatagramSocket) 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/api/system-current.txt b/api/system-current.txt
index 0adbc13d60b6..4296bf7d0943 100644
--- a/api/system-current.txt
+++ b/api/system-current.txt
@@ -24518,7 +24518,9 @@ package android.net {
     method public static void setThreadStatsTagBackup();
     method public static void setThreadStatsTagRestore();
     method public static void setThreadStatsUid(int);
+    method public static void tagDatagramSocket(java.net.DatagramSocket) throws java.net.SocketException;
     method public static void tagSocket(java.net.Socket) throws java.net.SocketException;
+    method public static void untagDatagramSocket(java.net.DatagramSocket) 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/api/test-current.txt b/api/test-current.txt
index 63e9a3ed576d..d248d4a9fac7 100644
--- a/api/test-current.txt
+++ b/api/test-current.txt
@@ -22931,7 +22931,9 @@ package android.net {
     method public static void incrementOperationCount(int);
     method public static void incrementOperationCount(int, int);
     method public static void setThreadStatsTag(int);
+    method public static void tagDatagramSocket(java.net.DatagramSocket) throws java.net.SocketException;
     method public static void tagSocket(java.net.Socket) throws java.net.SocketException;
+    method public static void untagDatagramSocket(java.net.DatagramSocket) 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 2137c3a30aa5..e82485da32f7 100644
--- a/core/java/android/net/TrafficStats.java
+++ b/core/java/android/net/TrafficStats.java
@@ -28,6 +28,7 @@ import com.android.server.NetworkManagementSocketTagger;
 
 import dalvik.system.SocketTagger;
 
+import java.net.DatagramSocket;
 import java.net.Socket;
 import java.net.SocketException;
 
@@ -226,6 +227,27 @@ public class TrafficStats {
     }
 
     /**
+     * Tag the given {@link DatagramSocket} with any statistics parameters
+     * active for the current thread. Subsequent calls always replace any
+     * existing parameters. When finished, call
+     * {@link #untagDatagramSocket(DatagramSocket)} to remove statistics
+     * parameters.
+     *
+     * @see #setThreadStatsTag(int)
+     * @see #setThreadStatsUid(int)
+     */
+    public static void tagDatagramSocket(DatagramSocket socket) throws SocketException {
+        SocketTagger.get().tag(socket);
+    }
+
+    /**
+     * Remove any statistics parameters from the given {@link DatagramSocket}.
+     */
+    public static void untagDatagramSocket(DatagramSocket socket) throws SocketException {
+        SocketTagger.get().untag(socket);
+    }
+
+    /**
      * Start profiling data usage for current UID. Only one profiling session
      * can be active at a time.
      *