From 8881325d71e3863b3ec89971b6254c3316017010 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Thu, 13 Jun 2019 13:54:46 -0700
Subject: [PATCH 01/30] Implement Firebase segmentation SDK device local cache

---
 .../firebase-segmentation.gradle              |  9 +-
 .../CustomInstallationIdMappingCache.java     | 98 +++++++++++++++++++
 2 files changed, 103 insertions(+), 4 deletions(-)
 create mode 100644 firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdMappingCache.java

diff --git a/firebase-segmentation/firebase-segmentation.gradle b/firebase-segmentation/firebase-segmentation.gradle
index 88ea12a1ed0..11a5046d6d3 100644
--- a/firebase-segmentation/firebase-segmentation.gradle
+++ b/firebase-segmentation/firebase-segmentation.gradle
@@ -57,7 +57,7 @@ android {
     compileSdkVersion project.targetSdkVersion
 
     defaultConfig {
-        minSdkVersion project.minSdkVersion
+        minSdkVersion 21
         targetSdkVersion project.targetSdkVersion
         multiDexEnabled true
         versionName version
@@ -95,8 +95,9 @@ dependencies {
     testImplementation 'junit:junit:4.12'
     testImplementation "org.robolectric:robolectric:$robolectricVersion"
 
-    androidTestImplementation 'androidx.annotation:annotation:1.1.0'
-    androidTestImplementation 'androidx.test:runner:1.2.0'
+    androidTestImplementation "androidx.annotation:annotation:1.1.0"
     androidTestImplementation 'androidx.test:rules:1.2.0'
-    androidTestImplementation 'androidx.test.ext:junit:1.1.1'
+    androidTestImplementation 'androidx.test:runner:1.2.0'
+    androidTestImplementation "com.google.truth:truth:$googleTruthVersion"
+    androidTestImplementation 'junit:junit:4.12'
 }
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdMappingCache.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdMappingCache.java
new file mode 100644
index 00000000000..9826ea9e23f
--- /dev/null
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdMappingCache.java
@@ -0,0 +1,98 @@
+// Copyright 2019 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.firebase.segmentation;
+
+import android.database.Cursor;
+import android.database.sqlite.SQLiteDatabase;
+import androidx.annotation.Nullable;
+import com.google.android.gms.common.internal.Preconditions;
+import com.google.firebase.FirebaseApp;
+
+class CustomInstallationIdMappingCache {
+
+  // Status of each cache entry
+  enum CacheStatus {
+    // Cache entry is synced to Firebase backend
+    SYNCED,
+    // Cache entry is waiting for Firebase backend response or pending internal retry for retryable
+    // errors.
+    PENDING,
+    // Cache entry is not accepted by Firebase backend.
+    ERROR
+  }
+
+  private static final String LOCAL_DB_NAME = "CustomInstallationIdCache";
+  private static final String TABLE_NAME = "InstallationIdMapping";
+
+  private static final String GMP_APP_ID_COLUMN_NAME = "GmpAppId";
+  private static final String FIREBASE_APP_NAME_COLUMN_NAME = "AppName";
+  private static final String INSTANCE_ID_COLUMN_NAME = "Iid";
+  private static final String CUSTOM_INSTALLATION_ID_COLUMN_NAME = "Cid";
+  private static final String CACHE_STATUS_COLUMN = "Status";
+
+  private static final String QUERY_WHERE_CLAUSE =
+      String.format(
+          "%s = ? " + "AND " + "%s = ?", GMP_APP_ID_COLUMN_NAME, FIREBASE_APP_NAME_COLUMN_NAME);
+
+  private final SQLiteDatabase localDb;
+
+  CustomInstallationIdMappingCache() {
+    // Since different FirebaseApp in the same Android application should have the same application
+    // context and same dir path, so that use the context of the default FirebaseApp to create/open
+    // the database.
+    localDb =
+        SQLiteDatabase.openOrCreateDatabase(
+            FirebaseApp.getInstance()
+                    .getApplicationContext()
+                    .getNoBackupFilesDir()
+                    .getAbsolutePath()
+                + "/"
+                + LOCAL_DB_NAME,
+            null);
+
+    localDb.execSQL(
+        String.format(
+            "CREATE TABLE IF NOT EXISTS %s(%s TEXT PRIMARY KEY, %s TEXT PRIMARY KEY, "
+                + "%s TEXT NOT NULL, %s TEXT NOT NULL, %s INTEGER NOT NULL);",
+            TABLE_NAME,
+            GMP_APP_ID_COLUMN_NAME,
+            FIREBASE_APP_NAME_COLUMN_NAME,
+            INSTANCE_ID_COLUMN_NAME,
+            CUSTOM_INSTALLATION_ID_COLUMN_NAME,
+            CACHE_STATUS_COLUMN));
+  }
+
+  @Nullable
+  String readIid(FirebaseApp firebaseApp) {
+    String gmpAppId = firebaseApp.getOptions().getApplicationId();
+    String appName = firebaseApp.getName();
+    Cursor cursor =
+        localDb.query(
+            TABLE_NAME,
+            new String[] {INSTANCE_ID_COLUMN_NAME},
+            QUERY_WHERE_CLAUSE,
+            new String[] {gmpAppId, appName},
+            null,
+            null,
+            null);
+    String iid = null;
+    while (cursor.moveToNext()) {
+      Preconditions.checkArgument(
+          iid == null, "Multiple iid found for " + "firebase app %s", appName);
+      iid = cursor.getString(cursor.getColumnIndex(INSTANCE_ID_COLUMN_NAME));
+    }
+    return iid;
+  }
+}

From 864748f4ac58f4662187be16d75084a16229e80f Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Fri, 14 Jun 2019 14:25:33 -0700
Subject: [PATCH 02/30] [Firebase Segmentation] Add custom installation id
 cache layer and tests for it.

---
 .../firebase-segmentation.gradle              |  4 +
 .../CustomInstallationIdCacheTest.java        | 77 +++++++++++++++++++
 ...he.java => CustomInstallationIdCache.java} | 62 +++++++++++----
 .../CustomInstallationIdCacheEntryValue.java  | 37 +++++++++
 4 files changed, 167 insertions(+), 13 deletions(-)
 create mode 100644 firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
 rename firebase-segmentation/src/main/java/com/google/firebase/segmentation/{CustomInstallationIdMappingCache.java => CustomInstallationIdCache.java} (58%)
 create mode 100644 firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCacheEntryValue.java

diff --git a/firebase-segmentation/firebase-segmentation.gradle b/firebase-segmentation/firebase-segmentation.gradle
index 11a5046d6d3..cc24fe30ced 100644
--- a/firebase-segmentation/firebase-segmentation.gradle
+++ b/firebase-segmentation/firebase-segmentation.gradle
@@ -91,11 +91,15 @@ dependencies {
     implementation 'androidx.multidex:multidex:2.0.0'
     implementation 'com.google.android.gms:play-services-tasks:16.0.1'
 
+    compileOnly "com.google.auto.value:auto-value-annotations:1.6.5"
+    annotationProcessor "com.google.auto.value:auto-value:1.6.2"
+
     testImplementation 'androidx.test:core:1.2.0'
     testImplementation 'junit:junit:4.12'
     testImplementation "org.robolectric:robolectric:$robolectricVersion"
 
     androidTestImplementation "androidx.annotation:annotation:1.1.0"
+    androidTestImplementation 'androidx.test.ext:junit:1.1.1'
     androidTestImplementation 'androidx.test:rules:1.2.0'
     androidTestImplementation 'androidx.test:runner:1.2.0'
     androidTestImplementation "com.google.truth:truth:$googleTruthVersion"
diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
new file mode 100644
index 00000000000..c6b7ce0cecb
--- /dev/null
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
@@ -0,0 +1,77 @@
+// Copyright 2019 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.firebase.segmentation;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+
+import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import com.google.firebase.FirebaseApp;
+import com.google.firebase.FirebaseOptions;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/** Instrumented tests for {@link CustomInstallationIdCache} */
+@RunWith(AndroidJUnit4.class)
+public class CustomInstallationIdCacheTest {
+
+  private FirebaseApp firebaseApp0;
+  private FirebaseApp firebaseApp1;
+  private CustomInstallationIdCache cache;
+
+  @Before
+  public void setUp() {
+    FirebaseApp.clearInstancesForTest();
+    firebaseApp0 =
+        FirebaseApp.initializeApp(
+            InstrumentationRegistry.getContext(),
+            new FirebaseOptions.Builder().setApplicationId("1:123456789:android:abcdef").build());
+    firebaseApp1 =
+        FirebaseApp.initializeApp(
+            InstrumentationRegistry.getContext(),
+            new FirebaseOptions.Builder().setApplicationId("1:987654321:android:abcdef").build(),
+            "firebase_app_1");
+    cache = new CustomInstallationIdCache();
+  }
+
+  @After
+  public void cleanUp() {
+    cache.clear();
+  }
+
+  @Test
+  public void testReadCacheEntry_Null() {
+    assertNull(cache.readCacheEntryValue(firebaseApp0));
+    assertNull(cache.readCacheEntryValue(firebaseApp1));
+  }
+
+  @Test
+  public void testUpdateAndReadCacheEntry() {
+    cache.insertOrUpdateCacheEntry(
+        firebaseApp0,
+        CustomInstallationIdCacheEntryValue.create(
+            "123456", "cAAAAAAAAAA", CustomInstallationIdCache.CacheStatus.SYNCED));
+    CustomInstallationIdCacheEntryValue entryValue = cache.readCacheEntryValue(firebaseApp0);
+    assertNotNull(entryValue);
+    assertThat(entryValue.getCustomInstallationId()).isEqualTo("123456");
+    assertThat(entryValue.getFirebaseInstanceId()).isEqualTo("cAAAAAAAAAA");
+    assertThat(entryValue.getCacheStatus()).isEqualTo(CustomInstallationIdCache.CacheStatus.SYNCED);
+    assertNull(cache.readCacheEntryValue(firebaseApp1));
+  }
+}
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdMappingCache.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
similarity index 58%
rename from firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdMappingCache.java
rename to firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
index 9826ea9e23f..afb1c86b5b0 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdMappingCache.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
@@ -17,12 +17,15 @@
 import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 import com.google.android.gms.common.internal.Preconditions;
 import com.google.firebase.FirebaseApp;
 
-class CustomInstallationIdMappingCache {
+class CustomInstallationIdCache {
 
   // Status of each cache entry
+  // NOTE: never change the ordinal of the enum values because the enum values are stored in cache
+  // as their ordinal numbers.
   enum CacheStatus {
     // Cache entry is synced to Firebase backend
     SYNCED,
@@ -38,8 +41,8 @@ enum CacheStatus {
 
   private static final String GMP_APP_ID_COLUMN_NAME = "GmpAppId";
   private static final String FIREBASE_APP_NAME_COLUMN_NAME = "AppName";
-  private static final String INSTANCE_ID_COLUMN_NAME = "Iid";
   private static final String CUSTOM_INSTALLATION_ID_COLUMN_NAME = "Cid";
+  private static final String INSTANCE_ID_COLUMN_NAME = "Iid";
   private static final String CACHE_STATUS_COLUMN = "Status";
 
   private static final String QUERY_WHERE_CLAUSE =
@@ -48,7 +51,7 @@ enum CacheStatus {
 
   private final SQLiteDatabase localDb;
 
-  CustomInstallationIdMappingCache() {
+  CustomInstallationIdCache() {
     // Since different FirebaseApp in the same Android application should have the same application
     // context and same dir path, so that use the context of the default FirebaseApp to create/open
     // the database.
@@ -64,35 +67,68 @@ enum CacheStatus {
 
     localDb.execSQL(
         String.format(
-            "CREATE TABLE IF NOT EXISTS %s(%s TEXT PRIMARY KEY, %s TEXT PRIMARY KEY, "
-                + "%s TEXT NOT NULL, %s TEXT NOT NULL, %s INTEGER NOT NULL);",
+            "CREATE TABLE IF NOT EXISTS %s(%s TEXT NOT NULL, %s TEXT NOT NULL, "
+                + "%s TEXT NOT NULL, %s TEXT NOT NULL, %s INTEGER NOT NULL, PRIMARY KEY (%s, %s));",
             TABLE_NAME,
             GMP_APP_ID_COLUMN_NAME,
             FIREBASE_APP_NAME_COLUMN_NAME,
-            INSTANCE_ID_COLUMN_NAME,
             CUSTOM_INSTALLATION_ID_COLUMN_NAME,
-            CACHE_STATUS_COLUMN));
+            INSTANCE_ID_COLUMN_NAME,
+            CACHE_STATUS_COLUMN,
+            GMP_APP_ID_COLUMN_NAME,
+            FIREBASE_APP_NAME_COLUMN_NAME));
   }
 
   @Nullable
-  String readIid(FirebaseApp firebaseApp) {
+  CustomInstallationIdCacheEntryValue readCacheEntryValue(FirebaseApp firebaseApp) {
     String gmpAppId = firebaseApp.getOptions().getApplicationId();
     String appName = firebaseApp.getName();
     Cursor cursor =
         localDb.query(
             TABLE_NAME,
-            new String[] {INSTANCE_ID_COLUMN_NAME},
+            new String[] {
+              CUSTOM_INSTALLATION_ID_COLUMN_NAME, INSTANCE_ID_COLUMN_NAME, CACHE_STATUS_COLUMN
+            },
             QUERY_WHERE_CLAUSE,
             new String[] {gmpAppId, appName},
             null,
             null,
             null);
-    String iid = null;
+    CustomInstallationIdCacheEntryValue value = null;
     while (cursor.moveToNext()) {
       Preconditions.checkArgument(
-          iid == null, "Multiple iid found for " + "firebase app %s", appName);
-      iid = cursor.getString(cursor.getColumnIndex(INSTANCE_ID_COLUMN_NAME));
+          value == null, "Multiple cache entries found for " + "firebase app %s", appName);
+      value =
+          CustomInstallationIdCacheEntryValue.create(
+              cursor.getString(cursor.getColumnIndex(CUSTOM_INSTALLATION_ID_COLUMN_NAME)),
+              cursor.getString(cursor.getColumnIndex(INSTANCE_ID_COLUMN_NAME)),
+              CacheStatus.values()[cursor.getInt(cursor.getColumnIndex(CACHE_STATUS_COLUMN))]);
     }
-    return iid;
+    return value;
+  }
+
+  void insertOrUpdateCacheEntry(
+      FirebaseApp firebaseApp, CustomInstallationIdCacheEntryValue entryValue) {
+    String gmpAppId = firebaseApp.getOptions().getApplicationId();
+    String appName = firebaseApp.getName();
+    localDb.execSQL(
+        String.format(
+            "INSERT OR REPLACE INTO %s(%s, %s, %s, %s, %s) VALUES(%s, %s, %s, %s, %s)",
+            TABLE_NAME,
+            GMP_APP_ID_COLUMN_NAME,
+            FIREBASE_APP_NAME_COLUMN_NAME,
+            CUSTOM_INSTALLATION_ID_COLUMN_NAME,
+            INSTANCE_ID_COLUMN_NAME,
+            CACHE_STATUS_COLUMN,
+            "\"" + gmpAppId + "\"",
+            "\"" + appName + "\"",
+            "\"" + entryValue.getCustomInstallationId() + "\"",
+            "\"" + entryValue.getFirebaseInstanceId() + "\"",
+            entryValue.getCacheStatus().ordinal()));
+  }
+
+  @VisibleForTesting
+  void clear() {
+    localDb.execSQL(String.format("DROP TABLE IF EXISTS %s", TABLE_NAME));
   }
 }
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCacheEntryValue.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCacheEntryValue.java
new file mode 100644
index 00000000000..c79d5f091a9
--- /dev/null
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCacheEntryValue.java
@@ -0,0 +1,37 @@
+// Copyright 2019 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.firebase.segmentation;
+
+import com.google.auto.value.AutoValue;
+import com.google.firebase.segmentation.CustomInstallationIdCache.CacheStatus;
+
+/**
+ * This class represents a cache entry value in {@link CustomInstallationIdCache}, which contains a
+ * Firebase instance id, a custom installation id and the cache status of this entry.
+ */
+@AutoValue
+abstract class CustomInstallationIdCacheEntryValue {
+  abstract String getCustomInstallationId();
+
+  abstract String getFirebaseInstanceId();
+
+  abstract CacheStatus getCacheStatus();
+
+  static CustomInstallationIdCacheEntryValue create(
+      String customInstallationId, String firebaseInstanceId, CacheStatus cacheStatus) {
+    return new AutoValue_CustomInstallationIdCacheEntryValue(
+        customInstallationId, firebaseInstanceId, cacheStatus);
+  }
+}

From 0a3ebf6a5b2aa44d36469caf016e938e9376255b Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Fri, 14 Jun 2019 14:38:05 -0700
Subject: [PATCH 03/30] Add test for updating cache

---
 .../CustomInstallationIdCacheTest.java           | 16 ++++++++++++----
 .../segmentation/CustomInstallationIdCache.java  |  6 ++----
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
index c6b7ce0cecb..0dd24398e32 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
@@ -15,7 +15,6 @@
 package com.google.firebase.segmentation;
 
 import static com.google.common.truth.Truth.assertThat;
-import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 
 import androidx.test.InstrumentationRegistry;
@@ -66,12 +65,21 @@ public void testUpdateAndReadCacheEntry() {
     cache.insertOrUpdateCacheEntry(
         firebaseApp0,
         CustomInstallationIdCacheEntryValue.create(
-            "123456", "cAAAAAAAAAA", CustomInstallationIdCache.CacheStatus.SYNCED));
+            "123456", "cAAAAAAAAAA", CustomInstallationIdCache.CacheStatus.PENDING));
     CustomInstallationIdCacheEntryValue entryValue = cache.readCacheEntryValue(firebaseApp0);
-    assertNotNull(entryValue);
     assertThat(entryValue.getCustomInstallationId()).isEqualTo("123456");
     assertThat(entryValue.getFirebaseInstanceId()).isEqualTo("cAAAAAAAAAA");
-    assertThat(entryValue.getCacheStatus()).isEqualTo(CustomInstallationIdCache.CacheStatus.SYNCED);
+    assertThat(entryValue.getCacheStatus())
+        .isEqualTo(CustomInstallationIdCache.CacheStatus.PENDING);
     assertNull(cache.readCacheEntryValue(firebaseApp1));
+
+    cache.insertOrUpdateCacheEntry(
+        firebaseApp0,
+        CustomInstallationIdCacheEntryValue.create(
+            "123456", "cAAAAAAAAAA", CustomInstallationIdCache.CacheStatus.SYNCED));
+    entryValue = cache.readCacheEntryValue(firebaseApp0);
+    assertThat(entryValue.getCustomInstallationId()).isEqualTo("123456");
+    assertThat(entryValue.getFirebaseInstanceId()).isEqualTo("cAAAAAAAAAA");
+    assertThat(entryValue.getCacheStatus()).isEqualTo(CustomInstallationIdCache.CacheStatus.SYNCED);
   }
 }
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
index afb1c86b5b0..94df707d3fe 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
@@ -109,8 +109,6 @@ CustomInstallationIdCacheEntryValue readCacheEntryValue(FirebaseApp firebaseApp)
 
   void insertOrUpdateCacheEntry(
       FirebaseApp firebaseApp, CustomInstallationIdCacheEntryValue entryValue) {
-    String gmpAppId = firebaseApp.getOptions().getApplicationId();
-    String appName = firebaseApp.getName();
     localDb.execSQL(
         String.format(
             "INSERT OR REPLACE INTO %s(%s, %s, %s, %s, %s) VALUES(%s, %s, %s, %s, %s)",
@@ -120,8 +118,8 @@ void insertOrUpdateCacheEntry(
             CUSTOM_INSTALLATION_ID_COLUMN_NAME,
             INSTANCE_ID_COLUMN_NAME,
             CACHE_STATUS_COLUMN,
-            "\"" + gmpAppId + "\"",
-            "\"" + appName + "\"",
+            "\"" + firebaseApp.getOptions().getApplicationId() + "\"",
+            "\"" + firebaseApp.getName() + "\"",
             "\"" + entryValue.getCustomInstallationId() + "\"",
             "\"" + entryValue.getFirebaseInstanceId() + "\"",
             entryValue.getCacheStatus().ordinal()));

From 2d158ed63a92b1130cca4f34286177de42e45e5e Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Fri, 14 Jun 2019 17:03:52 -0700
Subject: [PATCH 04/30] Switch to use SQLiteOpenHelper

---
 .../firebase-segmentation.gradle              |   2 +-
 .../CustomInstallationIdCache.java            | 176 +++++++++++++-----
 2 files changed, 127 insertions(+), 51 deletions(-)

diff --git a/firebase-segmentation/firebase-segmentation.gradle b/firebase-segmentation/firebase-segmentation.gradle
index cc24fe30ced..dc4606715c5 100644
--- a/firebase-segmentation/firebase-segmentation.gradle
+++ b/firebase-segmentation/firebase-segmentation.gradle
@@ -57,7 +57,7 @@ android {
     compileSdkVersion project.targetSdkVersion
 
     defaultConfig {
-        minSdkVersion 21
+        minSdkVersion project.minSdkVersion
         targetSdkVersion project.targetSdkVersion
         multiDexEnabled true
         versionName version
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
index 94df707d3fe..e2647f75cca 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
@@ -14,8 +14,11 @@
 
 package com.google.firebase.segmentation;
 
+import android.content.Context;
 import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
+import android.database.sqlite.SQLiteOpenHelper;
+import android.os.Build;
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import com.google.android.gms.common.internal.Preconditions;
@@ -49,34 +52,105 @@ enum CacheStatus {
       String.format(
           "%s = ? " + "AND " + "%s = ?", GMP_APP_ID_COLUMN_NAME, FIREBASE_APP_NAME_COLUMN_NAME);
 
-  private final SQLiteDatabase localDb;
+  /**
+   * A SQLiteOpenHelper that configures database connections just the way we like them, delegating
+   * to SQLiteSchema to actually do the work of migration.
+   *
+   * <p>The order of events when opening a new connection is as follows:
+   *
+   * <ol>
+   *   <li>New connection
+   *   <li>onConfigure (API 16 and above)
+   *   <li>onCreate / onUpgrade (optional; if version already matches these aren't called)
+   *   <li>onOpen
+   * </ol>
+   *
+   * <p>This OpenHelper attempts to obtain exclusive access to the database and attempts to do so as
+   * early as possible. On Jelly Bean devices and above (some 98% of devices at time of writing)
+   * this happens naturally during onConfigure. On pre-Jelly Bean devices all other methods ensure
+   * that the configuration is applied before any action is taken.
+   */
+  private static class OpenHelper extends SQLiteOpenHelper {
+    // TODO: when we do schema upgrades in the future we need to make sure both downgrades and
+    // upgrades work as expected, e.g. `up+down+up` is equivalent to `up`.
+    private static int SCHEMA_VERSION = 1;
+
+    private boolean configured = false;
+
+    private OpenHelper(Context context) {
+      super(context, LOCAL_DB_NAME, null, SCHEMA_VERSION);
+    }
+
+    @Override
+    public void onConfigure(SQLiteDatabase db) {
+      // Note that this is only called automatically by the SQLiteOpenHelper base class on Jelly
+      // Bean and above.
+      configured = true;
+
+      db.rawQuery("PRAGMA busy_timeout=0;", new String[0]).close();
+
+      if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.JELLY_BEAN) {
+        db.setForeignKeyConstraintsEnabled(true);
+      }
+    }
+
+    private void ensureConfigured(SQLiteDatabase db) {
+      if (!configured) {
+        onConfigure(db);
+      }
+    }
+
+    @Override
+    public void onCreate(SQLiteDatabase db) {
+      ensureConfigured(db);
+      // Create custom id mapping table.
+      db.execSQL(
+          String.format(
+              "CREATE TABLE IF NOT EXISTS %s(%s TEXT NOT NULL, %s TEXT NOT NULL, "
+                  + "%s TEXT NOT NULL, %s TEXT NOT NULL, %s INTEGER NOT NULL, PRIMARY KEY (%s, %s));",
+              TABLE_NAME,
+              GMP_APP_ID_COLUMN_NAME,
+              FIREBASE_APP_NAME_COLUMN_NAME,
+              CUSTOM_INSTALLATION_ID_COLUMN_NAME,
+              INSTANCE_ID_COLUMN_NAME,
+              CACHE_STATUS_COLUMN,
+              GMP_APP_ID_COLUMN_NAME,
+              FIREBASE_APP_NAME_COLUMN_NAME));
+    }
+
+    @Override
+    public void onUpgrade(SQLiteDatabase db, int oldVersion, int newVersion) {
+      ensureConfigured(db);
+    }
+
+    @Override
+    public void onDowngrade(SQLiteDatabase db, int oldVersion, int newVersion) {
+      ensureConfigured(db);
+    }
+
+    @Override
+    public void onOpen(SQLiteDatabase db) {
+      ensureConfigured(db);
+    }
+  }
+
+  private final OpenHelper openHelper;
 
   CustomInstallationIdCache() {
-    // Since different FirebaseApp in the same Android application should have the same application
-    // context and same dir path, so that use the context of the default FirebaseApp to create/open
+    // Since different FirebaseApp in the same Android application should
+    // have the same application
+    // context and same dir path, so that use the context of the default
+    // FirebaseApp to create/open
     // the database.
-    localDb =
-        SQLiteDatabase.openOrCreateDatabase(
-            FirebaseApp.getInstance()
-                    .getApplicationContext()
-                    .getNoBackupFilesDir()
-                    .getAbsolutePath()
-                + "/"
-                + LOCAL_DB_NAME,
-            null);
-
-    localDb.execSQL(
-        String.format(
-            "CREATE TABLE IF NOT EXISTS %s(%s TEXT NOT NULL, %s TEXT NOT NULL, "
-                + "%s TEXT NOT NULL, %s TEXT NOT NULL, %s INTEGER NOT NULL, PRIMARY KEY (%s, %s));",
-            TABLE_NAME,
-            GMP_APP_ID_COLUMN_NAME,
-            FIREBASE_APP_NAME_COLUMN_NAME,
-            CUSTOM_INSTALLATION_ID_COLUMN_NAME,
-            INSTANCE_ID_COLUMN_NAME,
-            CACHE_STATUS_COLUMN,
-            GMP_APP_ID_COLUMN_NAME,
-            FIREBASE_APP_NAME_COLUMN_NAME));
+    openHelper = new OpenHelper(FirebaseApp.getInstance().getApplicationContext());
+  }
+
+  private SQLiteDatabase getReadableDb() {
+    return openHelper.getReadableDatabase();
+  }
+
+  private SQLiteDatabase getWritableDb() {
+    return openHelper.getWritableDatabase();
   }
 
   @Nullable
@@ -84,16 +158,17 @@ CustomInstallationIdCacheEntryValue readCacheEntryValue(FirebaseApp firebaseApp)
     String gmpAppId = firebaseApp.getOptions().getApplicationId();
     String appName = firebaseApp.getName();
     Cursor cursor =
-        localDb.query(
-            TABLE_NAME,
-            new String[] {
-              CUSTOM_INSTALLATION_ID_COLUMN_NAME, INSTANCE_ID_COLUMN_NAME, CACHE_STATUS_COLUMN
-            },
-            QUERY_WHERE_CLAUSE,
-            new String[] {gmpAppId, appName},
-            null,
-            null,
-            null);
+        getReadableDb()
+            .query(
+                TABLE_NAME,
+                new String[] {
+                  CUSTOM_INSTALLATION_ID_COLUMN_NAME, INSTANCE_ID_COLUMN_NAME, CACHE_STATUS_COLUMN
+                },
+                QUERY_WHERE_CLAUSE,
+                new String[] {gmpAppId, appName},
+                null,
+                null,
+                null);
     CustomInstallationIdCacheEntryValue value = null;
     while (cursor.moveToNext()) {
       Preconditions.checkArgument(
@@ -109,24 +184,25 @@ CustomInstallationIdCacheEntryValue readCacheEntryValue(FirebaseApp firebaseApp)
 
   void insertOrUpdateCacheEntry(
       FirebaseApp firebaseApp, CustomInstallationIdCacheEntryValue entryValue) {
-    localDb.execSQL(
-        String.format(
-            "INSERT OR REPLACE INTO %s(%s, %s, %s, %s, %s) VALUES(%s, %s, %s, %s, %s)",
-            TABLE_NAME,
-            GMP_APP_ID_COLUMN_NAME,
-            FIREBASE_APP_NAME_COLUMN_NAME,
-            CUSTOM_INSTALLATION_ID_COLUMN_NAME,
-            INSTANCE_ID_COLUMN_NAME,
-            CACHE_STATUS_COLUMN,
-            "\"" + firebaseApp.getOptions().getApplicationId() + "\"",
-            "\"" + firebaseApp.getName() + "\"",
-            "\"" + entryValue.getCustomInstallationId() + "\"",
-            "\"" + entryValue.getFirebaseInstanceId() + "\"",
-            entryValue.getCacheStatus().ordinal()));
+    getWritableDb()
+        .execSQL(
+            String.format(
+                "INSERT OR REPLACE INTO %s(%s, %s, %s, %s, %s) VALUES(%s, %s, %s, %s, %s)",
+                TABLE_NAME,
+                GMP_APP_ID_COLUMN_NAME,
+                FIREBASE_APP_NAME_COLUMN_NAME,
+                CUSTOM_INSTALLATION_ID_COLUMN_NAME,
+                INSTANCE_ID_COLUMN_NAME,
+                CACHE_STATUS_COLUMN,
+                "\"" + firebaseApp.getOptions().getApplicationId() + "\"",
+                "\"" + firebaseApp.getName() + "\"",
+                "\"" + entryValue.getCustomInstallationId() + "\"",
+                "\"" + entryValue.getFirebaseInstanceId() + "\"",
+                entryValue.getCacheStatus().ordinal()));
   }
 
   @VisibleForTesting
   void clear() {
-    localDb.execSQL(String.format("DROP TABLE IF EXISTS %s", TABLE_NAME));
+    getWritableDb().execSQL(String.format("DELETE FROM %s", TABLE_NAME));
   }
 }

From f118d39bf6cef56330d37ce154afa246b3891269 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Mon, 17 Jun 2019 14:20:39 -0700
Subject: [PATCH 05/30] Switch to use SharedPreferences from SQLite.

---
 .../CustomInstallationIdCacheTest.java        |   2 +-
 .../CustomInstallationIdCache.java            | 207 +++++-------------
 2 files changed, 51 insertions(+), 158 deletions(-)

diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
index 0dd24398e32..3e085e32a22 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
@@ -51,7 +51,7 @@ public void setUp() {
 
   @After
   public void cleanUp() {
-    cache.clear();
+    cache.clearAll();
   }
 
   @Test
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
index e2647f75cca..cb50fb3891c 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
@@ -14,14 +14,10 @@
 
 package com.google.firebase.segmentation;
 
-import android.content.Context;
-import android.database.Cursor;
-import android.database.sqlite.SQLiteDatabase;
-import android.database.sqlite.SQLiteOpenHelper;
-import android.os.Build;
+import android.content.SharedPreferences;
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
-import com.google.android.gms.common.internal.Preconditions;
+import com.google.android.gms.common.util.Strings;
 import com.google.firebase.FirebaseApp;
 
 class CustomInstallationIdCache {
@@ -36,173 +32,70 @@ enum CacheStatus {
     // errors.
     PENDING,
     // Cache entry is not accepted by Firebase backend.
-    ERROR
+    ERROR,
   }
 
-  private static final String LOCAL_DB_NAME = "CustomInstallationIdCache";
-  private static final String TABLE_NAME = "InstallationIdMapping";
+  private static final String SHARED_PREFS_NAME = "CustomInstallationIdCache";
 
-  private static final String GMP_APP_ID_COLUMN_NAME = "GmpAppId";
-  private static final String FIREBASE_APP_NAME_COLUMN_NAME = "AppName";
-  private static final String CUSTOM_INSTALLATION_ID_COLUMN_NAME = "Cid";
-  private static final String INSTANCE_ID_COLUMN_NAME = "Iid";
-  private static final String CACHE_STATUS_COLUMN = "Status";
+  private static final String CUSTOM_INSTALLATION_ID_KEY = "Cid";
+  private static final String INSTANCE_ID_KEY = "Iid";
+  private static final String CACHE_STATUS_KEY = "Status";
 
-  private static final String QUERY_WHERE_CLAUSE =
-      String.format(
-          "%s = ? " + "AND " + "%s = ?", GMP_APP_ID_COLUMN_NAME, FIREBASE_APP_NAME_COLUMN_NAME);
-
-  /**
-   * A SQLiteOpenHelper that configures database connections just the way we like them, delegating
-   * to SQLiteSchema to actually do the work of migration.
-   *
-   * <p>The order of events when opening a new connection is as follows:
-   *
-   * <ol>
-   *   <li>New connection
-   *   <li>onConfigure (API 16 and above)
-   *   <li>onCreate / onUpgrade (optional; if version already matches these aren't called)
-   *   <li>onOpen
-   * </ol>
-   *
-   * <p>This OpenHelper attempts to obtain exclusive access to the database and attempts to do so as
-   * early as possible. On Jelly Bean devices and above (some 98% of devices at time of writing)
-   * this happens naturally during onConfigure. On pre-Jelly Bean devices all other methods ensure
-   * that the configuration is applied before any action is taken.
-   */
-  private static class OpenHelper extends SQLiteOpenHelper {
-    // TODO: when we do schema upgrades in the future we need to make sure both downgrades and
-    // upgrades work as expected, e.g. `up+down+up` is equivalent to `up`.
-    private static int SCHEMA_VERSION = 1;
-
-    private boolean configured = false;
-
-    private OpenHelper(Context context) {
-      super(context, LOCAL_DB_NAME, null, SCHEMA_VERSION);
-    }
-
-    @Override
-    public void onConfigure(SQLiteDatabase db) {
-      // Note that this is only called automatically by the SQLiteOpenHelper base class on Jelly
-      // Bean and above.
-      configured = true;
-
-      db.rawQuery("PRAGMA busy_timeout=0;", new String[0]).close();
-
-      if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.JELLY_BEAN) {
-        db.setForeignKeyConstraintsEnabled(true);
-      }
-    }
-
-    private void ensureConfigured(SQLiteDatabase db) {
-      if (!configured) {
-        onConfigure(db);
-      }
-    }
-
-    @Override
-    public void onCreate(SQLiteDatabase db) {
-      ensureConfigured(db);
-      // Create custom id mapping table.
-      db.execSQL(
-          String.format(
-              "CREATE TABLE IF NOT EXISTS %s(%s TEXT NOT NULL, %s TEXT NOT NULL, "
-                  + "%s TEXT NOT NULL, %s TEXT NOT NULL, %s INTEGER NOT NULL, PRIMARY KEY (%s, %s));",
-              TABLE_NAME,
-              GMP_APP_ID_COLUMN_NAME,
-              FIREBASE_APP_NAME_COLUMN_NAME,
-              CUSTOM_INSTALLATION_ID_COLUMN_NAME,
-              INSTANCE_ID_COLUMN_NAME,
-              CACHE_STATUS_COLUMN,
-              GMP_APP_ID_COLUMN_NAME,
-              FIREBASE_APP_NAME_COLUMN_NAME));
-    }
-
-    @Override
-    public void onUpgrade(SQLiteDatabase db, int oldVersion, int newVersion) {
-      ensureConfigured(db);
-    }
-
-    @Override
-    public void onDowngrade(SQLiteDatabase db, int oldVersion, int newVersion) {
-      ensureConfigured(db);
-    }
-
-    @Override
-    public void onOpen(SQLiteDatabase db) {
-      ensureConfigured(db);
-    }
-  }
-
-  private final OpenHelper openHelper;
+  private final SharedPreferences prefs;
 
   CustomInstallationIdCache() {
-    // Since different FirebaseApp in the same Android application should
-    // have the same application
-    // context and same dir path, so that use the context of the default
-    // FirebaseApp to create/open
-    // the database.
-    openHelper = new OpenHelper(FirebaseApp.getInstance().getApplicationContext());
+    // Since different FirebaseApp in the same Android application should have the same application
+    // context and same dir path, so that use the context of the default FirebaseApp to create the
+    // shared preferences.
+    prefs =
+        FirebaseApp.getInstance()
+            .getApplicationContext()
+            .getSharedPreferences(SHARED_PREFS_NAME, 0); // private mode
   }
 
-  private SQLiteDatabase getReadableDb() {
-    return openHelper.getReadableDatabase();
+  @Nullable
+  synchronized CustomInstallationIdCacheEntryValue readCacheEntryValue(FirebaseApp firebaseApp) {
+    String cid =
+        prefs.getString(getSharedPreferencesKey(firebaseApp, CUSTOM_INSTALLATION_ID_KEY), null);
+    String iid = prefs.getString(getSharedPreferencesKey(firebaseApp, INSTANCE_ID_KEY), null);
+    int status = prefs.getInt(getSharedPreferencesKey(firebaseApp, CACHE_STATUS_KEY), -1);
+
+    if (Strings.isEmptyOrWhitespace(cid) || Strings.isEmptyOrWhitespace(iid) || status == -1) {
+      return null;
+    }
+
+    return CustomInstallationIdCacheEntryValue.create(cid, iid, CacheStatus.values()[status]);
   }
 
-  private SQLiteDatabase getWritableDb() {
-    return openHelper.getWritableDatabase();
+  synchronized void insertOrUpdateCacheEntry(
+      FirebaseApp firebaseApp, CustomInstallationIdCacheEntryValue entryValue) {
+    SharedPreferences.Editor editor = prefs.edit();
+    editor.putString(
+        getSharedPreferencesKey(firebaseApp, CUSTOM_INSTALLATION_ID_KEY),
+        entryValue.getCustomInstallationId());
+    editor.putString(
+        getSharedPreferencesKey(firebaseApp, INSTANCE_ID_KEY), entryValue.getFirebaseInstanceId());
+    editor.putInt(
+        getSharedPreferencesKey(firebaseApp, CACHE_STATUS_KEY),
+        entryValue.getCacheStatus().ordinal());
+    editor.commit();
   }
 
-  @Nullable
-  CustomInstallationIdCacheEntryValue readCacheEntryValue(FirebaseApp firebaseApp) {
-    String gmpAppId = firebaseApp.getOptions().getApplicationId();
-    String appName = firebaseApp.getName();
-    Cursor cursor =
-        getReadableDb()
-            .query(
-                TABLE_NAME,
-                new String[] {
-                  CUSTOM_INSTALLATION_ID_COLUMN_NAME, INSTANCE_ID_COLUMN_NAME, CACHE_STATUS_COLUMN
-                },
-                QUERY_WHERE_CLAUSE,
-                new String[] {gmpAppId, appName},
-                null,
-                null,
-                null);
-    CustomInstallationIdCacheEntryValue value = null;
-    while (cursor.moveToNext()) {
-      Preconditions.checkArgument(
-          value == null, "Multiple cache entries found for " + "firebase app %s", appName);
-      value =
-          CustomInstallationIdCacheEntryValue.create(
-              cursor.getString(cursor.getColumnIndex(CUSTOM_INSTALLATION_ID_COLUMN_NAME)),
-              cursor.getString(cursor.getColumnIndex(INSTANCE_ID_COLUMN_NAME)),
-              CacheStatus.values()[cursor.getInt(cursor.getColumnIndex(CACHE_STATUS_COLUMN))]);
-    }
-    return value;
+  synchronized void clear(FirebaseApp firebaseApp) {
+    SharedPreferences.Editor editor = prefs.edit();
+    editor.remove(getSharedPreferencesKey(firebaseApp, CUSTOM_INSTALLATION_ID_KEY));
+    editor.remove(getSharedPreferencesKey(firebaseApp, INSTANCE_ID_KEY));
+    editor.remove(getSharedPreferencesKey(firebaseApp, CACHE_STATUS_KEY));
   }
 
-  void insertOrUpdateCacheEntry(
-      FirebaseApp firebaseApp, CustomInstallationIdCacheEntryValue entryValue) {
-    getWritableDb()
-        .execSQL(
-            String.format(
-                "INSERT OR REPLACE INTO %s(%s, %s, %s, %s, %s) VALUES(%s, %s, %s, %s, %s)",
-                TABLE_NAME,
-                GMP_APP_ID_COLUMN_NAME,
-                FIREBASE_APP_NAME_COLUMN_NAME,
-                CUSTOM_INSTALLATION_ID_COLUMN_NAME,
-                INSTANCE_ID_COLUMN_NAME,
-                CACHE_STATUS_COLUMN,
-                "\"" + firebaseApp.getOptions().getApplicationId() + "\"",
-                "\"" + firebaseApp.getName() + "\"",
-                "\"" + entryValue.getCustomInstallationId() + "\"",
-                "\"" + entryValue.getFirebaseInstanceId() + "\"",
-                entryValue.getCacheStatus().ordinal()));
+  private static String getSharedPreferencesKey(FirebaseApp firebaseApp, String key) {
+    return String.format("%s|%s", firebaseApp.getPersistenceKey(), key);
   }
 
   @VisibleForTesting
-  void clear() {
-    getWritableDb().execSQL(String.format("DELETE FROM %s", TABLE_NAME));
+  synchronized void clearAll() {
+    SharedPreferences.Editor editor = prefs.edit();
+    editor.clear();
+    editor.commit();
   }
 }

From 4da5d31f31d2160e6ed3702e72d305f04a90e0c5 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Mon, 17 Jun 2019 17:05:01 -0700
Subject: [PATCH 06/30] Change the cache class to be singleton

---
 .../segmentation/CustomInstallationIdCacheTest.java    |  2 +-
 .../segmentation/CustomInstallationIdCache.java        | 10 +++++++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
index 3e085e32a22..2645ab1571f 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
@@ -46,7 +46,7 @@ public void setUp() {
             InstrumentationRegistry.getContext(),
             new FirebaseOptions.Builder().setApplicationId("1:987654321:android:abcdef").build(),
             "firebase_app_1");
-    cache = new CustomInstallationIdCache();
+    cache = CustomInstallationIdCache.getInstance();
   }
 
   @After
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
index cb50fb3891c..1e48ca6d6c9 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
@@ -41,9 +41,17 @@ enum CacheStatus {
   private static final String INSTANCE_ID_KEY = "Iid";
   private static final String CACHE_STATUS_KEY = "Status";
 
+  private static CustomInstallationIdCache singleton = null;
   private final SharedPreferences prefs;
 
-  CustomInstallationIdCache() {
+  static CustomInstallationIdCache getInstance() {
+    if (singleton == null) {
+      singleton = new CustomInstallationIdCache();
+    }
+    return singleton;
+  }
+
+  private CustomInstallationIdCache() {
     // Since different FirebaseApp in the same Android application should have the same application
     // context and same dir path, so that use the context of the default FirebaseApp to create the
     // shared preferences.

From d1ff0ec0bcd7b111ea67d18950f8c5f455a759e9 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Tue, 18 Jun 2019 10:41:01 -0700
Subject: [PATCH 07/30] Wrap shared pref commit in a async task.

---
 .../CustomInstallationIdCacheTest.java        | 28 +++++++------
 .../CustomInstallationIdCache.java            | 39 ++++++++++++++-----
 2 files changed, 47 insertions(+), 20 deletions(-)

diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
index 2645ab1571f..5783294cfa3 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
@@ -16,9 +16,11 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 
 import androidx.test.InstrumentationRegistry;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
+import com.google.android.gms.tasks.Tasks;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.FirebaseOptions;
 import org.junit.After;
@@ -50,8 +52,8 @@ public void setUp() {
   }
 
   @After
-  public void cleanUp() {
-    cache.clearAll();
+  public void cleanUp() throws Exception {
+    Tasks.await(cache.clearAll());
   }
 
   @Test
@@ -61,11 +63,13 @@ public void testReadCacheEntry_Null() {
   }
 
   @Test
-  public void testUpdateAndReadCacheEntry() {
-    cache.insertOrUpdateCacheEntry(
-        firebaseApp0,
-        CustomInstallationIdCacheEntryValue.create(
-            "123456", "cAAAAAAAAAA", CustomInstallationIdCache.CacheStatus.PENDING));
+  public void testUpdateAndReadCacheEntry() throws Exception {
+    assertTrue(
+        Tasks.await(
+            cache.insertOrUpdateCacheEntry(
+                firebaseApp0,
+                CustomInstallationIdCacheEntryValue.create(
+                    "123456", "cAAAAAAAAAA", CustomInstallationIdCache.CacheStatus.PENDING))));
     CustomInstallationIdCacheEntryValue entryValue = cache.readCacheEntryValue(firebaseApp0);
     assertThat(entryValue.getCustomInstallationId()).isEqualTo("123456");
     assertThat(entryValue.getFirebaseInstanceId()).isEqualTo("cAAAAAAAAAA");
@@ -73,10 +77,12 @@ public void testUpdateAndReadCacheEntry() {
         .isEqualTo(CustomInstallationIdCache.CacheStatus.PENDING);
     assertNull(cache.readCacheEntryValue(firebaseApp1));
 
-    cache.insertOrUpdateCacheEntry(
-        firebaseApp0,
-        CustomInstallationIdCacheEntryValue.create(
-            "123456", "cAAAAAAAAAA", CustomInstallationIdCache.CacheStatus.SYNCED));
+    assertTrue(
+        Tasks.await(
+            cache.insertOrUpdateCacheEntry(
+                firebaseApp0,
+                CustomInstallationIdCacheEntryValue.create(
+                    "123456", "cAAAAAAAAAA", CustomInstallationIdCache.CacheStatus.SYNCED))));
     entryValue = cache.readCacheEntryValue(firebaseApp0);
     assertThat(entryValue.getCustomInstallationId()).isEqualTo("123456");
     assertThat(entryValue.getFirebaseInstanceId()).isEqualTo("cAAAAAAAAAA");
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
index 1e48ca6d6c9..1863b976d9d 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
@@ -14,11 +14,16 @@
 
 package com.google.firebase.segmentation;
 
+import android.content.Context;
 import android.content.SharedPreferences;
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import com.google.android.gms.common.util.Strings;
+import com.google.android.gms.tasks.Task;
+import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.firebase.FirebaseApp;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
 
 class CustomInstallationIdCache {
 
@@ -42,6 +47,7 @@ enum CacheStatus {
   private static final String CACHE_STATUS_KEY = "Status";
 
   private static CustomInstallationIdCache singleton = null;
+  private final Executor ioExecuter;
   private final SharedPreferences prefs;
 
   static CustomInstallationIdCache getInstance() {
@@ -58,7 +64,9 @@ private CustomInstallationIdCache() {
     prefs =
         FirebaseApp.getInstance()
             .getApplicationContext()
-            .getSharedPreferences(SHARED_PREFS_NAME, 0); // private mode
+            .getSharedPreferences(SHARED_PREFS_NAME, Context.MODE_PRIVATE);
+
+    ioExecuter = Executors.newFixedThreadPool(2);
   }
 
   @Nullable
@@ -75,7 +83,7 @@ synchronized CustomInstallationIdCacheEntryValue readCacheEntryValue(FirebaseApp
     return CustomInstallationIdCacheEntryValue.create(cid, iid, CacheStatus.values()[status]);
   }
 
-  synchronized void insertOrUpdateCacheEntry(
+  synchronized Task<Boolean> insertOrUpdateCacheEntry(
       FirebaseApp firebaseApp, CustomInstallationIdCacheEntryValue entryValue) {
     SharedPreferences.Editor editor = prefs.edit();
     editor.putString(
@@ -86,24 +94,37 @@ synchronized void insertOrUpdateCacheEntry(
     editor.putInt(
         getSharedPreferencesKey(firebaseApp, CACHE_STATUS_KEY),
         entryValue.getCacheStatus().ordinal());
-    editor.commit();
+    return commitSharedPreferencesEditAsync(editor);
   }
 
-  synchronized void clear(FirebaseApp firebaseApp) {
+  synchronized Task<Boolean> clear(FirebaseApp firebaseApp) {
     SharedPreferences.Editor editor = prefs.edit();
     editor.remove(getSharedPreferencesKey(firebaseApp, CUSTOM_INSTALLATION_ID_KEY));
     editor.remove(getSharedPreferencesKey(firebaseApp, INSTANCE_ID_KEY));
     editor.remove(getSharedPreferencesKey(firebaseApp, CACHE_STATUS_KEY));
+    return commitSharedPreferencesEditAsync(editor);
+  }
+
+  @VisibleForTesting
+  synchronized Task<Boolean> clearAll() {
+    SharedPreferences.Editor editor = prefs.edit();
+    editor.clear();
+    return commitSharedPreferencesEditAsync(editor);
   }
 
   private static String getSharedPreferencesKey(FirebaseApp firebaseApp, String key) {
     return String.format("%s|%s", firebaseApp.getPersistenceKey(), key);
   }
 
-  @VisibleForTesting
-  synchronized void clearAll() {
-    SharedPreferences.Editor editor = prefs.edit();
-    editor.clear();
-    editor.commit();
+  private Task<Boolean> commitSharedPreferencesEditAsync(SharedPreferences.Editor editor) {
+    TaskCompletionSource<Boolean> result = new TaskCompletionSource<Boolean>();
+    ioExecuter.execute(
+        new Runnable() {
+          @Override
+          public void run() {
+            result.setResult(editor.commit());
+          }
+        });
+    return result.getTask();
   }
 }

From 41fbfee9e518794f9de09ee7e47ec5716fc92ead Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Tue, 18 Jun 2019 11:02:46 -0700
Subject: [PATCH 08/30] Address comments

---
 .../firebase/segmentation/CustomInstallationIdCache.java     | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
index 1863b976d9d..2a7fb54d1e7 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
@@ -17,6 +17,7 @@
 import android.content.Context;
 import android.content.SharedPreferences;
 import androidx.annotation.Nullable;
+import androidx.annotation.RestrictTo;
 import androidx.annotation.VisibleForTesting;
 import com.google.android.gms.common.util.Strings;
 import com.google.android.gms.tasks.Task;
@@ -50,7 +51,7 @@ enum CacheStatus {
   private final Executor ioExecuter;
   private final SharedPreferences prefs;
 
-  static CustomInstallationIdCache getInstance() {
+  synchronized static CustomInstallationIdCache getInstance() {
     if (singleton == null) {
       singleton = new CustomInstallationIdCache();
     }
@@ -105,7 +106,7 @@ synchronized Task<Boolean> clear(FirebaseApp firebaseApp) {
     return commitSharedPreferencesEditAsync(editor);
   }
 
-  @VisibleForTesting
+  @RestrictTo(RestrictTo.Scope.TESTS)
   synchronized Task<Boolean> clearAll() {
     SharedPreferences.Editor editor = prefs.edit();
     editor.clear();

From 5fd2fa0d6f4b83e152cb9bceb1fa2467e74e58e0 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Tue, 18 Jun 2019 11:57:08 -0700
Subject: [PATCH 09/30] Google format fix

---
 .../firebase/segmentation/CustomInstallationIdCache.java       | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
index 2a7fb54d1e7..5096a265714 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
@@ -18,7 +18,6 @@
 import android.content.SharedPreferences;
 import androidx.annotation.Nullable;
 import androidx.annotation.RestrictTo;
-import androidx.annotation.VisibleForTesting;
 import com.google.android.gms.common.util.Strings;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
@@ -51,7 +50,7 @@ enum CacheStatus {
   private final Executor ioExecuter;
   private final SharedPreferences prefs;
 
-  synchronized static CustomInstallationIdCache getInstance() {
+  static synchronized CustomInstallationIdCache getInstance() {
     if (singleton == null) {
       singleton = new CustomInstallationIdCache();
     }

From dba0c0eb57aa655b60ebe88de91b121835fd9e00 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Tue, 18 Jun 2019 15:16:35 -0700
Subject: [PATCH 10/30] Replace some deprecated code.

---
 .../segmentation/CustomInstallationIdCacheTest.java         | 6 +++---
 .../segmentation/FirebaseSegmentationInstrumentedTest.java  | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
index 5783294cfa3..ced06450c3c 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
@@ -18,7 +18,7 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
-import androidx.test.InstrumentationRegistry;
+import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import com.google.android.gms.tasks.Tasks;
 import com.google.firebase.FirebaseApp;
@@ -41,11 +41,11 @@ public void setUp() {
     FirebaseApp.clearInstancesForTest();
     firebaseApp0 =
         FirebaseApp.initializeApp(
-            InstrumentationRegistry.getContext(),
+            ApplicationProvider.getApplicationContext(),
             new FirebaseOptions.Builder().setApplicationId("1:123456789:android:abcdef").build());
     firebaseApp1 =
         FirebaseApp.initializeApp(
-            InstrumentationRegistry.getContext(),
+            ApplicationProvider.getApplicationContext(),
             new FirebaseOptions.Builder().setApplicationId("1:987654321:android:abcdef").build(),
             "firebase_app_1");
     cache = CustomInstallationIdCache.getInstance();
diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
index db8fcb9b873..739df092564 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
@@ -16,7 +16,7 @@
 
 import static org.junit.Assert.assertNull;
 
-import androidx.test.InstrumentationRegistry;
+import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.FirebaseOptions;
@@ -39,7 +39,7 @@ public void setUp() {
     FirebaseApp.clearInstancesForTest();
     firebaseApp =
         FirebaseApp.initializeApp(
-            InstrumentationRegistry.getContext(),
+            ApplicationProvider.getApplicationContext(),
             new FirebaseOptions.Builder().setApplicationId("1:123456789:android:abcdef").build());
   }
 

From a9a43a45ac18db388618156c962fc55eeda75424 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Tue, 18 Jun 2019 16:02:04 -0700
Subject: [PATCH 11/30] Package refactor

---
 .../{ => local}/CustomInstallationIdCacheTest.java          | 6 ++++--
 .../segmentation/{ => local}/CustomInstallationIdCache.java | 2 +-
 .../{ => local}/CustomInstallationIdCacheEntryValue.java    | 4 ++--
 3 files changed, 7 insertions(+), 5 deletions(-)
 rename firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/{ => local}/CustomInstallationIdCacheTest.java (95%)
 rename firebase-segmentation/src/main/java/com/google/firebase/segmentation/{ => local}/CustomInstallationIdCache.java (99%)
 rename firebase-segmentation/src/main/java/com/google/firebase/segmentation/{ => local}/CustomInstallationIdCacheEntryValue.java (90%)

diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java
similarity index 95%
rename from firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
rename to firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java
index ced06450c3c..4a63177ba26 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/CustomInstallationIdCacheTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package com.google.firebase.segmentation;
+package com.google.firebase.segmentation.local;
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertNull;
@@ -28,7 +28,9 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-/** Instrumented tests for {@link CustomInstallationIdCache} */
+/**
+ * Instrumented tests for {@link com.google.firebase.segmentation.local.CustomInstallationIdCache}
+ */
 @RunWith(AndroidJUnit4.class)
 public class CustomInstallationIdCacheTest {
 
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java
similarity index 99%
rename from firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
rename to firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java
index 5096a265714..2b93231eb39 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCache.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package com.google.firebase.segmentation;
+package com.google.firebase.segmentation.local;
 
 import android.content.Context;
 import android.content.SharedPreferences;
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCacheEntryValue.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheEntryValue.java
similarity index 90%
rename from firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCacheEntryValue.java
rename to firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheEntryValue.java
index c79d5f091a9..2d3b5f3c3a6 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/CustomInstallationIdCacheEntryValue.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheEntryValue.java
@@ -12,10 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package com.google.firebase.segmentation;
+package com.google.firebase.segmentation.local;
 
 import com.google.auto.value.AutoValue;
-import com.google.firebase.segmentation.CustomInstallationIdCache.CacheStatus;
+import com.google.firebase.segmentation.local.CustomInstallationIdCache.CacheStatus;
 
 /**
  * This class represents a cache entry value in {@link CustomInstallationIdCache}, which contains a

From ca6dacfff2f9d9b29b0495b09927bf374892b96a Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Tue, 18 Jun 2019 16:04:34 -0700
Subject: [PATCH 12/30] nit

---
 .../segmentation/local/CustomInstallationIdCacheTest.java       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java
index 4a63177ba26..b7028f1e595 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java
@@ -29,7 +29,7 @@
 import org.junit.runner.RunWith;
 
 /**
- * Instrumented tests for {@link com.google.firebase.segmentation.local.CustomInstallationIdCache}
+ * Instrumented tests for {@link CustomInstallationIdCache}
  */
 @RunWith(AndroidJUnit4.class)
 public class CustomInstallationIdCacheTest {

From e7fff8152e3e912821622703565f8741497b6b9e Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Tue, 18 Jun 2019 16:05:32 -0700
Subject: [PATCH 13/30] nit

---
 .../segmentation/local/CustomInstallationIdCacheTest.java     | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java
index b7028f1e595..27ab706b3a8 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java
@@ -28,9 +28,7 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-/**
- * Instrumented tests for {@link CustomInstallationIdCache}
- */
+/** Instrumented tests for {@link CustomInstallationIdCache} */
 @RunWith(AndroidJUnit4.class)
 public class CustomInstallationIdCacheTest {
 

From b381889c8ef8e86c0f8f662538421ce1763aba5b Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Wed, 19 Jun 2019 16:47:20 -0700
Subject: [PATCH 14/30] Add the state machine of updating custom installation
 id in the local cache and update to Firebase Segmentation backend. CL also
 contains unit tests. (The http client is not implemented yet.)

---
 .../firebase-segmentation.gradle              |   7 +
 .../FirebaseSegmentationInstrumentedTest.java | 189 +++++++++++++++++-
 .../local/CustomInstallationIdCacheTest.java  |  31 +--
 .../segmentation/FirebaseSegmentation.java    | 184 ++++++++++++++++-
 .../SetCustomInstallationIdException.java     |  68 +++++++
 .../google/firebase/segmentation/Utils.java   |  33 +++
 .../local/CustomInstallationIdCache.java      |  87 ++++----
 .../CustomInstallationIdCacheEntryValue.java  |  10 +-
 .../remote/SegmentationServiceClient.java     |  56 ++++++
 .../FirebaseSegmentationRegistrarTest.java    |   3 -
 10 files changed, 588 insertions(+), 80 deletions(-)
 create mode 100644 firebase-segmentation/src/main/java/com/google/firebase/segmentation/SetCustomInstallationIdException.java
 create mode 100644 firebase-segmentation/src/main/java/com/google/firebase/segmentation/Utils.java
 create mode 100644 firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java

diff --git a/firebase-segmentation/firebase-segmentation.gradle b/firebase-segmentation/firebase-segmentation.gradle
index dc4606715c5..1f76b8b593a 100644
--- a/firebase-segmentation/firebase-segmentation.gradle
+++ b/firebase-segmentation/firebase-segmentation.gradle
@@ -83,13 +83,18 @@ android {
 
 dependencies {
     implementation project(':firebase-common')
+    implementation project(':protolite-well-known-types')
 
     implementation('com.google.firebase:firebase-iid:17.0.3') {
         exclude group: "com.google.firebase", module: "firebase-common"
     }
+    implementation 'io.grpc:grpc-stub:1.21.0'
+    implementation 'io.grpc:grpc-protobuf-lite:1.21.0'
+    implementation 'io.grpc:grpc-okhttp:1.21.0'
     implementation 'androidx.appcompat:appcompat:1.0.2'
     implementation 'androidx.multidex:multidex:2.0.0'
     implementation 'com.google.android.gms:play-services-tasks:16.0.1'
+    implementation 'com.squareup.okhttp:okhttp:2.7.5'
 
     compileOnly "com.google.auto.value:auto-value-annotations:1.6.5"
     annotationProcessor "com.google.auto.value:auto-value:1.6.2"
@@ -104,4 +109,6 @@ dependencies {
     androidTestImplementation 'androidx.test:runner:1.2.0'
     androidTestImplementation "com.google.truth:truth:$googleTruthVersion"
     androidTestImplementation 'junit:junit:4.12'
+    androidTestImplementation 'org.mockito:mockito-core:2.25.0'
+    androidTestImplementation 'org.mockito:mockito-android:2.25.0'
 }
diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
index 739df092564..289c448fd1f 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
@@ -14,15 +14,34 @@
 
 package com.google.firebase.segmentation;
 
+import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.fail;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.when;
 
+import androidx.annotation.NonNull;
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
+import com.google.android.gms.tasks.Tasks;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.FirebaseOptions;
+import com.google.firebase.iid.FirebaseInstanceId;
+import com.google.firebase.iid.InstanceIdResult;
+import com.google.firebase.segmentation.local.CustomInstallationIdCache;
+import com.google.firebase.segmentation.local.CustomInstallationIdCacheEntryValue;
+import com.google.firebase.segmentation.remote.SegmentationServiceClient;
+import java.util.concurrent.ExecutionException;
+import org.junit.After;
 import org.junit.Before;
+import org.junit.FixMethodOrder;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.junit.runners.MethodSorters;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 
 /**
  * Instrumented test, which will execute on an Android device.
@@ -30,21 +49,185 @@
  * @see <a href="http://d.android.com/tools/testing">Testing documentation</a>
  */
 @RunWith(AndroidJUnit4.class)
+@FixMethodOrder(MethodSorters.NAME_ASCENDING)
 public class FirebaseSegmentationInstrumentedTest {
 
+  private static final String CUSTOM_INSTALLATION_ID = "123";
+  private static final String FIREBASE_INSTANCE_ID = "cAAAAAAAAAA";
+
   private FirebaseApp firebaseApp;
+  @Mock private FirebaseInstanceId firebaseInstanceId;
+  @Mock private SegmentationServiceClient backendClientReturnsOk;
+  @Mock private SegmentationServiceClient backendClientReturnsError;
+  private CustomInstallationIdCache actualCache;
+  @Mock private CustomInstallationIdCache cacheReturnsError;
 
   @Before
   public void setUp() {
+    MockitoAnnotations.initMocks(this);
     FirebaseApp.clearInstancesForTest();
     firebaseApp =
         FirebaseApp.initializeApp(
             ApplicationProvider.getApplicationContext(),
-            new FirebaseOptions.Builder().setApplicationId("1:123456789:android:abcdef").build());
+            new FirebaseOptions.Builder()
+                .setApplicationId("1" + ":123456789:android:abcdef")
+                .build());
+    actualCache = new CustomInstallationIdCache(firebaseApp);
+
+    when(backendClientReturnsOk.updateCustomInstallationId(
+            anyLong(), anyString(), anyString(), anyString()))
+        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.OK));
+    when(backendClientReturnsOk.clearCustomInstallationId(anyLong(), anyString(), anyString()))
+        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.OK));
+    when(backendClientReturnsError.updateCustomInstallationId(
+            anyLong(), anyString(), anyString(), anyString()))
+        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.SERVER_INTERNAL_ERROR));
+    when(backendClientReturnsError.clearCustomInstallationId(anyLong(), anyString(), anyString()))
+        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.SERVER_INTERNAL_ERROR));
+    when(firebaseInstanceId.getInstanceId())
+        .thenReturn(
+            Tasks.forResult(
+                new InstanceIdResult() {
+                  @NonNull
+                  @Override
+                  public String getId() {
+                    return FIREBASE_INSTANCE_ID;
+                  }
+
+                  @NonNull
+                  @Override
+                  public String getToken() {
+                    return "iid_token";
+                  }
+                }));
+    when(cacheReturnsError.insertOrUpdateCacheEntry(any())).thenReturn(Tasks.forResult(false));
+    when(cacheReturnsError.readCacheEntryValue()).thenReturn(null);
+  }
+
+  @After
+  public void cleanUp() throws Exception {
+    Tasks.await(actualCache.clear());
+  }
+
+  @Test
+  public void testUpdateCustomInstallationId_CacheOk_BackendOk() throws Exception {
+    FirebaseSegmentation firebaseSegmentation =
+        new FirebaseSegmentation(
+            firebaseApp, firebaseInstanceId, actualCache, backendClientReturnsOk);
+
+    // No exception, means success.
+    assertNull(Tasks.await(firebaseSegmentation.setCustomInstallationId(CUSTOM_INSTALLATION_ID)));
+    CustomInstallationIdCacheEntryValue entryValue = actualCache.readCacheEntryValue();
+    assertThat(entryValue.getCustomInstallationId()).isEqualTo(CUSTOM_INSTALLATION_ID);
+    assertThat(entryValue.getFirebaseInstanceId()).isEqualTo(FIREBASE_INSTANCE_ID);
+    assertThat(entryValue.getCacheStatus()).isEqualTo(CustomInstallationIdCache.CacheStatus.SYNCED);
+  }
+
+  @Test
+  public void testUpdateCustomInstallationId_CacheOk_BackendError() throws InterruptedException {
+    FirebaseSegmentation firebaseSegmentation =
+        new FirebaseSegmentation(
+            firebaseApp, firebaseInstanceId, actualCache, backendClientReturnsError);
+
+    // Expect exception
+    try {
+      Tasks.await(firebaseSegmentation.setCustomInstallationId(CUSTOM_INSTALLATION_ID));
+      fail();
+    } catch (ExecutionException expected) {
+      Throwable cause = expected.getCause();
+      assertThat(cause).isInstanceOf(SetCustomInstallationIdException.class);
+      assertThat(((SetCustomInstallationIdException) cause).getStatus())
+          .isEqualTo(SetCustomInstallationIdException.Status.BACKEND_ERROR);
+    }
+
+    CustomInstallationIdCacheEntryValue entryValue = actualCache.readCacheEntryValue();
+    assertThat(entryValue.getCustomInstallationId()).isEqualTo(CUSTOM_INSTALLATION_ID);
+    assertThat(entryValue.getFirebaseInstanceId()).isEqualTo(FIREBASE_INSTANCE_ID);
+    assertThat(entryValue.getCacheStatus())
+        .isEqualTo(CustomInstallationIdCache.CacheStatus.PENDING_UPDATE);
   }
 
   @Test
-  public void useAppContext() {
-    assertNull(FirebaseSegmentation.getInstance().setCustomInstallationId("123123").getResult());
+  public void testUpdateCustomInstallationId_CacheError_BackendOk() throws InterruptedException {
+    FirebaseSegmentation firebaseSegmentation =
+        new FirebaseSegmentation(
+            firebaseApp, firebaseInstanceId, cacheReturnsError, backendClientReturnsOk);
+
+    // Expect exception
+    try {
+      Tasks.await(firebaseSegmentation.setCustomInstallationId(CUSTOM_INSTALLATION_ID));
+      fail();
+    } catch (ExecutionException expected) {
+      Throwable cause = expected.getCause();
+      assertThat(cause).isInstanceOf(SetCustomInstallationIdException.class);
+      assertThat(((SetCustomInstallationIdException) cause).getStatus())
+          .isEqualTo(SetCustomInstallationIdException.Status.CLIENT_ERROR);
+    }
+  }
+
+  @Test
+  public void testClearCustomInstallationId_CacheOk_BackendOk() throws Exception {
+    Tasks.await(
+        actualCache.insertOrUpdateCacheEntry(
+            CustomInstallationIdCacheEntryValue.create(
+                CUSTOM_INSTALLATION_ID,
+                FIREBASE_INSTANCE_ID,
+                CustomInstallationIdCache.CacheStatus.SYNCED)));
+    FirebaseSegmentation firebaseSegmentation =
+        new FirebaseSegmentation(
+            firebaseApp, firebaseInstanceId, actualCache, backendClientReturnsOk);
+
+    // No exception, means success.
+    assertNull(Tasks.await(firebaseSegmentation.setCustomInstallationId(null)));
+    CustomInstallationIdCacheEntryValue entryValue = actualCache.readCacheEntryValue();
+    assertNull(entryValue);
+  }
+
+  @Test
+  public void testClearCustomInstallationId_CacheOk_BackendError() throws Exception {
+    Tasks.await(
+        actualCache.insertOrUpdateCacheEntry(
+            CustomInstallationIdCacheEntryValue.create(
+                CUSTOM_INSTALLATION_ID,
+                FIREBASE_INSTANCE_ID,
+                CustomInstallationIdCache.CacheStatus.SYNCED)));
+    FirebaseSegmentation firebaseSegmentation =
+        new FirebaseSegmentation(
+            firebaseApp, firebaseInstanceId, actualCache, backendClientReturnsError);
+
+    // Expect exception
+    try {
+      Tasks.await(firebaseSegmentation.setCustomInstallationId(null));
+      fail();
+    } catch (ExecutionException expected) {
+      Throwable cause = expected.getCause();
+      assertThat(cause).isInstanceOf(SetCustomInstallationIdException.class);
+      assertThat(((SetCustomInstallationIdException) cause).getStatus())
+          .isEqualTo(SetCustomInstallationIdException.Status.BACKEND_ERROR);
+    }
+
+    CustomInstallationIdCacheEntryValue entryValue = actualCache.readCacheEntryValue();
+    assertThat(entryValue.getCustomInstallationId().isEmpty()).isTrue();
+    assertThat(entryValue.getFirebaseInstanceId()).isEqualTo(FIREBASE_INSTANCE_ID);
+    assertThat(entryValue.getCacheStatus())
+        .isEqualTo(CustomInstallationIdCache.CacheStatus.PENDING_CLEAR);
+  }
+
+  @Test
+  public void testClearCustomInstallationId_CacheError_BackendOk() throws InterruptedException {
+    FirebaseSegmentation firebaseSegmentation =
+        new FirebaseSegmentation(
+            firebaseApp, firebaseInstanceId, cacheReturnsError, backendClientReturnsOk);
+
+    // Expect exception
+    try {
+      Tasks.await(firebaseSegmentation.setCustomInstallationId(CUSTOM_INSTALLATION_ID));
+      fail();
+    } catch (ExecutionException expected) {
+      Throwable cause = expected.getCause();
+      assertThat(cause).isInstanceOf(SetCustomInstallationIdException.class);
+      assertThat(((SetCustomInstallationIdException) cause).getStatus())
+          .isEqualTo(SetCustomInstallationIdException.Status.CLIENT_ERROR);
+    }
   }
 }
diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java
index 27ab706b3a8..019a2b8ba08 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java
@@ -34,7 +34,8 @@ public class CustomInstallationIdCacheTest {
 
   private FirebaseApp firebaseApp0;
   private FirebaseApp firebaseApp1;
-  private CustomInstallationIdCache cache;
+  private CustomInstallationIdCache cache0;
+  private CustomInstallationIdCache cache1;
 
   @Before
   public void setUp() {
@@ -48,42 +49,44 @@ public void setUp() {
             ApplicationProvider.getApplicationContext(),
             new FirebaseOptions.Builder().setApplicationId("1:987654321:android:abcdef").build(),
             "firebase_app_1");
-    cache = CustomInstallationIdCache.getInstance();
+    cache0 = new CustomInstallationIdCache(firebaseApp0);
+    cache1 = new CustomInstallationIdCache(firebaseApp1);
   }
 
   @After
   public void cleanUp() throws Exception {
-    Tasks.await(cache.clearAll());
+    Tasks.await(cache0.clear());
+    Tasks.await(cache1.clear());
   }
 
   @Test
   public void testReadCacheEntry_Null() {
-    assertNull(cache.readCacheEntryValue(firebaseApp0));
-    assertNull(cache.readCacheEntryValue(firebaseApp1));
+    assertNull(cache0.readCacheEntryValue());
+    assertNull(cache1.readCacheEntryValue());
   }
 
   @Test
   public void testUpdateAndReadCacheEntry() throws Exception {
     assertTrue(
         Tasks.await(
-            cache.insertOrUpdateCacheEntry(
-                firebaseApp0,
+            cache0.insertOrUpdateCacheEntry(
                 CustomInstallationIdCacheEntryValue.create(
-                    "123456", "cAAAAAAAAAA", CustomInstallationIdCache.CacheStatus.PENDING))));
-    CustomInstallationIdCacheEntryValue entryValue = cache.readCacheEntryValue(firebaseApp0);
+                    "123456",
+                    "cAAAAAAAAAA",
+                    CustomInstallationIdCache.CacheStatus.PENDING_UPDATE))));
+    CustomInstallationIdCacheEntryValue entryValue = cache0.readCacheEntryValue();
     assertThat(entryValue.getCustomInstallationId()).isEqualTo("123456");
     assertThat(entryValue.getFirebaseInstanceId()).isEqualTo("cAAAAAAAAAA");
     assertThat(entryValue.getCacheStatus())
-        .isEqualTo(CustomInstallationIdCache.CacheStatus.PENDING);
-    assertNull(cache.readCacheEntryValue(firebaseApp1));
+        .isEqualTo(CustomInstallationIdCache.CacheStatus.PENDING_UPDATE);
+    assertNull(cache1.readCacheEntryValue());
 
     assertTrue(
         Tasks.await(
-            cache.insertOrUpdateCacheEntry(
-                firebaseApp0,
+            cache0.insertOrUpdateCacheEntry(
                 CustomInstallationIdCacheEntryValue.create(
                     "123456", "cAAAAAAAAAA", CustomInstallationIdCache.CacheStatus.SYNCED))));
-    entryValue = cache.readCacheEntryValue(firebaseApp0);
+    entryValue = cache0.readCacheEntryValue();
     assertThat(entryValue.getCustomInstallationId()).isEqualTo("123456");
     assertThat(entryValue.getFirebaseInstanceId()).isEqualTo("cAAAAAAAAAA");
     assertThat(entryValue.getCacheStatus()).isEqualTo(CustomInstallationIdCache.CacheStatus.SYNCED);
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
index eca517db1b3..5bed0660b8a 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
@@ -15,21 +15,45 @@
 package com.google.firebase.segmentation;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+import androidx.annotation.RestrictTo;
 import com.google.android.gms.common.internal.Preconditions;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.Tasks;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.iid.FirebaseInstanceId;
+import com.google.firebase.iid.InstanceIdResult;
+import com.google.firebase.segmentation.SetCustomInstallationIdException.Status;
+import com.google.firebase.segmentation.local.CustomInstallationIdCache;
+import com.google.firebase.segmentation.local.CustomInstallationIdCacheEntryValue;
+import com.google.firebase.segmentation.remote.SegmentationServiceClient;
+import com.google.firebase.segmentation.remote.SegmentationServiceClient.Code;
 
 /** Entry point of Firebase Segmentation SDK. */
 public class FirebaseSegmentation {
 
   private final FirebaseApp firebaseApp;
   private final FirebaseInstanceId firebaseInstanceId;
+  private final CustomInstallationIdCache localCache;
+  private final SegmentationServiceClient backendServiceClient;
 
   FirebaseSegmentation(FirebaseApp firebaseApp) {
     this.firebaseApp = firebaseApp;
     this.firebaseInstanceId = FirebaseInstanceId.getInstance(firebaseApp);
+    localCache = new CustomInstallationIdCache(firebaseApp);
+    backendServiceClient = new SegmentationServiceClient();
+  }
+
+  @RestrictTo(RestrictTo.Scope.TESTS)
+  FirebaseSegmentation(
+      FirebaseApp firebaseApp,
+      FirebaseInstanceId firebaseInstanceId,
+      CustomInstallationIdCache localCache,
+      SegmentationServiceClient backendServiceClient) {
+    this.firebaseApp = firebaseApp;
+    this.firebaseInstanceId = firebaseInstanceId;
+    this.localCache = localCache;
+    this.backendServiceClient = backendServiceClient;
   }
 
   /**
@@ -51,11 +75,165 @@ public static FirebaseSegmentation getInstance() {
    */
   @NonNull
   public static FirebaseSegmentation getInstance(@NonNull FirebaseApp app) {
-    Preconditions.checkArgument(app != null, "Null is not a valid value of FirebaseApp.");
+    Preconditions.checkArgument(app != null, "Null is not a valid value " + "of FirebaseApp.");
     return app.get(FirebaseSegmentation.class);
   }
 
-  Task<Void> setCustomInstallationId(String customInstallationId) {
-    return Tasks.forResult(null);
+  Task<Void> setCustomInstallationId(@Nullable String customInstallationId) {
+    if (customInstallationId == null) {
+      return clearCustomInstallationId();
+    }
+    return updateCustomInstallationId(customInstallationId);
+  }
+
+  /**
+   * Update custom installation id of the {@link FirebaseApp} on Firebase segmentation backend and
+   * client side cache.
+   *
+   * <pre>
+   *     The workflow is:
+   *         check diff against cache or cache status is not SYNCED
+   *                                 |
+   *                  get Firebase instance id and token
+   *                      |                       |
+   *                      |      update cache with cache status PENDING_UPDATE
+   *                      |                       |
+   *                    send http request to backend
+   *                                 |
+   *              on success: set cache entry status to SYNCED
+   *                                 |
+   *                               return
+   * </pre>
+   */
+  private Task<Void> updateCustomInstallationId(String customInstallationId) {
+    CustomInstallationIdCacheEntryValue cacheEntryValue = localCache.readCacheEntryValue();
+    if (cacheEntryValue != null
+        && cacheEntryValue.getCustomInstallationId().equals(customInstallationId)
+        && cacheEntryValue.getCacheStatus() == CustomInstallationIdCache.CacheStatus.SYNCED) {
+      // If the given custom installation id matches up the cached
+      // value, there's no need to update.
+      return Tasks.forResult(null);
+    }
+
+    Task<InstanceIdResult> instanceIdResultTask = firebaseInstanceId.getInstanceId();
+    Task<Boolean> firstUpdateCacheResultTask =
+        instanceIdResultTask.onSuccessTask(
+            instanceIdResult ->
+                localCache.insertOrUpdateCacheEntry(
+                    CustomInstallationIdCacheEntryValue.create(
+                        customInstallationId,
+                        instanceIdResult.getId(),
+                        CustomInstallationIdCache.CacheStatus.PENDING_UPDATE)));
+
+    // Start requesting backend when first cache update is done.
+    Task<Code> backendRequestResultTask =
+        firstUpdateCacheResultTask.onSuccessTask(
+            firstUpdateCacheResult -> {
+              if (firstUpdateCacheResult) {
+                String iid = instanceIdResultTask.getResult().getId();
+                String iidToken = instanceIdResultTask.getResult().getToken();
+                return backendServiceClient.updateCustomInstallationId(
+                    Utils.getProjectNumberFromAppId(firebaseApp.getOptions().getApplicationId()),
+                    customInstallationId,
+                    iid,
+                    iidToken);
+              } else {
+                throw new SetCustomInstallationIdException(
+                    "Failed to update client side cache", Status.CLIENT_ERROR);
+              }
+            });
+
+    Task<Boolean> finalUpdateCacheResultTask =
+        backendRequestResultTask.onSuccessTask(
+            backendRequestResult -> {
+              switch (backendRequestResult) {
+                case OK:
+                  return localCache.insertOrUpdateCacheEntry(
+                      CustomInstallationIdCacheEntryValue.create(
+                          customInstallationId,
+                          instanceIdResultTask.getResult().getId(),
+                          CustomInstallationIdCache.CacheStatus.SYNCED));
+                case ALREADY_EXISTS:
+                  throw new SetCustomInstallationIdException(
+                      Status.DUPLICATED_CUSTOM_INSTALLATION_ID);
+                default:
+                  throw new SetCustomInstallationIdException(Status.BACKEND_ERROR);
+              }
+            });
+
+    return finalUpdateCacheResultTask.onSuccessTask(
+        finalUpdateCacheResult -> {
+          if (finalUpdateCacheResult) {
+            return Tasks.forResult(null);
+          } else {
+            throw new SetCustomInstallationIdException(
+                "Failed to update client side cache", Status.CLIENT_ERROR);
+          }
+        });
+  }
+
+  /**
+   * Clear custom installation id of the {@link FirebaseApp} on Firebase segmentation backend and
+   * client side cache.
+   *
+   * <pre>
+   *     The workflow is:
+   *                  get Firebase instance id and token
+   *                      |                      |
+   *                      |    update cache with cache status PENDING_CLEAR
+   *                      |                      |
+   *                    send http request to backend
+   *                                  |
+   *                   on success: delete cache entry
+   *                                  |
+   *                               return
+   * </pre>
+   */
+  private Task<Void> clearCustomInstallationId() {
+    Task<InstanceIdResult> instanceIdResultTask = firebaseInstanceId.getInstanceId();
+    Task<Boolean> firstUpdateCacheResultTask =
+        instanceIdResultTask.onSuccessTask(
+            instanceIdResult ->
+                localCache.insertOrUpdateCacheEntry(
+                    CustomInstallationIdCacheEntryValue.create(
+                        "",
+                        instanceIdResult.getId(),
+                        CustomInstallationIdCache.CacheStatus.PENDING_CLEAR)));
+
+    Task<Code> backendRequestResultTask =
+        firstUpdateCacheResultTask.onSuccessTask(
+            firstUpdateCacheResult -> {
+              if (firstUpdateCacheResult) {
+                String iid = instanceIdResultTask.getResult().getId();
+                String iidToken = instanceIdResultTask.getResult().getToken();
+                return backendServiceClient.clearCustomInstallationId(
+                    Utils.getProjectNumberFromAppId(firebaseApp.getOptions().getApplicationId()),
+                    iid,
+                    iidToken);
+              } else {
+                throw new SetCustomInstallationIdException(
+                    "Failed to update client side cache", Status.CLIENT_ERROR);
+              }
+            });
+
+    Task<Boolean> finalUpdateCacheResultTask =
+        backendRequestResultTask.onSuccessTask(
+            backendRequestResult -> {
+              if (backendRequestResult == Code.OK) {
+                return localCache.clear();
+              } else {
+                throw new SetCustomInstallationIdException(Status.BACKEND_ERROR);
+              }
+            });
+
+    return finalUpdateCacheResultTask.onSuccessTask(
+        finalUpdateCacheResult -> {
+          if (finalUpdateCacheResult) {
+            return Tasks.forResult(null);
+          } else {
+            throw new SetCustomInstallationIdException(
+                "Failed to update client side cache", Status.CLIENT_ERROR);
+          }
+        });
   }
 }
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/SetCustomInstallationIdException.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/SetCustomInstallationIdException.java
new file mode 100644
index 00000000000..3c957ce3294
--- /dev/null
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/SetCustomInstallationIdException.java
@@ -0,0 +1,68 @@
+// Copyright 2019 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.firebase.segmentation;
+
+import androidx.annotation.NonNull;
+import com.google.firebase.FirebaseException;
+
+/** The class for all Exceptions thrown by {@link FirebaseSegmentation}. */
+public class SetCustomInstallationIdException extends FirebaseException {
+
+  public enum Status {
+    UNKOWN(0),
+
+    /** Error in Firebase SDK. */
+    CLIENT_ERROR(1),
+
+    /** Error when calling Firebase segmentation backend. */
+    BACKEND_ERROR(2),
+
+    /** The given custom installation is already tied to another Firebase installation. */
+    DUPLICATED_CUSTOM_INSTALLATION_ID(3);
+
+    private final int value;
+
+    Status(int value) {
+      this.value = value;
+    }
+  }
+
+  @NonNull private final Status status;
+
+  SetCustomInstallationIdException(@NonNull Status status) {
+    this.status = status;
+  }
+
+  SetCustomInstallationIdException(@NonNull String message, @NonNull Status status) {
+    super(message);
+    this.status = status;
+  }
+
+  SetCustomInstallationIdException(
+      @NonNull String message, @NonNull Status status, Throwable cause) {
+    super(message, cause);
+    this.status = status;
+  }
+
+  /**
+   * Gets the status code for the operation that failed.
+   *
+   * @return the code for the SetCustomInstallationIdException
+   */
+  @NonNull
+  public Status getStatus() {
+    return status;
+  }
+}
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/Utils.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/Utils.java
new file mode 100644
index 00000000000..ca231a89cb5
--- /dev/null
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/Utils.java
@@ -0,0 +1,33 @@
+// Copyright 2019 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.firebase.segmentation;
+
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+/** Util methods used for {@link FirebaseSegmentation} */
+class Utils {
+
+  private static final Pattern APP_ID_PATTERN =
+      Pattern.compile("^[^:]+:([0-9]+):(android|ios|web):([0-9a-f]+)");
+
+  static long getProjectNumberFromAppId(String appId) {
+    Matcher matcher = APP_ID_PATTERN.matcher(appId);
+    if (matcher.matches()) {
+      return Long.valueOf(matcher.group(1));
+    }
+    throw new IllegalArgumentException("Invalid app id " + appId);
+  }
+}
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java
index 2b93231eb39..f650c3463d8 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java
@@ -17,27 +17,30 @@
 import android.content.Context;
 import android.content.SharedPreferences;
 import androidx.annotation.Nullable;
-import androidx.annotation.RestrictTo;
-import com.google.android.gms.common.util.Strings;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.firebase.FirebaseApp;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 
-class CustomInstallationIdCache {
+/**
+ * A layer that locally caches a few Firebase Segmentation attributes on top the Segmentation
+ * backend API.
+ */
+public class CustomInstallationIdCache {
 
   // Status of each cache entry
   // NOTE: never change the ordinal of the enum values because the enum values are stored in cache
   // as their ordinal numbers.
-  enum CacheStatus {
+  public enum CacheStatus {
     // Cache entry is synced to Firebase backend
     SYNCED,
-    // Cache entry is waiting for Firebase backend response or pending internal retry for retryable
-    // errors.
-    PENDING,
-    // Cache entry is not accepted by Firebase backend.
-    ERROR,
+    // Cache entry is waiting for Firebase backend response or internal network retry (for update
+    // operation).
+    PENDING_UPDATE,
+    // Cache entry is waiting for Firebase backend response or internal network retry (for clear
+    // operation).
+    PENDING_CLEAR
   }
 
   private static final String SHARED_PREFS_NAME = "CustomInstallationIdCache";
@@ -46,78 +49,58 @@ enum CacheStatus {
   private static final String INSTANCE_ID_KEY = "Iid";
   private static final String CACHE_STATUS_KEY = "Status";
 
-  private static CustomInstallationIdCache singleton = null;
   private final Executor ioExecuter;
   private final SharedPreferences prefs;
+  private final String persistenceKey;
 
-  static synchronized CustomInstallationIdCache getInstance() {
-    if (singleton == null) {
-      singleton = new CustomInstallationIdCache();
-    }
-    return singleton;
-  }
-
-  private CustomInstallationIdCache() {
-    // Since different FirebaseApp in the same Android application should have the same application
-    // context and same dir path, so that use the context of the default FirebaseApp to create the
-    // shared preferences.
+  public CustomInstallationIdCache(FirebaseApp firebaseApp) {
+    // Different FirebaseApp in the same Android application should have the same application
+    // context and same dir path
     prefs =
-        FirebaseApp.getInstance()
+        firebaseApp
             .getApplicationContext()
             .getSharedPreferences(SHARED_PREFS_NAME, Context.MODE_PRIVATE);
-
+    persistenceKey = firebaseApp.getPersistenceKey();
     ioExecuter = Executors.newFixedThreadPool(2);
   }
 
   @Nullable
-  synchronized CustomInstallationIdCacheEntryValue readCacheEntryValue(FirebaseApp firebaseApp) {
-    String cid =
-        prefs.getString(getSharedPreferencesKey(firebaseApp, CUSTOM_INSTALLATION_ID_KEY), null);
-    String iid = prefs.getString(getSharedPreferencesKey(firebaseApp, INSTANCE_ID_KEY), null);
-    int status = prefs.getInt(getSharedPreferencesKey(firebaseApp, CACHE_STATUS_KEY), -1);
+  public synchronized CustomInstallationIdCacheEntryValue readCacheEntryValue() {
+    String cid = prefs.getString(getSharedPreferencesKey(CUSTOM_INSTALLATION_ID_KEY), null);
+    String iid = prefs.getString(getSharedPreferencesKey(INSTANCE_ID_KEY), null);
+    int status = prefs.getInt(getSharedPreferencesKey(CACHE_STATUS_KEY), -1);
 
-    if (Strings.isEmptyOrWhitespace(cid) || Strings.isEmptyOrWhitespace(iid) || status == -1) {
+    if (cid == null || iid == null || status == -1) {
       return null;
     }
 
     return CustomInstallationIdCacheEntryValue.create(cid, iid, CacheStatus.values()[status]);
   }
 
-  synchronized Task<Boolean> insertOrUpdateCacheEntry(
-      FirebaseApp firebaseApp, CustomInstallationIdCacheEntryValue entryValue) {
+  public synchronized Task<Boolean> insertOrUpdateCacheEntry(
+      CustomInstallationIdCacheEntryValue entryValue) {
     SharedPreferences.Editor editor = prefs.edit();
     editor.putString(
-        getSharedPreferencesKey(firebaseApp, CUSTOM_INSTALLATION_ID_KEY),
-        entryValue.getCustomInstallationId());
-    editor.putString(
-        getSharedPreferencesKey(firebaseApp, INSTANCE_ID_KEY), entryValue.getFirebaseInstanceId());
-    editor.putInt(
-        getSharedPreferencesKey(firebaseApp, CACHE_STATUS_KEY),
-        entryValue.getCacheStatus().ordinal());
-    return commitSharedPreferencesEditAsync(editor);
-  }
-
-  synchronized Task<Boolean> clear(FirebaseApp firebaseApp) {
-    SharedPreferences.Editor editor = prefs.edit();
-    editor.remove(getSharedPreferencesKey(firebaseApp, CUSTOM_INSTALLATION_ID_KEY));
-    editor.remove(getSharedPreferencesKey(firebaseApp, INSTANCE_ID_KEY));
-    editor.remove(getSharedPreferencesKey(firebaseApp, CACHE_STATUS_KEY));
+        getSharedPreferencesKey(CUSTOM_INSTALLATION_ID_KEY), entryValue.getCustomInstallationId());
+    editor.putString(getSharedPreferencesKey(INSTANCE_ID_KEY), entryValue.getFirebaseInstanceId());
+    editor.putInt(getSharedPreferencesKey(CACHE_STATUS_KEY), entryValue.getCacheStatus().ordinal());
     return commitSharedPreferencesEditAsync(editor);
   }
 
-  @RestrictTo(RestrictTo.Scope.TESTS)
-  synchronized Task<Boolean> clearAll() {
+  public synchronized Task<Boolean> clear() {
     SharedPreferences.Editor editor = prefs.edit();
-    editor.clear();
+    editor.remove(getSharedPreferencesKey(CUSTOM_INSTALLATION_ID_KEY));
+    editor.remove(getSharedPreferencesKey(INSTANCE_ID_KEY));
+    editor.remove(getSharedPreferencesKey(CACHE_STATUS_KEY));
     return commitSharedPreferencesEditAsync(editor);
   }
 
-  private static String getSharedPreferencesKey(FirebaseApp firebaseApp, String key) {
-    return String.format("%s|%s", firebaseApp.getPersistenceKey(), key);
+  private String getSharedPreferencesKey(String key) {
+    return String.format("%s|%s", persistenceKey, key);
   }
 
   private Task<Boolean> commitSharedPreferencesEditAsync(SharedPreferences.Editor editor) {
-    TaskCompletionSource<Boolean> result = new TaskCompletionSource<Boolean>();
+    TaskCompletionSource<Boolean> result = new TaskCompletionSource<>();
     ioExecuter.execute(
         new Runnable() {
           @Override
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheEntryValue.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheEntryValue.java
index 2d3b5f3c3a6..05528cd40f0 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheEntryValue.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheEntryValue.java
@@ -22,14 +22,14 @@
  * Firebase instance id, a custom installation id and the cache status of this entry.
  */
 @AutoValue
-abstract class CustomInstallationIdCacheEntryValue {
-  abstract String getCustomInstallationId();
+public abstract class CustomInstallationIdCacheEntryValue {
+  public abstract String getCustomInstallationId();
 
-  abstract String getFirebaseInstanceId();
+  public abstract String getFirebaseInstanceId();
 
-  abstract CacheStatus getCacheStatus();
+  public abstract CacheStatus getCacheStatus();
 
-  static CustomInstallationIdCacheEntryValue create(
+  public static CustomInstallationIdCacheEntryValue create(
       String customInstallationId, String firebaseInstanceId, CacheStatus cacheStatus) {
     return new AutoValue_CustomInstallationIdCacheEntryValue(
         customInstallationId, firebaseInstanceId, cacheStatus);
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
new file mode 100644
index 00000000000..59398b369d5
--- /dev/null
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
@@ -0,0 +1,56 @@
+// Copyright 2019 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.firebase.segmentation.remote;
+
+import com.google.android.gms.tasks.Task;
+import com.google.android.gms.tasks.Tasks;
+import com.squareup.okhttp.OkHttpClient;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
+
+/** Http client that sends request to Firebase Segmentation backend API. To be implemented */
+public class SegmentationServiceClient {
+
+  private final OkHttpClient httpClient;
+  private final Executor httpRequestExecutor;
+
+  public enum Code {
+    OK,
+
+    SERVER_INTERNAL_ERROR,
+
+    ALREADY_EXISTS,
+
+    PERMISSION_DENIED
+  }
+
+  public SegmentationServiceClient() {
+    httpClient = new OkHttpClient();
+    httpRequestExecutor = Executors.newFixedThreadPool(4);
+  }
+
+  public Task<Code> updateCustomInstallationId(
+      long projectNumber,
+      String customInstallationId,
+      String firebaseInstanceId,
+      String firebaseInstanceIdToken) {
+    return Tasks.forResult(Code.OK);
+  }
+
+  public Task<Code> clearCustomInstallationId(
+      long projectNumber, String firebaseInstanceId, String firebaseInstanceIdToken) {
+    return Tasks.forResult(Code.OK);
+  }
+}
diff --git a/firebase-segmentation/src/test/java/com/google/firebase/segmentation/FirebaseSegmentationRegistrarTest.java b/firebase-segmentation/src/test/java/com/google/firebase/segmentation/FirebaseSegmentationRegistrarTest.java
index 1f3c441808f..56b0d120eb0 100644
--- a/firebase-segmentation/src/test/java/com/google/firebase/segmentation/FirebaseSegmentationRegistrarTest.java
+++ b/firebase-segmentation/src/test/java/com/google/firebase/segmentation/FirebaseSegmentationRegistrarTest.java
@@ -15,7 +15,6 @@
 package com.google.firebase.segmentation;
 
 import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertNull;
 
 import androidx.test.core.app.ApplicationProvider;
 import com.google.firebase.FirebaseApp;
@@ -48,10 +47,8 @@ public void getFirebaseInstallationsInstance() {
 
     FirebaseSegmentation defaultSegmentation = FirebaseSegmentation.getInstance();
     assertNotNull(defaultSegmentation);
-    assertNull(defaultSegmentation.setCustomInstallationId("12345").getResult());
 
     FirebaseSegmentation anotherSegmentation = FirebaseSegmentation.getInstance(anotherApp);
     assertNotNull(anotherSegmentation);
-    assertNull(anotherSegmentation.setCustomInstallationId("ghdjaas").getResult());
   }
 }

From 1adcfbd864f3ed65f0898034bea59fefb0864156 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Wed, 19 Jun 2019 17:02:20 -0700
Subject: [PATCH 15/30] minor format fix

---
 .../segmentation/FirebaseSegmentationInstrumentedTest.java    | 4 +---
 .../google/firebase/segmentation/FirebaseSegmentation.java    | 2 +-
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
index 289c448fd1f..8519e641cb7 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
@@ -69,9 +69,7 @@ public void setUp() {
     firebaseApp =
         FirebaseApp.initializeApp(
             ApplicationProvider.getApplicationContext(),
-            new FirebaseOptions.Builder()
-                .setApplicationId("1" + ":123456789:android:abcdef")
-                .build());
+            new FirebaseOptions.Builder().setApplicationId("1:123456789:android:abcdef").build());
     actualCache = new CustomInstallationIdCache(firebaseApp);
 
     when(backendClientReturnsOk.updateCustomInstallationId(
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
index 5bed0660b8a..e7566cd965b 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
@@ -75,7 +75,7 @@ public static FirebaseSegmentation getInstance() {
    */
   @NonNull
   public static FirebaseSegmentation getInstance(@NonNull FirebaseApp app) {
-    Preconditions.checkArgument(app != null, "Null is not a valid value " + "of FirebaseApp.");
+    Preconditions.checkArgument(app != null, "Null is not a valid value of FirebaseApp.");
     return app.get(FirebaseSegmentation.class);
   }
 

From 6091f82032fc78512d64de8e9e5c2257b98cdeda Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Thu, 20 Jun 2019 13:46:10 -0700
Subject: [PATCH 16/30] Address comments #1

---
 .../firebase/segmentation/FirebaseSegmentation.java       | 3 ++-
 .../segmentation/local/CustomInstallationIdCache.java     | 8 +-------
 2 files changed, 3 insertions(+), 8 deletions(-)

diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
index e7566cd965b..34de68597c1 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
@@ -79,7 +79,8 @@ public static FirebaseSegmentation getInstance(@NonNull FirebaseApp app) {
     return app.get(FirebaseSegmentation.class);
   }
 
-  Task<Void> setCustomInstallationId(@Nullable String customInstallationId) {
+  @NonNull
+  public synchronized Task<Void> setCustomInstallationId(@Nullable String customInstallationId) {
     if (customInstallationId == null) {
       return clearCustomInstallationId();
     }
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java
index f650c3463d8..a1346062440 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java
@@ -101,13 +101,7 @@ private String getSharedPreferencesKey(String key) {
 
   private Task<Boolean> commitSharedPreferencesEditAsync(SharedPreferences.Editor editor) {
     TaskCompletionSource<Boolean> result = new TaskCompletionSource<>();
-    ioExecuter.execute(
-        new Runnable() {
-          @Override
-          public void run() {
-            result.setResult(editor.commit());
-          }
-        });
+    ioExecuter.execute(() -> result.setResult(editor.commit()));
     return result.getTask();
   }
 }

From 6a0f50241fe81025e2c6c48e803cdf78abb6e5f6 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Wed, 26 Jun 2019 13:21:20 -0700
Subject: [PATCH 17/30] Http client in Firebase Segmentation SDK to call
 backend service.

---
 .../firebase-segmentation.gradle              |  45 +-----
 .../segmentation/FirebaseSegmentation.java    |  17 +-
 .../remote/SegmentationServiceClient.java     | 153 +++++++++++++++++-
 3 files changed, 165 insertions(+), 50 deletions(-)

diff --git a/firebase-segmentation/firebase-segmentation.gradle b/firebase-segmentation/firebase-segmentation.gradle
index 1f76b8b593a..70b0f82fa14 100644
--- a/firebase-segmentation/firebase-segmentation.gradle
+++ b/firebase-segmentation/firebase-segmentation.gradle
@@ -14,45 +14,12 @@
 
 plugins {
     id 'firebase-library'
-    id 'com.google.protobuf'
 }
 
 firebaseLibrary {
     testLab.enabled = true
 }
 
-protobuf {
-    // Configure the protoc executable
-    protoc {
-        // Download from repositories
-        artifact = 'com.google.protobuf:protoc:3.4.0'
-    }
-    plugins {
-        grpc {
-            artifact = 'io.grpc:protoc-gen-grpc-java:1.12.0'
-        }
-        javalite {
-            // The codegen for lite comes as a separate artifact
-            artifact = 'com.google.protobuf:protoc-gen-javalite:3.0.0'
-        }
-    }
-    generateProtoTasks {
-        all().each { task ->
-            task.builtins {
-                // In most cases you don't need the full Java output
-                // if you use the lite output.
-                remove java
-            }
-            task.plugins {
-                grpc {
-                    option 'lite'
-                }
-                javalite {}
-            }
-        }
-    }
-}
-
 android {
     compileSdkVersion project.targetSdkVersion
 
@@ -65,8 +32,11 @@ android {
     }
     sourceSets {
         main {
-            proto {
-                srcDir 'src/main/proto'
+            java {
+            }
+        }
+        test {
+            java {
             }
         }
     }
@@ -83,14 +53,11 @@ android {
 
 dependencies {
     implementation project(':firebase-common')
-    implementation project(':protolite-well-known-types')
 
     implementation('com.google.firebase:firebase-iid:17.0.3') {
         exclude group: "com.google.firebase", module: "firebase-common"
     }
-    implementation 'io.grpc:grpc-stub:1.21.0'
-    implementation 'io.grpc:grpc-protobuf-lite:1.21.0'
-    implementation 'io.grpc:grpc-okhttp:1.21.0'
+
     implementation 'androidx.appcompat:appcompat:1.0.2'
     implementation 'androidx.multidex:multidex:2.0.0'
     implementation 'com.google.android.gms:play-services-tasks:16.0.1'
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
index 34de68597c1..6a7550cb5a4 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
@@ -135,6 +135,7 @@ private Task<Void> updateCustomInstallationId(String customInstallationId) {
                 String iidToken = instanceIdResultTask.getResult().getToken();
                 return backendServiceClient.updateCustomInstallationId(
                     Utils.getProjectNumberFromAppId(firebaseApp.getOptions().getApplicationId()),
+                    firebaseApp.getOptions().getApiKey(),
                     customInstallationId,
                     iid,
                     iidToken);
@@ -154,7 +155,9 @@ private Task<Void> updateCustomInstallationId(String customInstallationId) {
                           customInstallationId,
                           instanceIdResultTask.getResult().getId(),
                           CustomInstallationIdCache.CacheStatus.SYNCED));
-                case ALREADY_EXISTS:
+                case HTTP_CLIENT_ERROR:
+                  throw new SetCustomInstallationIdException(Status.CLIENT_ERROR);
+                case CONFLICT:
                   throw new SetCustomInstallationIdException(
                       Status.DUPLICATED_CUSTOM_INSTALLATION_ID);
                 default:
@@ -209,6 +212,7 @@ private Task<Void> clearCustomInstallationId() {
                 String iidToken = instanceIdResultTask.getResult().getToken();
                 return backendServiceClient.clearCustomInstallationId(
                     Utils.getProjectNumberFromAppId(firebaseApp.getOptions().getApplicationId()),
+                    firebaseApp.getOptions().getApiKey(),
                     iid,
                     iidToken);
               } else {
@@ -220,10 +224,13 @@ private Task<Void> clearCustomInstallationId() {
     Task<Boolean> finalUpdateCacheResultTask =
         backendRequestResultTask.onSuccessTask(
             backendRequestResult -> {
-              if (backendRequestResult == Code.OK) {
-                return localCache.clear();
-              } else {
-                throw new SetCustomInstallationIdException(Status.BACKEND_ERROR);
+              switch (backendRequestResult) {
+                case OK:
+                  return localCache.clear();
+                case HTTP_CLIENT_ERROR:
+                  throw new SetCustomInstallationIdException(Status.CLIENT_ERROR);
+                default:
+                  throw new SetCustomInstallationIdException(Status.BACKEND_ERROR);
               }
             });
 
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
index 59398b369d5..5292bc74180 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
@@ -15,25 +15,47 @@
 package com.google.firebase.segmentation.remote;
 
 import com.google.android.gms.tasks.Task;
+import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.android.gms.tasks.Tasks;
+import com.squareup.okhttp.MediaType;
 import com.squareup.okhttp.OkHttpClient;
+import com.squareup.okhttp.Request;
+import com.squareup.okhttp.RequestBody;
+import com.squareup.okhttp.Response;
+import java.io.IOException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
+import org.json.JSONException;
+import org.json.JSONObject;
 
 /** Http client that sends request to Firebase Segmentation backend API. To be implemented */
 public class SegmentationServiceClient {
 
+  private static final String FIREBASE_SEGMENTATION_API_DOMAIN =
+      "firebasesegmentation.googleapis.com";
+  private static final String UPDATE_REQUEST_RESOURCE_NAME_FORMAT =
+      "projects/%s/installations/%s/customSegmentationData";
+  private static final String CLEAR_REQUEST_RESOURCE_NAME_FORMAT =
+      "projects/%s/installations/%s/customSegmentationData:clear";
+  private static final String FIREBASE_SEGMENTATION_API_VERSION = "alpha1";
+
+  private static final MediaType JSON = MediaType.parse("application/json; charset=utf-8");
+
   private final OkHttpClient httpClient;
   private final Executor httpRequestExecutor;
 
   public enum Code {
     OK,
 
-    SERVER_INTERNAL_ERROR,
+    HTTP_CLIENT_ERROR,
+
+    CONFLICT,
 
-    ALREADY_EXISTS,
+    NETWORK_ERROR,
 
-    PERMISSION_DENIED
+    SERVER_ERROR,
+
+    UNAUTHORIZED,
   }
 
   public SegmentationServiceClient() {
@@ -43,14 +65,133 @@ public SegmentationServiceClient() {
 
   public Task<Code> updateCustomInstallationId(
       long projectNumber,
+      String apiKey,
       String customInstallationId,
       String firebaseInstanceId,
       String firebaseInstanceIdToken) {
-    return Tasks.forResult(Code.OK);
+    String resourceName =
+        String.format(UPDATE_REQUEST_RESOURCE_NAME_FORMAT, projectNumber, firebaseInstanceId);
+
+    RequestBody requestBody;
+    try {
+      requestBody =
+          RequestBody.create(
+              JSON,
+              buildUpdateCustomSegmentationDataRequestBody(resourceName, customInstallationId)
+                  .toString());
+    } catch (JSONException e) {
+      return Tasks.forException(e);
+    }
+
+    Request request =
+        new Request.Builder()
+            .url(
+                String.format(
+                    "https://%s/%s/%s",
+                    FIREBASE_SEGMENTATION_API_DOMAIN,
+                    FIREBASE_SEGMENTATION_API_VERSION,
+                    resourceName))
+            .header("X-Goog-Api-Key", apiKey)
+            .header("Authorization", "FIREBASE_INSTALLATIONS_AUTH " + firebaseInstanceIdToken)
+            .header("Content-Type", "application/json")
+            .patch(requestBody)
+            .build();
+
+    TaskCompletionSource<Code> taskCompletionSource = new TaskCompletionSource<>();
+    httpRequestExecutor.execute(
+        () -> {
+          try {
+            Response response = httpClient.newCall(request).execute();
+            switch (response.code()) {
+              case 200:
+                taskCompletionSource.setResult(Code.OK);
+                break;
+              case 401:
+                taskCompletionSource.setResult(Code.UNAUTHORIZED);
+                break;
+              case 409:
+                taskCompletionSource.setResult(Code.CONFLICT);
+                break;
+              default:
+                taskCompletionSource.setResult(Code.SERVER_ERROR);
+                break;
+            }
+          } catch (IOException e) {
+            taskCompletionSource.setResult(Code.NETWORK_ERROR);
+          }
+        });
+    return taskCompletionSource.getTask();
+  }
+
+  private static JSONObject buildUpdateCustomSegmentationDataRequestBody(
+      String resourceName, String customInstallationId) throws JSONException {
+    JSONObject rlt = new JSONObject();
+    rlt.put(
+        "update_mask",
+        "custom_segmentation_data.name,custom_segmentation_data.custom_installation_id");
+    JSONObject customSegmentationData = new JSONObject();
+    customSegmentationData.put("name", resourceName);
+    customSegmentationData.put("custom_installation_id", customInstallationId);
+    rlt.put("custom_segmentation_data", customSegmentationData);
+    return rlt;
   }
 
   public Task<Code> clearCustomInstallationId(
-      long projectNumber, String firebaseInstanceId, String firebaseInstanceIdToken) {
-    return Tasks.forResult(Code.OK);
+      long projectNumber,
+      String apiKey,
+      String firebaseInstanceId,
+      String firebaseInstanceIdToken) {
+    String resourceName =
+        String.format(CLEAR_REQUEST_RESOURCE_NAME_FORMAT, projectNumber, firebaseInstanceId);
+
+    RequestBody requestBody;
+    try {
+      requestBody =
+          RequestBody.create(
+              JSON, buildClearCustomSegmentationDataRequestBody(resourceName).toString());
+    } catch (JSONException e) {
+      return Tasks.forException(e);
+    }
+
+    Request request =
+        new Request.Builder()
+            .url(
+                String.format(
+                    "https://%s/%s/%s",
+                    FIREBASE_SEGMENTATION_API_DOMAIN,
+                    FIREBASE_SEGMENTATION_API_VERSION,
+                    resourceName))
+            .header("X-Goog-Api-Key", apiKey)
+            .header("Authorization", "FIREBASE_INSTALLATIONS_AUTH " + firebaseInstanceIdToken)
+            .header("Content-Type", "application/json")
+            .post(requestBody)
+            .build();
+
+    TaskCompletionSource<Code> taskCompletionSource = new TaskCompletionSource<>();
+    httpRequestExecutor.execute(
+        () -> {
+          try {
+            Response response = httpClient.newCall(request).execute();
+            switch (response.code()) {
+              case 200:
+                taskCompletionSource.setResult(Code.OK);
+                break;
+              case 401:
+                taskCompletionSource.setResult(Code.UNAUTHORIZED);
+                break;
+              default:
+                taskCompletionSource.setResult(Code.SERVER_ERROR);
+                break;
+            }
+          } catch (IOException e) {
+            taskCompletionSource.setResult(Code.NETWORK_ERROR);
+          }
+        });
+    return taskCompletionSource.getTask();
+  }
+
+  private static JSONObject buildClearCustomSegmentationDataRequestBody(String resourceName)
+      throws JSONException {
+    return new JSONObject().put("name", resourceName);
   }
 }

From 7f4097836d1787bc1f267c2170dc318738423479 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Wed, 26 Jun 2019 13:26:07 -0700
Subject: [PATCH 18/30] Revert unintentional change

---
 subprojects.cfg | 1 -
 1 file changed, 1 deletion(-)

diff --git a/subprojects.cfg b/subprojects.cfg
index e27229cee5b..6a09b6d4a45 100644
--- a/subprojects.cfg
+++ b/subprojects.cfg
@@ -15,7 +15,6 @@ firebase-inappmessaging-display
 firebase-storage
 firebase-storage:test-app
 firebase-segmentation
-
 protolite-well-known-types
 
 tools:errorprone

From b2fc302762cb4562646587466a08eb7e244c5928 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Wed, 26 Jun 2019 13:39:30 -0700
Subject: [PATCH 19/30] Fix connected device test

---
 .../FirebaseSegmentationInstrumentedTest.java   | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
index 8519e641cb7..3bac1321da6 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
@@ -69,19 +69,24 @@ public void setUp() {
     firebaseApp =
         FirebaseApp.initializeApp(
             ApplicationProvider.getApplicationContext(),
-            new FirebaseOptions.Builder().setApplicationId("1:123456789:android:abcdef").build());
+            new FirebaseOptions.Builder()
+                .setApplicationId("1" + ":123456789:android:abcdef")
+                .setApiKey("api_key")
+                .build());
     actualCache = new CustomInstallationIdCache(firebaseApp);
 
     when(backendClientReturnsOk.updateCustomInstallationId(
-            anyLong(), anyString(), anyString(), anyString()))
+            anyLong(), anyString(), anyString(), anyString(), anyString()))
         .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.OK));
-    when(backendClientReturnsOk.clearCustomInstallationId(anyLong(), anyString(), anyString()))
+    when(backendClientReturnsOk.clearCustomInstallationId(
+            anyLong(), anyString(), anyString(), anyString()))
         .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.OK));
     when(backendClientReturnsError.updateCustomInstallationId(
+            anyLong(), anyString(), anyString(), anyString(), anyString()))
+        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.SERVER_ERROR));
+    when(backendClientReturnsError.clearCustomInstallationId(
             anyLong(), anyString(), anyString(), anyString()))
-        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.SERVER_INTERNAL_ERROR));
-    when(backendClientReturnsError.clearCustomInstallationId(anyLong(), anyString(), anyString()))
-        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.SERVER_INTERNAL_ERROR));
+        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.SERVER_ERROR));
     when(firebaseInstanceId.getInstanceId())
         .thenReturn(
             Tasks.forResult(

From 1f2ab34e45e6414d8663a70099650d0feac2b8f2 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Wed, 26 Jun 2019 13:39:30 -0700
Subject: [PATCH 20/30] Fix connected device test

---
 .../FirebaseSegmentationInstrumentedTest.java   | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
index 8519e641cb7..a2ac5f52b90 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
@@ -69,19 +69,24 @@ public void setUp() {
     firebaseApp =
         FirebaseApp.initializeApp(
             ApplicationProvider.getApplicationContext(),
-            new FirebaseOptions.Builder().setApplicationId("1:123456789:android:abcdef").build());
+            new FirebaseOptions.Builder()
+                .setApplicationId("1:123456789:android:abcdef")
+                .setApiKey("api_key")
+                .build());
     actualCache = new CustomInstallationIdCache(firebaseApp);
 
     when(backendClientReturnsOk.updateCustomInstallationId(
-            anyLong(), anyString(), anyString(), anyString()))
+            anyLong(), anyString(), anyString(), anyString(), anyString()))
         .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.OK));
-    when(backendClientReturnsOk.clearCustomInstallationId(anyLong(), anyString(), anyString()))
+    when(backendClientReturnsOk.clearCustomInstallationId(
+            anyLong(), anyString(), anyString(), anyString()))
         .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.OK));
     when(backendClientReturnsError.updateCustomInstallationId(
+            anyLong(), anyString(), anyString(), anyString(), anyString()))
+        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.SERVER_ERROR));
+    when(backendClientReturnsError.clearCustomInstallationId(
             anyLong(), anyString(), anyString(), anyString()))
-        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.SERVER_INTERNAL_ERROR));
-    when(backendClientReturnsError.clearCustomInstallationId(anyLong(), anyString(), anyString()))
-        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.SERVER_INTERNAL_ERROR));
+        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.SERVER_ERROR));
     when(firebaseInstanceId.getInstanceId())
         .thenReturn(
             Tasks.forResult(

From 9880624380da20818bcc43a9902cd6bf9b2bba73 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Mon, 8 Jul 2019 12:46:15 -0700
Subject: [PATCH 21/30] 1. Add a few annotations to make java code Kotlin
 friendly 2. Some fixes for the http request format

---
 .../firebase-segmentation.gradle              |  2 +-
 .../segmentation/FirebaseSegmentation.java    |  1 +
 .../FirebaseSegmentationRegistrar.java        |  2 ++
 .../local/CustomInstallationIdCache.java      |  7 ++--
 .../CustomInstallationIdCacheEntryValue.java  |  9 ++++-
 .../remote/SegmentationServiceClient.java     | 36 +++++++++----------
 6 files changed, 34 insertions(+), 23 deletions(-)

diff --git a/firebase-segmentation/firebase-segmentation.gradle b/firebase-segmentation/firebase-segmentation.gradle
index 70b0f82fa14..fa48a4ac58e 100644
--- a/firebase-segmentation/firebase-segmentation.gradle
+++ b/firebase-segmentation/firebase-segmentation.gradle
@@ -70,7 +70,7 @@ dependencies {
     testImplementation 'junit:junit:4.12'
     testImplementation "org.robolectric:robolectric:$robolectricVersion"
 
-    androidTestImplementation "androidx.annotation:annotation:1.1.0"
+    androidTestImplementation "androidx.annotation:annotation:1.0.0"
     androidTestImplementation 'androidx.test.ext:junit:1.1.1'
     androidTestImplementation 'androidx.test:rules:1.2.0'
     androidTestImplementation 'androidx.test:runner:1.2.0'
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
index 6a7550cb5a4..26af8c6a817 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
@@ -17,6 +17,7 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.annotation.RestrictTo;
+
 import com.google.android.gms.common.internal.Preconditions;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.Tasks;
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentationRegistrar.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentationRegistrar.java
index 7d1d5fcfaab..ca7f688d60a 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentationRegistrar.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentationRegistrar.java
@@ -14,6 +14,7 @@
 
 package com.google.firebase.segmentation;
 
+import androidx.annotation.NonNull;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.components.Component;
 import com.google.firebase.components.ComponentRegistrar;
@@ -25,6 +26,7 @@
 public class FirebaseSegmentationRegistrar implements ComponentRegistrar {
 
   @Override
+  @NonNull
   public List<Component<?>> getComponents() {
     return Arrays.asList(
         Component.builder(FirebaseSegmentation.class)
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java
index a1346062440..05f652d4723 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java
@@ -16,6 +16,7 @@
 
 import android.content.Context;
 import android.content.SharedPreferences;
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
@@ -53,7 +54,7 @@ public enum CacheStatus {
   private final SharedPreferences prefs;
   private final String persistenceKey;
 
-  public CustomInstallationIdCache(FirebaseApp firebaseApp) {
+  public CustomInstallationIdCache(@NonNull FirebaseApp firebaseApp) {
     // Different FirebaseApp in the same Android application should have the same application
     // context and same dir path
     prefs =
@@ -77,8 +78,9 @@ public synchronized CustomInstallationIdCacheEntryValue readCacheEntryValue() {
     return CustomInstallationIdCacheEntryValue.create(cid, iid, CacheStatus.values()[status]);
   }
 
+  @NonNull
   public synchronized Task<Boolean> insertOrUpdateCacheEntry(
-      CustomInstallationIdCacheEntryValue entryValue) {
+      @NonNull CustomInstallationIdCacheEntryValue entryValue) {
     SharedPreferences.Editor editor = prefs.edit();
     editor.putString(
         getSharedPreferencesKey(CUSTOM_INSTALLATION_ID_KEY), entryValue.getCustomInstallationId());
@@ -87,6 +89,7 @@ public synchronized Task<Boolean> insertOrUpdateCacheEntry(
     return commitSharedPreferencesEditAsync(editor);
   }
 
+  @NonNull
   public synchronized Task<Boolean> clear() {
     SharedPreferences.Editor editor = prefs.edit();
     editor.remove(getSharedPreferencesKey(CUSTOM_INSTALLATION_ID_KEY));
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheEntryValue.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheEntryValue.java
index 05528cd40f0..5e2c1944278 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheEntryValue.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheEntryValue.java
@@ -14,6 +14,7 @@
 
 package com.google.firebase.segmentation.local;
 
+import androidx.annotation.NonNull;
 import com.google.auto.value.AutoValue;
 import com.google.firebase.segmentation.local.CustomInstallationIdCache.CacheStatus;
 
@@ -23,14 +24,20 @@
  */
 @AutoValue
 public abstract class CustomInstallationIdCacheEntryValue {
+  @NonNull
   public abstract String getCustomInstallationId();
 
+  @NonNull
   public abstract String getFirebaseInstanceId();
 
+  @NonNull
   public abstract CacheStatus getCacheStatus();
 
+  @NonNull
   public static CustomInstallationIdCacheEntryValue create(
-      String customInstallationId, String firebaseInstanceId, CacheStatus cacheStatus) {
+      @NonNull String customInstallationId,
+      @NonNull String firebaseInstanceId,
+      @NonNull CacheStatus cacheStatus) {
     return new AutoValue_CustomInstallationIdCacheEntryValue(
         customInstallationId, firebaseInstanceId, cacheStatus);
   }
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
index 5292bc74180..71dfb19c6c9 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
@@ -14,6 +14,7 @@
 
 package com.google.firebase.segmentation.remote;
 
+import androidx.annotation.NonNull;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.android.gms.tasks.Tasks;
@@ -63,12 +64,13 @@ public SegmentationServiceClient() {
     httpRequestExecutor = Executors.newFixedThreadPool(4);
   }
 
+  @NonNull
   public Task<Code> updateCustomInstallationId(
       long projectNumber,
-      String apiKey,
-      String customInstallationId,
-      String firebaseInstanceId,
-      String firebaseInstanceIdToken) {
+      @NonNull String apiKey,
+      @NonNull String customInstallationId,
+      @NonNull String firebaseInstanceId,
+      @NonNull String firebaseInstanceIdToken) {
     String resourceName =
         String.format(UPDATE_REQUEST_RESOURCE_NAME_FORMAT, projectNumber, firebaseInstanceId);
 
@@ -87,11 +89,11 @@ public Task<Code> updateCustomInstallationId(
         new Request.Builder()
             .url(
                 String.format(
-                    "https://%s/%s/%s",
+                    "https://%s/%s/%s?key=%s",
                     FIREBASE_SEGMENTATION_API_DOMAIN,
                     FIREBASE_SEGMENTATION_API_VERSION,
-                    resourceName))
-            .header("X-Goog-Api-Key", apiKey)
+                    resourceName,
+                    apiKey))
             .header("Authorization", "FIREBASE_INSTALLATIONS_AUTH " + firebaseInstanceIdToken)
             .header("Content-Type", "application/json")
             .patch(requestBody)
@@ -125,22 +127,18 @@ public Task<Code> updateCustomInstallationId(
 
   private static JSONObject buildUpdateCustomSegmentationDataRequestBody(
       String resourceName, String customInstallationId) throws JSONException {
-    JSONObject rlt = new JSONObject();
-    rlt.put(
-        "update_mask",
-        "custom_segmentation_data.name,custom_segmentation_data.custom_installation_id");
     JSONObject customSegmentationData = new JSONObject();
     customSegmentationData.put("name", resourceName);
     customSegmentationData.put("custom_installation_id", customInstallationId);
-    rlt.put("custom_segmentation_data", customSegmentationData);
-    return rlt;
+    return customSegmentationData;
   }
 
+  @NonNull
   public Task<Code> clearCustomInstallationId(
       long projectNumber,
-      String apiKey,
-      String firebaseInstanceId,
-      String firebaseInstanceIdToken) {
+      @NonNull String apiKey,
+      @NonNull String firebaseInstanceId,
+      @NonNull String firebaseInstanceIdToken) {
     String resourceName =
         String.format(CLEAR_REQUEST_RESOURCE_NAME_FORMAT, projectNumber, firebaseInstanceId);
 
@@ -157,11 +155,11 @@ public Task<Code> clearCustomInstallationId(
         new Request.Builder()
             .url(
                 String.format(
-                    "https://%s/%s/%s",
+                    "https://%s/%s/%s?key=%s",
                     FIREBASE_SEGMENTATION_API_DOMAIN,
                     FIREBASE_SEGMENTATION_API_VERSION,
-                    resourceName))
-            .header("X-Goog-Api-Key", apiKey)
+                    resourceName,
+                    apiKey))
             .header("Authorization", "FIREBASE_INSTALLATIONS_AUTH " + firebaseInstanceIdToken)
             .header("Content-Type", "application/json")
             .post(requestBody)

From 8ffb5bb6e35e4b753c0eb6376a66f964ad37db4e Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Mon, 8 Jul 2019 13:26:45 -0700
Subject: [PATCH 22/30] Fix java format

---
 .../com/google/firebase/segmentation/FirebaseSegmentation.java   | 1 -
 1 file changed, 1 deletion(-)

diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
index 26af8c6a817..6a7550cb5a4 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
@@ -17,7 +17,6 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.annotation.RestrictTo;
-
 import com.google.android.gms.common.internal.Preconditions;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.Tasks;

From daf4698905e6fb54bd4f9fd190c2b2dfc8e51c83 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Mon, 8 Jul 2019 13:46:38 -0700
Subject: [PATCH 23/30] Fix API version

---
 .../firebase/segmentation/remote/SegmentationServiceClient.java | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
index 71dfb19c6c9..e258b84ad06 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
@@ -38,7 +38,7 @@ public class SegmentationServiceClient {
       "projects/%s/installations/%s/customSegmentationData";
   private static final String CLEAR_REQUEST_RESOURCE_NAME_FORMAT =
       "projects/%s/installations/%s/customSegmentationData:clear";
-  private static final String FIREBASE_SEGMENTATION_API_VERSION = "alpha1";
+  private static final String FIREBASE_SEGMENTATION_API_VERSION = "v1alpha";
 
   private static final MediaType JSON = MediaType.parse("application/json; charset=utf-8");
 

From 7985e146526d6b648d322ba81b3da5a752c22df5 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Tue, 9 Jul 2019 10:55:56 -0700
Subject: [PATCH 24/30] Change the segmentation API implementation to
 synchronous and put the entire synchronous code block in async task.

---
 .../firebase-segmentation.gradle              |  14 +-
 .../FirebaseSegmentationInstrumentedTest.java |  34 ++-
 .../local/CustomInstallationIdCacheTest.java  |  21 +-
 .../segmentation/FirebaseSegmentation.java    | 229 ++++++++++--------
 .../local/CustomInstallationIdCache.java      |  20 +-
 .../remote/SegmentationServiceClient.java     |  90 +++----
 6 files changed, 186 insertions(+), 222 deletions(-)

diff --git a/firebase-segmentation/firebase-segmentation.gradle b/firebase-segmentation/firebase-segmentation.gradle
index fa48a4ac58e..d16e46b402b 100644
--- a/firebase-segmentation/firebase-segmentation.gradle
+++ b/firebase-segmentation/firebase-segmentation.gradle
@@ -30,16 +30,6 @@ android {
         versionName version
         testInstrumentationRunner "androidx.test.runner.AndroidJUnitRunner"
     }
-    sourceSets {
-        main {
-            java {
-            }
-        }
-        test {
-            java {
-            }
-        }
-    }
     compileOptions {
         sourceCompatibility JavaVersion.VERSION_1_8
         targetCompatibility JavaVersion.VERSION_1_8
@@ -59,8 +49,8 @@ dependencies {
     }
 
     implementation 'androidx.appcompat:appcompat:1.0.2'
-    implementation 'androidx.multidex:multidex:2.0.0'
-    implementation 'com.google.android.gms:play-services-tasks:16.0.1'
+    implementation 'androidx.multidex:multidex:2.0.1'
+    implementation 'com.google.android.gms:play-services-tasks:17.0.0'
     implementation 'com.squareup.okhttp:okhttp:2.7.5'
 
     compileOnly "com.google.auto.value:auto-value-annotations:1.6.5"
diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
index a2ac5f52b90..19498782a78 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
@@ -77,16 +77,16 @@ public void setUp() {
 
     when(backendClientReturnsOk.updateCustomInstallationId(
             anyLong(), anyString(), anyString(), anyString(), anyString()))
-        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.OK));
+        .thenReturn(SegmentationServiceClient.Code.OK);
     when(backendClientReturnsOk.clearCustomInstallationId(
             anyLong(), anyString(), anyString(), anyString()))
-        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.OK));
+        .thenReturn(SegmentationServiceClient.Code.OK);
     when(backendClientReturnsError.updateCustomInstallationId(
             anyLong(), anyString(), anyString(), anyString(), anyString()))
-        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.SERVER_ERROR));
+        .thenReturn(SegmentationServiceClient.Code.SERVER_ERROR);
     when(backendClientReturnsError.clearCustomInstallationId(
             anyLong(), anyString(), anyString(), anyString()))
-        .thenReturn(Tasks.forResult(SegmentationServiceClient.Code.SERVER_ERROR));
+        .thenReturn(SegmentationServiceClient.Code.SERVER_ERROR);
     when(firebaseInstanceId.getInstanceId())
         .thenReturn(
             Tasks.forResult(
@@ -103,13 +103,13 @@ public String getToken() {
                     return "iid_token";
                   }
                 }));
-    when(cacheReturnsError.insertOrUpdateCacheEntry(any())).thenReturn(Tasks.forResult(false));
+    when(cacheReturnsError.insertOrUpdateCacheEntry(any())).thenReturn(false);
     when(cacheReturnsError.readCacheEntryValue()).thenReturn(null);
   }
 
   @After
   public void cleanUp() throws Exception {
-    Tasks.await(actualCache.clear());
+    actualCache.clear();
   }
 
   @Test
@@ -170,12 +170,11 @@ public void testUpdateCustomInstallationId_CacheError_BackendOk() throws Interru
 
   @Test
   public void testClearCustomInstallationId_CacheOk_BackendOk() throws Exception {
-    Tasks.await(
-        actualCache.insertOrUpdateCacheEntry(
-            CustomInstallationIdCacheEntryValue.create(
-                CUSTOM_INSTALLATION_ID,
-                FIREBASE_INSTANCE_ID,
-                CustomInstallationIdCache.CacheStatus.SYNCED)));
+    actualCache.insertOrUpdateCacheEntry(
+        CustomInstallationIdCacheEntryValue.create(
+            CUSTOM_INSTALLATION_ID,
+            FIREBASE_INSTANCE_ID,
+            CustomInstallationIdCache.CacheStatus.SYNCED));
     FirebaseSegmentation firebaseSegmentation =
         new FirebaseSegmentation(
             firebaseApp, firebaseInstanceId, actualCache, backendClientReturnsOk);
@@ -188,12 +187,11 @@ public void testClearCustomInstallationId_CacheOk_BackendOk() throws Exception {
 
   @Test
   public void testClearCustomInstallationId_CacheOk_BackendError() throws Exception {
-    Tasks.await(
-        actualCache.insertOrUpdateCacheEntry(
-            CustomInstallationIdCacheEntryValue.create(
-                CUSTOM_INSTALLATION_ID,
-                FIREBASE_INSTANCE_ID,
-                CustomInstallationIdCache.CacheStatus.SYNCED)));
+    actualCache.insertOrUpdateCacheEntry(
+        CustomInstallationIdCacheEntryValue.create(
+            CUSTOM_INSTALLATION_ID,
+            FIREBASE_INSTANCE_ID,
+            CustomInstallationIdCache.CacheStatus.SYNCED));
     FirebaseSegmentation firebaseSegmentation =
         new FirebaseSegmentation(
             firebaseApp, firebaseInstanceId, actualCache, backendClientReturnsError);
diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java
index 019a2b8ba08..9e22e522d2b 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/local/CustomInstallationIdCacheTest.java
@@ -20,7 +20,6 @@
 
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
-import com.google.android.gms.tasks.Tasks;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.FirebaseOptions;
 import org.junit.After;
@@ -55,8 +54,8 @@ public void setUp() {
 
   @After
   public void cleanUp() throws Exception {
-    Tasks.await(cache0.clear());
-    Tasks.await(cache1.clear());
+    cache0.clear();
+    cache1.clear();
   }
 
   @Test
@@ -68,12 +67,9 @@ public void testReadCacheEntry_Null() {
   @Test
   public void testUpdateAndReadCacheEntry() throws Exception {
     assertTrue(
-        Tasks.await(
-            cache0.insertOrUpdateCacheEntry(
-                CustomInstallationIdCacheEntryValue.create(
-                    "123456",
-                    "cAAAAAAAAAA",
-                    CustomInstallationIdCache.CacheStatus.PENDING_UPDATE))));
+        cache0.insertOrUpdateCacheEntry(
+            CustomInstallationIdCacheEntryValue.create(
+                "123456", "cAAAAAAAAAA", CustomInstallationIdCache.CacheStatus.PENDING_UPDATE)));
     CustomInstallationIdCacheEntryValue entryValue = cache0.readCacheEntryValue();
     assertThat(entryValue.getCustomInstallationId()).isEqualTo("123456");
     assertThat(entryValue.getFirebaseInstanceId()).isEqualTo("cAAAAAAAAAA");
@@ -82,10 +78,9 @@ public void testUpdateAndReadCacheEntry() throws Exception {
     assertNull(cache1.readCacheEntryValue());
 
     assertTrue(
-        Tasks.await(
-            cache0.insertOrUpdateCacheEntry(
-                CustomInstallationIdCacheEntryValue.create(
-                    "123456", "cAAAAAAAAAA", CustomInstallationIdCache.CacheStatus.SYNCED))));
+        cache0.insertOrUpdateCacheEntry(
+            CustomInstallationIdCacheEntryValue.create(
+                "123456", "cAAAAAAAAAA", CustomInstallationIdCache.CacheStatus.SYNCED)));
     entryValue = cache0.readCacheEntryValue();
     assertThat(entryValue.getCustomInstallationId()).isEqualTo("123456");
     assertThat(entryValue.getFirebaseInstanceId()).isEqualTo("cAAAAAAAAAA");
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
index 6a7550cb5a4..7a6d879d669 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
@@ -19,7 +19,7 @@
 import androidx.annotation.RestrictTo;
 import com.google.android.gms.common.internal.Preconditions;
 import com.google.android.gms.tasks.Task;
-import com.google.android.gms.tasks.Tasks;
+import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.iid.FirebaseInstanceId;
 import com.google.firebase.iid.InstanceIdResult;
@@ -28,6 +28,8 @@
 import com.google.firebase.segmentation.local.CustomInstallationIdCacheEntryValue;
 import com.google.firebase.segmentation.remote.SegmentationServiceClient;
 import com.google.firebase.segmentation.remote.SegmentationServiceClient.Code;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
 
 /** Entry point of Firebase Segmentation SDK. */
 public class FirebaseSegmentation {
@@ -36,12 +38,14 @@ public class FirebaseSegmentation {
   private final FirebaseInstanceId firebaseInstanceId;
   private final CustomInstallationIdCache localCache;
   private final SegmentationServiceClient backendServiceClient;
+  private final Executor executor;
 
   FirebaseSegmentation(FirebaseApp firebaseApp) {
     this.firebaseApp = firebaseApp;
     this.firebaseInstanceId = FirebaseInstanceId.getInstance(firebaseApp);
-    localCache = new CustomInstallationIdCache(firebaseApp);
-    backendServiceClient = new SegmentationServiceClient();
+    this.localCache = new CustomInstallationIdCache(firebaseApp);
+    this.backendServiceClient = new SegmentationServiceClient();
+    this.executor = Executors.newFixedThreadPool(4);
   }
 
   @RestrictTo(RestrictTo.Scope.TESTS)
@@ -54,6 +58,7 @@ public class FirebaseSegmentation {
     this.firebaseInstanceId = firebaseInstanceId;
     this.localCache = localCache;
     this.backendServiceClient = backendServiceClient;
+    this.executor = Executors.newFixedThreadPool(4);
   }
 
   /**
@@ -107,73 +112,79 @@ public synchronized Task<Void> setCustomInstallationId(@Nullable String customIn
    * </pre>
    */
   private Task<Void> updateCustomInstallationId(String customInstallationId) {
-    CustomInstallationIdCacheEntryValue cacheEntryValue = localCache.readCacheEntryValue();
-    if (cacheEntryValue != null
-        && cacheEntryValue.getCustomInstallationId().equals(customInstallationId)
-        && cacheEntryValue.getCacheStatus() == CustomInstallationIdCache.CacheStatus.SYNCED) {
-      // If the given custom installation id matches up the cached
-      // value, there's no need to update.
-      return Tasks.forResult(null);
-    }
+    TaskCompletionSource<Void> taskCompletionSource = new TaskCompletionSource<>();
+
+    executor.execute(
+        () -> {
+          CustomInstallationIdCacheEntryValue cacheEntryValue = localCache.readCacheEntryValue();
+          if (cacheEntryValue != null
+              && cacheEntryValue.getCustomInstallationId().equals(customInstallationId)
+              && cacheEntryValue.getCacheStatus() == CustomInstallationIdCache.CacheStatus.SYNCED) {
+            // If the given custom installation id matches up the cached
+            // value, there's no need to update.
+            taskCompletionSource.setResult(null);
+            return;
+          }
 
-    Task<InstanceIdResult> instanceIdResultTask = firebaseInstanceId.getInstanceId();
-    Task<Boolean> firstUpdateCacheResultTask =
-        instanceIdResultTask.onSuccessTask(
-            instanceIdResult ->
-                localCache.insertOrUpdateCacheEntry(
-                    CustomInstallationIdCacheEntryValue.create(
-                        customInstallationId,
-                        instanceIdResult.getId(),
-                        CustomInstallationIdCache.CacheStatus.PENDING_UPDATE)));
-
-    // Start requesting backend when first cache update is done.
-    Task<Code> backendRequestResultTask =
-        firstUpdateCacheResultTask.onSuccessTask(
-            firstUpdateCacheResult -> {
-              if (firstUpdateCacheResult) {
-                String iid = instanceIdResultTask.getResult().getId();
-                String iidToken = instanceIdResultTask.getResult().getToken();
-                return backendServiceClient.updateCustomInstallationId(
-                    Utils.getProjectNumberFromAppId(firebaseApp.getOptions().getApplicationId()),
-                    firebaseApp.getOptions().getApiKey(),
-                    customInstallationId,
-                    iid,
-                    iidToken);
-              } else {
-                throw new SetCustomInstallationIdException(
-                    "Failed to update client side cache", Status.CLIENT_ERROR);
-              }
-            });
-
-    Task<Boolean> finalUpdateCacheResultTask =
-        backendRequestResultTask.onSuccessTask(
-            backendRequestResult -> {
-              switch (backendRequestResult) {
-                case OK:
-                  return localCache.insertOrUpdateCacheEntry(
+          InstanceIdResult instanceIdResult = firebaseInstanceId.getInstanceId().getResult();
+          boolean firstUpdateCacheResult =
+              localCache.insertOrUpdateCacheEntry(
+                  CustomInstallationIdCacheEntryValue.create(
+                      customInstallationId,
+                      instanceIdResult.getId(),
+                      CustomInstallationIdCache.CacheStatus.PENDING_UPDATE));
+
+          if (!firstUpdateCacheResult) {
+            taskCompletionSource.setException(
+                new SetCustomInstallationIdException(
+                    "Failed to update client side cache", Status.CLIENT_ERROR));
+            return;
+          }
+
+          // Start requesting backend when first cache updae is done.
+          String iid = instanceIdResult.getId();
+          String iidToken = instanceIdResult.getToken();
+          Code backendRequestResult =
+              backendServiceClient.updateCustomInstallationId(
+                  Utils.getProjectNumberFromAppId(firebaseApp.getOptions().getApplicationId()),
+                  firebaseApp.getOptions().getApiKey(),
+                  customInstallationId,
+                  iid,
+                  iidToken);
+
+          boolean finalUpdateCacheResult;
+          switch (backendRequestResult) {
+            case OK:
+              finalUpdateCacheResult =
+                  localCache.insertOrUpdateCacheEntry(
                       CustomInstallationIdCacheEntryValue.create(
                           customInstallationId,
-                          instanceIdResultTask.getResult().getId(),
+                          instanceIdResult.getId(),
                           CustomInstallationIdCache.CacheStatus.SYNCED));
-                case HTTP_CLIENT_ERROR:
-                  throw new SetCustomInstallationIdException(Status.CLIENT_ERROR);
-                case CONFLICT:
-                  throw new SetCustomInstallationIdException(
-                      Status.DUPLICATED_CUSTOM_INSTALLATION_ID);
-                default:
-                  throw new SetCustomInstallationIdException(Status.BACKEND_ERROR);
-              }
-            });
-
-    return finalUpdateCacheResultTask.onSuccessTask(
-        finalUpdateCacheResult -> {
+              break;
+            case HTTP_CLIENT_ERROR:
+              taskCompletionSource.setException(
+                  new SetCustomInstallationIdException(Status.CLIENT_ERROR));
+              return;
+            case CONFLICT:
+              taskCompletionSource.setException(
+                  new SetCustomInstallationIdException(Status.DUPLICATED_CUSTOM_INSTALLATION_ID));
+              return;
+            default:
+              taskCompletionSource.setException(
+                  new SetCustomInstallationIdException(Status.BACKEND_ERROR));
+              return;
+          }
+
           if (finalUpdateCacheResult) {
-            return Tasks.forResult(null);
+            taskCompletionSource.setResult(null);
           } else {
-            throw new SetCustomInstallationIdException(
-                "Failed to update client side cache", Status.CLIENT_ERROR);
+            taskCompletionSource.setException(
+                new SetCustomInstallationIdException(
+                    "Failed to update client side cache", Status.CLIENT_ERROR));
           }
         });
+    return taskCompletionSource.getTask();
   }
 
   /**
@@ -194,54 +205,58 @@ private Task<Void> updateCustomInstallationId(String customInstallationId) {
    * </pre>
    */
   private Task<Void> clearCustomInstallationId() {
-    Task<InstanceIdResult> instanceIdResultTask = firebaseInstanceId.getInstanceId();
-    Task<Boolean> firstUpdateCacheResultTask =
-        instanceIdResultTask.onSuccessTask(
-            instanceIdResult ->
-                localCache.insertOrUpdateCacheEntry(
-                    CustomInstallationIdCacheEntryValue.create(
-                        "",
-                        instanceIdResult.getId(),
-                        CustomInstallationIdCache.CacheStatus.PENDING_CLEAR)));
-
-    Task<Code> backendRequestResultTask =
-        firstUpdateCacheResultTask.onSuccessTask(
-            firstUpdateCacheResult -> {
-              if (firstUpdateCacheResult) {
-                String iid = instanceIdResultTask.getResult().getId();
-                String iidToken = instanceIdResultTask.getResult().getToken();
-                return backendServiceClient.clearCustomInstallationId(
-                    Utils.getProjectNumberFromAppId(firebaseApp.getOptions().getApplicationId()),
-                    firebaseApp.getOptions().getApiKey(),
-                    iid,
-                    iidToken);
-              } else {
-                throw new SetCustomInstallationIdException(
-                    "Failed to update client side cache", Status.CLIENT_ERROR);
-              }
-            });
-
-    Task<Boolean> finalUpdateCacheResultTask =
-        backendRequestResultTask.onSuccessTask(
-            backendRequestResult -> {
-              switch (backendRequestResult) {
-                case OK:
-                  return localCache.clear();
-                case HTTP_CLIENT_ERROR:
-                  throw new SetCustomInstallationIdException(Status.CLIENT_ERROR);
-                default:
-                  throw new SetCustomInstallationIdException(Status.BACKEND_ERROR);
-              }
-            });
-
-    return finalUpdateCacheResultTask.onSuccessTask(
-        finalUpdateCacheResult -> {
+
+    TaskCompletionSource<Void> taskCompletionSource = new TaskCompletionSource<>();
+
+    executor.execute(
+        () -> {
+          InstanceIdResult instanceIdResult = firebaseInstanceId.getInstanceId().getResult();
+          boolean firstUpdateCacheResult =
+              localCache.insertOrUpdateCacheEntry(
+                  CustomInstallationIdCacheEntryValue.create(
+                      "",
+                      instanceIdResult.getId(),
+                      CustomInstallationIdCache.CacheStatus.PENDING_CLEAR));
+
+          if (!firstUpdateCacheResult) {
+            taskCompletionSource.setException(
+                new SetCustomInstallationIdException(
+                    "Failed to update client side cache", Status.CLIENT_ERROR));
+            return;
+          }
+
+          String iid = instanceIdResult.getId();
+          String iidToken = instanceIdResult.getToken();
+          Code backendRequestResult =
+              backendServiceClient.clearCustomInstallationId(
+                  Utils.getProjectNumberFromAppId(firebaseApp.getOptions().getApplicationId()),
+                  firebaseApp.getOptions().getApiKey(),
+                  iid,
+                  iidToken);
+
+          boolean finalUpdateCacheResult;
+          switch (backendRequestResult) {
+            case OK:
+              finalUpdateCacheResult = localCache.clear();
+              break;
+            case HTTP_CLIENT_ERROR:
+              taskCompletionSource.setException(
+                  new SetCustomInstallationIdException(Status.CLIENT_ERROR));
+              return;
+            default:
+              taskCompletionSource.setException(
+                  new SetCustomInstallationIdException(Status.BACKEND_ERROR));
+              return;
+          }
+
           if (finalUpdateCacheResult) {
-            return Tasks.forResult(null);
+            taskCompletionSource.setResult(null);
           } else {
-            throw new SetCustomInstallationIdException(
-                "Failed to update client side cache", Status.CLIENT_ERROR);
+            taskCompletionSource.setException(
+                new SetCustomInstallationIdException(
+                    "Failed to update client side cache", Status.CLIENT_ERROR));
           }
         });
+    return taskCompletionSource.getTask();
   }
 }
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java
index 05f652d4723..307d5d49923 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/local/CustomInstallationIdCache.java
@@ -18,11 +18,7 @@
 import android.content.SharedPreferences;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
-import com.google.android.gms.tasks.Task;
-import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.firebase.FirebaseApp;
-import java.util.concurrent.Executor;
-import java.util.concurrent.Executors;
 
 /**
  * A layer that locally caches a few Firebase Segmentation attributes on top the Segmentation
@@ -50,7 +46,6 @@ public enum CacheStatus {
   private static final String INSTANCE_ID_KEY = "Iid";
   private static final String CACHE_STATUS_KEY = "Status";
 
-  private final Executor ioExecuter;
   private final SharedPreferences prefs;
   private final String persistenceKey;
 
@@ -62,7 +57,6 @@ public CustomInstallationIdCache(@NonNull FirebaseApp firebaseApp) {
             .getApplicationContext()
             .getSharedPreferences(SHARED_PREFS_NAME, Context.MODE_PRIVATE);
     persistenceKey = firebaseApp.getPersistenceKey();
-    ioExecuter = Executors.newFixedThreadPool(2);
   }
 
   @Nullable
@@ -79,32 +73,26 @@ public synchronized CustomInstallationIdCacheEntryValue readCacheEntryValue() {
   }
 
   @NonNull
-  public synchronized Task<Boolean> insertOrUpdateCacheEntry(
+  public synchronized boolean insertOrUpdateCacheEntry(
       @NonNull CustomInstallationIdCacheEntryValue entryValue) {
     SharedPreferences.Editor editor = prefs.edit();
     editor.putString(
         getSharedPreferencesKey(CUSTOM_INSTALLATION_ID_KEY), entryValue.getCustomInstallationId());
     editor.putString(getSharedPreferencesKey(INSTANCE_ID_KEY), entryValue.getFirebaseInstanceId());
     editor.putInt(getSharedPreferencesKey(CACHE_STATUS_KEY), entryValue.getCacheStatus().ordinal());
-    return commitSharedPreferencesEditAsync(editor);
+    return editor.commit();
   }
 
   @NonNull
-  public synchronized Task<Boolean> clear() {
+  public synchronized boolean clear() {
     SharedPreferences.Editor editor = prefs.edit();
     editor.remove(getSharedPreferencesKey(CUSTOM_INSTALLATION_ID_KEY));
     editor.remove(getSharedPreferencesKey(INSTANCE_ID_KEY));
     editor.remove(getSharedPreferencesKey(CACHE_STATUS_KEY));
-    return commitSharedPreferencesEditAsync(editor);
+    return editor.commit();
   }
 
   private String getSharedPreferencesKey(String key) {
     return String.format("%s|%s", persistenceKey, key);
   }
-
-  private Task<Boolean> commitSharedPreferencesEditAsync(SharedPreferences.Editor editor) {
-    TaskCompletionSource<Boolean> result = new TaskCompletionSource<>();
-    ioExecuter.execute(() -> result.setResult(editor.commit()));
-    return result.getTask();
-  }
 }
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
index e258b84ad06..a88ecd7b662 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
@@ -15,17 +15,12 @@
 package com.google.firebase.segmentation.remote;
 
 import androidx.annotation.NonNull;
-import com.google.android.gms.tasks.Task;
-import com.google.android.gms.tasks.TaskCompletionSource;
-import com.google.android.gms.tasks.Tasks;
 import com.squareup.okhttp.MediaType;
 import com.squareup.okhttp.OkHttpClient;
 import com.squareup.okhttp.Request;
 import com.squareup.okhttp.RequestBody;
 import com.squareup.okhttp.Response;
 import java.io.IOException;
-import java.util.concurrent.Executor;
-import java.util.concurrent.Executors;
 import org.json.JSONException;
 import org.json.JSONObject;
 
@@ -43,7 +38,6 @@ public class SegmentationServiceClient {
   private static final MediaType JSON = MediaType.parse("application/json; charset=utf-8");
 
   private final OkHttpClient httpClient;
-  private final Executor httpRequestExecutor;
 
   public enum Code {
     OK,
@@ -61,11 +55,10 @@ public enum Code {
 
   public SegmentationServiceClient() {
     httpClient = new OkHttpClient();
-    httpRequestExecutor = Executors.newFixedThreadPool(4);
   }
 
   @NonNull
-  public Task<Code> updateCustomInstallationId(
+  public Code updateCustomInstallationId(
       long projectNumber,
       @NonNull String apiKey,
       @NonNull String customInstallationId,
@@ -82,7 +75,8 @@ public Task<Code> updateCustomInstallationId(
               buildUpdateCustomSegmentationDataRequestBody(resourceName, customInstallationId)
                   .toString());
     } catch (JSONException e) {
-      return Tasks.forException(e);
+      // This should never happen
+      throw new IllegalStateException(e);
     }
 
     Request request =
@@ -99,30 +93,21 @@ public Task<Code> updateCustomInstallationId(
             .patch(requestBody)
             .build();
 
-    TaskCompletionSource<Code> taskCompletionSource = new TaskCompletionSource<>();
-    httpRequestExecutor.execute(
-        () -> {
-          try {
-            Response response = httpClient.newCall(request).execute();
-            switch (response.code()) {
-              case 200:
-                taskCompletionSource.setResult(Code.OK);
-                break;
-              case 401:
-                taskCompletionSource.setResult(Code.UNAUTHORIZED);
-                break;
-              case 409:
-                taskCompletionSource.setResult(Code.CONFLICT);
-                break;
-              default:
-                taskCompletionSource.setResult(Code.SERVER_ERROR);
-                break;
-            }
-          } catch (IOException e) {
-            taskCompletionSource.setResult(Code.NETWORK_ERROR);
-          }
-        });
-    return taskCompletionSource.getTask();
+    try {
+      Response response = httpClient.newCall(request).execute();
+      switch (response.code()) {
+        case 200:
+          return Code.OK;
+        case 401:
+          return Code.UNAUTHORIZED;
+        case 409:
+          return Code.CONFLICT;
+        default:
+          return Code.SERVER_ERROR;
+      }
+    } catch (IOException e) {
+      return Code.NETWORK_ERROR;
+    }
   }
 
   private static JSONObject buildUpdateCustomSegmentationDataRequestBody(
@@ -134,7 +119,7 @@ private static JSONObject buildUpdateCustomSegmentationDataRequestBody(
   }
 
   @NonNull
-  public Task<Code> clearCustomInstallationId(
+  public Code clearCustomInstallationId(
       long projectNumber,
       @NonNull String apiKey,
       @NonNull String firebaseInstanceId,
@@ -148,7 +133,8 @@ public Task<Code> clearCustomInstallationId(
           RequestBody.create(
               JSON, buildClearCustomSegmentationDataRequestBody(resourceName).toString());
     } catch (JSONException e) {
-      return Tasks.forException(e);
+      // This should never happen
+      throw new IllegalStateException(e);
     }
 
     Request request =
@@ -165,27 +151,19 @@ public Task<Code> clearCustomInstallationId(
             .post(requestBody)
             .build();
 
-    TaskCompletionSource<Code> taskCompletionSource = new TaskCompletionSource<>();
-    httpRequestExecutor.execute(
-        () -> {
-          try {
-            Response response = httpClient.newCall(request).execute();
-            switch (response.code()) {
-              case 200:
-                taskCompletionSource.setResult(Code.OK);
-                break;
-              case 401:
-                taskCompletionSource.setResult(Code.UNAUTHORIZED);
-                break;
-              default:
-                taskCompletionSource.setResult(Code.SERVER_ERROR);
-                break;
-            }
-          } catch (IOException e) {
-            taskCompletionSource.setResult(Code.NETWORK_ERROR);
-          }
-        });
-    return taskCompletionSource.getTask();
+    try {
+      Response response = httpClient.newCall(request).execute();
+      switch (response.code()) {
+        case 200:
+          return Code.OK;
+        case 401:
+          return Code.UNAUTHORIZED;
+        default:
+          return Code.SERVER_ERROR;
+      }
+    } catch (IOException e) {
+      return Code.NETWORK_ERROR;
+    }
   }
 
   private static JSONObject buildClearCustomSegmentationDataRequestBody(String resourceName)

From 9f36d3523761ddfe9bfb15af527f0be223134048 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Tue, 9 Jul 2019 11:08:01 -0700
Subject: [PATCH 25/30] Fix a async getResult race issue.

---
 .../segmentation/FirebaseSegmentation.java    | 23 +++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
index 7a6d879d669..12b5ff52b88 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
@@ -20,6 +20,7 @@
 import com.google.android.gms.common.internal.Preconditions;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
+import com.google.android.gms.tasks.Tasks;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.iid.FirebaseInstanceId;
 import com.google.firebase.iid.InstanceIdResult;
@@ -126,7 +127,16 @@ private Task<Void> updateCustomInstallationId(String customInstallationId) {
             return;
           }
 
-          InstanceIdResult instanceIdResult = firebaseInstanceId.getInstanceId().getResult();
+          InstanceIdResult instanceIdResult;
+          try {
+            instanceIdResult = Tasks.await(firebaseInstanceId.getInstanceId());
+          } catch (Exception e) {
+            taskCompletionSource.setException(
+                new SetCustomInstallationIdException(
+                    "Failed to get Firebase instance id", Status.CLIENT_ERROR));
+            return;
+          }
+
           boolean firstUpdateCacheResult =
               localCache.insertOrUpdateCacheEntry(
                   CustomInstallationIdCacheEntryValue.create(
@@ -210,7 +220,16 @@ private Task<Void> clearCustomInstallationId() {
 
     executor.execute(
         () -> {
-          InstanceIdResult instanceIdResult = firebaseInstanceId.getInstanceId().getResult();
+          InstanceIdResult instanceIdResult;
+          try {
+            instanceIdResult = Tasks.await(firebaseInstanceId.getInstanceId());
+          } catch (Exception e) {
+            taskCompletionSource.setException(
+                new SetCustomInstallationIdException(
+                    "Failed to get Firebase instance id", Status.CLIENT_ERROR));
+            return;
+          }
+
           boolean firstUpdateCacheResult =
               localCache.insertOrUpdateCacheEntry(
                   CustomInstallationIdCacheEntryValue.create(

From 047c0af81eff97623b823400b0d9293d1a1d449f Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Wed, 10 Jul 2019 09:56:56 -0700
Subject: [PATCH 26/30] OkHttpClient -> HttpsUrlConnection

---
 .../firebase-segmentation.gradle              |   1 -
 .../remote/SegmentationServiceClient.java     | 122 ++++++++----------
 2 files changed, 55 insertions(+), 68 deletions(-)

diff --git a/firebase-segmentation/firebase-segmentation.gradle b/firebase-segmentation/firebase-segmentation.gradle
index d16e46b402b..c11085bd161 100644
--- a/firebase-segmentation/firebase-segmentation.gradle
+++ b/firebase-segmentation/firebase-segmentation.gradle
@@ -51,7 +51,6 @@ dependencies {
     implementation 'androidx.appcompat:appcompat:1.0.2'
     implementation 'androidx.multidex:multidex:2.0.1'
     implementation 'com.google.android.gms:play-services-tasks:17.0.0'
-    implementation 'com.squareup.okhttp:okhttp:2.7.5'
 
     compileOnly "com.google.auto.value:auto-value-annotations:1.6.5"
     annotationProcessor "com.google.auto.value:auto-value:1.6.2"
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
index a88ecd7b662..042afef6fc2 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
@@ -15,12 +15,10 @@
 package com.google.firebase.segmentation.remote;
 
 import androidx.annotation.NonNull;
-import com.squareup.okhttp.MediaType;
-import com.squareup.okhttp.OkHttpClient;
-import com.squareup.okhttp.Request;
-import com.squareup.okhttp.RequestBody;
-import com.squareup.okhttp.Response;
 import java.io.IOException;
+import java.io.OutputStream;
+import java.net.URL;
+import javax.net.ssl.HttpsURLConnection;
 import org.json.JSONException;
 import org.json.JSONObject;
 
@@ -35,10 +33,6 @@ public class SegmentationServiceClient {
       "projects/%s/installations/%s/customSegmentationData:clear";
   private static final String FIREBASE_SEGMENTATION_API_VERSION = "v1alpha";
 
-  private static final MediaType JSON = MediaType.parse("application/json; charset=utf-8");
-
-  private final OkHttpClient httpClient;
-
   public enum Code {
     OK,
 
@@ -53,10 +47,6 @@ public enum Code {
     UNAUTHORIZED,
   }
 
-  public SegmentationServiceClient() {
-    httpClient = new OkHttpClient();
-  }
-
   @NonNull
   public Code updateCustomInstallationId(
       long projectNumber,
@@ -66,36 +56,35 @@ public Code updateCustomInstallationId(
       @NonNull String firebaseInstanceIdToken) {
     String resourceName =
         String.format(UPDATE_REQUEST_RESOURCE_NAME_FORMAT, projectNumber, firebaseInstanceId);
-
-    RequestBody requestBody;
     try {
-      requestBody =
-          RequestBody.create(
-              JSON,
-              buildUpdateCustomSegmentationDataRequestBody(resourceName, customInstallationId)
-                  .toString());
-    } catch (JSONException e) {
-      // This should never happen
-      throw new IllegalStateException(e);
-    }
-
-    Request request =
-        new Request.Builder()
-            .url(
-                String.format(
-                    "https://%s/%s/%s?key=%s",
-                    FIREBASE_SEGMENTATION_API_DOMAIN,
-                    FIREBASE_SEGMENTATION_API_VERSION,
-                    resourceName,
-                    apiKey))
-            .header("Authorization", "FIREBASE_INSTALLATIONS_AUTH " + firebaseInstanceIdToken)
-            .header("Content-Type", "application/json")
-            .patch(requestBody)
-            .build();
+      URL url =
+          new URL(
+              String.format(
+                  "https://%s/%s/%s?key=%s",
+                  FIREBASE_SEGMENTATION_API_DOMAIN,
+                  FIREBASE_SEGMENTATION_API_VERSION,
+                  resourceName,
+                  apiKey));
+
+      HttpsURLConnection httpsURLConnection = (HttpsURLConnection) url.openConnection();
+      httpsURLConnection.setDoOutput(true);
+      httpsURLConnection.setRequestMethod("PATCH");
+      httpsURLConnection.addRequestProperty(
+          "Authorization", "FIREBASE_INSTALLATIONS_AUTH " + firebaseInstanceIdToken);
+      httpsURLConnection.addRequestProperty("Content-Type", "application/json");
+      OutputStream os = httpsURLConnection.getOutputStream();
+      try {
+        os.write(
+            buildUpdateCustomSegmentationDataRequestBody(resourceName, customInstallationId)
+                .toString()
+                .getBytes("UTF-8"));
+      } catch (JSONException e) {
+        throw new IllegalStateException(e);
+      }
+      httpsURLConnection.connect();
 
-    try {
-      Response response = httpClient.newCall(request).execute();
-      switch (response.code()) {
+      int httpResponseCode = httpsURLConnection.getResponseCode();
+      switch (httpResponseCode) {
         case 200:
           return Code.OK;
         case 401:
@@ -126,34 +115,33 @@ public Code clearCustomInstallationId(
       @NonNull String firebaseInstanceIdToken) {
     String resourceName =
         String.format(CLEAR_REQUEST_RESOURCE_NAME_FORMAT, projectNumber, firebaseInstanceId);
-
-    RequestBody requestBody;
     try {
-      requestBody =
-          RequestBody.create(
-              JSON, buildClearCustomSegmentationDataRequestBody(resourceName).toString());
-    } catch (JSONException e) {
-      // This should never happen
-      throw new IllegalStateException(e);
-    }
-
-    Request request =
-        new Request.Builder()
-            .url(
-                String.format(
-                    "https://%s/%s/%s?key=%s",
-                    FIREBASE_SEGMENTATION_API_DOMAIN,
-                    FIREBASE_SEGMENTATION_API_VERSION,
-                    resourceName,
-                    apiKey))
-            .header("Authorization", "FIREBASE_INSTALLATIONS_AUTH " + firebaseInstanceIdToken)
-            .header("Content-Type", "application/json")
-            .post(requestBody)
-            .build();
+      URL url =
+          new URL(
+              String.format(
+                  "https://%s/%s/%s?key=%s",
+                  FIREBASE_SEGMENTATION_API_DOMAIN,
+                  FIREBASE_SEGMENTATION_API_VERSION,
+                  resourceName,
+                  apiKey));
+
+      HttpsURLConnection httpsURLConnection = (HttpsURLConnection) url.openConnection();
+      httpsURLConnection.setDoOutput(true);
+      httpsURLConnection.setRequestMethod("POST");
+      httpsURLConnection.addRequestProperty(
+          "Authorization", "FIREBASE_INSTALLATIONS_AUTH " + firebaseInstanceIdToken);
+      httpsURLConnection.addRequestProperty("Content-Type", "application/json");
+      OutputStream os = httpsURLConnection.getOutputStream();
+      try {
+        os.write(
+            buildClearCustomSegmentationDataRequestBody(resourceName).toString().getBytes("UTF-8"));
+      } catch (JSONException e) {
+        throw new IllegalStateException(e);
+      }
+      httpsURLConnection.connect();
 
-    try {
-      Response response = httpClient.newCall(request).execute();
-      switch (response.code()) {
+      int httpResponseCode = httpsURLConnection.getResponseCode();
+      switch (httpResponseCode) {
         case 200:
           return Code.OK;
         case 401:

From 8b39c3170189b33d92c3646f89cb0dd62de713a0 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Thu, 18 Jul 2019 12:04:22 -0700
Subject: [PATCH 27/30] Use gzip for compressing content and fix ourput stream
 memory leak risk.

---
 .../remote/SegmentationServiceClient.java     | 28 +++++++++++++------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
index 042afef6fc2..03956abd5e8 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
@@ -16,8 +16,8 @@
 
 import androidx.annotation.NonNull;
 import java.io.IOException;
-import java.io.OutputStream;
 import java.net.URL;
+import java.util.zip.GZIPOutputStream;
 import javax.net.ssl.HttpsURLConnection;
 import org.json.JSONException;
 import org.json.JSONObject;
@@ -33,6 +33,11 @@ public class SegmentationServiceClient {
       "projects/%s/installations/%s/customSegmentationData:clear";
   private static final String FIREBASE_SEGMENTATION_API_VERSION = "v1alpha";
 
+  private static final String CONTENT_TYPE_HEADER_KEY = "Content-Type";
+  private static final String JSON_CONTENT_TYPE = "application/json";
+  private static final String CONTENT_ENCODING_HEADER_KEY = "Content-Encoding";
+  private static final String GZIP_CONTENT_ENCODING = "gzip";
+
   public enum Code {
     OK,
 
@@ -71,18 +76,21 @@ public Code updateCustomInstallationId(
       httpsURLConnection.setRequestMethod("PATCH");
       httpsURLConnection.addRequestProperty(
           "Authorization", "FIREBASE_INSTALLATIONS_AUTH " + firebaseInstanceIdToken);
-      httpsURLConnection.addRequestProperty("Content-Type", "application/json");
-      OutputStream os = httpsURLConnection.getOutputStream();
+      httpsURLConnection.addRequestProperty(CONTENT_TYPE_HEADER_KEY, JSON_CONTENT_TYPE);
+      httpsURLConnection.addRequestProperty(CONTENT_ENCODING_HEADER_KEY, GZIP_CONTENT_ENCODING);
+      GZIPOutputStream gzipOutputStream =
+          new GZIPOutputStream(httpsURLConnection.getOutputStream());
       try {
-        os.write(
+        gzipOutputStream.write(
             buildUpdateCustomSegmentationDataRequestBody(resourceName, customInstallationId)
                 .toString()
                 .getBytes("UTF-8"));
       } catch (JSONException e) {
         throw new IllegalStateException(e);
+      } finally {
+        gzipOutputStream.close();
       }
       httpsURLConnection.connect();
-
       int httpResponseCode = httpsURLConnection.getResponseCode();
       switch (httpResponseCode) {
         case 200:
@@ -130,13 +138,17 @@ public Code clearCustomInstallationId(
       httpsURLConnection.setRequestMethod("POST");
       httpsURLConnection.addRequestProperty(
           "Authorization", "FIREBASE_INSTALLATIONS_AUTH " + firebaseInstanceIdToken);
-      httpsURLConnection.addRequestProperty("Content-Type", "application/json");
-      OutputStream os = httpsURLConnection.getOutputStream();
+      httpsURLConnection.addRequestProperty(CONTENT_TYPE_HEADER_KEY, JSON_CONTENT_TYPE);
+      httpsURLConnection.addRequestProperty(CONTENT_ENCODING_HEADER_KEY, GZIP_CONTENT_ENCODING);
+      GZIPOutputStream gzipOutputStream =
+          new GZIPOutputStream(httpsURLConnection.getOutputStream());
       try {
-        os.write(
+        gzipOutputStream.write(
             buildClearCustomSegmentationDataRequestBody(resourceName).toString().getBytes("UTF-8"));
       } catch (JSONException e) {
         throw new IllegalStateException(e);
+      } finally {
+        gzipOutputStream.close();
       }
       httpsURLConnection.connect();
 

From dc1a63c5f7af6a383fbc6c164bb3745c8e1fbbaf Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Mon, 22 Jul 2019 11:42:46 -0700
Subject: [PATCH 28/30] Addressed a few comments

---
 .../segmentation/FirebaseSegmentation.java    | 260 ++++++++----------
 .../remote/SegmentationServiceClient.java     |   3 +-
 2 files changed, 114 insertions(+), 149 deletions(-)

diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
index 12b5ff52b88..2008ea09359 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
@@ -16,10 +16,9 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
-import androidx.annotation.RestrictTo;
+import androidx.annotation.WorkerThread;
 import com.google.android.gms.common.internal.Preconditions;
 import com.google.android.gms.tasks.Task;
-import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.android.gms.tasks.Tasks;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.iid.FirebaseInstanceId;
@@ -29,6 +28,7 @@
 import com.google.firebase.segmentation.local.CustomInstallationIdCacheEntryValue;
 import com.google.firebase.segmentation.remote.SegmentationServiceClient;
 import com.google.firebase.segmentation.remote.SegmentationServiceClient.Code;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 
@@ -42,14 +42,13 @@ public class FirebaseSegmentation {
   private final Executor executor;
 
   FirebaseSegmentation(FirebaseApp firebaseApp) {
-    this.firebaseApp = firebaseApp;
-    this.firebaseInstanceId = FirebaseInstanceId.getInstance(firebaseApp);
-    this.localCache = new CustomInstallationIdCache(firebaseApp);
-    this.backendServiceClient = new SegmentationServiceClient();
-    this.executor = Executors.newFixedThreadPool(4);
+    this(
+        firebaseApp,
+        FirebaseInstanceId.getInstance(firebaseApp),
+        new CustomInstallationIdCache(firebaseApp),
+        new SegmentationServiceClient());
   }
 
-  @RestrictTo(RestrictTo.Scope.TESTS)
   FirebaseSegmentation(
       FirebaseApp firebaseApp,
       FirebaseInstanceId firebaseInstanceId,
@@ -88,9 +87,9 @@ public static FirebaseSegmentation getInstance(@NonNull FirebaseApp app) {
   @NonNull
   public synchronized Task<Void> setCustomInstallationId(@Nullable String customInstallationId) {
     if (customInstallationId == null) {
-      return clearCustomInstallationId();
+      return Tasks.call(executor, () -> clearCustomInstallationId());
     }
-    return updateCustomInstallationId(customInstallationId);
+    return Tasks.call(executor, () -> updateCustomInstallationId(customInstallationId));
   }
 
   /**
@@ -112,89 +111,73 @@ public synchronized Task<Void> setCustomInstallationId(@Nullable String customIn
    *                               return
    * </pre>
    */
-  private Task<Void> updateCustomInstallationId(String customInstallationId) {
-    TaskCompletionSource<Void> taskCompletionSource = new TaskCompletionSource<>();
-
-    executor.execute(
-        () -> {
-          CustomInstallationIdCacheEntryValue cacheEntryValue = localCache.readCacheEntryValue();
-          if (cacheEntryValue != null
-              && cacheEntryValue.getCustomInstallationId().equals(customInstallationId)
-              && cacheEntryValue.getCacheStatus() == CustomInstallationIdCache.CacheStatus.SYNCED) {
-            // If the given custom installation id matches up the cached
-            // value, there's no need to update.
-            taskCompletionSource.setResult(null);
-            return;
-          }
-
-          InstanceIdResult instanceIdResult;
-          try {
-            instanceIdResult = Tasks.await(firebaseInstanceId.getInstanceId());
-          } catch (Exception e) {
-            taskCompletionSource.setException(
-                new SetCustomInstallationIdException(
-                    "Failed to get Firebase instance id", Status.CLIENT_ERROR));
-            return;
-          }
+  @WorkerThread
+  private Void updateCustomInstallationId(String customInstallationId)
+      throws SetCustomInstallationIdException {
+    CustomInstallationIdCacheEntryValue cacheEntryValue = localCache.readCacheEntryValue();
+    if (cacheEntryValue != null
+        && cacheEntryValue.getCustomInstallationId().equals(customInstallationId)
+        && cacheEntryValue.getCacheStatus() == CustomInstallationIdCache.CacheStatus.SYNCED) {
+      // If the given custom installation id matches up the cached
+      // value, there's no need to update.
+      return null;
+    }
 
-          boolean firstUpdateCacheResult =
-              localCache.insertOrUpdateCacheEntry(
-                  CustomInstallationIdCacheEntryValue.create(
-                      customInstallationId,
-                      instanceIdResult.getId(),
-                      CustomInstallationIdCache.CacheStatus.PENDING_UPDATE));
+    InstanceIdResult instanceIdResult;
+    try {
+      instanceIdResult = Tasks.await(firebaseInstanceId.getInstanceId());
+    } catch (ExecutionException | InterruptedException e) {
+      throw new SetCustomInstallationIdException(
+          "Failed to get Firebase instance id", Status.CLIENT_ERROR);
+    }
 
-          if (!firstUpdateCacheResult) {
-            taskCompletionSource.setException(
-                new SetCustomInstallationIdException(
-                    "Failed to update client side cache", Status.CLIENT_ERROR));
-            return;
-          }
+    boolean firstUpdateCacheResult =
+        localCache.insertOrUpdateCacheEntry(
+            CustomInstallationIdCacheEntryValue.create(
+                customInstallationId,
+                instanceIdResult.getId(),
+                CustomInstallationIdCache.CacheStatus.PENDING_UPDATE));
 
-          // Start requesting backend when first cache updae is done.
-          String iid = instanceIdResult.getId();
-          String iidToken = instanceIdResult.getToken();
-          Code backendRequestResult =
-              backendServiceClient.updateCustomInstallationId(
-                  Utils.getProjectNumberFromAppId(firebaseApp.getOptions().getApplicationId()),
-                  firebaseApp.getOptions().getApiKey(),
-                  customInstallationId,
-                  iid,
-                  iidToken);
+    if (!firstUpdateCacheResult) {
+      throw new SetCustomInstallationIdException(
+          "Failed to update client side cache", Status.CLIENT_ERROR);
+    }
 
-          boolean finalUpdateCacheResult;
-          switch (backendRequestResult) {
-            case OK:
-              finalUpdateCacheResult =
-                  localCache.insertOrUpdateCacheEntry(
-                      CustomInstallationIdCacheEntryValue.create(
-                          customInstallationId,
-                          instanceIdResult.getId(),
-                          CustomInstallationIdCache.CacheStatus.SYNCED));
-              break;
-            case HTTP_CLIENT_ERROR:
-              taskCompletionSource.setException(
-                  new SetCustomInstallationIdException(Status.CLIENT_ERROR));
-              return;
-            case CONFLICT:
-              taskCompletionSource.setException(
-                  new SetCustomInstallationIdException(Status.DUPLICATED_CUSTOM_INSTALLATION_ID));
-              return;
-            default:
-              taskCompletionSource.setException(
-                  new SetCustomInstallationIdException(Status.BACKEND_ERROR));
-              return;
-          }
+    // Start requesting backend when first cache updae is done.
+    String iid = instanceIdResult.getId();
+    String iidToken = instanceIdResult.getToken();
+    Code backendRequestResult =
+        backendServiceClient.updateCustomInstallationId(
+            Utils.getProjectNumberFromAppId(firebaseApp.getOptions().getApplicationId()),
+            firebaseApp.getOptions().getApiKey(),
+            customInstallationId,
+            iid,
+            iidToken);
+
+    boolean finalUpdateCacheResult;
+    switch (backendRequestResult) {
+      case OK:
+        finalUpdateCacheResult =
+            localCache.insertOrUpdateCacheEntry(
+                CustomInstallationIdCacheEntryValue.create(
+                    customInstallationId,
+                    instanceIdResult.getId(),
+                    CustomInstallationIdCache.CacheStatus.SYNCED));
+        break;
+      case HTTP_CLIENT_ERROR:
+        throw new SetCustomInstallationIdException(Status.CLIENT_ERROR);
+      case CONFLICT:
+        throw new SetCustomInstallationIdException(Status.DUPLICATED_CUSTOM_INSTALLATION_ID);
+      default:
+        throw new SetCustomInstallationIdException(Status.BACKEND_ERROR);
+    }
 
-          if (finalUpdateCacheResult) {
-            taskCompletionSource.setResult(null);
-          } else {
-            taskCompletionSource.setException(
-                new SetCustomInstallationIdException(
-                    "Failed to update client side cache", Status.CLIENT_ERROR));
-          }
-        });
-    return taskCompletionSource.getTask();
+    if (finalUpdateCacheResult) {
+      return null;
+    } else {
+      throw new SetCustomInstallationIdException(
+          "Failed to update client side cache", Status.CLIENT_ERROR);
+    }
   }
 
   /**
@@ -214,68 +197,51 @@ private Task<Void> updateCustomInstallationId(String customInstallationId) {
    *                               return
    * </pre>
    */
-  private Task<Void> clearCustomInstallationId() {
-
-    TaskCompletionSource<Void> taskCompletionSource = new TaskCompletionSource<>();
-
-    executor.execute(
-        () -> {
-          InstanceIdResult instanceIdResult;
-          try {
-            instanceIdResult = Tasks.await(firebaseInstanceId.getInstanceId());
-          } catch (Exception e) {
-            taskCompletionSource.setException(
-                new SetCustomInstallationIdException(
-                    "Failed to get Firebase instance id", Status.CLIENT_ERROR));
-            return;
-          }
-
-          boolean firstUpdateCacheResult =
-              localCache.insertOrUpdateCacheEntry(
-                  CustomInstallationIdCacheEntryValue.create(
-                      "",
-                      instanceIdResult.getId(),
-                      CustomInstallationIdCache.CacheStatus.PENDING_CLEAR));
+  @WorkerThread
+  private Void clearCustomInstallationId() throws SetCustomInstallationIdException {
+    InstanceIdResult instanceIdResult;
+    try {
+      instanceIdResult = Tasks.await(firebaseInstanceId.getInstanceId());
+    } catch (ExecutionException | InterruptedException e) {
+      throw new SetCustomInstallationIdException(
+          "Failed to get Firebase instance id", Status.CLIENT_ERROR);
+    }
 
-          if (!firstUpdateCacheResult) {
-            taskCompletionSource.setException(
-                new SetCustomInstallationIdException(
-                    "Failed to update client side cache", Status.CLIENT_ERROR));
-            return;
-          }
+    boolean firstUpdateCacheResult =
+        localCache.insertOrUpdateCacheEntry(
+            CustomInstallationIdCacheEntryValue.create(
+                "", instanceIdResult.getId(), CustomInstallationIdCache.CacheStatus.PENDING_CLEAR));
 
-          String iid = instanceIdResult.getId();
-          String iidToken = instanceIdResult.getToken();
-          Code backendRequestResult =
-              backendServiceClient.clearCustomInstallationId(
-                  Utils.getProjectNumberFromAppId(firebaseApp.getOptions().getApplicationId()),
-                  firebaseApp.getOptions().getApiKey(),
-                  iid,
-                  iidToken);
+    if (!firstUpdateCacheResult) {
+      throw new SetCustomInstallationIdException(
+          "Failed to update client side cache", Status.CLIENT_ERROR);
+    }
 
-          boolean finalUpdateCacheResult;
-          switch (backendRequestResult) {
-            case OK:
-              finalUpdateCacheResult = localCache.clear();
-              break;
-            case HTTP_CLIENT_ERROR:
-              taskCompletionSource.setException(
-                  new SetCustomInstallationIdException(Status.CLIENT_ERROR));
-              return;
-            default:
-              taskCompletionSource.setException(
-                  new SetCustomInstallationIdException(Status.BACKEND_ERROR));
-              return;
-          }
+    String iid = instanceIdResult.getId();
+    String iidToken = instanceIdResult.getToken();
+    Code backendRequestResult =
+        backendServiceClient.clearCustomInstallationId(
+            Utils.getProjectNumberFromAppId(firebaseApp.getOptions().getApplicationId()),
+            firebaseApp.getOptions().getApiKey(),
+            iid,
+            iidToken);
+
+    boolean finalUpdateCacheResult;
+    switch (backendRequestResult) {
+      case OK:
+        finalUpdateCacheResult = localCache.clear();
+        break;
+      case HTTP_CLIENT_ERROR:
+        throw new SetCustomInstallationIdException(Status.CLIENT_ERROR);
+      default:
+        throw new SetCustomInstallationIdException(Status.BACKEND_ERROR);
+    }
 
-          if (finalUpdateCacheResult) {
-            taskCompletionSource.setResult(null);
-          } else {
-            taskCompletionSource.setException(
-                new SetCustomInstallationIdException(
-                    "Failed to update client side cache", Status.CLIENT_ERROR));
-          }
-        });
-    return taskCompletionSource.getTask();
+    if (finalUpdateCacheResult) {
+      return null;
+    } else {
+      throw new SetCustomInstallationIdException(
+          "Failed to update client side cache", Status.CLIENT_ERROR);
+    }
   }
 }
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
index 03956abd5e8..86b5f945205 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
@@ -90,7 +90,7 @@ public Code updateCustomInstallationId(
       } finally {
         gzipOutputStream.close();
       }
-      httpsURLConnection.connect();
+
       int httpResponseCode = httpsURLConnection.getResponseCode();
       switch (httpResponseCode) {
         case 200:
@@ -150,7 +150,6 @@ public Code clearCustomInstallationId(
       } finally {
         gzipOutputStream.close();
       }
-      httpsURLConnection.connect();
 
       int httpResponseCode = httpsURLConnection.getResponseCode();
       switch (httpResponseCode) {

From 143ed7451b48dbbf8ca510017bca201546f30bde Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Tue, 30 Jul 2019 11:10:48 -0700
Subject: [PATCH 29/30] FirebaseSegmentation SDK 1. Clean up http client
 response code. 2. When updateCustomInstallationId is called, on non-retryable
 server errors, the SDK should clean up the local cache. Instead, for
 retryable errors, SDK can keep the local cache for retrying update later.

---
 .../FirebaseSegmentationInstrumentedTest.java | 28 +++++++++++++-
 .../segmentation/FirebaseSegmentation.java    | 37 ++++++++++++++-----
 .../SetCustomInstallationIdException.java     |  4 +-
 .../remote/SegmentationServiceClient.java     | 14 +++++--
 4 files changed, 66 insertions(+), 17 deletions(-)

diff --git a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
index 19498782a78..034ee1eb339 100644
--- a/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
+++ b/firebase-segmentation/src/androidTest/java/com/google/firebase/segmentation/FirebaseSegmentationInstrumentedTest.java
@@ -127,7 +127,8 @@ public void testUpdateCustomInstallationId_CacheOk_BackendOk() throws Exception
   }
 
   @Test
-  public void testUpdateCustomInstallationId_CacheOk_BackendError() throws InterruptedException {
+  public void testUpdateCustomInstallationId_CacheOk_BackendError_Retryable()
+      throws InterruptedException {
     FirebaseSegmentation firebaseSegmentation =
         new FirebaseSegmentation(
             firebaseApp, firebaseInstanceId, actualCache, backendClientReturnsError);
@@ -150,6 +151,31 @@ public void testUpdateCustomInstallationId_CacheOk_BackendError() throws Interru
         .isEqualTo(CustomInstallationIdCache.CacheStatus.PENDING_UPDATE);
   }
 
+  @Test
+  public void testUpdateCustomInstallationId_CacheOk_BackendError_NotRetryable()
+      throws InterruptedException {
+    when(backendClientReturnsError.updateCustomInstallationId(
+            anyLong(), anyString(), anyString(), anyString(), anyString()))
+        .thenReturn(SegmentationServiceClient.Code.CONFLICT);
+    FirebaseSegmentation firebaseSegmentation =
+        new FirebaseSegmentation(
+            firebaseApp, firebaseInstanceId, actualCache, backendClientReturnsError);
+
+    // Expect exception
+    try {
+      Tasks.await(firebaseSegmentation.setCustomInstallationId(CUSTOM_INSTALLATION_ID));
+      fail();
+    } catch (ExecutionException expected) {
+      Throwable cause = expected.getCause();
+      assertThat(cause).isInstanceOf(SetCustomInstallationIdException.class);
+      assertThat(((SetCustomInstallationIdException) cause).getStatus())
+          .isEqualTo(SetCustomInstallationIdException.Status.DUPLICATED_CUSTOM_INSTALLATION_ID);
+    }
+
+    CustomInstallationIdCacheEntryValue entryValue = actualCache.readCacheEntryValue();
+    assertThat(entryValue).isNull();
+  }
+
   @Test
   public void testUpdateCustomInstallationId_CacheError_BackendOk() throws InterruptedException {
     FirebaseSegmentation firebaseSegmentation =
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
index 2008ea09359..85adf2fb708 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
@@ -128,7 +128,7 @@ private Void updateCustomInstallationId(String customInstallationId)
       instanceIdResult = Tasks.await(firebaseInstanceId.getInstanceId());
     } catch (ExecutionException | InterruptedException e) {
       throw new SetCustomInstallationIdException(
-          "Failed to get Firebase instance id", Status.CLIENT_ERROR);
+          Status.CLIENT_ERROR, "Failed to get Firebase instance id");
     }
 
     boolean firstUpdateCacheResult =
@@ -140,7 +140,7 @@ private Void updateCustomInstallationId(String customInstallationId)
 
     if (!firstUpdateCacheResult) {
       throw new SetCustomInstallationIdException(
-          "Failed to update client side cache", Status.CLIENT_ERROR);
+          Status.CLIENT_ERROR, "Failed to update client side cache");
     }
 
     // Start requesting backend when first cache updae is done.
@@ -164,11 +164,22 @@ private Void updateCustomInstallationId(String customInstallationId)
                     instanceIdResult.getId(),
                     CustomInstallationIdCache.CacheStatus.SYNCED));
         break;
-      case HTTP_CLIENT_ERROR:
-        throw new SetCustomInstallationIdException(Status.CLIENT_ERROR);
+      case UNAUTHORIZED:
+        localCache.clear();
+        throw new SetCustomInstallationIdException(
+            Status.CLIENT_ERROR, "Instance id token is invalid.");
       case CONFLICT:
-        throw new SetCustomInstallationIdException(Status.DUPLICATED_CUSTOM_INSTALLATION_ID);
+        localCache.clear();
+        throw new SetCustomInstallationIdException(
+            Status.DUPLICATED_CUSTOM_INSTALLATION_ID,
+            "The custom installation id is used by another Firebase installation in your project.");
+      case HTTP_CLIENT_ERROR:
+        localCache.clear();
+        throw new SetCustomInstallationIdException(Status.CLIENT_ERROR, "Http client error(4xx)");
+      case NETWORK_ERROR:
+      case SERVER_ERROR:
       default:
+        // These are considered retryable errors, so not to clean up the cache.
         throw new SetCustomInstallationIdException(Status.BACKEND_ERROR);
     }
 
@@ -176,7 +187,7 @@ private Void updateCustomInstallationId(String customInstallationId)
       return null;
     } else {
       throw new SetCustomInstallationIdException(
-          "Failed to update client side cache", Status.CLIENT_ERROR);
+          Status.CLIENT_ERROR, "Failed to update client side cache");
     }
   }
 
@@ -204,7 +215,7 @@ private Void clearCustomInstallationId() throws SetCustomInstallationIdException
       instanceIdResult = Tasks.await(firebaseInstanceId.getInstanceId());
     } catch (ExecutionException | InterruptedException e) {
       throw new SetCustomInstallationIdException(
-          "Failed to get Firebase instance id", Status.CLIENT_ERROR);
+          Status.CLIENT_ERROR, "Failed to get Firebase instance id");
     }
 
     boolean firstUpdateCacheResult =
@@ -214,7 +225,7 @@ private Void clearCustomInstallationId() throws SetCustomInstallationIdException
 
     if (!firstUpdateCacheResult) {
       throw new SetCustomInstallationIdException(
-          "Failed to update client side cache", Status.CLIENT_ERROR);
+          Status.CLIENT_ERROR, "Failed to update client side cache");
     }
 
     String iid = instanceIdResult.getId();
@@ -231,9 +242,15 @@ private Void clearCustomInstallationId() throws SetCustomInstallationIdException
       case OK:
         finalUpdateCacheResult = localCache.clear();
         break;
+      case UNAUTHORIZED:
+        throw new SetCustomInstallationIdException(
+            Status.CLIENT_ERROR, "Instance id token is invalid.");
       case HTTP_CLIENT_ERROR:
-        throw new SetCustomInstallationIdException(Status.CLIENT_ERROR);
+        throw new SetCustomInstallationIdException(Status.CLIENT_ERROR, "Http client error(4xx)");
+      case NETWORK_ERROR:
+      case SERVER_ERROR:
       default:
+        // These are considered retryable errors, so not to clean up the cache.
         throw new SetCustomInstallationIdException(Status.BACKEND_ERROR);
     }
 
@@ -241,7 +258,7 @@ private Void clearCustomInstallationId() throws SetCustomInstallationIdException
       return null;
     } else {
       throw new SetCustomInstallationIdException(
-          "Failed to update client side cache", Status.CLIENT_ERROR);
+          Status.CLIENT_ERROR, "Failed to update client side cache");
     }
   }
 }
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/SetCustomInstallationIdException.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/SetCustomInstallationIdException.java
index 3c957ce3294..2291c078a03 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/SetCustomInstallationIdException.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/SetCustomInstallationIdException.java
@@ -45,13 +45,13 @@ public enum Status {
     this.status = status;
   }
 
-  SetCustomInstallationIdException(@NonNull String message, @NonNull Status status) {
+  SetCustomInstallationIdException(@NonNull Status status, @NonNull String message) {
     super(message);
     this.status = status;
   }
 
   SetCustomInstallationIdException(
-      @NonNull String message, @NonNull Status status, Throwable cause) {
+      @NonNull Status status, @NonNull String message, Throwable cause) {
     super(message, cause);
     this.status = status;
   }
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
index 86b5f945205..6d0436d8b24 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
@@ -41,15 +41,15 @@ public class SegmentationServiceClient {
   public enum Code {
     OK,
 
-    HTTP_CLIENT_ERROR,
-
     CONFLICT,
 
+    UNAUTHORIZED,
+
     NETWORK_ERROR,
 
-    SERVER_ERROR,
+    HTTP_CLIENT_ERROR,
 
-    UNAUTHORIZED,
+    SERVER_ERROR,
   }
 
   @NonNull
@@ -100,6 +100,9 @@ public Code updateCustomInstallationId(
         case 409:
           return Code.CONFLICT;
         default:
+          if (httpResponseCode / 100 == 4) {
+            return Code.HTTP_CLIENT_ERROR;
+          }
           return Code.SERVER_ERROR;
       }
     } catch (IOException e) {
@@ -158,6 +161,9 @@ public Code clearCustomInstallationId(
         case 401:
           return Code.UNAUTHORIZED;
         default:
+          if (httpResponseCode / 100 == 4) {
+            return Code.HTTP_CLIENT_ERROR;
+          }
           return Code.SERVER_ERROR;
       }
     } catch (IOException e) {

From d467379ab17f48719cfbafe4eafbbf6d86af8fc2 Mon Sep 17 00:00:00 2001
From: Di Wu <arete@google.com>
Date: Tue, 6 Aug 2019 11:22:26 -0700
Subject: [PATCH 30/30] Restrict Firebase API key to Android app package name.

---
 .../segmentation/FirebaseSegmentation.java    |  4 +-
 .../remote/SegmentationServiceClient.java     | 41 +++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
index 85adf2fb708..b1d5e6e6742 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/FirebaseSegmentation.java
@@ -35,6 +35,8 @@
 /** Entry point of Firebase Segmentation SDK. */
 public class FirebaseSegmentation {
 
+  public static final String TAG = "FirebaseSegmentation";
+
   private final FirebaseApp firebaseApp;
   private final FirebaseInstanceId firebaseInstanceId;
   private final CustomInstallationIdCache localCache;
@@ -46,7 +48,7 @@ public class FirebaseSegmentation {
         firebaseApp,
         FirebaseInstanceId.getInstance(firebaseApp),
         new CustomInstallationIdCache(firebaseApp),
-        new SegmentationServiceClient());
+        new SegmentationServiceClient(firebaseApp.getApplicationContext()));
   }
 
   FirebaseSegmentation(
diff --git a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
index 6d0436d8b24..65058bb97da 100644
--- a/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
+++ b/firebase-segmentation/src/main/java/com/google/firebase/segmentation/remote/SegmentationServiceClient.java
@@ -14,7 +14,14 @@
 
 package com.google.firebase.segmentation.remote;
 
+import static com.google.firebase.segmentation.FirebaseSegmentation.TAG;
+
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.util.Log;
 import androidx.annotation.NonNull;
+import com.google.android.gms.common.util.AndroidUtilsLight;
+import com.google.android.gms.common.util.Hex;
 import java.io.IOException;
 import java.net.URL;
 import java.util.zip.GZIPOutputStream;
@@ -37,6 +44,14 @@ public class SegmentationServiceClient {
   private static final String JSON_CONTENT_TYPE = "application/json";
   private static final String CONTENT_ENCODING_HEADER_KEY = "Content-Encoding";
   private static final String GZIP_CONTENT_ENCODING = "gzip";
+  private static final String X_ANDROID_PACKAGE_HEADER_KEY = "X-Android-Package";
+  private static final String X_ANDROID_CERT_HEADER_KEY = "X-Android-Cert";
+
+  private final Context context;
+
+  public SegmentationServiceClient(@NonNull Context context) {
+    this.context = context;
+  }
 
   public enum Code {
     OK,
@@ -78,6 +93,9 @@ public Code updateCustomInstallationId(
           "Authorization", "FIREBASE_INSTALLATIONS_AUTH " + firebaseInstanceIdToken);
       httpsURLConnection.addRequestProperty(CONTENT_TYPE_HEADER_KEY, JSON_CONTENT_TYPE);
       httpsURLConnection.addRequestProperty(CONTENT_ENCODING_HEADER_KEY, GZIP_CONTENT_ENCODING);
+      httpsURLConnection.addRequestProperty(X_ANDROID_PACKAGE_HEADER_KEY, context.getPackageName());
+      httpsURLConnection.addRequestProperty(
+          X_ANDROID_CERT_HEADER_KEY, getFingerprintHashForPackage());
       GZIPOutputStream gzipOutputStream =
           new GZIPOutputStream(httpsURLConnection.getOutputStream());
       try {
@@ -143,6 +161,9 @@ public Code clearCustomInstallationId(
           "Authorization", "FIREBASE_INSTALLATIONS_AUTH " + firebaseInstanceIdToken);
       httpsURLConnection.addRequestProperty(CONTENT_TYPE_HEADER_KEY, JSON_CONTENT_TYPE);
       httpsURLConnection.addRequestProperty(CONTENT_ENCODING_HEADER_KEY, GZIP_CONTENT_ENCODING);
+      httpsURLConnection.addRequestProperty(X_ANDROID_PACKAGE_HEADER_KEY, context.getPackageName());
+      httpsURLConnection.addRequestProperty(
+          X_ANDROID_CERT_HEADER_KEY, getFingerprintHashForPackage());
       GZIPOutputStream gzipOutputStream =
           new GZIPOutputStream(httpsURLConnection.getOutputStream());
       try {
@@ -175,4 +196,24 @@ private static JSONObject buildClearCustomSegmentationDataRequestBody(String res
       throws JSONException {
     return new JSONObject().put("name", resourceName);
   }
+
+  /** Gets the Android package's SHA-1 fingerprint. */
+  private String getFingerprintHashForPackage() {
+    byte[] hash;
+
+    try {
+      hash = AndroidUtilsLight.getPackageCertificateHashBytes(context, context.getPackageName());
+
+      if (hash == null) {
+        Log.e(TAG, "Could not get fingerprint hash for package: " + context.getPackageName());
+        return null;
+      } else {
+        String cert = Hex.bytesToStringUppercase(hash, /* zeroTerminated= */ false);
+        return Hex.bytesToStringUppercase(hash, /* zeroTerminated= */ false);
+      }
+    } catch (PackageManager.NameNotFoundException e) {
+      Log.e(TAG, "No such package: " + context.getPackageName(), e);
+      return null;
+    }
+  }
 }