Set TrafficStats tag in FirebaseInstallationServiceClient.createFirebaseInstallation (#2270)

* Set TrafficStats tag in FirebaseInstallationServiceClient.createFirebaseInstallation

Fixes #2227

Tested:
I verified that com.google.googlequicksearchbox doesn't have the violation anymore in the e2e test.

Signed-off-by: Régis Décamps <[email protected]>

* Add TrafficStats tags for Create isntallation, Delete installation and Generate Token.

* Rename tags to be prefixed with TRAFFIC_STATS

* Update copyright header for 2021.

Author: regisd

firebase-installations/src/main/java/com/google/firebase/installations/remote/FirebaseInstallationServiceClient.java

@@ -1,4 +1,4 @@
-// Copyright 2019 Google LLC
+// Copyright 2019-2021 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -20,6 +20,7 @@
 
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.net.TrafficStats;
 import android.text.TextUtils;
 import android.util.JsonReader;
 import android.util.Log;
@@ -56,6 +57,16 @@
  * @hide
  */
 public class FirebaseInstallationServiceClient {
+
+  // TrafficStats tags must be kept in sync with java/com/google/android/gms/libs/punchclock
+  private static final int TRAFFIC_STATS_FIREBASE_INSTALLATIONS_TAG = 0x00008000;
+  private static final int TRAFFIC_STATS_CREATE_INSTALLATION_TAG =
+      TRAFFIC_STATS_FIREBASE_INSTALLATIONS_TAG | 0x1;
+  private static final int TRAFFIC_STATS_DELETE_INSTALLATION_TAG =
+      TRAFFIC_STATS_FIREBASE_INSTALLATIONS_TAG | 0x2;
+  private static final int TRAFFIC_STATS_GENERATE_AUTH_TOKEN_TAG =
+      TRAFFIC_STATS_FIREBASE_INSTALLATIONS_TAG | 0x3;
+
   private static final String FIREBASE_INSTALLATIONS_API_DOMAIN =
       "firebaseinstallations.googleapis.com";
   private static final String CREATE_REQUEST_RESOURCE_NAME_FORMAT = "projects/%s/installations";
@@ -153,6 +164,7 @@ public InstallationResponse createFirebaseInstallation(
     URL url = getFullyQualifiedRequestUri(resourceName);
     for (int retryCount = 0; retryCount <= MAX_RETRIES; retryCount++) {
 
+      TrafficStats.setThreadStatsTag(TRAFFIC_STATS_CREATE_INSTALLATION_TAG);
       HttpURLConnection httpURLConnection = openHttpURLConnection(url, apiKey);
 
       try {
@@ -194,6 +206,7 @@ public InstallationResponse createFirebaseInstallation(
         continue;
       } finally {
         httpURLConnection.disconnect();
+        TrafficStats.clearThreadStatsTag();
       }
     }
 
@@ -309,6 +322,7 @@ public void deleteFirebaseInstallation(
 
     int retryCount = 0;
     while (retryCount <= MAX_RETRIES) {
+      TrafficStats.setThreadStatsTag(TRAFFIC_STATS_DELETE_INSTALLATION_TAG);
       HttpURLConnection httpURLConnection = openHttpURLConnection(url, apiKey);
       try {
         httpURLConnection.setRequestMethod("DELETE");
@@ -335,6 +349,7 @@ public void deleteFirebaseInstallation(
         retryCount++;
       } finally {
         httpURLConnection.disconnect();
+        TrafficStats.clearThreadStatsTag();
       }
     }
 
@@ -390,6 +405,7 @@ public TokenResult generateAuthToken(
     URL url = getFullyQualifiedRequestUri(resourceName);
     for (int retryCount = 0; retryCount <= MAX_RETRIES; retryCount++) {
 
+      TrafficStats.setThreadStatsTag(TRAFFIC_STATS_GENERATE_AUTH_TOKEN_TAG);
       HttpURLConnection httpURLConnection = openHttpURLConnection(url, apiKey);
       try {
         httpURLConnection.setRequestMethod("POST");
@@ -430,6 +446,7 @@ public TokenResult generateAuthToken(
         continue;
       } finally {
         httpURLConnection.disconnect();
+        TrafficStats.clearThreadStatsTag();
       }
     }
     throw new FirebaseInstallationsException(
@@ -619,7 +636,8 @@ private static String readErrorResponse(HttpURLConnection conn) {
         response.append(input).append('\n');
       }
       return String.format(
-          "Error when communicating with the Firebase Installations server API. HTTP response: [%d %s: %s]",
+          "Error when communicating with the Firebase Installations server API. HTTP response: [%d"
+              + " %s: %s]",
           conn.getResponseCode(), conn.getResponseMessage(), response);
     } catch (IOException ignored) {
       return null;