Introduce CDM AssociationStore

Add public (within system_server) AssociationStore interface that
provides APIs for retreiving and monitoring changes to the CDM
associations.
Add AssociationStoreImpl - implementation of the AssociationStore which
additionally adds adds methods for adding, removing and updating
associations to be used by CdmService.

Bug: 211398735
Test: atest CtsCompanionDeviceManagerCoreTestCases
Test: atest CtsCompanionDeviceManagerUiAutomationTestCases
Change-Id: I82bc948c9949c3034b539530ef9641baf220601c
diff --git a/core/java/android/companion/AssociationInfo.java b/core/java/android/companion/AssociationInfo.java
index 3f02aa2..78d5137 100644
--- a/core/java/android/companion/AssociationInfo.java
+++ b/core/java/android/companion/AssociationInfo.java
@@ -197,6 +197,20 @@
         return macAddress.equals(mDeviceMacAddress);
     }
 
+    /** @hide */
+    public @NonNull String toShortString() {
+        final StringBuilder sb = new StringBuilder();
+        sb.append("id=").append(mId);
+        if (mDeviceMacAddress != null) {
+            sb.append(", addr=").append(getDeviceMacAddressAsString());
+        }
+        if (mSelfManaged) {
+            sb.append(", self-managed");
+        }
+        sb.append(", pkg=u").append(mUserId).append('/').append(mPackageName);
+        return sb.toString();
+    }
+
     @Override
     public String toString() {
         return "Association{"
diff --git a/services/companion/java/com/android/server/companion/AssociationRequestsProcessor.java b/services/companion/java/com/android/server/companion/AssociationRequestsProcessor.java
index 637994f..1914164 100644
--- a/services/companion/java/com/android/server/companion/AssociationRequestsProcessor.java
+++ b/services/companion/java/com/android/server/companion/AssociationRequestsProcessor.java
@@ -22,7 +22,6 @@
 import static android.companion.CompanionDeviceManager.COMPANION_DEVICE_DISCOVERY_PACKAGE_NAME;
 import static android.content.ComponentName.createRelative;
 
-import static com.android.internal.util.CollectionUtils.filter;
 import static com.android.server.companion.CompanionDeviceManagerService.DEBUG;
 import static com.android.server.companion.CompanionDeviceManagerService.LOG_TAG;
 import static com.android.server.companion.PermissionsUtils.enforcePermissionsForAssociation;
@@ -57,6 +56,7 @@
 
 import java.util.Arrays;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Set;
 
 /**
@@ -124,14 +124,17 @@
     private static final int ASSOCIATE_WITHOUT_PROMPT_MAX_PER_TIME_WINDOW = 5;
     private static final long ASSOCIATE_WITHOUT_PROMPT_WINDOW_MS = 60 * 60 * 1000; // 60 min;
 
-    private final Context mContext;
-    private final CompanionDeviceManagerService mService;
-    private final PackageManagerInternal mPackageManager;
+    private final @NonNull Context mContext;
+    private final @NonNull CompanionDeviceManagerService mService;
+    private final @NonNull PackageManagerInternal mPackageManager;
+    private final @NonNull AssociationStore mAssociationStore;
 
-    AssociationRequestsProcessor(CompanionDeviceManagerService service) {
+    AssociationRequestsProcessor(@NonNull CompanionDeviceManagerService service,
+            @NonNull AssociationStore associationStore) {
         mContext = service.getContext();
         mService = service;
         mPackageManager = service.mPackageManagerInternal;
+        mAssociationStore = associationStore;
     }
 
     /**
@@ -330,18 +333,24 @@
         }
 
         // Throttle frequent associations
-        long now = System.currentTimeMillis();
-        Set<AssociationInfo> recentAssociations = filter(
-                mService.getAssociations(userId, packageName),
-                a -> now - a.getTimeApprovedMs() < ASSOCIATE_WITHOUT_PROMPT_WINDOW_MS);
-
-        if (recentAssociations.size() >= ASSOCIATE_WITHOUT_PROMPT_MAX_PER_TIME_WINDOW) {
-            Slog.w(TAG, "Too many associations. " + packageName
-                    + " already associated " + recentAssociations.size()
-                    + " devices within the last " + ASSOCIATE_WITHOUT_PROMPT_WINDOW_MS
-                    + "ms: " + recentAssociations);
-            return false;
+        final long now = System.currentTimeMillis();
+        final List<AssociationInfo> associationForPackage =
+                mAssociationStore.getAssociationsForPackage(userId, packageName);
+        // Number of "recent" associations.
+        int recent = 0;
+        for (AssociationInfo association : associationForPackage) {
+            final boolean isRecent =
+                    now - association.getTimeApprovedMs() < ASSOCIATE_WITHOUT_PROMPT_WINDOW_MS;
+            if (isRecent) {
+                if (++recent >= ASSOCIATE_WITHOUT_PROMPT_MAX_PER_TIME_WINDOW) {
+                    Slog.w(TAG, "Too many associations: " + packageName + " already "
+                            + "associated " + recent + " devices within the last "
+                            + ASSOCIATE_WITHOUT_PROMPT_WINDOW_MS + "ms");
+                    return false;
+                }
+            }
         }
+
         String[] sameOemCerts = mContext.getResources()
                 .getStringArray(com.android.internal.R.array.config_companionDeviceCerts);
 
diff --git a/services/companion/java/com/android/server/companion/AssociationStore.java b/services/companion/java/com/android/server/companion/AssociationStore.java
new file mode 100644
index 0000000..58fc8f7
--- /dev/null
+++ b/services/companion/java/com/android/server/companion/AssociationStore.java
@@ -0,0 +1,126 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * 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.android.server.companion;
+
+import android.annotation.IntDef;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.annotation.UserIdInt;
+import android.companion.AssociationInfo;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.util.Collection;
+import java.util.List;
+
+/**
+ * Interface for a store of {@link AssociationInfo}-s.
+ */
+public interface AssociationStore {
+
+    @IntDef(prefix = { "CHANGE_TYPE_" }, value = {
+            CHANGE_TYPE_ADDED,
+            CHANGE_TYPE_REMOVED,
+            CHANGE_TYPE_UPDATED_ADDRESS_CHANGED,
+            CHANGE_TYPE_UPDATED_ADDRESS_UNCHANGED,
+    })
+    @Retention(RetentionPolicy.SOURCE)
+    @interface ChangeType {}
+
+    int CHANGE_TYPE_ADDED = 0;
+    int CHANGE_TYPE_REMOVED = 1;
+    int CHANGE_TYPE_UPDATED_ADDRESS_CHANGED = 2;
+    int CHANGE_TYPE_UPDATED_ADDRESS_UNCHANGED = 3;
+
+    /**  Listener for any changes to {@link AssociationInfo}-s. */
+    interface OnChangeListener {
+        default void onAssociationChanged(
+                @ChangeType int changeType, AssociationInfo association) {}
+
+        default void onAssociationAdded(AssociationInfo association) {}
+
+        default void onAssociationRemoved(AssociationInfo association) {}
+
+        default void onAssociationUpdated(AssociationInfo association, boolean addressChanged) {}
+    }
+
+    /**
+     * @return all CDM associations.
+     */
+    @NonNull
+    Collection<AssociationInfo> getAssociations();
+
+    /**
+     * @return a {@link List} of associations that belong to the user.
+     */
+    @NonNull
+    List<AssociationInfo> getAssociationsForUser(@UserIdInt int userId);
+
+    /**
+     * @return a {@link List} of association that belong to the package.
+     */
+    @NonNull
+    List<AssociationInfo> getAssociationsForPackage(
+            @UserIdInt int userId, @NonNull String packageName);
+
+    /**
+     * @return an association with the given address that belong to the given package if such an
+     * association exists, otherwise {@code null}.
+     */
+    @Nullable
+    AssociationInfo getAssociationsForPackageWithAddress(
+            @UserIdInt int userId, @NonNull String packageName, @NonNull String macAddress);
+
+    /**
+     * @return an association with the given id if such an association exists, otherwise
+     * {@code null}.
+     */
+    @Nullable
+    AssociationInfo getAssociationById(int id);
+
+    /**
+     * @return all associations with the given MAc address.
+     */
+    @NonNull
+    List<AssociationInfo> getAssociationsByAddress(@NonNull String macAddress);
+
+    /** Register a {@link OnChangeListener} */
+    void registerListener(@NonNull OnChangeListener listener);
+
+    /** Un-register a previously registered {@link OnChangeListener} */
+    void unregisterListener(@NonNull OnChangeListener listener);
+
+    /** @hide */
+    static String changeTypeToString(@ChangeType int changeType) {
+        switch (changeType) {
+            case CHANGE_TYPE_ADDED:
+                return "ASSOCIATION_ADDED";
+
+            case CHANGE_TYPE_REMOVED:
+                return "ASSOCIATION_REMOVED";
+
+            case CHANGE_TYPE_UPDATED_ADDRESS_CHANGED:
+                return "ASSOCIATION_UPDATED";
+
+            case CHANGE_TYPE_UPDATED_ADDRESS_UNCHANGED:
+                return "ASSOCIATION_UPDATED_ADDRESS_UNCHANGED";
+
+            default:
+                return "Unknown (" + changeType + ")";
+        }
+    }
+}
diff --git a/services/companion/java/com/android/server/companion/AssociationStoreImpl.java b/services/companion/java/com/android/server/companion/AssociationStoreImpl.java
new file mode 100644
index 0000000..3f0200e
--- /dev/null
+++ b/services/companion/java/com/android/server/companion/AssociationStoreImpl.java
@@ -0,0 +1,302 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * 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.android.server.companion;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.annotation.UserIdInt;
+import android.companion.AssociationInfo;
+import android.net.MacAddress;
+import android.util.Log;
+import android.util.SparseArray;
+
+import com.android.internal.annotations.GuardedBy;
+import com.android.internal.util.CollectionUtils;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+
+/**
+ * Implementation of the {@link AssociationStore}, with addition of the methods for modification.
+ * <ul>
+ * <li> {@link #addAssociation(AssociationInfo)}
+ * <li> {@link #removeAssociation(int)}
+ * <li> {@link #updateAssociation(AssociationInfo)}
+ * </ul>
+ *
+ * The class has package-private access level, and instances of the class should only be created by
+ * the {@link CompanionDeviceManagerService}.
+ * Other system component (both inside and outside if the com.android.server.companion package)
+ * should use public {@link AssociationStore} interface.
+ */
+class AssociationStoreImpl implements AssociationStore {
+    private static final boolean DEBUG = false;
+    private static final String TAG = "AssociationStore";
+
+    private final Object mLock = new Object();
+
+    @GuardedBy("mLock")
+    private final Map<Integer, AssociationInfo> mIdMap;
+    @GuardedBy("mLock")
+    private final Map<MacAddress, Set<Integer>> mAddressMap;
+    @GuardedBy("mLock")
+    private final SparseArray<List<AssociationInfo>> mCachedPerUser = new SparseArray<>();
+
+    @GuardedBy("mListeners")
+    private final Set<OnChangeListener> mListeners = new LinkedHashSet<>();
+
+    AssociationStoreImpl(Collection<AssociationInfo> associations) {
+        synchronized (mLock) {
+            final int size = associations.size();
+            mIdMap = new HashMap<>(size);
+            mAddressMap = new HashMap<>(size);
+
+            for (AssociationInfo association : associations) {
+                final int id = association.getId();
+                mIdMap.put(id, association);
+
+                final MacAddress address = association.getDeviceMacAddress();
+                if (address != null) {
+                    mAddressMap.computeIfAbsent(address, it -> new HashSet<>()).add(id);
+                }
+            }
+        }
+    }
+
+    void addAssociation(@NonNull AssociationInfo association) {
+        final int id = association.getId();
+
+        if (DEBUG) {
+            Log.i(TAG, "addAssociation() " + association.toShortString());
+            Log.d(TAG, "  association=" + association);
+        }
+
+        synchronized (mLock) {
+            if (mIdMap.containsKey(id)) {
+                if (DEBUG) Log.w(TAG, "Association already stored.");
+                return;
+            }
+            mIdMap.put(id, association);
+
+            final MacAddress address = association.getDeviceMacAddress();
+            if (address != null) {
+                mAddressMap.computeIfAbsent(address, it -> new HashSet<>()).add(id);
+            }
+
+            invalidateCacheForUserLocked(association.getUserId());
+        }
+
+        broadcastChange(CHANGE_TYPE_ADDED, association);
+    }
+
+    void updateAssociation(@NonNull AssociationInfo updated) {
+        final int id = updated.getId();
+
+        if (DEBUG) {
+            Log.i(TAG, "updateAssociation() " + updated.toShortString());
+            Log.d(TAG, "  updated=" + updated);
+        }
+
+        final AssociationInfo current;
+        final boolean macAddressChanged;
+        synchronized (mLock) {
+            current = mIdMap.get(id);
+            if (current == null) {
+                if (DEBUG) Log.w(TAG, "Association with id " + id + " does not exist.");
+                return;
+            }
+            if (DEBUG) Log.d(TAG, "  current=" + current);
+
+            if (current.equals(updated)) {
+                if (DEBUG) Log.w(TAG, "  No changes.");
+                return;
+            }
+
+            // Update the ID-to-Association map.
+            mIdMap.put(id, updated);
+
+            // Update the MacAddress-to-List<Association> map if needed.
+            final MacAddress updatedAddress = updated.getDeviceMacAddress();
+            final MacAddress currentAddress = current.getDeviceMacAddress();
+            macAddressChanged = Objects.equals(
+                    current.getDeviceMacAddress(), updated.getDeviceMacAddress());
+            if (macAddressChanged) {
+                if (currentAddress != null) {
+                    mAddressMap.get(currentAddress).remove(id);
+                }
+                if (updatedAddress != null) {
+                    mAddressMap.computeIfAbsent(updatedAddress, it -> new HashSet<>()).add(id);
+                }
+            }
+        }
+
+        final int changeType = macAddressChanged ? CHANGE_TYPE_UPDATED_ADDRESS_CHANGED
+                : CHANGE_TYPE_UPDATED_ADDRESS_UNCHANGED;
+        broadcastChange(changeType, updated);
+    }
+
+    void removeAssociation(int id) {
+        if (DEBUG) Log.i(TAG, "removeAssociation() id=" + id);
+
+        final AssociationInfo association;
+        synchronized (mLock) {
+            association = mIdMap.remove(id);
+
+            if (association == null) {
+                if (DEBUG) Log.w(TAG, "Association with id " + id + " is not stored.");
+                return;
+            } else {
+                if (DEBUG) {
+                    Log.i(TAG, "removed " + association.toShortString());
+                    Log.d(TAG, "  association=" + association);
+                }
+            }
+
+            final MacAddress macAddress = association.getDeviceMacAddress();
+            if (macAddress != null) {
+                mAddressMap.get(macAddress).remove(id);
+            }
+
+            invalidateCacheForUserLocked(association.getUserId());
+        }
+
+        broadcastChange(CHANGE_TYPE_REMOVED, association);
+    }
+
+    public @NonNull Collection<AssociationInfo> getAssociations() {
+        final Collection<AssociationInfo> allAssociations;
+        synchronized (mLock) {
+            allAssociations = mIdMap.values();
+        }
+        return Collections.unmodifiableCollection(allAssociations);
+    }
+
+    public @NonNull List<AssociationInfo> getAssociationsForUser(@UserIdInt int userId) {
+        synchronized (mLock) {
+            return getAssociationsForUserLocked(userId);
+        }
+    }
+
+    public @NonNull List<AssociationInfo> getAssociationsForPackage(
+            @UserIdInt int userId, @NonNull String packageName) {
+        final List<AssociationInfo> associationsForUser = getAssociationsForUser(userId);
+        final List<AssociationInfo> associationsForPackage =
+                CollectionUtils.filter(associationsForUser,
+                        it -> it.getPackageName().equals(packageName));
+        return Collections.unmodifiableList(associationsForPackage);
+    }
+
+    public @Nullable AssociationInfo getAssociationsForPackageWithAddress(
+            @UserIdInt int userId, @NonNull String packageName, @NonNull String macAddress) {
+        final List<AssociationInfo> associations = getAssociationsByAddress(macAddress);
+        return CollectionUtils.find(associations,
+                it -> it.belongsToPackage(userId, packageName));
+    }
+
+    public @Nullable AssociationInfo getAssociationById(int id) {
+        synchronized (mLock) {
+            return mIdMap.get(id);
+        }
+    }
+
+    public @NonNull List<AssociationInfo> getAssociationsByAddress(@NonNull String macAddress) {
+        final MacAddress address = MacAddress.fromString(macAddress);
+
+        synchronized (mLock) {
+            final Set<Integer> ids = mAddressMap.get(address);
+            if (ids == null) return Collections.emptyList();
+
+            final List<AssociationInfo> associations = new ArrayList<>();
+            for (AssociationInfo association : mIdMap.values()) {
+                if (address.equals(association.getDeviceMacAddress())) {
+                    associations.add(association);
+                }
+            }
+
+            return Collections.unmodifiableList(associations);
+        }
+    }
+
+    @GuardedBy("mLock")
+    private @NonNull List<AssociationInfo> getAssociationsForUserLocked(@UserIdInt int userId) {
+        final List<AssociationInfo> cached = mCachedPerUser.get(userId);
+        if (cached != null) {
+            return cached;
+        }
+
+        final List<AssociationInfo> associationsForUser = new ArrayList<>();
+        for (AssociationInfo association : mIdMap.values()) {
+            if (association.getUserId() == userId) {
+                associationsForUser.add(association);
+            }
+        }
+        final List<AssociationInfo> set = Collections.unmodifiableList(associationsForUser);
+        mCachedPerUser.set(userId, set);
+        return set;
+    }
+
+    @GuardedBy("mLock")
+    private void invalidateCacheForUserLocked(@UserIdInt int userId) {
+        mCachedPerUser.delete(userId);
+    }
+
+    public void registerListener(@NonNull OnChangeListener listener) {
+        synchronized (mListeners) {
+            mListeners.add(listener);
+        }
+    }
+
+    public void unregisterListener(@NonNull OnChangeListener listener) {
+        synchronized (mListeners) {
+            mListeners.remove(listener);
+        }
+    }
+
+    private void broadcastChange(@ChangeType int changeType, AssociationInfo association) {
+        synchronized (mListeners) {
+            for (OnChangeListener listener : mListeners) {
+                listener.onAssociationChanged(changeType, association);
+
+                switch (changeType) {
+                    case CHANGE_TYPE_ADDED:
+                        listener.onAssociationAdded(association);
+                        break;
+
+                    case CHANGE_TYPE_REMOVED:
+                        listener.onAssociationRemoved(association);
+                        break;
+
+                    case CHANGE_TYPE_UPDATED_ADDRESS_CHANGED:
+                        listener.onAssociationUpdated(association, true);
+                        break;
+
+                    case CHANGE_TYPE_UPDATED_ADDRESS_UNCHANGED:
+                        listener.onAssociationUpdated(association, false);
+                        break;
+                }
+            }
+        }
+    }
+}
diff --git a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
index a06672b..5aa1c93 100644
--- a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
+++ b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
@@ -27,15 +27,12 @@
 import static android.os.Process.SYSTEM_UID;
 import static android.os.UserHandle.getCallingUserId;
 
-import static com.android.internal.util.CollectionUtils.add;
 import static com.android.internal.util.CollectionUtils.any;
-import static com.android.internal.util.CollectionUtils.filter;
 import static com.android.internal.util.CollectionUtils.find;
-import static com.android.internal.util.CollectionUtils.forEach;
-import static com.android.internal.util.CollectionUtils.map;
 import static com.android.internal.util.Preconditions.checkState;
 import static com.android.internal.util.function.pooled.PooledLambda.obtainMessage;
 import static com.android.internal.util.function.pooled.PooledLambda.obtainRunnable;
+import static com.android.server.companion.AssociationStore.CHANGE_TYPE_UPDATED_ADDRESS_UNCHANGED;
 import static com.android.server.companion.PermissionsUtils.checkCallerCanManageAssociationsForPackage;
 import static com.android.server.companion.PermissionsUtils.checkCallerCanManageCompanionDevice;
 import static com.android.server.companion.PermissionsUtils.enforceCallerCanInteractWithUserId;
@@ -45,8 +42,6 @@
 import static com.android.server.companion.RolesUtils.addRoleHolderForAssociation;
 import static com.android.server.companion.RolesUtils.removeRoleHolderForAssociation;
 
-import static java.util.Collections.emptySet;
-import static java.util.Collections.unmodifiableSet;
 import static java.util.Objects.requireNonNull;
 import static java.util.concurrent.TimeUnit.MINUTES;
 
@@ -82,7 +77,6 @@
 import android.content.pm.PackageItemInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.PackageManagerInternal;
-import android.content.pm.UserInfo;
 import android.net.MacAddress;
 import android.net.NetworkPolicyManager;
 import android.os.Binder;
@@ -112,9 +106,7 @@
 import com.android.internal.notification.NotificationAccessConfirmationActivityContract;
 import com.android.internal.os.BackgroundThread;
 import com.android.internal.util.ArrayUtils;
-import com.android.internal.util.CollectionUtils;
 import com.android.internal.util.DumpUtils;
-import com.android.internal.util.function.pooled.PooledLambda;
 import com.android.server.FgThread;
 import com.android.server.LocalServices;
 import com.android.server.SystemService;
@@ -127,6 +119,7 @@
 import java.text.DateFormat;
 import java.text.SimpleDateFormat;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Date;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -135,12 +128,11 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.TimeZone;
-import java.util.function.Function;
-import java.util.function.Predicate;
 
 /** @hide */
 @SuppressLint("LongLogTag")
-public class CompanionDeviceManagerService extends SystemService {
+public class CompanionDeviceManagerService extends SystemService
+        implements AssociationStore.OnChangeListener {
     static final String LOG_TAG = "CompanionDeviceManagerService";
     static final boolean DEBUG = false;
 
@@ -162,10 +154,11 @@
         sDateFormat.setTimeZone(TimeZone.getDefault());
     }
 
-    private final CompanionDeviceManagerImpl mImpl;
     // Persistent data store for all Associations.
-    private final PersistentDataStore mPersistentDataStore;
-    private final AssociationRequestsProcessor mAssociationRequestsProcessor;
+    private PersistentDataStore mPersistentStore;
+    private AssociationStoreImpl mAssociationStore;
+    private AssociationRequestsProcessor mAssociationRequestsProcessor;
+
     private PowerWhitelistManager mPowerWhitelistManager;
     private IAppOpsService mAppOpsManager;
     private BluetoothAdapter mBluetoothAdapter;
@@ -187,20 +180,16 @@
     private final RemoteCallbackList<IOnAssociationsChangedListener> mListeners =
             new RemoteCallbackList<>();
 
-    final Object mLock = new Object();
     final Handler mMainHandler = Handler.getMain();
     private CompanionDevicePresenceController mCompanionDevicePresenceController;
 
-    /** Maps a {@link UserIdInt} to a set of associations for the user. */
-    @GuardedBy("mLock")
-    private final SparseArray<Set<AssociationInfo>> mCachedAssociations = new SparseArray<>();
     /**
      * A structure that consist of two nested maps, and effectively maps (userId + packageName) to
      * a list of IDs that have been previously assigned to associations for that package.
      * We maintain this structure so that we never re-use association IDs for the same package
      * (until it's uninstalled).
      */
-    @GuardedBy("mLock")
+    @GuardedBy("mPreviouslyUsedIds")
     private final SparseArray<Map<String, Set<Integer>>> mPreviouslyUsedIds = new SparseArray<>();
 
     ActivityTaskManagerInternal mAtmInternal;
@@ -209,8 +198,6 @@
 
     public CompanionDeviceManagerService(Context context) {
         super(context);
-        mImpl = new CompanionDeviceManagerImpl();
-        mPersistentDataStore = new PersistentDataStore();
 
         mPowerWhitelistManager = context.getSystemService(PowerWhitelistManager.class);
         mAppOpsManager = IAppOpsService.Stub.asInterface(
@@ -221,49 +208,36 @@
         mPermissionControllerManager = requireNonNull(
                 context.getSystemService(PermissionControllerManager.class));
         mUserManager = context.getSystemService(UserManager.class);
-        mCompanionDevicePresenceController = new CompanionDevicePresenceController(this);
-        mAssociationRequestsProcessor = new AssociationRequestsProcessor(this);
-
-        registerPackageMonitor();
-    }
-
-    private void registerPackageMonitor() {
-        new PackageMonitor() {
-            @Override
-            public void onPackageRemoved(String packageName, int uid) {
-                final int userId = getChangingUserId();
-                Slog.i(LOG_TAG, "onPackageRemoved() u" + userId + "/" + packageName);
-
-                clearAssociationForPackage(userId, packageName);
-            }
-
-            @Override
-            public void onPackageDataCleared(String packageName, int uid) {
-                final int userId = getChangingUserId();
-                Slog.i(LOG_TAG, "onPackageDataCleared() u" + userId + "/" + packageName);
-
-                clearAssociationForPackage(userId, packageName);
-            }
-
-            @Override
-            public void onPackageModified(String packageName) {
-                final int userId = getChangingUserId();
-                Slog.i(LOG_TAG, "onPackageModified() u" + userId + "/" + packageName);
-
-                forEach(getAssociations(userId, packageName), association ->
-                        updateSpecialAccessPermissionForAssociatedPackage(association));
-            }
-        }.register(getContext(), FgThread.get().getLooper(), UserHandle.ALL, true);
     }
 
     @Override
     public void onStart() {
-        publishBinderService(Context.COMPANION_DEVICE_SERVICE, mImpl);
+        mPersistentStore = new PersistentDataStore();
+        final Set<AssociationInfo> allAssociations = new ArraySet<>();
+
+        synchronized (mPreviouslyUsedIds) {
+            // The data is stored in DE directories, so we can read the data for all users now
+            // (which would not be possible if the data was stored to CE directories).
+            mPersistentStore.readStateForUsers(
+                    mUserManager.getAliveUsers(), allAssociations, mPreviouslyUsedIds);
+        }
+
+        mAssociationStore = new AssociationStoreImpl(allAssociations);
+        mAssociationStore.registerListener(this);
+
+        mCompanionDevicePresenceController = new CompanionDevicePresenceController(this);
+        mAssociationRequestsProcessor = new AssociationRequestsProcessor(this, mAssociationStore);
+
+        // Publish "binder service"
+        final CompanionDeviceManagerImpl impl = new CompanionDeviceManagerImpl();
+        publishBinderService(Context.COMPANION_DEVICE_SERVICE, impl);
     }
 
     @Override
     public void onBootPhase(int phase) {
         if (phase == SystemService.PHASE_SYSTEM_SERVICES_READY) {
+            registerPackageMonitor();
+
             // Init Bluetooth
             mBluetoothAdapter = BluetoothAdapter.getDefaultAdapter();
             if (mBluetoothAdapter != null) {
@@ -282,7 +256,7 @@
     @Override
     public void onUserUnlocking(@NonNull TargetUser user) {
         final int userId = user.getUserIdentifier();
-        final Set<AssociationInfo> associations = getAllAssociationsForUser(userId);
+        final List<AssociationInfo> associations = mAssociationStore.getAssociationsForUser(userId);
 
         if (associations.isEmpty()) return;
 
@@ -293,42 +267,18 @@
                 MINUTES.toMillis(10));
     }
 
-    @NonNull
-    Set<AssociationInfo> getAllAssociationsForUser(@UserIdInt int userId) {
-        synchronized (mLock) {
-            readPersistedStateForUserIfNeededLocked(userId);
-            // This returns non-null, because the readAssociationsInfoForUserIfNeededLocked() method
-            // we just called adds an empty set, if there was no previously saved data.
-            return mCachedAssociations.get(userId);
-        }
-    }
-
-    @NonNull
-    Set<AssociationInfo> getAssociations(@UserIdInt int userId, @NonNull String packageName) {
-        return filter(getAllAssociationsForUser(userId),
-                a -> a.belongsToPackage(userId, packageName));
-    }
-
-    @Nullable
-    private AssociationInfo getAssociation(int associationId) {
-        return find(getAllAssociations(), association -> association.getId() == associationId);
-    }
-
-    @Nullable
-    AssociationInfo getAssociation(
-            @UserIdInt int userId, @NonNull String packageName, @NonNull String macAddress) {
-        return find(getAssociations(userId, packageName), a -> a.isLinkedTo(macAddress));
-    }
-
     @Nullable
     AssociationInfo getAssociationWithCallerChecks(
             @UserIdInt int userId, @NonNull String packageName, @NonNull String macAddress) {
-        return sanitizeWithCallerChecks(getAssociation(userId, packageName, macAddress));
+        final AssociationInfo association = mAssociationStore.getAssociationsForPackageWithAddress(
+                userId, packageName, macAddress);
+        return sanitizeWithCallerChecks(association);
     }
 
     @Nullable
     AssociationInfo getAssociationWithCallerChecks(int associationId) {
-        return sanitizeWithCallerChecks(getAssociation(associationId));
+        final AssociationInfo association = mAssociationStore.getAssociationById(associationId);
+        return sanitizeWithCallerChecks(association);
     }
 
     @Nullable
@@ -344,19 +294,6 @@
         return association;
     }
 
-    private Set<AssociationInfo> getAllAssociations() {
-        final long identity = Binder.clearCallingIdentity();
-        try {
-            final Set<AssociationInfo> result = new ArraySet<>();
-            for (UserInfo user : mUserManager.getAliveUsers()) {
-                result.addAll(getAllAssociationsForUser(user.id));
-            }
-            return result;
-        } finally {
-            Binder.restoreCallingIdentity(identity);
-        }
-    }
-
     void maybeGrantAutoRevokeExemptions() {
         Slog.d(LOG_TAG, "maybeGrantAutoRevokeExemptions()");
         PackageManager pm = getContext().getPackageManager();
@@ -369,10 +306,8 @@
             }
 
             try {
-                Set<AssociationInfo> associations = getAllAssociationsForUser(userId);
-                if (associations == null) {
-                    continue;
-                }
+                final List<AssociationInfo> associations =
+                        mAssociationStore.getAssociationsForUser(userId);
                 for (AssociationInfo a : associations) {
                     try {
                         int uid = pm.getPackageUidAsUser(a.getPackageName(), userId);
@@ -387,6 +322,61 @@
         }
     }
 
+    @Override
+    public void onAssociationChanged(
+            @AssociationStore.ChangeType int changeType, AssociationInfo association) {
+        final int id = association.getId();
+        final int userId = association.getUserId();
+        final String packageName = association.getPackageName();
+
+        if (changeType == AssociationStore.CHANGE_TYPE_REMOVED) {
+            markIdAsPreviouslyUsedForPackage(id, userId, packageName);
+        }
+
+        final List<AssociationInfo> updatedAssociations =
+                mAssociationStore.getAssociationsForUser(userId);
+        final Map<String, Set<Integer>> usedIdsForUser = getPreviouslyUsedIdsForUser(userId);
+        BackgroundThread.getHandler().post(() ->
+                mPersistentStore.persistStateForUser(userId, updatedAssociations, usedIdsForUser));
+
+        // Notify listeners if ADDED, REMOVED or UPDATED_ADDRESS_CHANGED.
+        // Do NOT notify when UPDATED_ADDRESS_UNCHANGED, which means a minor tweak in association's
+        // configs, which "listeners" won't (and shouldn't) be able to see.
+        if (changeType != CHANGE_TYPE_UPDATED_ADDRESS_UNCHANGED) {
+            notifyListeners(userId, updatedAssociations);
+        }
+        updateAtm(userId, updatedAssociations);
+
+        restartBleScan();
+    }
+
+    private void notifyListeners(
+            @UserIdInt int userId, @NonNull List<AssociationInfo> associations) {
+        mListeners.broadcast((listener, callbackUserId) -> {
+            if ((int) callbackUserId == userId) {
+                try {
+                    listener.onAssociationsChanged(associations);
+                } catch (RemoteException ignored) {
+                }
+            }
+        });
+    }
+
+    private void markIdAsPreviouslyUsedForPackage(
+            int associationId, @UserIdInt int userId, @NonNull String packageName) {
+        synchronized (mPreviouslyUsedIds) {
+            Map<String, Set<Integer>> usedIdsForUser = mPreviouslyUsedIds.get(userId);
+            if (usedIdsForUser == null) {
+                usedIdsForUser = new HashMap<>();
+                mPreviouslyUsedIds.put(userId, usedIdsForUser);
+            }
+
+            final Set<Integer> usedIdsForPackage =
+                    usedIdsForUser.computeIfAbsent(packageName, it -> new HashSet<>());
+            usedIdsForPackage.add(associationId);
+        }
+    }
+
     class CompanionDeviceManagerImpl extends ICompanionDeviceManager.Stub {
 
         @Override
@@ -424,8 +414,7 @@
                 checkUsesFeature(packageName, getCallingUserId());
             }
 
-            return new ArrayList<>(
-                    CompanionDeviceManagerService.this.getAssociations(userId, packageName));
+            return mAssociationStore.getAssociationsForPackage(userId, packageName);
         }
 
         @Override
@@ -433,8 +422,7 @@
             enforceCallerCanInteractWithUserId(getContext(), userId);
             enforceCallerCanManageCompanionDevice(getContext(), "getAllAssociationsForUser");
 
-            return new ArrayList<>(
-                    CompanionDeviceManagerService.this.getAllAssociationsForUser(userId));
+            return mAssociationStore.getAssociationsForUser(userId);
         }
 
         @Override
@@ -470,7 +458,7 @@
                         + "(ie. it belongs to a different package or a different user).");
             }
 
-            disassociateInternal(userId, association.getId());
+            disassociateInternal(association.getId());
         }
 
         @Override
