From 7b124cb8e6f0ad6de77001f55b4da24b709caa2f Mon Sep 17 00:00:00 2001
From: Jeff Sharkey <jsharkey@android.com>
Date: Fri, 4 Dec 2015 15:21:52 -0700
Subject: [PATCH] Add tagging variants for DatagramSocket.

Oops, we missed these along the way.

Bug: 25799174
Change-Id: I4ad368f9faee4b1996d605534dce4c2b23dbe200
---
 core/java/android/net/TrafficStats.java | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/core/java/android/net/TrafficStats.java b/core/java/android/net/TrafficStats.java
index 2137c3a30a..e82485da32 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;
 
@@ -225,6 +226,27 @@ public class TrafficStats {
         SocketTagger.get().untag(socket);
     }
 
+    /**
+     * 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.