@@ -483,7 +471,7 @@
                         + "or belongs to a different user");
             }
 
-            disassociateInternal(association.getUserId(), associationId);
+            disassociateInternal(associationId);
         }
 
         @Override
@@ -540,7 +528,7 @@
                 return true;
             }
 
-            return any(CompanionDeviceManagerService.this.getAssociations(userId, packageName),
+            return any(mAssociationStore.getAssociationsForPackage(userId, packageName),
                     a -> a.isLinkedTo(macAddress));
         }
 
@@ -623,25 +611,18 @@
             final int userId = getCallingUserId();
             enforceCallerIsSystemOr(userId, packageName);
 
-            Set<AssociationInfo> deviceAssociations = filter(
-                    CompanionDeviceManagerService.this.getAssociations(userId, packageName),
-                    a -> a.isLinkedTo(deviceAddress));
+            final AssociationInfo association =
+                    mAssociationStore.getAssociationsForPackageWithAddress(
+                            userId, packageName, deviceAddress);
 
-            if (deviceAssociations.isEmpty()) {
+            if (association == null) {
                 throw new RemoteException(new DeviceNotAssociatedException("App " + packageName
                         + " is not associated with device " + deviceAddress
                         + " for user " + userId));
             }
 
-            updateAssociations(associations -> map(associations, association -> {
-                if (association.belongsToPackage(userId, packageName)
-                        && association.isLinkedTo(deviceAddress)) {
-                    association.setNotifyOnDeviceNearby(active);
-                }
-                return association;
-            }), userId);
-
-            restartBleScan();
+            association.setNotifyOnDeviceNearby(active);
+            mAssociationStore.updateAssociation(association);
         }
 
         @Override
@@ -664,14 +645,16 @@
             enforceCallerIsSystemOr(userId, callingPackage);
 
             checkState(!ArrayUtils.isEmpty(
-                    CompanionDeviceManagerService.this.getAssociations(userId, callingPackage)),
+                    mAssociationStore.getAssociationsForPackage(userId, callingPackage)),
                     "App must have an association before calling this API");
             checkUsesFeature(callingPackage, userId);
         }
 
         @Override
         public boolean canPairWithoutPrompt(String packageName, String macAddress, int userId) {
-            final AssociationInfo association = getAssociation(userId, packageName, macAddress);
+            final AssociationInfo association =
+                    mAssociationStore.getAssociationsForPackageWithAddress(
+                            userId, packageName, macAddress);
             if (association == null) {
                 return false;
             }
@@ -684,7 +667,8 @@
                 String[] args, ShellCallback callback, ResultReceiver resultReceiver)
                 throws RemoteException {
             enforceCallerCanManageCompanionDevice(getContext(), "onShellCommand");
-            new CompanionDeviceShellCommand(CompanionDeviceManagerService.this)
+            new CompanionDeviceShellCommand(
+                    CompanionDeviceManagerService.this, mAssociationStore)
                     .exec(this, in, out, err, args, callback, resultReceiver);
         }
 
@@ -697,13 +681,8 @@
             }
 
             fout.append("Companion Device Associations:").append('\n');
-            synchronized (mLock) {
-                for (UserInfo user : getAllUsers()) {
-                    forEach(mCachedAssociations.get(user.id), a -> {
-                        fout.append("  ").append(a.toString()).append('\n');
-                    });
-                }
-
+            for (AssociationInfo a : mAssociationStore.getAssociations()) {
+                fout.append("  ").append(a.toString()).append('\n');
             }
 
             fout.append("Currently Connected Devices:").append('\n');
@@ -754,31 +733,56 @@
             @Nullable String deviceProfile, boolean selfManaged) {
         final int id = getNewAssociationIdForPackage(userId, packageName);
         final long timestamp = System.currentTimeMillis();
+
         final AssociationInfo association = new AssociationInfo(id, userId, packageName,
                 macAddress, displayName, deviceProfile, selfManaged, false, timestamp);
+        Slog.i(LOG_TAG, "New CDM association created=" + association);
+        mAssociationStore.addAssociation(association);
 
         updateSpecialAccessPermissionForAssociatedPackage(association);
-        recordAssociation(association, userId);
 
         return association;
     }
 
-    @GuardedBy("mLock")
+    @NonNull
+    private Map<String, Set<Integer>> getPreviouslyUsedIdsForUser(@UserIdInt int userId) {
+        synchronized (mPreviouslyUsedIds) {
+            return getPreviouslyUsedIdsForUserLocked(userId);
+        }
+    }
+
+    @GuardedBy("mPreviouslyUsedIds")
+    @NonNull
+    private Map<String, Set<Integer>> getPreviouslyUsedIdsForUserLocked(@UserIdInt int userId) {
+        final Map<String, Set<Integer>> usedIdsForUser = mPreviouslyUsedIds.get(userId);
+        if (usedIdsForUser == null) {
+            return Collections.emptyMap();
+        }
+        return deepUnmodifiableCopy(usedIdsForUser);
+    }
+
+    @GuardedBy("mPreviouslyUsedIds")
     @NonNull
     private Set<Integer> getPreviouslyUsedIdsForPackageLocked(
             @UserIdInt int userId, @NonNull String packageName) {
-        final Set<Integer> previouslyUsedIds = mPreviouslyUsedIds.get(userId).get(packageName);
-        if (previouslyUsedIds != null) return previouslyUsedIds;
-        return emptySet();
+        // "Deeply unmodifiable" map: the map itself and the Set<Integer> values it contains are all
+        // unmodifiable.
+        final Map<String, Set<Integer>> usedIdsForUser = getPreviouslyUsedIdsForUserLocked(userId);
+        final Set<Integer> usedIdsForPackage = usedIdsForUser.get(packageName);
+
+        if (usedIdsForPackage == null) {
+            return Collections.emptySet();
+        }
+
+        //The set is already unmodifiable.
+        return usedIdsForPackage;
     }
 
     private int getNewAssociationIdForPackage(@UserIdInt int userId, @NonNull String packageName) {
-        synchronized (mLock) {
-            readPersistedStateForUserIfNeededLocked(userId);
-
+        synchronized (mPreviouslyUsedIds) {
             // First: collect all IDs currently in use for this user's Associations.
             final SparseBooleanArray usedIds = new SparseBooleanArray();
-            for (AssociationInfo it : getAllAssociationsForUser(userId)) {
+            for (AssociationInfo it : mAssociationStore.getAssociationsForUser(userId)) {
                 usedIds.put(it.getId(), true);
             }
 
@@ -804,41 +808,14 @@
         }
     }
 
-    //TODO also revoke notification access
-    void disassociateInternal(@UserIdInt int userId, int associationId) {
-        updateAssociations(associations ->
-                filterOut(associations, it -> {
-                    if (it.getId() != associationId) return false;
-
-                    onAssociationPreRemove(it);
-                    markIdAsPreviouslyUsedForPackage(
-                            it.getId(), it.getUserId(), it.getPackageName());
-                    return true;
-                }), userId);
-
-        restartBleScan();
+    //TODO: also revoke notification access
+    void disassociateInternal(int associationId) {
+        onAssociationPreRemove(associationId);
+        mAssociationStore.removeAssociation(associationId);
     }
 
-    void clearAssociationForPackage(@UserIdInt int userId, @NonNull String packageName) {
-        if (DEBUG) Slog.d(LOG_TAG, "clearAssociationForPackage() u" + userId + "/" + packageName);
-
-        mCompanionDevicePresenceController.unbindDevicePresenceListener(packageName, userId);
-        updateAssociations(set -> filterOut(set, it -> it.belongsToPackage(userId, packageName)),
-                userId);
-    }
-
-    private void markIdAsPreviouslyUsedForPackage(
-            int associationId, @UserIdInt int userId, @NonNull String packageName) {
-        synchronized (mLock) {
-            // Mark as previously used.
-            readPersistedStateForUserIfNeededLocked(userId);
-            mPreviouslyUsedIds.get(userId)
-                    .computeIfAbsent(packageName, it -> new HashSet<>())
-                    .add(associationId);
-        }
-    }
-
-    void onAssociationPreRemove(AssociationInfo association) {
+    void onAssociationPreRemove(int associationId) {
+        final AssociationInfo association = mAssociationStore.getAssociationById(associationId);
         if (association.isNotifyOnDeviceNearby()
                 || (association.isSelfManaged()
                 && mPresentSelfManagedDevices.contains(association.getId()))) {
@@ -849,7 +826,7 @@
         String deviceProfile = association.getDeviceProfile();
         if (deviceProfile != null) {
             AssociationInfo otherAssociationWithDeviceProfile = find(
-                    getAllAssociationsForUser(association.getUserId()),
+                    mAssociationStore.getAssociationsForUser(association.getUserId()),
                     a -> !a.equals(association) && deviceProfile.equals(a.getDeviceProfile()));
             if (otherAssociationWithDeviceProfile != null) {
                 Slog.i(LOG_TAG, "Not revoking " + deviceProfile
@@ -941,49 +918,7 @@
                         .getPackageInfoAsUser(packageName, flags , userId));
     }
 
-    private void recordAssociation(AssociationInfo association, int userId) {
-        Slog.i(LOG_TAG, "recordAssociation(" + association + ")");
-        updateAssociations(associations -> add(associations, association), userId);
-    }
-
-    private void updateAssociations(Function<Set<AssociationInfo>, Set<AssociationInfo>> update,
-            int userId) {
-        final List<AssociationInfo> associationList;
-        synchronized (mLock) {
-            if (DEBUG) Slog.d(LOG_TAG, "Updating Associations set...");
-
-            final Set<AssociationInfo> prevAssociations = getAllAssociationsForUser(userId);
-            if (DEBUG) Slog.d(LOG_TAG, "  > Before : " + prevAssociations + "...");
-
-            final Set<AssociationInfo> updatedAssociations = update.apply(
-                    new ArraySet<>(prevAssociations));
-            if (DEBUG) Slog.d(LOG_TAG, "  > After: " + updatedAssociations);
-
-            associationList = new ArrayList<>(updatedAssociations);
-
-            mCachedAssociations.put(userId, unmodifiableSet(updatedAssociations));
-
-            BackgroundThread.getHandler().sendMessage(
-                    PooledLambda.obtainMessage(
-                            (associations, usedIds) ->
-                                    mPersistentDataStore
-                                            .persistStateForUser(userId, associations, usedIds),
-                            updatedAssociations, deepCopy(mPreviouslyUsedIds.get(userId))));
-
-            updateAtm(userId, updatedAssociations);
-        }
-
-        mListeners.broadcast((listener, callbackUserId) -> {
-            if ((int) callbackUserId == userId) {
-                try {
-                    listener.onAssociationsChanged(associationList);
-                } catch (RemoteException ignored) {
-                }
-            }
-        });
-    }
-
-    private void updateAtm(int userId, Set<AssociationInfo> associations) {
+    private void updateAtm(int userId, List<AssociationInfo> associations) {
         final Set<Integer> companionAppUids = new ArraySet<>();
         for (AssociationInfo association : associations) {
             final int uid = mPackageManagerInternal.getPackageUid(association.getPackageName(),
@@ -998,50 +933,18 @@
         }
     }
 
-    @GuardedBy("mLock")
-    private void readPersistedStateForUserIfNeededLocked(@UserIdInt int userId) {
-        if (mCachedAssociations.get(userId) != null) return;
-
-        Slog.i(LOG_TAG, "Reading state for user " + userId + "  from the disk");
-
-        final Set<AssociationInfo> associations = new ArraySet<>();
-        final Map<String, Set<Integer>> previouslyUsedIds = new ArrayMap<>();
-        mPersistentDataStore.readStateForUser(userId, associations, previouslyUsedIds);
-
-        if (DEBUG) {
-            Slog.d(LOG_TAG, "  > associations=" + associations + "\n"
-                    + "  > previouslyUsedIds=" + previouslyUsedIds);
-        }
-
-        mCachedAssociations.put(userId, unmodifiableSet(associations));
-        mPreviouslyUsedIds.append(userId, previouslyUsedIds);
-    }
-
-    private List<UserInfo> getAllUsers() {
-        final long identity = Binder.clearCallingIdentity();
-        try {
-            return mUserManager.getUsers();
-        } finally {
-            Binder.restoreCallingIdentity(identity);
-        }
-    }
-
     void onDeviceConnected(String address) {
         Slog.d(LOG_TAG, "onDeviceConnected(address = " + address + ")");
 
         mCurrentlyConnectedDevices.add(address);
 
-        for (UserInfo user : getAllUsers()) {
-            for (AssociationInfo association : getAllAssociationsForUser(user.id)) {
-                if (association.isLinkedTo(address)) {
-                    if (association.getDeviceProfile() != null) {
-                        Slog.i(LOG_TAG, "Granting role " + association.getDeviceProfile()
-                                + " to " + association.getPackageName()
-                                + " due to device connected: " + association.getDeviceMacAddress());
+        for (AssociationInfo association : mAssociationStore.getAssociationsByAddress(address)) {
+            if (association.getDeviceProfile() != null) {
+                Slog.i(LOG_TAG, "Granting role " + association.getDeviceProfile()
+                        + " to " + association.getPackageName()
+                        + " due to device connected: " + association.getDeviceMacAddress());
 
-                        addRoleHolderForAssociation(getContext(), association);
-                    }
-                }
+                addRoleHolderForAssociation(getContext(), association);
             }
         }
 
@@ -1134,7 +1037,9 @@
                 Date lastNearby = mDevicesLastNearby.valueAt(i);
 
                 if (isDeviceDisappeared(lastNearby)) {
-                    for (AssociationInfo association : getAllAssociations(address)) {
+                    final List<AssociationInfo> associations =
+                            mAssociationStore.getAssociationsByAddress(address);
+                    for (AssociationInfo association : associations) {
                         if (association.isNotifyOnDeviceNearby()) {
                             mCompanionDevicePresenceController.unbindDevicePresenceListener(
                                     association.getPackageName(), association.getUserId());
@@ -1176,20 +1081,6 @@
         }
     }
 
-    private Set<AssociationInfo> getAllAssociations(String deviceAddress) {
-        List<UserInfo> aliveUsers = mUserManager.getAliveUsers();
-        Set<AssociationInfo> result = new ArraySet<>();
-        for (int i = 0, size = aliveUsers.size(); i < size; i++) {
-            UserInfo user = aliveUsers.get(i);
-            for (AssociationInfo association : getAllAssociationsForUser(user.id)) {
-                if (association.isLinkedTo(deviceAddress)) {
-                    result.add(association);
-                }
-            }
-        }
-        return result;
-    }
-
     private void onDeviceNearby(String address) {
         Date timestamp = new Date();
         Date oldTimestamp = mDevicesLastNearby.put(address, timestamp);
@@ -1205,7 +1096,9 @@
                 || timestamp.getTime() - oldTimestamp.getTime() >= DEVICE_DISAPPEARED_TIMEOUT_MS;
         if (justAppeared) {
             Slog.i(LOG_TAG, "onDeviceNearby(justAppeared, address = " + address + ")");
-            for (AssociationInfo association : getAllAssociations(address)) {
+            final List<AssociationInfo> associations =
+                    mAssociationStore.getAssociationsByAddress(address);
+            for (AssociationInfo association : associations) {
                 if (association.isNotifyOnDeviceNearby()) {
                     mCompanionDevicePresenceController.onDeviceNotifyAppeared(association,
                             getContext(), mMainHandler);
@@ -1218,7 +1111,9 @@
         Slog.i(LOG_TAG, "onDeviceDisappeared(address = " + address + ")");
 
         boolean hasDeviceListeners = false;
-        for (AssociationInfo association : getAllAssociations(address)) {
+        final List<AssociationInfo> associations =
+                mAssociationStore.getAssociationsByAddress(address);
+        for (AssociationInfo association : associations) {
             if (association.isNotifyOnDeviceNearby()) {
                 mCompanionDevicePresenceController.onDeviceNotifyDisappeared(
                         association, getContext(), mMainHandler);
@@ -1291,7 +1186,7 @@
     private List<ScanFilter> getBleScanFilters() {
         ArrayList<ScanFilter> result = new ArrayList<>();
         ArraySet<String> addressesSeen = new ArraySet<>();
-        for (AssociationInfo association : getAllAssociations()) {
+        for (AssociationInfo association : mAssociationStore.getAssociations()) {
             if (association.isSelfManaged()) {
                 continue;
             }
@@ -1331,17 +1226,6 @@
         }
     }
 
-    private static @NonNull <T> Set<T> filterOut(
-            @NonNull Set<T> set, @NonNull Predicate<? super T> predicate) {
-        return CollectionUtils.filter(set, predicate.negate());
-    }
-
-    private Map<String, Set<Integer>> deepCopy(Map<String, Set<Integer>> orig) {
-        final Map<String, Set<Integer>> copy = new HashMap<>(orig.size(), 1f);
-        forEach(orig, (key, value) -> copy.put(key, new ArraySet<>(value)));
-        return copy;
-    }
-
     void checkUsesFeature(@NonNull String pkg, @UserIdInt int userId) {
         if (getCallingUid() == SYSTEM_UID) return;
 
@@ -1356,4 +1240,61 @@
                 + FEATURE_COMPANION_DEVICE_SETUP
                 + " in manifest to use this API");
     }
+
+    private void registerPackageMonitor() {
+        new PackageMonitor() {
+            @Override
+            public void onPackageRemoved(String packageName, int uid) {
+                final int userId = getChangingUserId();
+                Slog.i(LOG_TAG, "onPackageRemoved() u" + userId + "/" + packageName);
+
+                clearAssociationForPackage(userId, packageName);
+            }
+
+            @Override
+            public void onPackageDataCleared(String packageName, int uid) {
+                final int userId = getChangingUserId();
+                Slog.i(LOG_TAG, "onPackageDataCleared() u" + userId + "/" + packageName);
+
+                clearAssociationForPackage(userId, packageName);
+            }
+
+            @Override
+            public void onPackageModified(String packageName) {
+                final int userId = getChangingUserId();
+                Slog.i(LOG_TAG, "onPackageModified() u" + userId + "/" + packageName);
+
+                final List<AssociationInfo> associationsForPackage =
+                        mAssociationStore.getAssociationsForPackage(userId, packageName);
+                for (AssociationInfo association : associationsForPackage) {
+                    updateSpecialAccessPermissionForAssociatedPackage(association);
+                }
+            }
+        }.register(getContext(), FgThread.get().getLooper(), UserHandle.ALL, true);
+    }
+
+    private void clearAssociationForPackage(@UserIdInt int userId, @NonNull String packageName) {
+        if (DEBUG) Slog.d(LOG_TAG, "clearAssociationForPackage() u" + userId + "/" + packageName);
+
+        // First, unbind CompanionService if needed.
+        mCompanionDevicePresenceController.unbindDevicePresenceListener(packageName, userId);
+
+        // Clear associations.
+        final List<AssociationInfo> associationsForPackage =
+                mAssociationStore.getAssociationsForPackage(userId, packageName);
+        for (AssociationInfo association : associationsForPackage) {
+            mAssociationStore.removeAssociation(association.getId());
+        }
+    }
+
+    private static Map<String, Set<Integer>> deepUnmodifiableCopy(Map<String, Set<Integer>> orig) {
+        final Map<String, Set<Integer>> copy = new HashMap<>();
+
+        for (Map.Entry<String, Set<Integer>> entry : orig.entrySet()) {
+            final Set<Integer> valueCopy = new HashSet<>(entry.getValue());
+            copy.put(entry.getKey(), Collections.unmodifiableSet(valueCopy));
+        }
+
+        return Collections.unmodifiableMap(copy);
+    }
 }
diff --git a/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java b/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java
index 5cb3079..5c0571d 100644
--- a/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java
+++ b/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java
@@ -16,7 +16,6 @@
 
 package com.android.server.companion;
 
-import static com.android.internal.util.CollectionUtils.forEach;
 import static com.android.server.companion.CompanionDeviceManagerService.LOG_TAG;
 
 import android.companion.AssociationInfo;
@@ -24,24 +23,33 @@
 import android.util.Slog;
 
 import java.io.PrintWriter;
+import java.util.List;
 
 class CompanionDeviceShellCommand extends android.os.ShellCommand {
     private final CompanionDeviceManagerService mService;
+    private final AssociationStore mAssociationStore;
 
-    CompanionDeviceShellCommand(CompanionDeviceManagerService service) {
+    CompanionDeviceShellCommand(CompanionDeviceManagerService service,
+            AssociationStore associationStore) {
         mService = service;
+        mAssociationStore = associationStore;
     }
 
     @Override
     public int onCommand(String cmd) {
+        final PrintWriter out = getOutPrintWriter();
         try {
             switch (cmd) {
                 case "list": {
-                    forEach(
-                            mService.getAllAssociationsForUser(getNextArgInt()),
-                            a -> getOutPrintWriter()
-                                    .println(a.getPackageName() + " "
-                                            + a.getDeviceMacAddress()));
+                    final int userId = getNextArgInt();
+                    final List<AssociationInfo> associationsForUser =
+                            mAssociationStore.getAssociationsForUser(userId);
+                    for (AssociationInfo association : associationsForUser) {
+                        // TODO(b/212535524): use AssociationInfo.toShortString(), once it's not
+                        //  longer referenced in tests.
+                        out.println(association.getPackageName() + " "
+                                + association.getDeviceMacAddress());
+                    }
                 }
                 break;
 
@@ -60,7 +68,7 @@
                     final AssociationInfo association =
                             mService.getAssociationWithCallerChecks(userId, packageName, address);
                     if (association != null) {
-                        mService.disassociateInternal(userId, association.getId());
+                        mService.disassociateInternal(association.getId());
                     }
                 }
                 break;
diff --git a/services/companion/java/com/android/server/companion/PersistentDataStore.java b/services/companion/java/com/android/server/companion/PersistentDataStore.java
index 87558df..e2a814b 100644
--- a/services/companion/java/com/android/server/companion/PersistentDataStore.java
+++ b/services/companion/java/com/android/server/companion/PersistentDataStore.java
@@ -33,11 +33,14 @@
 import android.annotation.Nullable;
 import android.annotation.UserIdInt;
 import android.companion.AssociationInfo;
+import android.content.pm.UserInfo;
 import android.net.MacAddress;
 import android.os.Environment;
+import android.util.ArrayMap;
 import android.util.AtomicFile;
 import android.util.ExceptionUtils;
 import android.util.Slog;
+import android.util.SparseArray;
 import android.util.TypedXmlPullParser;
 import android.util.TypedXmlSerializer;
 import android.util.Xml;
@@ -51,7 +54,9 @@
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
+import java.util.Collection;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
@@ -68,8 +73,8 @@
  *
  * Before Android T the data was stored using the v0 schema.
  *
- * @see #readAssociationsV0(TypedXmlPullParser, int, Set)
- * @see #readAssociationV0(TypedXmlPullParser, int, int, Set)
+ * @see #readAssociationsV0(TypedXmlPullParser, int, Collection)
+ * @see #readAssociationV0(TypedXmlPullParser, int, int, Collection)
  *
  * The following snippet is a sample of a the file that is using v0 schema.
  * <pre>{@code
@@ -100,8 +105,8 @@
  * optional.
  *
  * @see #CURRENT_PERSISTENCE_VERSION
- * @see #readAssociationsV1(TypedXmlPullParser, int, Set)
- * @see #readAssociationV1(TypedXmlPullParser, int, Set)
+ * @see #readAssociationsV1(TypedXmlPullParser, int, Collection)
+ * @see #readAssociationV1(TypedXmlPullParser, int, Collection)
  * @see #readPreviouslyUsedIdsV1(TypedXmlPullParser, Map)
  *
  * The following snippet is a sample of a the file that is using v0 schema.
@@ -168,6 +173,23 @@
     private final @NonNull ConcurrentMap<Integer, AtomicFile> mUserIdToStorageFile =
             new ConcurrentHashMap<>();
 
+    void readStateForUsers(@NonNull List<UserInfo> users,
+            @NonNull Set<AssociationInfo> allAssociationsOut,
+            @NonNull SparseArray<Map<String, Set<Integer>>> previouslyUsedIdsPerUserOut) {
+        for (UserInfo user : users) {
+            final int userId = user.id;
+            // Previously used IDs are stored in the "out" collection per-user.
+            final Map<String, Set<Integer>> previouslyUsedIds = new ArrayMap<>();
+
+            // Associations for all users are stored in a single "flat" set: so we read directly
+            // into it.
+            readStateForUser(userId, allAssociationsOut, previouslyUsedIds);
+
+            // Save previously used IDs for this user into the "out" structure.
+            previouslyUsedIdsPerUserOut.append(userId, previouslyUsedIds);
+        }
+    }
+
     /**
      * Reads previously persisted data for the given user "into" the provided containers.
      *
@@ -176,7 +198,7 @@
      * @param previouslyUsedIdsPerPackageOut a container to read the used IDs "into".
      */
     void readStateForUser(@UserIdInt int userId,
-            @NonNull Set<AssociationInfo> associationsOut,
+            @NonNull Collection<AssociationInfo> associationsOut,
             @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackageOut) {
         Slog.i(LOG_TAG, "Reading associations for user " + userId + " from disk");
         final AtomicFile file = getStorageFileForUser(userId);
@@ -237,7 +259,8 @@
      * @param associations a set of user's associations.
      * @param previouslyUsedIdsPerPackage a set previously used Association IDs for the user.
      */
-    void persistStateForUser(@UserIdInt int userId, @NonNull Set<AssociationInfo> associations,
+    void persistStateForUser(@UserIdInt int userId,
+            @NonNull Collection<AssociationInfo> associations,
             @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackage) {
         Slog.i(LOG_TAG, "Writing associations for user " + userId + " to disk");
         if (DEBUG) Slog.d(LOG_TAG, "  > " + associations);
@@ -250,7 +273,7 @@
     }
 
     private int readStateFromFileLocked(@UserIdInt int userId, @NonNull AtomicFile file,
-            @NonNull String rootTag, @Nullable Set<AssociationInfo> associationsOut,
+            @NonNull String rootTag, @Nullable Collection<AssociationInfo> associationsOut,
             @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackageOut) {
         try (FileInputStream in = file.openRead()) {
             final TypedXmlPullParser parser = Xml.resolvePullParser(in);
@@ -282,7 +305,7 @@
     }
 
     private void persistStateToFileLocked(@NonNull AtomicFile file,
-            @Nullable Set<AssociationInfo> associations,
+            @Nullable Collection<AssociationInfo> associations,
             @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackage) {
         file.write(out -> {
             try {
@@ -321,7 +344,7 @@
     }
 
     private static void readAssociationsV0(@NonNull TypedXmlPullParser parser,
-            @UserIdInt int userId, @NonNull Set<AssociationInfo> out)
+            @UserIdInt int userId, @NonNull Collection<AssociationInfo> out)
             throws XmlPullParserException, IOException {
         requireStartOfTag(parser, XML_TAG_ASSOCIATIONS);
 
@@ -342,7 +365,8 @@
     }
 
     private static void readAssociationV0(@NonNull TypedXmlPullParser parser, @UserIdInt int userId,
-            int associationId, @NonNull Set<AssociationInfo> out) throws XmlPullParserException {
+            int associationId, @NonNull Collection<AssociationInfo> out)
+            throws XmlPullParserException {
         requireStartOfTag(parser, XML_TAG_ASSOCIATION);
 
         final String appPackage = readStringAttribute(parser, XML_ATTR_PACKAGE);
@@ -360,7 +384,7 @@
     }
 
     private static void readAssociationsV1(@NonNull TypedXmlPullParser parser,
-            @UserIdInt int userId, @NonNull Set<AssociationInfo> out)
+            @UserIdInt int userId, @NonNull Collection<AssociationInfo> out)
             throws XmlPullParserException, IOException {
         requireStartOfTag(parser, XML_TAG_ASSOCIATIONS);
 
@@ -374,7 +398,7 @@
     }
 
     private static void readAssociationV1(@NonNull TypedXmlPullParser parser, @UserIdInt int userId,
-            @NonNull Set<AssociationInfo> out) throws XmlPullParserException, IOException {
+            @NonNull Collection<AssociationInfo> out) throws XmlPullParserException, IOException {
         requireStartOfTag(parser, XML_TAG_ASSOCIATION);
 
         final int associationId = readIntAttribute(parser, XML_ATTR_ID);
@@ -421,9 +445,11 @@
     }
 
     private static void writeAssociations(@NonNull XmlSerializer parent,
-            @Nullable Set<AssociationInfo> associations) throws IOException {
+            @Nullable Collection<AssociationInfo> associations) throws IOException {
         final XmlSerializer serializer = parent.startTag(null, XML_TAG_ASSOCIATIONS);
-        forEach(associations, it -> writeAssociation(serializer, it));
+        for (AssociationInfo association : associations) {
+            writeAssociation(serializer, association);
+        }
         serializer.endTag(null, XML_TAG_ASSOCIATIONS);
     }