Merge changes from topic "am-5ad49821c2a945fcbcb805cc2067110c" into main

* changes:
  [automerger skipped] Merge "Import translations. DO NOT MERGE ANYWHERE" into 24D1-dev am: b0079067e1 -s ours
  [automerger skipped] Import translations. DO NOT MERGE ANYWHERE am: e6bef82914 -s ours
  [automerger skipped] Merge "Import translations. DO NOT MERGE ANYWHERE" into 24D1-dev am: 98ff433765 -s ours
  [automerger skipped] Import translations. DO NOT MERGE ANYWHERE am: bd2c35b294 -s ours
  [automerger skipped] Merge "Import translations. DO NOT MERGE ANYWHERE" into 24D1-dev am: 5e4d313eab -s ours
  [automerger skipped] Import translations. DO NOT MERGE ANYWHERE am: 553fb49b2b -s ours
  [automerger skipped] Merge "Import translations. DO NOT MERGE ANYWHERE" into 24D1-dev am: 841a6b2ee1 -s ours
  [automerger skipped] Import translations. DO NOT MERGE ANYWHERE am: fcb1a9cab3 -s ours
  [automerger skipped] Merge "Import translations. DO NOT MERGE ANYWHERE" into 24D1-dev am: 7ea0b6c2e7 -s ours
  [automerger skipped] Import translations. DO NOT MERGE ANYWHERE am: 29f260ffb3 -s ours
diff --git a/core/java/android/companion/AssociationInfo.java b/core/java/android/companion/AssociationInfo.java
index 843158c..b4b96e2 100644
--- a/core/java/android/companion/AssociationInfo.java
+++ b/core/java/android/companion/AssociationInfo.java
@@ -252,6 +252,14 @@
     }
 
     /**
+     * @return true if the association is not revoked nor pending
+     * @hide
+     */
+    public boolean isActive() {
+        return !mRevoked && !mPending;
+    }
+
+    /**
      * @return the last time self reported disconnected for selfManaged only.
      * @hide
      */
diff --git a/services/companion/java/com/android/server/companion/BackupRestoreProcessor.java b/services/companion/java/com/android/server/companion/BackupRestoreProcessor.java
index f2409fb..5e52e06 100644
--- a/services/companion/java/com/android/server/companion/BackupRestoreProcessor.java
+++ b/services/companion/java/com/android/server/companion/BackupRestoreProcessor.java
@@ -18,7 +18,8 @@
 
 import static android.os.UserHandle.getCallingUserId;
 
-import static com.android.server.companion.CompanionDeviceManagerService.PerUserAssociationSet;
+import static com.android.server.companion.association.AssociationDiskStore.readAssociationsFromPayload;
+import static com.android.server.companion.utils.RolesUtils.addRoleHolderForAssociation;
 
 import android.annotation.NonNull;
 import android.annotation.SuppressLint;
@@ -26,62 +27,50 @@
 import android.companion.AssociationInfo;
 import android.companion.Flags;
 import android.companion.datatransfer.SystemDataTransferRequest;
+import android.content.Context;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManagerInternal;
-import android.util.ArraySet;
-import android.util.Log;
 import android.util.Slog;
 
-import com.android.internal.annotations.GuardedBy;
 import com.android.internal.util.CollectionUtils;
 import com.android.server.companion.association.AssociationDiskStore;
 import com.android.server.companion.association.AssociationRequestsProcessor;
 import com.android.server.companion.association.AssociationStore;
+import com.android.server.companion.association.Associations;
 import com.android.server.companion.datatransfer.SystemDataTransferRequestStore;
 
 import java.nio.ByteBuffer;
-import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
-import java.util.Set;
 import java.util.function.Predicate;
 
 @SuppressLint("LongLogTag")
 class BackupRestoreProcessor {
-    static final String TAG = "CDM_BackupRestoreProcessor";
+    private static final String TAG = "CDM_BackupRestoreProcessor";
     private static final int BACKUP_AND_RESTORE_VERSION = 0;
 
+    private final Context mContext;
     @NonNull
-    private final CompanionDeviceManagerService mService;
-    @NonNull
-    private final PackageManagerInternal mPackageManager;
+    private final PackageManagerInternal mPackageManagerInternal;
     @NonNull
     private final AssociationStore mAssociationStore;
     @NonNull
-    private final AssociationDiskStore mPersistentStore;
+    private final AssociationDiskStore mAssociationDiskStore;
     @NonNull
     private final SystemDataTransferRequestStore mSystemDataTransferRequestStore;
     @NonNull
     private final AssociationRequestsProcessor mAssociationRequestsProcessor;
 
-    /**
-     * A structure that consists of a set of restored associations that are pending corresponding
-     * companion app to be installed.
-     */
-    @GuardedBy("mAssociationsPendingAppInstall")
-    private final PerUserAssociationSet mAssociationsPendingAppInstall =
-            new PerUserAssociationSet();
-
-    BackupRestoreProcessor(@NonNull CompanionDeviceManagerService service,
+    BackupRestoreProcessor(@NonNull Context context,
+                           @NonNull PackageManagerInternal packageManagerInternal,
                            @NonNull AssociationStore associationStore,
-                           @NonNull AssociationDiskStore persistentStore,
+                           @NonNull AssociationDiskStore associationDiskStore,
                            @NonNull SystemDataTransferRequestStore systemDataTransferRequestStore,
                            @NonNull AssociationRequestsProcessor associationRequestsProcessor) {
-        mService = service;
-        mPackageManager = service.mPackageManagerInternal;
+        mContext = context;
+        mPackageManagerInternal = packageManagerInternal;
         mAssociationStore = associationStore;
-        mPersistentStore = persistentStore;
+        mAssociationDiskStore = associationDiskStore;
         mSystemDataTransferRequestStore = systemDataTransferRequestStore;
         mAssociationRequestsProcessor = associationRequestsProcessor;
     }
@@ -93,9 +82,9 @@
      * | (4) SystemDataTransferRequest length | SystemDataTransferRequest XML (without userId)|
      */
     byte[] getBackupPayload(int userId) {
-        // Persist state first to generate an up-to-date XML file
-        mService.persistStateForUser(userId);
-        byte[] associationsPayload = mPersistentStore.getBackupPayload(userId);
+        Slog.i(TAG, "getBackupPayload() userId=[" + userId + "].");
+
+        byte[] associationsPayload = mAssociationDiskStore.getBackupPayload(userId);
         int associationsPayloadLength = associationsPayload.length;
 
         // System data transfer requests are persisted up-to-date already
@@ -119,6 +108,9 @@
      * Create new associations and system data transfer request consents using backed up payload.
      */
     void applyRestoredPayload(byte[] payload, int userId) {
+        Slog.i(TAG, "applyRestoredPayload() userId=[" + userId + "], payload size=["
+                + payload.length + "].");
+
         ByteBuffer buffer = ByteBuffer.wrap(payload);
 
         // Make sure that payload version matches current version to ensure proper deserialization
@@ -131,9 +123,8 @@
         // Read the bytes containing backed-up associations
         byte[] associationsPayload = new byte[buffer.getInt()];
         buffer.get(associationsPayload);
-        final Set<AssociationInfo> restoredAssociations = new HashSet<>();
-        mPersistentStore.readStateFromPayload(associationsPayload, userId,
-                restoredAssociations, new HashMap<>());
+        final Associations restoredAssociations = readAssociationsFromPayload(
+                associationsPayload, userId);
 
         // Read the bytes containing backed-up system data transfer requests user consent
         byte[] requestsPayload = new byte[buffer.getInt()];
@@ -142,13 +133,13 @@
                 mSystemDataTransferRequestStore.readRequestsFromPayload(requestsPayload, userId);
 
         // Get a list of installed packages ahead of time.
-        List<ApplicationInfo> installedApps = mPackageManager.getInstalledApplications(
+        List<ApplicationInfo> installedApps = mPackageManagerInternal.getInstalledApplications(
                 0, userId, getCallingUserId());
 
         // Restored device may have a different user ID than the backed-up user's user-ID. Since
         // association ID is dependent on the user ID, restored associations must account for
         // this potential difference on their association IDs.
-        for (AssociationInfo restored : restoredAssociations) {
+        for (AssociationInfo restored : restoredAssociations.getAssociations()) {
             // Don't restore a revoked association. Since they weren't added to the device being
             // restored in the first place, there is no need to worry about revoking a role that
             // was never granted either.
@@ -168,10 +159,9 @@
 
             // Create a new association reassigned to this user and a valid association ID
             final String packageName = restored.getPackageName();
-            final int newId = mService.getNewAssociationIdForPackage(userId, packageName);
-            AssociationInfo newAssociation =
-                    new AssociationInfo.Builder(newId, userId, packageName, restored)
-                            .build();
+            final int newId = mAssociationStore.getNextId(userId);
+            AssociationInfo newAssociation = new AssociationInfo.Builder(newId, userId, packageName,
+                    restored).build();
 
             // Check if the companion app for this association is already installed, then do one
             // of the following:
@@ -179,13 +169,15 @@
             // the role attached to this association to the app.
             // (2) If the app isn't yet installed, then add this association to the list of pending
             // associations to be added when the package is installed in the future.
-            boolean isPackageInstalled = installedApps.stream()
-                    .anyMatch(app -> packageName.equals(app.packageName));
+            boolean isPackageInstalled = installedApps.stream().anyMatch(
+                    app -> packageName.equals(app.packageName));
             if (isPackageInstalled) {
                 mAssociationRequestsProcessor.maybeGrantRoleAndStoreAssociation(newAssociation,
                         null, null);
             } else {
-                addToPendingAppInstall(newAssociation);
+                newAssociation = (new AssociationInfo.Builder(newAssociation)).setPending(true)
+                        .build();
+                mAssociationStore.addAssociation(newAssociation);
             }
 
             // Re-map restored system data transfer requests to newly created associations
@@ -195,32 +187,27 @@
                 mSystemDataTransferRequestStore.writeRequest(userId, newRequest);
             }
         }
-
-        // Persist restored state.
-        mService.persistStateForUser(userId);
     }
 
-    void addToPendingAppInstall(@NonNull AssociationInfo association) {
-        association = (new AssociationInfo.Builder(association))
-                .setPending(true)
-                .build();
-
-        synchronized (mAssociationsPendingAppInstall) {
-            mAssociationsPendingAppInstall.forUser(association.getUserId()).add(association);
+    public void restorePendingAssociations(int userId, String packageName) {
+        List<AssociationInfo> pendingAssociations = mAssociationStore.getPendingAssociations(userId,
+                packageName);
+        if (!pendingAssociations.isEmpty()) {
+            Slog.i(TAG, "Found pending associations for package=[" + packageName
+                    + "]. Restoring...");
         }
-    }
-
-    void removeFromPendingAppInstall(@NonNull AssociationInfo association) {
-        synchronized (mAssociationsPendingAppInstall) {
-            mAssociationsPendingAppInstall.forUser(association.getUserId()).remove(association);
-        }
-    }
-
-    @NonNull
-    Set<AssociationInfo> getAssociationsPendingAppInstallForUser(@UserIdInt int userId) {
-        synchronized (mAssociationsPendingAppInstall) {
-            // Return a copy.
-            return new ArraySet<>(mAssociationsPendingAppInstall.forUser(userId));
+        for (AssociationInfo association : pendingAssociations) {
+            AssociationInfo newAssociation = new AssociationInfo.Builder(association)
+                    .setPending(false)
+                    .build();
+            addRoleHolderForAssociation(mContext, newAssociation, success -> {
+                if (success) {
+                    mAssociationStore.updateAssociation(newAssociation);
+                    Slog.i(TAG, "Association=[" + association + "] is restored.");
+                } else {
+                    Slog.e(TAG, "Failed to restore association=[" + association + "].");
+                }
+            });
         }
     }
 
@@ -231,7 +218,7 @@
     private boolean handleCollision(@UserIdInt int userId,
             AssociationInfo restored,
             List<SystemDataTransferRequest> restoredRequests) {
-        List<AssociationInfo> localAssociations = mAssociationStore.getAssociationsForPackage(
+        List<AssociationInfo> localAssociations = mAssociationStore.getActiveAssociationsByPackage(
                 restored.getUserId(), restored.getPackageName());
         Predicate<AssociationInfo> isSameDevice = associationInfo -> {
             boolean matchesMacAddress = Objects.equals(
@@ -248,7 +235,7 @@
             return false;
         }
 
-        Log.d(TAG, "Conflict detected with association id=" + local.getId()
+        Slog.d(TAG, "Conflict detected with association id=" + local.getId()
                 + " while restoring CDM backup. Keeping local association.");
 
         List<SystemDataTransferRequest> localRequests = mSystemDataTransferRequestStore
@@ -266,8 +253,8 @@
                 continue;
             }
 
-            Log.d(TAG, "Restoring " + restoredRequest.getClass().getSimpleName()
-                    + " to an existing association id=" + local.getId() + ".");
+            Slog.d(TAG, "Restoring " + restoredRequest.getClass().getSimpleName()
+                    + " to an existing association id=[" + local.getId() + "].");
 
             SystemDataTransferRequest newRequest =
                     restoredRequest.copyWithNewId(local.getId());
diff --git a/services/companion/java/com/android/server/companion/CompanionApplicationController.java b/services/companion/java/com/android/server/companion/CompanionApplicationController.java
index c801489..0a41485 100644
--- a/services/companion/java/com/android/server/companion/CompanionApplicationController.java
+++ b/services/companion/java/com/android/server/companion/CompanionApplicationController.java
@@ -397,7 +397,7 @@
         // First, disable hint mode for Auto profile and mark not BOUND for primary service ONLY.
         if (isPrimary) {
             final List<AssociationInfo> associations =
-                    mAssociationStore.getAssociationsForPackage(userId, packageName);
+                    mAssociationStore.getActiveAssociationsByPackage(userId, packageName);
 
             for (AssociationInfo association : associations) {
                 final String deviceProfile = association.getDeviceProfile();
@@ -442,7 +442,7 @@
                 mObservableUuidStore.getObservableUuidsForPackage(userId, packageName);
 
         for (AssociationInfo ai :
-                mAssociationStore.getAssociationsForPackage(userId, packageName)) {
+                mAssociationStore.getActiveAssociationsByPackage(userId, packageName)) {
             final int associationId = ai.getId();
             stillAssociated = true;
             if (ai.isSelfManaged()) {
diff --git a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
index 3846e98..73ebbc7 100644
--- a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
+++ b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
@@ -37,12 +37,9 @@
 import static com.android.internal.util.CollectionUtils.any;
 import static com.android.internal.util.Preconditions.checkState;
 import static com.android.internal.util.function.pooled.PooledLambda.obtainMessage;
-import static com.android.server.companion.association.AssociationStore.CHANGE_TYPE_UPDATED_ADDRESS_UNCHANGED;
-import static com.android.server.companion.utils.AssociationUtils.getFirstAssociationIdForUser;
-import static com.android.server.companion.utils.AssociationUtils.getLastAssociationIdForUser;
-import static com.android.server.companion.utils.PackageUtils.isRestrictedSettingsAllowed;
 import static com.android.server.companion.utils.PackageUtils.enforceUsesCompanionDeviceFeature;
 import static com.android.server.companion.utils.PackageUtils.getPackageInfo;
+import static com.android.server.companion.utils.PackageUtils.isRestrictedSettingsAllowed;
 import static com.android.server.companion.utils.PermissionsUtils.checkCallerCanManageCompanionDevice;
 import static com.android.server.companion.utils.PermissionsUtils.enforceCallerCanManageAssociationsForPackage;
 import static com.android.server.companion.utils.PermissionsUtils.enforceCallerCanObservingDevicePresenceByUuid;
@@ -82,20 +79,16 @@
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.PackageManagerInternal;
-import android.content.pm.UserInfo;
 import android.hardware.power.Mode;
 import android.net.MacAddress;
 import android.net.NetworkPolicyManager;
 import android.os.Binder;
 import android.os.Environment;
-import android.os.Handler;
-import android.os.Message;
 import android.os.Parcel;
 import android.os.ParcelFileDescriptor;
 import android.os.ParcelUuid;
+import android.os.PowerExemptionManager;
 import android.os.PowerManagerInternal;
-import android.os.PowerWhitelistManager;
-import android.os.RemoteCallbackList;
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.os.SystemProperties;
@@ -105,13 +98,9 @@
 import android.util.ExceptionUtils;
 import android.util.Log;
 import android.util.Slog;
-import android.util.SparseArray;
-import android.util.SparseBooleanArray;
 
-import com.android.internal.annotations.GuardedBy;
 import com.android.internal.app.IAppOpsService;
 import com.android.internal.content.PackageMonitor;
-import com.android.internal.infra.PerUser;
 import com.android.internal.notification.NotificationAccessConfirmationActivityContract;
 import com.android.internal.os.BackgroundThread;
 import com.android.internal.util.ArrayUtils;
@@ -121,8 +110,8 @@
 import com.android.server.SystemService;
 import com.android.server.companion.association.AssociationDiskStore;
 import com.android.server.companion.association.AssociationRequestsProcessor;
-import com.android.server.companion.association.AssociationRevokeProcessor;
 import com.android.server.companion.association.AssociationStore;
+import com.android.server.companion.association.DisassociationProcessor;
 import com.android.server.companion.association.InactiveAssociationsRemovalService;
 import com.android.server.companion.datatransfer.SystemDataTransferProcessor;
 import com.android.server.companion.datatransfer.SystemDataTransferRequestStore;
@@ -139,7 +128,6 @@
 import java.io.File;
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
@@ -164,80 +152,51 @@
     private static final long ASSOCIATION_REMOVAL_TIME_WINDOW_DEFAULT = DAYS.toMillis(90);
     private static final int MAX_CN_LENGTH = 500;
 
-    private final ActivityManager mActivityManager;
-    private AssociationDiskStore mAssociationDiskStore;
-    private final PersistUserStateHandler mUserPersistenceHandler;
-
-    private final AssociationStore mAssociationStore;
-    private final SystemDataTransferRequestStore mSystemDataTransferRequestStore;
-    private AssociationRequestsProcessor mAssociationRequestsProcessor;
-    private SystemDataTransferProcessor mSystemDataTransferProcessor;
-    private BackupRestoreProcessor mBackupRestoreProcessor;
-    private CompanionDevicePresenceMonitor mDevicePresenceMonitor;
-    private CompanionApplicationController mCompanionAppController;
-    private CompanionTransportManager mTransportManager;
-    private AssociationRevokeProcessor mAssociationRevokeProcessor;
-
     private final ActivityTaskManagerInternal mAtmInternal;
     private final ActivityManagerInternal mAmInternal;
     private final IAppOpsService mAppOpsManager;
-    private final PowerWhitelistManager mPowerWhitelistManager;
-    private final UserManager mUserManager;
-    public final PackageManagerInternal mPackageManagerInternal;
+    private final PowerExemptionManager mPowerExemptionManager;
+    private final PackageManagerInternal mPackageManagerInternal;
     private final PowerManagerInternal mPowerManagerInternal;
 
-    /**
-     * A structure that consists 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("mPreviouslyUsedIds")
-    private final SparseArray<Map<String, Set<Integer>>> mPreviouslyUsedIds = new SparseArray<>();
-
-    private final RemoteCallbackList<IOnAssociationsChangedListener> mListeners =
-            new RemoteCallbackList<>();
-
-    private CrossDeviceSyncController mCrossDeviceSyncController;
-
-    private ObservableUuidStore mObservableUuidStore;
+    private final AssociationStore mAssociationStore;
+    private final SystemDataTransferRequestStore mSystemDataTransferRequestStore;
+    private final ObservableUuidStore mObservableUuidStore;
+    private final AssociationRequestsProcessor mAssociationRequestsProcessor;
+    private final SystemDataTransferProcessor mSystemDataTransferProcessor;
+    private final BackupRestoreProcessor mBackupRestoreProcessor;
+    private final CompanionDevicePresenceMonitor mDevicePresenceMonitor;
+    private final CompanionApplicationController mCompanionAppController;
+    private final CompanionTransportManager mTransportManager;
+    private final DisassociationProcessor mDisassociationProcessor;
+    private final CrossDeviceSyncController mCrossDeviceSyncController;
 
     public CompanionDeviceManagerService(Context context) {
         super(context);
 
-        mActivityManager = context.getSystemService(ActivityManager.class);
-        mPowerWhitelistManager = context.getSystemService(PowerWhitelistManager.class);
+        final ActivityManager activityManager = context.getSystemService(ActivityManager.class);
+        mPowerExemptionManager = context.getSystemService(PowerExemptionManager.class);
         mAppOpsManager = IAppOpsService.Stub.asInterface(
                 ServiceManager.getService(Context.APP_OPS_SERVICE));
         mAtmInternal = LocalServices.getService(ActivityTaskManagerInternal.class);
         mAmInternal = LocalServices.getService(ActivityManagerInternal.class);
         mPackageManagerInternal = LocalServices.getService(PackageManagerInternal.class);
-        mUserManager = context.getSystemService(UserManager.class);
-
-        mUserPersistenceHandler = new PersistUserStateHandler();
-        mAssociationStore = new AssociationStore();
-        mSystemDataTransferRequestStore = new SystemDataTransferRequestStore();
-
+        final UserManager userManager = context.getSystemService(UserManager.class);
         mPowerManagerInternal = LocalServices.getService(PowerManagerInternal.class);
+
+        final AssociationDiskStore associationDiskStore = new AssociationDiskStore();
+        mAssociationStore = new AssociationStore(userManager, associationDiskStore);
+        mSystemDataTransferRequestStore = new SystemDataTransferRequestStore();
         mObservableUuidStore = new ObservableUuidStore();
-    }
 
-    @Override
-    public void onStart() {
-        final Context context = getContext();
+        // Init processors
+        mAssociationRequestsProcessor = new AssociationRequestsProcessor(context,
+                mPackageManagerInternal, mAssociationStore);
+        mBackupRestoreProcessor = new BackupRestoreProcessor(context, mPackageManagerInternal,
+                mAssociationStore, associationDiskStore, mSystemDataTransferRequestStore,
+                mAssociationRequestsProcessor);
 
-        mAssociationDiskStore = new AssociationDiskStore();
-        mAssociationRequestsProcessor = new AssociationRequestsProcessor(
-                /* cdmService */ this, mAssociationStore);
-        mBackupRestoreProcessor = new BackupRestoreProcessor(
-                /* cdmService */ this, mAssociationStore, mAssociationDiskStore,
-                mSystemDataTransferRequestStore, mAssociationRequestsProcessor);
-
-        mObservableUuidStore.getObservableUuidsForUser(getContext().getUserId());
-
-        mAssociationStore.registerListener(mAssociationStoreChangeListener);
-
-        mDevicePresenceMonitor = new CompanionDevicePresenceMonitor(mUserManager,
+        mDevicePresenceMonitor = new CompanionDevicePresenceMonitor(userManager,
                 mAssociationStore, mObservableUuidStore, mDevicePresenceCallback);
 
         mCompanionAppController = new CompanionApplicationController(
@@ -246,11 +205,9 @@
 
         mTransportManager = new CompanionTransportManager(context, mAssociationStore);
 
-        mAssociationRevokeProcessor = new AssociationRevokeProcessor(this, mAssociationStore,
-                mPackageManagerInternal, mDevicePresenceMonitor, mCompanionAppController,
-                mSystemDataTransferRequestStore, mTransportManager);
-
-        loadAssociationsFromDisk();
+        mDisassociationProcessor = new DisassociationProcessor(context, activityManager,
+                mAssociationStore, mPackageManagerInternal, mDevicePresenceMonitor,
+                mCompanionAppController, mSystemDataTransferRequestStore, mTransportManager);
 
         mSystemDataTransferProcessor = new SystemDataTransferProcessor(this,
                 mPackageManagerInternal, mAssociationStore,
@@ -258,6 +215,16 @@
 
         // TODO(b/279663946): move context sync to a dedicated system service
         mCrossDeviceSyncController = new CrossDeviceSyncController(getContext(), mTransportManager);
+    }
+
+    @Override
+    public void onStart() {
+        // Init association stores
+        mAssociationStore.refreshCache();
+        mAssociationStore.registerLocalListener(mAssociationStoreChangeListener);
+
+        // Init UUID store
+        mObservableUuidStore.getObservableUuidsForUser(getContext().getUserId());
 
         // Publish "binder" service.
         final CompanionDeviceManagerImpl impl = new CompanionDeviceManagerImpl();
@@ -267,50 +234,6 @@
         LocalServices.addService(CompanionDeviceManagerServiceInternal.class, new LocalService());
     }
 
-    void loadAssociationsFromDisk() {
-        final Set<AssociationInfo> allAssociations = new ArraySet<>();
-        synchronized (mPreviouslyUsedIds) {
-            List<Integer> userIds = new ArrayList<>();
-            for (UserInfo user : mUserManager.getAliveUsers()) {
-                userIds.add(user.id);
-            }
-            // 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).
-            mAssociationDiskStore.readStateForUsers(userIds, allAssociations, mPreviouslyUsedIds);
-        }
-
-        final Set<AssociationInfo> activeAssociations =
-                new ArraySet<>(/* capacity */ allAssociations.size());
-        // A set contains the userIds that need to persist state after remove the app
-        // from the list of role holders.
-        final Set<Integer> usersToPersistStateFor = new ArraySet<>();
-
-        for (AssociationInfo association : allAssociations) {
-            if (association.isPending()) {
-                mBackupRestoreProcessor.addToPendingAppInstall(association);
-            } else if (!association.isRevoked()) {
-                activeAssociations.add(association);
-            } else if (mAssociationRevokeProcessor.maybeRemoveRoleHolderForAssociation(
-                    association)) {
-                // Nothing more to do here, but we'll need to persist all the associations to the
-                // disk afterwards.
-                usersToPersistStateFor.add(association.getUserId());
-            } else {
-                mAssociationRevokeProcessor.addToPendingRoleHolderRemoval(association);
-            }
-        }
-
-        mAssociationStore.setAssociationsToCache(activeAssociations);
-
-        // IMPORTANT: only do this AFTER mAssociationStore.setAssociations(), because
-        // persistStateForUser() queries AssociationStore.
-        // (If persistStateForUser() is invoked before mAssociationStore.setAssociations() it
-        // would effectively just clear-out all the persisted associations).
-        for (int userId : usersToPersistStateFor) {
-            persistStateForUser(userId);
-        }
-    }
-
     @Override
     public void onBootPhase(int phase) {
         final Context context = getContext();
@@ -329,8 +252,10 @@
 
     @Override
     public void onUserUnlocking(@NonNull TargetUser user) {
+        Slog.d(TAG, "onUserUnlocking...");
         final int userId = user.getUserIdentifier();
-        final List<AssociationInfo> associations = mAssociationStore.getAssociationsForUser(userId);
+        final List<AssociationInfo> associations = mAssociationStore.getActiveAssociationsByUser(
+                userId);
 
         if (associations.isEmpty()) return;
 
@@ -359,7 +284,8 @@
                         ? Collections.emptyList() : Arrays.asList(bluetoothDeviceUuids);
 
                 for (AssociationInfo ai :
-                        mAssociationStore.getAssociationsByAddress(bluetoothDevice.getAddress())) {
+                        mAssociationStore.getActiveAssociationsByAddress(
+                                bluetoothDevice.getAddress())) {
                     Slog.i(TAG, "onUserUnlocked, device id( " + ai.getId() + " ) is connected");
                     mDevicePresenceMonitor.onBluetoothCompanionDeviceConnected(ai.getId());
                 }
@@ -379,7 +305,7 @@
     @NonNull
     AssociationInfo getAssociationWithCallerChecks(
             @UserIdInt int userId, @NonNull String packageName, @NonNull String macAddress) {
-        AssociationInfo association = mAssociationStore.getAssociationsForPackageWithAddress(
+        AssociationInfo association = mAssociationStore.getFirstAssociationByAddress(
                 userId, packageName, macAddress);
         association = sanitizeWithCallerChecks(getContext(), association);
         if (association != null) {
@@ -533,7 +459,7 @@
      */
     private boolean shouldBindPackage(@UserIdInt int userId, @NonNull String packageName) {
         final List<AssociationInfo> packageAssociations =
-                mAssociationStore.getAssociationsForPackage(userId, packageName);
+                mAssociationStore.getActiveAssociationsByPackage(userId, packageName);
         final List<ObservableUuid> observableUuids =
                 mObservableUuidStore.getObservableUuidsForPackage(userId, packageName);
 
@@ -551,77 +477,6 @@
         return false;
     }
 
-    private void onAssociationChangedInternal(
-            @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);
-
-        mUserPersistenceHandler.postPersistUserState(userId);
-
-        // 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);
-    }
-
-    void persistStateForUser(@UserIdInt int userId) {
-        // We want to store both active associations and the revoked (removed) association that we
-        // are keeping around for the final clean-up (delayed role holder removal).
-        final List<AssociationInfo> allAssociations;
-        // Start with the active associations - these we can get from the AssociationStore.
-        allAssociations = new ArrayList<>(
-                mAssociationStore.getAssociationsForUser(userId));
-        // ... and add the revoked (removed) association, that are yet to be permanently removed.
-        allAssociations.addAll(
-                mAssociationRevokeProcessor.getPendingRoleHolderRemovalAssociationsForUser(userId));
-        // ... and add the restored associations that are pending missing package installation.
-        allAssociations.addAll(mBackupRestoreProcessor
-                .getAssociationsPendingAppInstallForUser(userId));
-
-        final Map<String, Set<Integer>> usedIdsForUser = getPreviouslyUsedIdsForUser(userId);
-
-        mAssociationDiskStore.persistStateForUser(userId, allAssociations, usedIdsForUser);
-    }
-
-    private void notifyListeners(
-            @UserIdInt int userId, @NonNull List<AssociationInfo> associations) {
-        mListeners.broadcast((listener, callbackUserId) -> {
-            int listenerUserId = (int) callbackUserId;
-            if (listenerUserId == userId || listenerUserId == UserHandle.USER_ALL) {
-                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);
-        }
-    }
-
     private void onPackageRemoveOrDataClearedInternal(
             @UserIdInt int userId, @NonNull String packageName) {
         if (DEBUG) {
@@ -629,19 +484,20 @@
                     + packageName);
         }
 
-        // Clear associations.
+        // Clear all associations for the package.
         final List<AssociationInfo> associationsForPackage =
-                mAssociationStore.getAssociationsForPackage(userId, packageName);
+                mAssociationStore.getAssociationsByPackage(userId, packageName);
+        if (!associationsForPackage.isEmpty()) {
+            Slog.i(TAG, "Package removed or data cleared for user=[" + userId + "], package=["
+                    + packageName + "]. Cleaning up CDM data...");
+        }
+        for (AssociationInfo association : associationsForPackage) {
+            mDisassociationProcessor.disassociate(association.getId());
+        }
+
+        // Clear observable UUIDs for the package.
         final List<ObservableUuid> uuidsTobeObserved =
                 mObservableUuidStore.getObservableUuidsForPackage(userId, packageName);
-        for (AssociationInfo association : associationsForPackage) {
-            mAssociationStore.removeAssociation(association.getId());
-        }
-        // Clear role holders
-        for (AssociationInfo association : associationsForPackage) {
-            mAssociationRevokeProcessor.maybeRemoveRoleHolderForAssociation(association);
-        }
-        // Clear the uuids to be observed.
         for (ObservableUuid uuid : uuidsTobeObserved) {
             mObservableUuidStore.removeObservableUuid(userId, uuid.getUuid(), packageName);
         }
@@ -652,31 +508,13 @@
     private void onPackageModifiedInternal(@UserIdInt int userId, @NonNull String packageName) {
         if (DEBUG) Log.i(TAG, "onPackageModified() u" + userId + "/" + packageName);
 
-        final List<AssociationInfo> associationsForPackage =
-                mAssociationStore.getAssociationsForPackage(userId, packageName);
-        for (AssociationInfo association : associationsForPackage) {
-            updateSpecialAccessPermissionForAssociatedPackage(association.getUserId(),
-                    association.getPackageName());
-        }
+        updateSpecialAccessPermissionForAssociatedPackage(userId, packageName);
 
         mCompanionAppController.onPackagesChanged(userId);
     }
 
     private void onPackageAddedInternal(@UserIdInt int userId, @NonNull String packageName) {
-        if (DEBUG) Log.i(TAG, "onPackageAddedInternal() u" + userId + "/" + packageName);
-
-        Set<AssociationInfo> associationsPendingAppInstall = mBackupRestoreProcessor
-                .getAssociationsPendingAppInstallForUser(userId);
-        for (AssociationInfo association : associationsPendingAppInstall) {
-            if (!packageName.equals(association.getPackageName())) continue;
-
-            AssociationInfo newAssociation = new AssociationInfo.Builder(association)
-                    .setPending(false)
-                    .build();
-            mAssociationRequestsProcessor.maybeGrantRoleAndStoreAssociation(newAssociation,
-                    null, null);
-            mBackupRestoreProcessor.removeFromPendingAppInstall(association);
-        }
+        mBackupRestoreProcessor.restorePendingAssociations(userId, packageName);
     }
 
     // Revoke associations if the selfManaged companion device does not connect for 3 months.
@@ -698,7 +536,7 @@
             final int id = association.getId();
 
             Slog.i(TAG, "Removing inactive self-managed association id=" + id);
-            mAssociationRevokeProcessor.disassociateInternal(id);
+            mDisassociationProcessor.disassociate(id);
         }
     }
 
@@ -750,7 +588,7 @@
                 enforceUsesCompanionDeviceFeature(getContext(), userId, packageName);
             }
 
-            return mAssociationStore.getAssociationsForPackage(userId, packageName);
+            return mAssociationStore.getActiveAssociationsByPackage(userId, packageName);
         }
 
         @Override
@@ -761,9 +599,9 @@
             enforceCallerIsSystemOrCanInteractWithUserId(getContext(), userId);
 
             if (userId == UserHandle.USER_ALL) {
-                return List.copyOf(mAssociationStore.getAssociations());
+                return mAssociationStore.getActiveAssociations();
             }
-            return mAssociationStore.getAssociationsForUser(userId);
+            return mAssociationStore.getActiveAssociationsByUser(userId);
         }
 
         @Override
@@ -773,7 +611,8 @@
             addOnAssociationsChangedListener_enforcePermission();
 
             enforceCallerIsSystemOrCanInteractWithUserId(getContext(), userId);
-            mListeners.register(listener, userId);
+
+            mAssociationStore.registerRemoteListener(listener, userId);
         }
 
         @Override
@@ -784,7 +623,7 @@
 
             enforceCallerIsSystemOrCanInteractWithUserId(getContext(), userId);
 
-            mListeners.unregister(listener);
+            mAssociationStore.unregisterRemoteListener(listener);
         }
 
         @Override
@@ -843,16 +682,16 @@
 
             final AssociationInfo association =
                     getAssociationWithCallerChecks(userId, packageName, deviceMacAddress);
-            mAssociationRevokeProcessor.disassociateInternal(association.getId());
+            mDisassociationProcessor.disassociate(association.getId());
         }
 
         @Override
         public void disassociate(int associationId) {
-            Log.i(TAG, "disassociate() associationId=" + associationId);
+            Slog.i(TAG, "disassociate() associationId=" + associationId);
 
             final AssociationInfo association =
                     getAssociationWithCallerChecks(associationId);
-            mAssociationRevokeProcessor.disassociateInternal(association.getId());
+            mDisassociationProcessor.disassociate(association.getId());
         }
 
         @Override
@@ -867,8 +706,7 @@
                 throw new IllegalArgumentException("Component name is too long.");
             }
 
-            final long identity = Binder.clearCallingIdentity();
-            try {
+            return Binder.withCleanCallingIdentity(() -> {
                 if (!isRestrictedSettingsAllowed(getContext(), callingPackage, callingUid)) {
                     Slog.e(TAG, "Side loaded app must enable restricted "
                             + "setting before request the notification access");
@@ -882,9 +720,7 @@
                                 | PendingIntent.FLAG_CANCEL_CURRENT,
                         null /* options */,
                         new UserHandle(userId));
-            } finally {
-                Binder.restoreCallingIdentity(identity);
-            }
+            });
         }
 
         /**
@@ -912,7 +748,7 @@
                 return true;
             }
 
-            return any(mAssociationStore.getAssociationsForPackage(userId, packageName),
+            return any(mAssociationStore.getActiveAssociationsByPackage(userId, packageName),
                     a -> a.isLinkedTo(macAddress));
         }
 
@@ -1166,7 +1002,7 @@
             final int userId = getCallingUserId();
             enforceCallerIsSystemOr(userId, packageName);
 
-            AssociationInfo association = mAssociationStore.getAssociationsForPackageWithAddress(
+            AssociationInfo association = mAssociationStore.getFirstAssociationByAddress(
                     userId, packageName, deviceAddress);
 
             if (association == null) {
@@ -1239,14 +1075,15 @@
 
             enforceUsesCompanionDeviceFeature(getContext(), userId, callingPackage);
             checkState(!ArrayUtils.isEmpty(
-                            mAssociationStore.getAssociationsForPackage(userId, callingPackage)),
+                            mAssociationStore.getActiveAssociationsByPackage(userId,
+                                    callingPackage)),
                     "App must have an association before calling this API");
         }
 
         @Override
         public boolean canPairWithoutPrompt(String packageName, String macAddress, int userId) {
             final AssociationInfo association =
-                    mAssociationStore.getAssociationsForPackageWithAddress(
+                    mAssociationStore.getFirstAssociationByAddress(
                             userId, packageName, macAddress);
             if (association == null) {
                 return false;
@@ -1269,13 +1106,11 @@
 
         @Override
         public byte[] getBackupPayload(int userId) {
-            Log.i(TAG, "getBackupPayload() userId=" + userId);
             return mBackupRestoreProcessor.getBackupPayload(userId);
         }
 
         @Override
         public void applyRestoredPayload(byte[] payload, int userId) {
-            Log.i(TAG, "applyRestoredPayload() userId=" + userId);
             mBackupRestoreProcessor.applyRestoredPayload(payload, userId);
         }
 
@@ -1286,7 +1121,7 @@
             return new CompanionDeviceShellCommand(CompanionDeviceManagerService.this,
                     mAssociationStore, mDevicePresenceMonitor, mTransportManager,
                     mSystemDataTransferProcessor, mAssociationRequestsProcessor,
-                    mBackupRestoreProcessor, mAssociationRevokeProcessor)
+                    mBackupRestoreProcessor, mDisassociationProcessor)
                     .exec(this, in.getFileDescriptor(), out.getFileDescriptor(),
                             err.getFileDescriptor(), args);
         }
@@ -1314,88 +1149,6 @@
                 /* callback */ null, /* resultReceiver */ null);
     }
 
-    @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) {
-        // "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;
-    }
-
-    /**
-     * Get a new association id for the package.
-     */
-    public int getNewAssociationIdForPackage(@UserIdInt int userId, @NonNull String packageName) {
-        synchronized (mPreviouslyUsedIds) {
-            // First: collect all IDs currently in use for this user's Associations.
-            final SparseBooleanArray usedIds = new SparseBooleanArray();
-
-            // We should really only be checking associations for the given user (i.e.:
-            // mAssociationStore.getAssociationsForUser(userId)), BUT in the past we've got in a
-            // state where association IDs were not assigned correctly in regard to
-            // user-to-association-ids-range (e.g. associations with IDs from 1 to 100,000 should
-            // always belong to u0), so let's check all the associations.
-            for (AssociationInfo it : mAssociationStore.getAssociations()) {
-                usedIds.put(it.getId(), true);
-            }
-
-            // Some IDs may be reserved by associations that aren't stored yet due to missing
-            // package after a backup restoration. We don't want the ID to have been taken by
-            // another association by the time when it is activated from the package installation.
-            final Set<AssociationInfo> pendingAssociations = mBackupRestoreProcessor
-                    .getAssociationsPendingAppInstallForUser(userId);
-            for (AssociationInfo it : pendingAssociations) {
-                usedIds.put(it.getId(), true);
-            }
-
-            // Second: collect all IDs that have been previously used for this package (and user).
-            final Set<Integer> previouslyUsedIds =
-                    getPreviouslyUsedIdsForPackageLocked(userId, packageName);
-
-            int id = getFirstAssociationIdForUser(userId);
-            final int lastAvailableIdForUser = getLastAssociationIdForUser(userId);
-
-            // Find first ID that isn't used now AND has never been used for the given package.
-            while (usedIds.get(id) || previouslyUsedIds.contains(id)) {
-                // Increment and try again
-                id++;
-                // ... but first check if the ID is valid (within the range allocated to the user).
-                if (id > lastAvailableIdForUser) {
-                    throw new RuntimeException("Cannot create a new Association ID for "
-                            + packageName + " for user " + userId);
-                }
-            }
-
-            return id;
-        }
-    }
-
     /**
      * Update special access for the association's package
      */
@@ -1403,20 +1156,27 @@
         final PackageInfo packageInfo =
                 getPackageInfo(getContext(), userId, packageName);
 
-        Binder.withCleanCallingIdentity(() -> updateSpecialAccessPermissionAsSystem(packageInfo));
+        Binder.withCleanCallingIdentity(() -> updateSpecialAccessPermissionAsSystem(packageInfo,
+                userId, packageName));
     }
 
-    private void updateSpecialAccessPermissionAsSystem(PackageInfo packageInfo) {
+    private void updateSpecialAccessPermissionAsSystem(PackageInfo packageInfo, int userId,
+            String packageName) {
         if (packageInfo == null) {
             return;
         }
+
+        List<AssociationInfo> associations = mAssociationStore.getActiveAssociationsByPackage(
+                userId, packageName);
+
         if (containsEither(packageInfo.requestedPermissions,
                 android.Manifest.permission.RUN_IN_BACKGROUND,
-                android.Manifest.permission.REQUEST_COMPANION_RUN_IN_BACKGROUND)) {
-            mPowerWhitelistManager.addToWhitelist(packageInfo.packageName);
+                android.Manifest.permission.REQUEST_COMPANION_RUN_IN_BACKGROUND)
+                && !associations.isEmpty()) {
+            mPowerExemptionManager.addToPermanentAllowList(packageInfo.packageName);
         } else {
             try {
-                mPowerWhitelistManager.removeFromWhitelist(packageInfo.packageName);
+                mPowerExemptionManager.removeFromPermanentAllowList(packageInfo.packageName);
             } catch (UnsupportedOperationException e) {
                 Slog.w(TAG, packageInfo.packageName + " can't be removed from power save"
                         + " whitelist. It might due to the package is whitelisted by the system.");
@@ -1427,7 +1187,8 @@
         try {
             if (containsEither(packageInfo.requestedPermissions,
                     android.Manifest.permission.USE_DATA_IN_BACKGROUND,
-                    android.Manifest.permission.REQUEST_COMPANION_USE_DATA_IN_BACKGROUND)) {
+                    android.Manifest.permission.REQUEST_COMPANION_USE_DATA_IN_BACKGROUND)
+                    && !associations.isEmpty()) {
                 networkPolicyManager.addUidPolicy(
                         packageInfo.applicationInfo.uid,
                         NetworkPolicyManager.POLICY_ALLOW_METERED_BACKGROUND);
@@ -1487,7 +1248,7 @@
 
             try {
                 final List<AssociationInfo> associations =
-                        mAssociationStore.getAssociationsForUser(userId);
+                        mAssociationStore.getActiveAssociationsByUser(userId);
                 for (AssociationInfo a : associations) {
                     try {
                         int uid = pm.getPackageUidAsUser(a.getPackageName(), userId);
@@ -1506,7 +1267,16 @@
             new AssociationStore.OnChangeListener() {
                 @Override
                 public void onAssociationChanged(int changeType, AssociationInfo association) {
-                    onAssociationChangedInternal(changeType, association);
+                    Slog.d(TAG, "onAssociationChanged changeType=[" + changeType
+                            + "], association=[" + association);
+
+                    final int userId = association.getUserId();
+                    final List<AssociationInfo> updatedAssociations =
+                            mAssociationStore.getActiveAssociationsByUser(userId);
+
+                    updateAtm(userId, updatedAssociations);
+                    updateSpecialAccessPermissionForAssociatedPackage(association.getUserId(),
+                            association.getPackageName());
                 }
             };
 
@@ -1634,64 +1404,4 @@
             }
         }
     }
-
-    /**
-     * This method must only be called from {@link CompanionDeviceShellCommand} for testing
-     * purposes only!
-     */
-    void persistState() {
-        mUserPersistenceHandler.clearMessages();
-        for (UserInfo user : mUserManager.getAliveUsers()) {
-            persistStateForUser(user.id);
-        }
-    }
-
-    /**
-     * This class is dedicated to handling requests to persist user state.
-     */
-    @SuppressLint("HandlerLeak")
-    private class PersistUserStateHandler extends Handler {
-        PersistUserStateHandler() {
-            super(BackgroundThread.get().getLooper());
-        }
-
-        /**
-         * Persists user state unless there is already an outstanding request for the given user.
-         */
-        synchronized void postPersistUserState(@UserIdInt int userId) {
-            if (!hasMessages(userId)) {
-                sendMessage(obtainMessage(userId));
-            }
-        }
-
-        /**
-         * Clears *ALL* outstanding persist requests for *ALL* users.
-         */
-        synchronized void clearMessages() {
-            removeCallbacksAndMessages(null);
-        }
-
-        @Override
-        public void handleMessage(@NonNull Message msg) {
-            final int userId = msg.what;
-            persistStateForUser(userId);
-        }
-    }
-
-    /**
-     * Persist associations
-     */
-    public void postPersistUserState(@UserIdInt int userId) {
-        mUserPersistenceHandler.postPersistUserState(userId);
-    }
-
-    /**
-     * Set to store associations
-     */
-    public static class PerUserAssociationSet extends PerUser<Set<AssociationInfo>> {
-        @Override
-        protected @NonNull Set<AssociationInfo> create(int userId) {
-            return new ArraySet<>();
-        }
-    }
 }
diff --git a/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java b/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java
index 16877dc..a7a73cb 100644
--- a/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java
+++ b/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java
@@ -33,8 +33,8 @@
 import android.util.proto.ProtoOutputStream;
 
 import com.android.server.companion.association.AssociationRequestsProcessor;
-import com.android.server.companion.association.AssociationRevokeProcessor;
 import com.android.server.companion.association.AssociationStore;
+import com.android.server.companion.association.DisassociationProcessor;
 import com.android.server.companion.datatransfer.SystemDataTransferProcessor;
 import com.android.server.companion.datatransfer.contextsync.BitmapUtils;
 import com.android.server.companion.datatransfer.contextsync.CrossDeviceSyncController;
@@ -49,7 +49,7 @@
     private static final String TAG = "CDM_CompanionDeviceShellCommand";
 
     private final CompanionDeviceManagerService mService;
-    private final AssociationRevokeProcessor mRevokeProcessor;
+    private final DisassociationProcessor mDisassociationProcessor;
     private final AssociationStore mAssociationStore;
     private final CompanionDevicePresenceMonitor mDevicePresenceMonitor;
     private final CompanionTransportManager mTransportManager;
@@ -65,7 +65,7 @@
             SystemDataTransferProcessor systemDataTransferProcessor,
             AssociationRequestsProcessor associationRequestsProcessor,
             BackupRestoreProcessor backupRestoreProcessor,
-            AssociationRevokeProcessor revokeProcessor) {
+            DisassociationProcessor disassociationProcessor) {
         mService = service;
         mAssociationStore = associationStore;
         mDevicePresenceMonitor = devicePresenceMonitor;
@@ -73,7 +73,7 @@
         mSystemDataTransferProcessor = systemDataTransferProcessor;
         mAssociationRequestsProcessor = associationRequestsProcessor;
         mBackupRestoreProcessor = backupRestoreProcessor;
-        mRevokeProcessor = revokeProcessor;
+        mDisassociationProcessor = disassociationProcessor;
     }
 
     @Override
@@ -105,12 +105,15 @@
                 case "list": {
                     final int userId = getNextIntArgRequired();
                     final List<AssociationInfo> associationsForUser =
-                            mAssociationStore.getAssociationsForUser(userId);
+                            mAssociationStore.getActiveAssociationsByUser(userId);
+                    final int maxId = mAssociationStore.getMaxId(userId);
+                    out.println("Max ID: " + maxId);
+                    out.println("Association ID | Package Name | Mac Address");
                     for (AssociationInfo association : associationsForUser) {
                         // TODO(b/212535524): use AssociationInfo.toShortString(), once it's not
                         //  longer referenced in tests.
-                        out.println(association.getPackageName() + " "
-                                + association.getDeviceMacAddress() + " " + association.getId());
+                        out.println(association.getId() + " | " + association.getPackageName()
+                                + " | " + association.getDeviceMacAddress());
                     }
                 }
                 break;
@@ -132,28 +135,24 @@
                     final String address = getNextArgRequired();
                     final AssociationInfo association =
                             mService.getAssociationWithCallerChecks(userId, packageName, address);
-                    if (association != null) {
-                        mRevokeProcessor.disassociateInternal(association.getId());
-                    }
+                    mDisassociationProcessor.disassociate(association.getId());
                 }
                 break;
 
                 case "disassociate-all": {
                     final int userId = getNextIntArgRequired();
-                    final String packageName = getNextArgRequired();
                     final List<AssociationInfo> userAssociations =
-                            mAssociationStore.getAssociationsForPackage(userId, packageName);
+                            mAssociationStore.getAssociationsByUser(userId);
                     for (AssociationInfo association : userAssociations) {
                         if (sanitizeWithCallerChecks(mService.getContext(), association) != null) {
-                            mRevokeProcessor.disassociateInternal(association.getId());
+                            mDisassociationProcessor.disassociate(association.getId());
                         }
                     }
                 }
                 break;
 
-                case "clear-association-memory-cache":
-                    mService.persistState();
-                    mService.loadAssociationsFromDisk();
+                case "refresh-cache":
+                    mAssociationStore.refreshCache();
                     break;
 
                 case "simulate-device-appeared":
diff --git a/services/companion/java/com/android/server/companion/association/AssociationDiskStore.java b/services/companion/java/com/android/server/companion/association/AssociationDiskStore.java
index 75cb120..46d60f9 100644
--- a/services/companion/java/com/android/server/companion/association/AssociationDiskStore.java
+++ b/services/companion/java/com/android/server/companion/association/AssociationDiskStore.java
@@ -16,7 +16,6 @@
 
 package com.android.server.companion.association;
 
-import static com.android.internal.util.CollectionUtils.forEach;
 import static com.android.internal.util.XmlUtils.readBooleanAttribute;
 import static com.android.internal.util.XmlUtils.readIntAttribute;
 import static com.android.internal.util.XmlUtils.readLongAttribute;
@@ -26,7 +25,6 @@
 import static com.android.internal.util.XmlUtils.writeLongAttribute;
 import static com.android.internal.util.XmlUtils.writeStringAttribute;
 import static com.android.server.companion.utils.AssociationUtils.getFirstAssociationIdForUser;
-import static com.android.server.companion.utils.AssociationUtils.getLastAssociationIdForUser;
 import static com.android.server.companion.utils.DataStoreUtils.createStorageFileForUser;
 import static com.android.server.companion.utils.DataStoreUtils.fileToByteArray;
 import static com.android.server.companion.utils.DataStoreUtils.isEndOfTag;
@@ -40,10 +38,8 @@
 import android.companion.AssociationInfo;
 import android.net.MacAddress;
 import android.os.Environment;
-import android.util.ArrayMap;
 import android.util.AtomicFile;
 import android.util.Slog;
-import android.util.SparseArray;
 import android.util.Xml;
 
 import com.android.internal.util.XmlUtils;
@@ -59,11 +55,9 @@
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
-import java.util.Collection;
-import java.util.HashSet;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
@@ -82,8 +76,8 @@
  * <p>
  * Before Android T the data was stored using the v0 schema. See:
  * <ul>
- * <li>{@link #readAssociationsV0(TypedXmlPullParser, int, Collection) readAssociationsV0()}.
- * <li>{@link #readAssociationV0(TypedXmlPullParser, int, int, Collection) readAssociationV0()}.
+ * <li>{@link #readAssociationsV0(TypedXmlPullParser, int) readAssociationsV0()}.
+ * <li>{@link #readAssociationV0(TypedXmlPullParser, int, int) readAssociationV0()}.
  * </ul>
  *
  * The following snippet is a sample of a file that is using v0 schema.
@@ -116,15 +110,14 @@
  * optional.
  * <ul>
  * <li> {@link #CURRENT_PERSISTENCE_VERSION}
- * <li> {@link #readAssociationsV1(TypedXmlPullParser, int, Collection) readAssociationsV1()}
- * <li> {@link #readAssociationV1(TypedXmlPullParser, int, Collection) readAssociationV1()}
- * <li> {@link #readPreviouslyUsedIdsV1(TypedXmlPullParser, Map) readPreviouslyUsedIdsV1()}
+ * <li> {@link #readAssociationsV1(TypedXmlPullParser, int) readAssociationsV1()}
+ * <li> {@link #readAssociationV1(TypedXmlPullParser, int) readAssociationV1()}
  * </ul>
  *
  * The following snippet is a sample of a file that is using v1 schema.
  * <pre>{@code
  * <state persistence-version="1">
- *     <associations>
+ *     <associations max-id="3">
  *         <association
  *             id="1"
  *             package="com.sample.companion.app"
@@ -148,18 +141,12 @@
  *             time_approved="1634641160229"
  *             system_data_sync_flags="1"/>
  *     </associations>
- *
- *     <previously-used-ids>
- *         <package package_name="com.sample.companion.app">
- *             <id>2</id>
- *         </package>
- *     </previously-used-ids>
  * </state>
  * }</pre>
  */
 @SuppressLint("LongLogTag")
 public final class AssociationDiskStore {
-    private static final String TAG = "CompanionDevice_AssociationDiskStore";
+    private static final String TAG = "CDM_AssociationDiskStore";
 
     private static final int CURRENT_PERSISTENCE_VERSION = 1;
 
@@ -169,16 +156,11 @@
     private static final String XML_TAG_STATE = "state";
     private static final String XML_TAG_ASSOCIATIONS = "associations";
     private static final String XML_TAG_ASSOCIATION = "association";
-    private static final String XML_TAG_PREVIOUSLY_USED_IDS = "previously-used-ids";
-    private static final String XML_TAG_PACKAGE = "package";
     private static final String XML_TAG_TAG = "tag";
-    private static final String XML_TAG_ID = "id";
 
     private static final String XML_ATTR_PERSISTENCE_VERSION = "persistence-version";
+    private static final String XML_ATTR_MAX_ID = "max-id";
     private static final String XML_ATTR_ID = "id";
-    // Used in <package> elements, nested within <previously-used-ids> elements.
-    private static final String XML_ATTR_PACKAGE_NAME = "package_name";
-    // Used in <association> elements, nested within <associations> elements.
     private static final String XML_ATTR_PACKAGE = "package";
     private static final String XML_ATTR_MAC_ADDRESS = "mac_address";
     private static final String XML_ATTR_DISPLAY_NAME = "display_name";
@@ -199,38 +181,12 @@
     /**
      * Read all associations for given users
      */
-    public void readStateForUsers(@NonNull List<Integer> userIds,
-            @NonNull Set<AssociationInfo> allAssociationsOut,
-            @NonNull SparseArray<Map<String, Set<Integer>>> previouslyUsedIdsPerUserOut) {
+    public Map<Integer, Associations> readAssociationsByUsers(@NonNull List<Integer> userIds) {
+        Map<Integer, Associations> userToAssociationsMap = new HashMap<>();
         for (int userId : userIds) {
-            // 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.
-            final Set<AssociationInfo> associationsForUser = new HashSet<>();
-            readStateForUser(userId, associationsForUser, previouslyUsedIds);
-
-            // Go through all the associations for the user and check if their IDs are within
-            // the allowed range (for the user).
-            final int firstAllowedId = getFirstAssociationIdForUser(userId);
-            final int lastAllowedId = getLastAssociationIdForUser(userId);
-            for (AssociationInfo association : associationsForUser) {
-                final int id = association.getId();
-                if (id < firstAllowedId || id > lastAllowedId) {
-                    Slog.e(TAG, "Wrong association ID assignment: " + id + ". "
-                            + "Association belongs to u" + userId + " and thus its ID should be "
-                            + "within [" + firstAllowedId + ", " + lastAllowedId + "] range.");
-                    // TODO(b/224736262): try fixing (re-assigning) the ID?
-                }
-            }
-
-            // Add user's association to the "output" set.
-            allAssociationsOut.addAll(associationsForUser);
-
-            // Save previously used IDs for this user into the "out" structure.
-            previouslyUsedIdsPerUserOut.append(userId, previouslyUsedIds);
+            userToAssociationsMap.put(userId, readAssociationsByUser(userId));
         }
+        return userToAssociationsMap;
     }
 
     /**
@@ -240,16 +196,12 @@
      * retrieval from this datastore because it is not persisted (by design). This means that
      * persisted data is not guaranteed to be identical to the initial data that was stored at the
      * time of association.
-     *
-     * @param userId Android UserID
-     * @param associationsOut a container to read the {@link AssociationInfo}s "into".
-     * @param previouslyUsedIdsPerPackageOut a container to read the used IDs "into".
      */
-    private void readStateForUser(@UserIdInt int userId,
-            @NonNull Collection<AssociationInfo> associationsOut,
-            @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackageOut) {
-        Slog.i(TAG, "Reading associations for user " + userId + " from disk");
+    @NonNull
+    private Associations readAssociationsByUser(@UserIdInt int userId) {
+        Slog.i(TAG, "Reading associations for user " + userId + " from disk.");
         final AtomicFile file = getStorageFileForUser(userId);
+        Associations associations;
 
         // getStorageFileForUser() ALWAYS returns the SAME OBJECT, which allows us to synchronize
         // accesses to the file on the file system using this AtomicFile object.
@@ -260,7 +212,7 @@
             if (!file.getBaseFile().exists()) {
                 legacyBaseFile = getBaseLegacyStorageFileForUser(userId);
                 if (!legacyBaseFile.exists()) {
-                    return;
+                    return new Associations();
                 }
 
                 readFrom = new AtomicFile(legacyBaseFile);
@@ -270,13 +222,12 @@
                 rootTag = XML_TAG_STATE;
             }
 
-            final int version = readStateFromFileLocked(userId, readFrom, rootTag,
-                    associationsOut, previouslyUsedIdsPerPackageOut);
+            associations = readAssociationsFromFile(userId, readFrom, rootTag);
 
-            if (legacyBaseFile != null || version < CURRENT_PERSISTENCE_VERSION) {
+            if (legacyBaseFile != null || associations.getVersion() < CURRENT_PERSISTENCE_VERSION) {
                 // The data is either in the legacy file or in the legacy format, or both.
                 // Save the data to right file in using the current format.
-                persistStateToFileLocked(file, associationsOut, previouslyUsedIdsPerPackageOut);
+                writeAssociationsToFile(file, associations);
 
                 if (legacyBaseFile != null) {
                     // We saved the data to the right file, can delete the old file now.
@@ -284,89 +235,75 @@
                 }
             }
         }
+        return associations;
     }
 
     /**
-     * Persisted data to the disk.
-     *
-     * Note that associatedDevice field in {@link AssociationInfo} is not persisted by this
-     * datastore implementation.
-     *
-     * @param userId Android UserID
-     * @param associations a set of user's associations.
-     * @param previouslyUsedIdsPerPackage a set previously used Association IDs for the user.
+     * Write associations to disk for the user.
      */
-    public void persistStateForUser(@UserIdInt int userId,
-            @NonNull Collection<AssociationInfo> associations,
-            @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackage) {
+    public void writeAssociationsForUser(@UserIdInt int userId,
+            @NonNull Associations associations) {
         Slog.i(TAG, "Writing associations for user " + userId + " to disk");
 
         final AtomicFile file = getStorageFileForUser(userId);
         // getStorageFileForUser() ALWAYS returns the SAME OBJECT, which allows us to synchronize
         // accesses to the file on the file system using this AtomicFile object.
         synchronized (file) {
-            persistStateToFileLocked(file, associations, previouslyUsedIdsPerPackage);
+            writeAssociationsToFile(file, associations);
         }
     }
 
-    private int readStateFromFileLocked(@UserIdInt int userId, @NonNull AtomicFile file,
-            @NonNull String rootTag, @Nullable Collection<AssociationInfo> associationsOut,
-            @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackageOut) {
+    @NonNull
+    private static Associations readAssociationsFromFile(@UserIdInt int userId,
+            @NonNull AtomicFile file, @NonNull String rootTag) {
         try (FileInputStream in = file.openRead()) {
-            return readStateFromInputStream(userId, in, rootTag, associationsOut,
-                    previouslyUsedIdsPerPackageOut);
+            return readAssociationsFromInputStream(userId, in, rootTag);
         } catch (XmlPullParserException | IOException e) {
             Slog.e(TAG, "Error while reading associations file", e);
-            return -1;
+            return new Associations();
         }
     }
 
-    private int readStateFromInputStream(@UserIdInt int userId, @NonNull InputStream in,
-            @NonNull String rootTag, @Nullable Collection<AssociationInfo> associationsOut,
-            @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackageOut)
+    @NonNull
+    private static Associations readAssociationsFromInputStream(@UserIdInt int userId,
+            @NonNull InputStream in, @NonNull String rootTag)
             throws XmlPullParserException, IOException {
         final TypedXmlPullParser parser = Xml.resolvePullParser(in);
-
         XmlUtils.beginDocument(parser, rootTag);
+
         final int version = readIntAttribute(parser, XML_ATTR_PERSISTENCE_VERSION, 0);
+        Associations associations = new Associations();
+
         switch (version) {
             case 0:
-                readAssociationsV0(parser, userId, associationsOut);
+                associations = readAssociationsV0(parser, userId);
                 break;
             case 1:
                 while (true) {
                     parser.nextTag();
                     if (isStartOfTag(parser, XML_TAG_ASSOCIATIONS)) {
-                        readAssociationsV1(parser, userId, associationsOut);
-                    } else if (isStartOfTag(parser, XML_TAG_PREVIOUSLY_USED_IDS)) {
-                        readPreviouslyUsedIdsV1(parser, previouslyUsedIdsPerPackageOut);
+                        associations = readAssociationsV1(parser, userId);
                     } else if (isEndOfTag(parser, rootTag)) {
                         break;
                     }
                 }
                 break;
         }
-        return version;
+        return associations;
     }
 
-    private void persistStateToFileLocked(@NonNull AtomicFile file,
-            @Nullable Collection<AssociationInfo> associations,
-            @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackage) {
+    private void writeAssociationsToFile(@NonNull AtomicFile file,
+            @NonNull Associations associations) {
         // Writing to file could fail, for example, if the user has been recently removed and so was
         // their DE (/data/system_de/<user-id>/) directory.
         writeToFileSafely(file, out -> {
             final TypedXmlSerializer serializer = Xml.resolveSerializer(out);
-            serializer.setFeature(
-                    "http://xmlpull.org/v1/doc/features.html#indent-output", true);
-
+            serializer.setFeature("http://xmlpull.org/v1/doc/features.html#indent-output", true);
             serializer.startDocument(null, true);
             serializer.startTag(null, XML_TAG_STATE);
             writeIntAttribute(serializer,
                     XML_ATTR_PERSISTENCE_VERSION, CURRENT_PERSISTENCE_VERSION);
-
             writeAssociations(serializer, associations);
-            writePreviouslyUsedIds(serializer, previouslyUsedIdsPerPackage);
-
             serializer.endTag(null, XML_TAG_STATE);
             serializer.endDocument();
         });
@@ -379,7 +316,8 @@
      * IMPORTANT: the method will ALWAYS return the same {@link AtomicFile} object, which makes it
      * possible to synchronize reads and writes to the file using the returned object.
      */
-    private @NonNull AtomicFile getStorageFileForUser(@UserIdInt int userId) {
+    @NonNull
+    private AtomicFile getStorageFileForUser(@UserIdInt int userId) {
         return mUserIdToStorageFile.computeIfAbsent(userId,
                 u -> createStorageFileForUser(userId, FILE_NAME));
     }
@@ -399,14 +337,12 @@
     /**
      * Convert payload to a set of associations
      */
-    public void readStateFromPayload(byte[] payload, @UserIdInt int userId,
-                              @NonNull Set<AssociationInfo> associationsOut,
-                              @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackageOut) {
+    public static Associations readAssociationsFromPayload(byte[] payload, @UserIdInt int userId) {
         try (ByteArrayInputStream in = new ByteArrayInputStream(payload)) {
-            readStateFromInputStream(userId, in, XML_TAG_STATE, associationsOut,
-                    previouslyUsedIdsPerPackageOut);
+            return readAssociationsFromInputStream(userId, in, XML_TAG_STATE);
         } catch (XmlPullParserException | IOException e) {
             Slog.e(TAG, "Error while reading associations file", e);
+            return new Associations();
         }
     }
 
@@ -414,8 +350,8 @@
         return new File(Environment.getUserSystemDirectory(userId), FILE_NAME_LEGACY);
     }
 
-    private static void readAssociationsV0(@NonNull TypedXmlPullParser parser,
-            @UserIdInt int userId, @NonNull Collection<AssociationInfo> out)
+    private static Associations readAssociationsV0(@NonNull TypedXmlPullParser parser,
+            @UserIdInt int userId)
             throws XmlPullParserException, IOException {
         requireStartOfTag(parser, XML_TAG_ASSOCIATIONS);
 
@@ -426,52 +362,70 @@
         // means that CDM hasn't assigned any IDs yet, so we can just start from the first available
         // id for each user (eg. 1 for user 0; 100 001 - for user 1; 200 001 - for user 2; etc).
         int associationId = getFirstAssociationIdForUser(userId);
+        Associations associations = new Associations();
+        associations.setVersion(0);
+
         while (true) {
             parser.nextTag();
             if (isEndOfTag(parser, XML_TAG_ASSOCIATIONS)) break;
             if (!isStartOfTag(parser, XML_TAG_ASSOCIATION)) continue;
 
-            readAssociationV0(parser, userId, associationId++, out);
+            associations.addAssociation(readAssociationV0(parser, userId, associationId++));
         }
+
+        associations.setMaxId(associationId - 1);
+
+        return associations;
     }
 
-    private static void readAssociationV0(@NonNull TypedXmlPullParser parser, @UserIdInt int userId,
-            int associationId, @NonNull Collection<AssociationInfo> out)
+    private static AssociationInfo readAssociationV0(@NonNull TypedXmlPullParser parser,
+            @UserIdInt int userId, int associationId)
             throws XmlPullParserException {
         requireStartOfTag(parser, XML_TAG_ASSOCIATION);
 
         final String appPackage = readStringAttribute(parser, XML_ATTR_PACKAGE);
         final String tag = readStringAttribute(parser, XML_TAG_TAG);
         final String deviceAddress = readStringAttribute(parser, LEGACY_XML_ATTR_DEVICE);
-
-        if (appPackage == null || deviceAddress == null) return;
-
         final String profile = readStringAttribute(parser, XML_ATTR_PROFILE);
         final boolean notify = readBooleanAttribute(parser, XML_ATTR_NOTIFY_DEVICE_NEARBY);
         final long timeApproved = readLongAttribute(parser, XML_ATTR_TIME_APPROVED, 0L);
 
-        out.add(new AssociationInfo(associationId, userId, appPackage, tag,
+        return new AssociationInfo(associationId, userId, appPackage, tag,
                 MacAddress.fromString(deviceAddress), null, profile, null,
                 /* managedByCompanionApp */ false, notify, /* revoked */ false, /* pending */ false,
-                timeApproved, Long.MAX_VALUE, /* systemDataSyncFlags */ 0));
+                timeApproved, Long.MAX_VALUE, /* systemDataSyncFlags */ 0);
     }
 
-    private static void readAssociationsV1(@NonNull TypedXmlPullParser parser,
-            @UserIdInt int userId, @NonNull Collection<AssociationInfo> out)
+    private static Associations readAssociationsV1(@NonNull TypedXmlPullParser parser,
+            @UserIdInt int userId)
             throws XmlPullParserException, IOException {
         requireStartOfTag(parser, XML_TAG_ASSOCIATIONS);
 
+        // For old builds that don't have max-id attr,
+        // default maxId to 0 and get the maxId out of all association ids.
+        int maxId = readIntAttribute(parser, XML_ATTR_MAX_ID, 0);
+        Associations associations = new Associations();
+        associations.setVersion(1);
+
         while (true) {
             parser.nextTag();
             if (isEndOfTag(parser, XML_TAG_ASSOCIATIONS)) break;
             if (!isStartOfTag(parser, XML_TAG_ASSOCIATION)) continue;
 
-            readAssociationV1(parser, userId, out);
+            AssociationInfo association = readAssociationV1(parser, userId);
+            associations.addAssociation(association);
+
+            maxId = Math.max(maxId, association.getId());
         }
+
+        associations.setMaxId(maxId);
+
+        return associations;
     }
 
-    private static void readAssociationV1(@NonNull TypedXmlPullParser parser, @UserIdInt int userId,
-            @NonNull Collection<AssociationInfo> out) throws XmlPullParserException, IOException {
+    private static AssociationInfo readAssociationV1(@NonNull TypedXmlPullParser parser,
+            @UserIdInt int userId)
+            throws XmlPullParserException, IOException {
         requireStartOfTag(parser, XML_TAG_ASSOCIATION);
 
         final int associationId = readIntAttribute(parser, XML_ATTR_ID);
@@ -491,46 +445,19 @@
         final int systemDataSyncFlags = readIntAttribute(parser,
                 XML_ATTR_SYSTEM_DATA_SYNC_FLAGS, 0);
 
-        final AssociationInfo associationInfo = createAssociationInfoNoThrow(associationId, userId,
-                appPackage, tag, macAddress, displayName, profile, selfManaged, notify, revoked,
-                pending, timeApproved, lastTimeConnected, systemDataSyncFlags);
-        if (associationInfo != null) {
-            out.add(associationInfo);
-        }
-    }
-
-    private static void readPreviouslyUsedIdsV1(@NonNull TypedXmlPullParser parser,
-            @NonNull Map<String, Set<Integer>> out) throws XmlPullParserException, IOException {
-        requireStartOfTag(parser, XML_TAG_PREVIOUSLY_USED_IDS);
-
-        while (true) {
-            parser.nextTag();
-            if (isEndOfTag(parser, XML_TAG_PREVIOUSLY_USED_IDS)) break;
-            if (!isStartOfTag(parser, XML_TAG_PACKAGE)) continue;
-
-            final String packageName = readStringAttribute(parser, XML_ATTR_PACKAGE_NAME);
-            final Set<Integer> usedIds = new HashSet<>();
-
-            while (true) {
-                parser.nextTag();
-                if (isEndOfTag(parser, XML_TAG_PACKAGE)) break;
-                if (!isStartOfTag(parser, XML_TAG_ID)) continue;
-
-                parser.nextToken();
-                final int id = Integer.parseInt(parser.getText());
-                usedIds.add(id);
-            }
-
-            out.put(packageName, usedIds);
-        }
+        return new AssociationInfo(associationId, userId, appPackage, tag, macAddress, displayName,
+                profile, null, selfManaged, notify, revoked, pending, timeApproved,
+                lastTimeConnected, systemDataSyncFlags);
     }
 
     private static void writeAssociations(@NonNull XmlSerializer parent,
-            @Nullable Collection<AssociationInfo> associations) throws IOException {
+            @NonNull Associations associations)
+            throws IOException {
         final XmlSerializer serializer = parent.startTag(null, XML_TAG_ASSOCIATIONS);
-        for (AssociationInfo association : associations) {
+        for (AssociationInfo association : associations.getAssociations()) {
             writeAssociation(serializer, association);
         }
+        writeIntAttribute(serializer, XML_ATTR_MAX_ID, associations.getMaxId());
         serializer.endTag(null, XML_TAG_ASSOCIATIONS);
     }
 
@@ -557,26 +484,6 @@
         serializer.endTag(null, XML_TAG_ASSOCIATION);
     }
 
-    private static void writePreviouslyUsedIds(@NonNull XmlSerializer parent,
-            @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackage) throws IOException {
-        final XmlSerializer serializer = parent.startTag(null, XML_TAG_PREVIOUSLY_USED_IDS);
-        for (Map.Entry<String, Set<Integer>> entry : previouslyUsedIdsPerPackage.entrySet()) {
-            writePreviouslyUsedIdsForPackage(serializer, entry.getKey(), entry.getValue());
-        }
-        serializer.endTag(null, XML_TAG_PREVIOUSLY_USED_IDS);
-    }
-
-    private static void writePreviouslyUsedIdsForPackage(@NonNull XmlSerializer parent,
-            @NonNull String packageName, @NonNull Set<Integer> previouslyUsedIds)
-            throws IOException {
-        final XmlSerializer serializer = parent.startTag(null, XML_TAG_PACKAGE);
-        writeStringAttribute(serializer, XML_ATTR_PACKAGE_NAME, packageName);
-        forEach(previouslyUsedIds, id -> serializer.startTag(null, XML_TAG_ID)
-                .text(Integer.toString(id))
-                .endTag(null, XML_TAG_ID));
-        serializer.endTag(null, XML_TAG_PACKAGE);
-    }
-
     private static void requireStartOfTag(@NonNull XmlPullParser parser, @NonNull String tag)
             throws XmlPullParserException {
         if (isStartOfTag(parser, tag)) return;
@@ -587,22 +494,4 @@
     private static @Nullable MacAddress stringToMacAddress(@Nullable String address) {
         return address != null ? MacAddress.fromString(address) : null;
     }
-
-    private static AssociationInfo createAssociationInfoNoThrow(int associationId,
-            @UserIdInt int userId, @NonNull String appPackage, @Nullable String tag,
-            @Nullable MacAddress macAddress, @Nullable CharSequence displayName,
-            @Nullable String profile, boolean selfManaged, boolean notify, boolean revoked,
-            boolean pending, long timeApproved, long lastTimeConnected, int systemDataSyncFlags) {
-        AssociationInfo associationInfo = null;
-        try {
-            // We do not persist AssociatedDevice, which means that AssociationInfo retrieved from
-            // datastore is not guaranteed to be identical to the one from initial association.
-            associationInfo = new AssociationInfo(associationId, userId, appPackage, tag,
-                    macAddress, displayName, profile, null, selfManaged, notify,
-                    revoked, pending, timeApproved, lastTimeConnected, systemDataSyncFlags);
-        } catch (Exception e) {
-            Slog.e(TAG, "Could not create AssociationInfo", e);
-        }
-        return associationInfo;
-    }
 }
diff --git a/services/companion/java/com/android/server/companion/association/AssociationRequestsProcessor.java b/services/companion/java/com/android/server/companion/association/AssociationRequestsProcessor.java
index 29ec7c2..a02d9f9 100644
--- a/services/companion/java/com/android/server/companion/association/AssociationRequestsProcessor.java
+++ b/services/companion/java/com/android/server/companion/association/AssociationRequestsProcessor.java
@@ -24,7 +24,6 @@
 import static android.content.ComponentName.createRelative;
 import static android.content.pm.PackageManager.FEATURE_WATCH;
 
-import static com.android.server.companion.utils.MetricUtils.logCreateAssociation;
 import static com.android.server.companion.utils.PackageUtils.enforceUsesCompanionDeviceFeature;
 import static com.android.server.companion.utils.PermissionsUtils.enforcePermissionForCreatingAssociation;
 import static com.android.server.companion.utils.RolesUtils.addRoleHolderForAssociation;
@@ -128,17 +127,16 @@
     private static final long ASSOCIATE_WITHOUT_PROMPT_WINDOW_MS = 60 * 60 * 1000; // 60 min;
 
     private final @NonNull Context mContext;
-    private final @NonNull CompanionDeviceManagerService mService;
-    private final @NonNull PackageManagerInternal mPackageManager;
+    private final @NonNull PackageManagerInternal mPackageManagerInternal;
     private final @NonNull AssociationStore mAssociationStore;
     @NonNull
     private final ComponentName mCompanionDeviceActivity;
 
-    public AssociationRequestsProcessor(@NonNull CompanionDeviceManagerService service,
+    public AssociationRequestsProcessor(@NonNull Context context,
+            @NonNull PackageManagerInternal packageManagerInternal,
             @NonNull AssociationStore associationStore) {
-        mContext = service.getContext();
-        mService = service;
-        mPackageManager = service.mPackageManagerInternal;
+        mContext = context;
+        mPackageManagerInternal = packageManagerInternal;
         mAssociationStore = associationStore;
         mCompanionDeviceActivity = createRelative(
                 mContext.getString(R.string.config_companionDeviceManagerPackage),
@@ -160,7 +158,7 @@
         requireNonNull(packageName, "Package name MUST NOT be null");
         requireNonNull(callback, "Callback MUST NOT be null");
 
-        final int packageUid = mPackageManager.getPackageUid(packageName, 0, userId);
+        final int packageUid = mPackageManagerInternal.getPackageUid(packageName, 0, userId);
         Slog.d(TAG, "processNewAssociationRequest() " + "request=" + request + ", " + "package=u"
                 + userId + "/" + packageName + " (uid=" + packageUid + ")");
 
@@ -226,7 +224,7 @@
 
         enforceUsesCompanionDeviceFeature(mContext, userId, packageName);
 
-        final int packageUid = mPackageManager.getPackageUid(packageName, 0, userId);
+        final int packageUid = mPackageManagerInternal.getPackageUid(packageName, 0, userId);
 
         final Bundle extras = new Bundle();
         extras.putBoolean(EXTRA_FORCE_CANCEL_CONFIRMATION, true);
@@ -243,7 +241,7 @@
             @NonNull ResultReceiver resultReceiver, @Nullable MacAddress macAddress) {
         final String packageName = request.getPackageName();
         final int userId = request.getUserId();
-        final int packageUid = mPackageManager.getPackageUid(packageName, 0, userId);
+        final int packageUid = mPackageManagerInternal.getPackageUid(packageName, 0, userId);
 
         // 1. Need to check permissions again in case something changed, since we first received
         // this request.
@@ -267,15 +265,12 @@
             @NonNull AssociationRequest request, @NonNull String packageName, @UserIdInt int userId,
             @Nullable MacAddress macAddress, @NonNull IAssociationRequestCallback callback,
             @NonNull ResultReceiver resultReceiver) {
-        final long callingIdentity = Binder.clearCallingIdentity();
-        try {
+        Binder.withCleanCallingIdentity(() -> {
             createAssociation(userId, packageName, macAddress, request.getDisplayName(),
                     request.getDeviceProfile(), request.getAssociatedDevice(),
                     request.isSelfManaged(),
                     callback, resultReceiver);
-        } finally {
-            Binder.restoreCallingIdentity(callingIdentity);
-        }
+        });
     }
 
     /**
@@ -286,7 +281,7 @@
             @Nullable String deviceProfile, @Nullable AssociatedDevice associatedDevice,
             boolean selfManaged, @Nullable IAssociationRequestCallback callback,
             @Nullable ResultReceiver resultReceiver) {
-        final int id = mService.getNewAssociationIdForPackage(userId, packageName);
+        final int id = mAssociationStore.getNextId(userId);
         final long timestamp = System.currentTimeMillis();
 
         final AssociationInfo association = new AssociationInfo(id, userId, packageName,
@@ -296,10 +291,6 @@
 
         // Add role holder for association (if specified) and add new association to store.
         maybeGrantRoleAndStoreAssociation(association, callback, resultReceiver);
-
-        // Don't need to update the mRevokedAssociationsPendingRoleHolderRemoval since
-        // maybeRemoveRoleHolderForAssociation in PackageInactivityListener will handle the case
-        // that there are other devices with the same profile, so the role holder won't be removed.
     }
 
     /**
@@ -311,12 +302,12 @@
         // If the "Device Profile" is specified, make the companion application a holder of the
         // corresponding role.
         // If it is null, then the operation will succeed without granting any role.
-        addRoleHolderForAssociation(mService.getContext(), association, success -> {
+        addRoleHolderForAssociation(mContext, association, success -> {
             if (success) {
                 Slog.i(TAG, "Added " + association.getDeviceProfile() + " role to userId="
                         + association.getUserId() + ", packageName="
                         + association.getPackageName());
-                addAssociationToStore(association);
+                mAssociationStore.addAssociation(association);
                 sendCallbackAndFinish(association, callback, resultReceiver);
             } else {
                 Slog.e(TAG, "Failed to add u" + association.getUserId()
@@ -347,17 +338,6 @@
         mAssociationStore.updateAssociation(updated);
     }
 
-    private void addAssociationToStore(@NonNull AssociationInfo association) {
-        Slog.i(TAG, "New CDM association created=" + association);
-
-        mAssociationStore.addAssociation(association);
-
-        mService.updateSpecialAccessPermissionForAssociatedPackage(association.getUserId(),
-                association.getPackageName());
-
-        logCreateAssociation(association.getDeviceProfile());
-    }
-
     private void sendCallbackAndFinish(@Nullable AssociationInfo association,
             @Nullable IAssociationRequestCallback callback,
             @Nullable ResultReceiver resultReceiver) {
@@ -409,27 +389,22 @@
 
     private PendingIntent createPendingIntent(int packageUid, Intent intent) {
         final PendingIntent pendingIntent;
-        final long token = Binder.clearCallingIdentity();
 
         // Using uid of the application that will own the association (usually the same
         // application that sent the request) allows us to have multiple "pending" association
         // requests at the same time.
         // If the application already has a pending association request, that PendingIntent
         // will be cancelled except application wants to cancel the request by the system.
-        try {
-            pendingIntent = PendingIntent.getActivityAsUser(
+        return Binder.withCleanCallingIdentity(() ->
+            PendingIntent.getActivityAsUser(
                     mContext, /*requestCode */ packageUid, intent,
                     FLAG_ONE_SHOT | FLAG_CANCEL_CURRENT | FLAG_IMMUTABLE,
                     ActivityOptions.makeBasic()
                             .setPendingIntentCreatorBackgroundActivityStartMode(
                                     ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED)
                             .toBundle(),
-                    UserHandle.CURRENT);
-        } finally {
-            Binder.restoreCallingIdentity(token);
-        }
-
-        return pendingIntent;
+                    UserHandle.CURRENT)
+        );
     }
 
     private final ResultReceiver mOnRequestConfirmationReceiver =
@@ -470,7 +445,7 @@
         // Throttle frequent associations
         final long now = System.currentTimeMillis();
         final List<AssociationInfo> associationForPackage =
-                mAssociationStore.getAssociationsForPackage(userId, packageName);
+                mAssociationStore.getActiveAssociationsByPackage(userId, packageName);
         // Number of "recent" associations.
         int recent = 0;
         for (AssociationInfo association : associationForPackage) {
@@ -486,6 +461,6 @@
             }
         }
 
-        return PackageUtils.isPackageAllowlisted(mContext, mPackageManager, packageName);
+        return PackageUtils.isPackageAllowlisted(mContext, mPackageManagerInternal, packageName);
     }
 }
diff --git a/services/companion/java/com/android/server/companion/association/AssociationRevokeProcessor.java b/services/companion/java/com/android/server/companion/association/AssociationRevokeProcessor.java
deleted file mode 100644
index d1efbbc..0000000
--- a/services/companion/java/com/android/server/companion/association/AssociationRevokeProcessor.java
+++ /dev/null
@@ -1,383 +0,0 @@
-/*
- * Copyright (C) 2024 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.association;
-
-import static android.app.ActivityManager.RunningAppProcessInfo.IMPORTANCE_VISIBLE;
-import static android.companion.AssociationRequest.DEVICE_PROFILE_AUTOMOTIVE_PROJECTION;
-
-import static com.android.internal.util.CollectionUtils.any;
-import static com.android.server.companion.utils.MetricUtils.logRemoveAssociation;
-import static com.android.server.companion.utils.RolesUtils.removeRoleHolderForAssociation;
-import static com.android.server.companion.CompanionDeviceManagerService.PerUserAssociationSet;
-
-import android.annotation.NonNull;
-import android.annotation.SuppressLint;
-import android.annotation.UserIdInt;
-import android.app.ActivityManager;
-import android.companion.AssociationInfo;
-import android.content.Context;
-import android.content.pm.PackageManagerInternal;
-import android.os.Binder;
-import android.os.UserHandle;
-import android.util.ArraySet;
-import android.util.Log;
-import android.util.Slog;
-
-import com.android.internal.annotations.GuardedBy;
-import com.android.server.companion.CompanionApplicationController;
-import com.android.server.companion.CompanionDeviceManagerService;
-import com.android.server.companion.datatransfer.SystemDataTransferRequestStore;
-import com.android.server.companion.presence.CompanionDevicePresenceMonitor;
-import com.android.server.companion.transport.CompanionTransportManager;
-
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Set;
-
-/**
- * A class response for Association removal.
- */
-@SuppressLint("LongLogTag")
-public class AssociationRevokeProcessor {
-
-    private static final String TAG = "CDM_AssociationRevokeProcessor";
-    private static final boolean DEBUG = false;
-    private final @NonNull Context mContext;
-    private final @NonNull CompanionDeviceManagerService mService;
-    private final @NonNull AssociationStore mAssociationStore;
-    private final @NonNull PackageManagerInternal mPackageManagerInternal;
-    private final @NonNull CompanionDevicePresenceMonitor mDevicePresenceMonitor;
-    private final @NonNull SystemDataTransferRequestStore mSystemDataTransferRequestStore;
-    private final @NonNull CompanionApplicationController mCompanionAppController;
-    private final @NonNull CompanionTransportManager mTransportManager;
-    private final OnPackageVisibilityChangeListener mOnPackageVisibilityChangeListener;
-    private final ActivityManager mActivityManager;
-
-    /**
-     * A structure that consists of a set of revoked associations that pending for role holder
-     * removal per each user.
-     *
-     * @see #maybeRemoveRoleHolderForAssociation(AssociationInfo)
-     * @see #addToPendingRoleHolderRemoval(AssociationInfo)
-     * @see #removeFromPendingRoleHolderRemoval(AssociationInfo)
-     * @see #getPendingRoleHolderRemovalAssociationsForUser(int)
-     */
-    @GuardedBy("mRevokedAssociationsPendingRoleHolderRemoval")
-    private final PerUserAssociationSet mRevokedAssociationsPendingRoleHolderRemoval =
-            new PerUserAssociationSet();
-    /**
-     * Contains uid-s of packages pending to be removed from the role holder list (after
-     * revocation of an association), which will happen one the package is no longer visible to the
-     * user.
-     * For quicker uid -> (userId, packageName) look-up this is not a {@code Set<Integer>} but
-     * a {@code Map<Integer, String>} which maps uid-s to packageName-s (userId-s can be derived
-     * from uid-s using {@link UserHandle#getUserId(int)}).
-     *
-     * @see #maybeRemoveRoleHolderForAssociation(AssociationInfo)
-     * @see #addToPendingRoleHolderRemoval(AssociationInfo)
-     * @see #removeFromPendingRoleHolderRemoval(AssociationInfo)
-     */
-    @GuardedBy("mRevokedAssociationsPendingRoleHolderRemoval")
-    private final Map<Integer, String> mUidsPendingRoleHolderRemoval = new HashMap<>();
-
-    public AssociationRevokeProcessor(@NonNull CompanionDeviceManagerService service,
-            @NonNull AssociationStore associationStore,
-            @NonNull PackageManagerInternal packageManager,
-            @NonNull CompanionDevicePresenceMonitor devicePresenceMonitor,
-            @NonNull CompanionApplicationController applicationController,
-            @NonNull SystemDataTransferRequestStore systemDataTransferRequestStore,
-            @NonNull CompanionTransportManager companionTransportManager) {
-        mService = service;
-        mContext = service.getContext();
-        mActivityManager = mContext.getSystemService(ActivityManager.class);
-        mAssociationStore = associationStore;
-        mPackageManagerInternal = packageManager;
-        mOnPackageVisibilityChangeListener =
-                new OnPackageVisibilityChangeListener(mActivityManager);
-        mDevicePresenceMonitor = devicePresenceMonitor;
-        mCompanionAppController = applicationController;
-        mSystemDataTransferRequestStore = systemDataTransferRequestStore;
-        mTransportManager = companionTransportManager;
-    }
-
-    /**
-     * Disassociate an association
-     */
-    // TODO: also revoke notification access
-    public void disassociateInternal(int associationId) {
-        final AssociationInfo association = mAssociationStore.getAssociationById(associationId);
-        final int userId = association.getUserId();
-        final String packageName = association.getPackageName();
-        final String deviceProfile = association.getDeviceProfile();
-
-        // Detach transport if exists
-        mTransportManager.detachSystemDataTransport(packageName, userId, associationId);
-
-        if (!maybeRemoveRoleHolderForAssociation(association)) {
-            // Need to remove the app from list of the role holders, but will have to do it later
-            // (the app is in foreground at the moment).
-            addToPendingRoleHolderRemoval(association);
-        }
-
-        // Need to check if device still present now because CompanionDevicePresenceMonitor will
-        // remove current connected device after mAssociationStore.removeAssociation
-        final boolean wasPresent = mDevicePresenceMonitor.isDevicePresent(associationId);
-
-        // Removing the association.
-        mAssociationStore.removeAssociation(associationId);
-        // Do not need to persistUserState since CompanionDeviceManagerService will get callback
-        // from #onAssociationChanged, and it will handle the persistUserState which including
-        // active and revoked association.
-        logRemoveAssociation(deviceProfile);
-
-        // Remove all the system data transfer requests for the association.
-        mSystemDataTransferRequestStore.removeRequestsByAssociationId(userId, associationId);
-
-        if (!wasPresent || !association.isNotifyOnDeviceNearby()) return;
-        // The device was connected and the app was notified: check if we need to unbind the app
-        // now.
-        final boolean shouldStayBound = any(
-                mAssociationStore.getAssociationsForPackage(userId, packageName),
-                it -> it.isNotifyOnDeviceNearby()
-                        && mDevicePresenceMonitor.isDevicePresent(it.getId()));
-        if (shouldStayBound) return;
-        mCompanionAppController.unbindCompanionApplication(userId, packageName);
-    }
-
-    /**
-     * First, checks if the companion application should be removed from the list role holders when
-     * upon association's removal, i.e.: association's profile (matches the role) is not null,
-     * the application does not have other associations with the same profile, etc.
-     *
-     * <p>
-     * Then, if establishes that the application indeed has to be removed from the list of the role
-     * holders, checks if it could be done right now -
-     * {@link android.app.role.RoleManager#removeRoleHolderAsUser(String, String, int, UserHandle, java.util.concurrent.Executor, java.util.function.Consumer) RoleManager#removeRoleHolderAsUser()}
-     * will kill the application's process, which leads poor user experience if the application was
-     * in foreground when this happened, to avoid this CDMS delays invoking
-     * {@code RoleManager.removeRoleHolderAsUser()} until the app is no longer in foreground.
-     *
-     * @return {@code true} if the application does NOT need be removed from the list of the role
-     *         holders OR if the application was successfully removed from the list of role holders.
-     *         I.e.: from the role-management perspective the association is done with.
-     *         {@code false} if the application needs to be removed from the list of role the role
-     *         holders, BUT it CDMS would prefer to do it later.
-     *         I.e.: application is in the foreground at the moment, but invoking
-     *         {@code RoleManager.removeRoleHolderAsUser()} will kill the application's process,
-     *         which would lead to the poor UX, hence need to try later.
-     */
-    public boolean maybeRemoveRoleHolderForAssociation(@NonNull AssociationInfo association) {
-        if (DEBUG) Log.d(TAG, "maybeRemoveRoleHolderForAssociation() association=" + association);
-        final String deviceProfile = association.getDeviceProfile();
-
-        if (deviceProfile == null) {
-            // No role was granted to for this association, there is nothing else we need to here.
-            return true;
-        }
-        // Do not need to remove the system role since it was pre-granted by the system.
-        if (deviceProfile.equals(DEVICE_PROFILE_AUTOMOTIVE_PROJECTION)) {
-            return true;
-        }
-
-        // Check if the applications is associated with another devices with the profile. If so,
-        // it should remain the role holder.
-        final int id = association.getId();
-        final int userId = association.getUserId();
-        final String packageName = association.getPackageName();
-        final boolean roleStillInUse = any(
-                mAssociationStore.getAssociationsForPackage(userId, packageName),
-                it -> deviceProfile.equals(it.getDeviceProfile()) && id != it.getId());
-        if (roleStillInUse) {
-            // Application should remain a role holder, there is nothing else we need to here.
-            return true;
-        }
-
-        final int packageProcessImportance = getPackageProcessImportance(userId, packageName);
-        if (packageProcessImportance <= IMPORTANCE_VISIBLE) {
-            // Need to remove the app from the list of role holders, but the process is visible to
-            // the user at the moment, so we'll need to it later: log and return false.
-            Slog.i(TAG, "Cannot remove role holder for the removed association id=" + id
-                    + " now - process is visible.");
-            return false;
-        }
-
-        removeRoleHolderForAssociation(mContext, association.getUserId(),
-                association.getPackageName(), association.getDeviceProfile());
-        return true;
-    }
-
-    /**
-     * Set revoked flag for active association and add the revoked association and the uid into
-     * the caches.
-     *
-     * @see #mRevokedAssociationsPendingRoleHolderRemoval
-     * @see #mUidsPendingRoleHolderRemoval
-     * @see OnPackageVisibilityChangeListener
-     */
-    public void addToPendingRoleHolderRemoval(@NonNull AssociationInfo association) {
-        // First: set revoked flag
-        association = (new AssociationInfo.Builder(association)).setRevoked(true).build();
-        final String packageName = association.getPackageName();
-        final int userId = association.getUserId();
-        final int uid = mPackageManagerInternal.getPackageUid(packageName, /* flags */0, userId);
-        // Second: add to the set.
-        synchronized (mRevokedAssociationsPendingRoleHolderRemoval) {
-            mRevokedAssociationsPendingRoleHolderRemoval.forUser(association.getUserId())
-                    .add(association);
-            if (!mUidsPendingRoleHolderRemoval.containsKey(uid)) {
-                mUidsPendingRoleHolderRemoval.put(uid, packageName);
-
-                if (mUidsPendingRoleHolderRemoval.size() == 1) {
-                    // Just added first uid: start the listener
-                    mOnPackageVisibilityChangeListener.startListening();
-                }
-            }
-        }
-    }
-
-    /**
-     * @return a copy of the revoked associations set (safeguarding against
-     *         {@code ConcurrentModificationException}-s).
-     */
-    @NonNull
-    public Set<AssociationInfo> getPendingRoleHolderRemovalAssociationsForUser(
-            @UserIdInt int userId) {
-        synchronized (mRevokedAssociationsPendingRoleHolderRemoval) {
-            // Return a copy.
-            return new ArraySet<>(mRevokedAssociationsPendingRoleHolderRemoval.forUser(userId));
-        }
-    }
-
-    @SuppressLint("MissingPermission")
-    private int  getPackageProcessImportance(@UserIdInt int userId, @NonNull String packageName) {
-        return Binder.withCleanCallingIdentity(() -> {
-            final int uid =
-                    mPackageManagerInternal.getPackageUid(packageName, /* flags */0, userId);
-            return mActivityManager.getUidImportance(uid);
-        });
-    }
-
-    /**
-     * Remove the revoked association from the cache and also remove the uid from the map if
-     * there are other associations with the same package still pending for role holder removal.
-     *
-     * @see #mRevokedAssociationsPendingRoleHolderRemoval
-     * @see #mUidsPendingRoleHolderRemoval
-     * @see OnPackageVisibilityChangeListener
-     */
-    private void removeFromPendingRoleHolderRemoval(@NonNull AssociationInfo association) {
-        final String packageName = association.getPackageName();
-        final int userId = association.getUserId();
-        final int uid = mPackageManagerInternal.getPackageUid(packageName, /* flags */  0, userId);
-
-        synchronized (mRevokedAssociationsPendingRoleHolderRemoval) {
-            mRevokedAssociationsPendingRoleHolderRemoval.forUser(userId)
-                    .remove(association);
-
-            final boolean shouldKeepUidForRemoval = any(
-                    getPendingRoleHolderRemovalAssociationsForUser(userId),
-                    ai -> packageName.equals(ai.getPackageName()));
-            // Do not remove the uid from the map since other associations with
-            // the same packageName still pending for role holder removal.
-            if (!shouldKeepUidForRemoval) {
-                mUidsPendingRoleHolderRemoval.remove(uid);
-            }
-
-            if (mUidsPendingRoleHolderRemoval.isEmpty()) {
-                // The set is empty now - can "turn off" the listener.
-                mOnPackageVisibilityChangeListener.stopListening();
-            }
-        }
-    }
-
-    private String getPackageNameByUid(int uid) {
-        synchronized (mRevokedAssociationsPendingRoleHolderRemoval) {
-            return mUidsPendingRoleHolderRemoval.get(uid);
-        }
-    }
-
-    /**
-     * An OnUidImportanceListener class which watches the importance of the packages.
-     * In this class, we ONLY interested in the importance of the running process is greater than
-     * {@link ActivityManager.RunningAppProcessInfo#IMPORTANCE_VISIBLE} for the uids have been added
-     * into the {@link #mUidsPendingRoleHolderRemoval}. Lastly remove the role holder for the
-     * revoked associations for the same packages.
-     *
-     * @see #maybeRemoveRoleHolderForAssociation(AssociationInfo)
-     * @see #removeFromPendingRoleHolderRemoval(AssociationInfo)
-     * @see #getPendingRoleHolderRemovalAssociationsForUser(int)
-     */
-    private class OnPackageVisibilityChangeListener implements
-            ActivityManager.OnUidImportanceListener {
-        final @NonNull ActivityManager mAm;
-
-        OnPackageVisibilityChangeListener(@NonNull ActivityManager am) {
-            this.mAm = am;
-        }
-
-        @SuppressLint("MissingPermission")
-        void startListening() {
-            Binder.withCleanCallingIdentity(
-                    () -> mAm.addOnUidImportanceListener(
-                            /* listener */ OnPackageVisibilityChangeListener.this,
-                            ActivityManager.RunningAppProcessInfo.IMPORTANCE_VISIBLE));
-        }
-
-        @SuppressLint("MissingPermission")
-        void stopListening() {
-            Binder.withCleanCallingIdentity(
-                    () -> mAm.removeOnUidImportanceListener(
-                            /* listener */ OnPackageVisibilityChangeListener.this));
-        }
-
-        @Override
-        public void onUidImportance(int uid, int importance) {
-            if (importance <= ActivityManager.RunningAppProcessInfo.IMPORTANCE_VISIBLE) {
-                // The lower the importance value the more "important" the process is.
-                // We are only interested when the process ceases to be visible.
-                return;
-            }
-
-            final String packageName = getPackageNameByUid(uid);
-            if (packageName == null) {
-                // Not interested in this uid.
-                return;
-            }
-
-            final int userId = UserHandle.getUserId(uid);
-
-            boolean needToPersistStateForUser = false;
-
-            for (AssociationInfo association :
-                    getPendingRoleHolderRemovalAssociationsForUser(userId)) {
-                if (!packageName.equals(association.getPackageName())) continue;
-
-                if (!maybeRemoveRoleHolderForAssociation(association)) {
-                    // Did not remove the role holder, will have to try again later.
-                    continue;
-                }
-
-                removeFromPendingRoleHolderRemoval(association);
-                needToPersistStateForUser = true;
-            }
-
-            if (needToPersistStateForUser) {
-                mService.postPersistUserState(userId);
-            }
-        }
-    }
-}
diff --git a/services/companion/java/com/android/server/companion/association/AssociationStore.java b/services/companion/java/com/android/server/companion/association/AssociationStore.java
index 2f94bde..edebb55 100644
--- a/services/companion/java/com/android/server/companion/association/AssociationStore.java
+++ b/services/companion/java/com/android/server/companion/association/AssociationStore.java
@@ -16,15 +16,24 @@
 
 package com.android.server.companion.association;
 
+import static com.android.server.companion.utils.MetricUtils.logCreateAssociation;
+import static com.android.server.companion.utils.MetricUtils.logRemoveAssociation;
+
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SuppressLint;
 import android.annotation.UserIdInt;
 import android.companion.AssociationInfo;
+import android.companion.IOnAssociationsChangedListener;
+import android.content.pm.UserInfo;
 import android.net.MacAddress;
+import android.os.Binder;
+import android.os.RemoteCallbackList;
+import android.os.RemoteException;
+import android.os.UserHandle;
+import android.os.UserManager;
 import android.util.Slog;
-import android.util.SparseArray;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.util.CollectionUtils;
@@ -33,15 +42,14 @@
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 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;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 
 /**
  * Association store for CRUD.
@@ -109,63 +117,124 @@
 
     private final Object mLock = new Object();
 
-    @GuardedBy("mLock")
-    private final Map<Integer, AssociationInfo> mIdMap = new HashMap<>();
-    @GuardedBy("mLock")
-    private final Map<MacAddress, Set<Integer>> mAddressMap = new HashMap<>();
-    @GuardedBy("mLock")
-    private final SparseArray<List<AssociationInfo>> mCachedPerUser = new SparseArray<>();
+    private final ExecutorService mExecutor;
 
-    @GuardedBy("mListeners")
-    private final Set<OnChangeListener> mListeners = new LinkedHashSet<>();
+    @GuardedBy("mLock")
+    private boolean mPersisted = false;
+    @GuardedBy("mLock")
+    private final Map<Integer, AssociationInfo> mIdToAssociationMap = new HashMap<>();
+    @GuardedBy("mLock")
+    private final Map<Integer, Integer> mUserToMaxId = new HashMap<>();
+
+    @GuardedBy("mLocalListeners")
+    private final Set<OnChangeListener> mLocalListeners = new LinkedHashSet<>();
+    @GuardedBy("mRemoteListeners")
+    private final RemoteCallbackList<IOnAssociationsChangedListener> mRemoteListeners =
+            new RemoteCallbackList<>();
+
+    private final UserManager mUserManager;
+    private final AssociationDiskStore mDiskStore;
+
+    public AssociationStore(UserManager userManager, AssociationDiskStore diskStore) {
+        mUserManager = userManager;
+        mDiskStore = diskStore;
+        mExecutor = Executors.newSingleThreadExecutor();
+    }
+
+    /**
+     * Load all alive users' associations from disk to cache.
+     */
+    public void refreshCache() {
+        Binder.withCleanCallingIdentity(() -> {
+            List<Integer> userIds = new ArrayList<>();
+            for (UserInfo user : mUserManager.getAliveUsers()) {
+                userIds.add(user.id);
+            }
+
+            synchronized (mLock) {
+                mPersisted = false;
+
+                mIdToAssociationMap.clear();
+                mUserToMaxId.clear();
+
+                // 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).
+                Map<Integer, Associations> userToAssociationsMap =
+                        mDiskStore.readAssociationsByUsers(userIds);
+                for (Map.Entry<Integer, Associations> entry : userToAssociationsMap.entrySet()) {
+                    for (AssociationInfo association : entry.getValue().getAssociations()) {
+                        mIdToAssociationMap.put(association.getId(), association);
+                    }
+                    mUserToMaxId.put(entry.getKey(), entry.getValue().getMaxId());
+                }
+
+                mPersisted = true;
+            }
+        });
+    }
+
+    /**
+     * Get the current max association id.
+     */
+    public int getMaxId(int userId) {
+        synchronized (mLock) {
+            return mUserToMaxId.getOrDefault(userId, 0);
+        }
+    }
+
+    /**
+     * Get the next available association id.
+     */
+    public int getNextId(int userId) {
+        synchronized (mLock) {
+            return getMaxId(userId) + 1;
+        }
+    }
 
     /**
      * Add an association.
      */
     public void addAssociation(@NonNull AssociationInfo association) {
-        Slog.i(TAG, "Adding new association=" + association);
-
-        // Validity check first.
-        checkNotRevoked(association);
+        Slog.i(TAG, "Adding new association=[" + association + "]...");
 
         final int id = association.getId();
+        final int userId = association.getUserId();
 
         synchronized (mLock) {
-            if (mIdMap.containsKey(id)) {
-                Slog.e(TAG, "Association with id " + id + " already exists.");
+            if (mIdToAssociationMap.containsKey(id)) {
+                Slog.e(TAG, "Association with id=[" + id + "] already exists.");
                 return;
             }
-            mIdMap.put(id, association);
 
-            final MacAddress address = association.getDeviceMacAddress();
-            if (address != null) {
-                mAddressMap.computeIfAbsent(address, it -> new HashSet<>()).add(id);
-            }
+            mIdToAssociationMap.put(id, association);
+            mUserToMaxId.put(userId, Math.max(mUserToMaxId.getOrDefault(userId, 0), id));
 
-            invalidateCacheForUserLocked(association.getUserId());
+            writeCacheToDisk(userId);
 
             Slog.i(TAG, "Done adding new association.");
         }
 
-        broadcastChange(CHANGE_TYPE_ADDED, association);
+        logCreateAssociation(association.getDeviceProfile());
+
+        if (association.isActive()) {
+            broadcastChange(CHANGE_TYPE_ADDED, association);
+        }
     }
 
     /**
      * Update an association.
      */
     public void updateAssociation(@NonNull AssociationInfo updated) {
-        Slog.i(TAG, "Updating new association=" + updated);
-        // Validity check first.
-        checkNotRevoked(updated);
+        Slog.i(TAG, "Updating new association=[" + updated + "]...");
 
         final int id = updated.getId();
-
         final AssociationInfo current;
         final boolean macAddressChanged;
+
         synchronized (mLock) {
-            current = mIdMap.get(id);
+            current = mIdToAssociationMap.get(id);
             if (current == null) {
-                Slog.w(TAG, "Can't update association. It does not exist.");
+                Slog.w(TAG, "Can't update association id=[" + id + "]. It does not exist.");
                 return;
             }
 
@@ -174,174 +243,245 @@
                 return;
             }
 
-            // Update the ID-to-Association map.
-            mIdMap.put(id, updated);
-            // Invalidate the corresponding user cache entry.
-            invalidateCacheForUserLocked(current.getUserId());
+            mIdToAssociationMap.put(id, updated);
 
-            // Update the MacAddress-to-List<Association> map if needed.
+            writeCacheToDisk(updated.getUserId());
+        }
+
+        Slog.i(TAG, "Done updating association.");
+
+        if (current.isActive() && !updated.isActive()) {
+            broadcastChange(CHANGE_TYPE_REMOVED, updated);
+            return;
+        }
+
+        if (updated.isActive()) {
+            // Check if the MacAddress has changed.
             final MacAddress updatedAddress = updated.getDeviceMacAddress();
             final MacAddress currentAddress = current.getDeviceMacAddress();
             macAddressChanged = !Objects.equals(currentAddress, updatedAddress);
-            if (macAddressChanged) {
-                if (currentAddress != null) {
-                    mAddressMap.get(currentAddress).remove(id);
-                }
-                if (updatedAddress != null) {
-                    mAddressMap.computeIfAbsent(updatedAddress, it -> new HashSet<>()).add(id);
-                }
-            }
-            Slog.i(TAG, "Done updating association.");
-        }
 
-        final int changeType = macAddressChanged ? CHANGE_TYPE_UPDATED_ADDRESS_CHANGED
-                : CHANGE_TYPE_UPDATED_ADDRESS_UNCHANGED;
-        broadcastChange(changeType, updated);
+            broadcastChange(macAddressChanged ? CHANGE_TYPE_UPDATED_ADDRESS_CHANGED
+                    : CHANGE_TYPE_UPDATED_ADDRESS_UNCHANGED, updated);
+        }
     }
 
     /**
-     * Remove an association
+     * Remove an association.
      */
     public void removeAssociation(int id) {
-        Slog.i(TAG, "Removing association id=" + id);
+        Slog.i(TAG, "Removing association id=[" + id + "]...");
 
         final AssociationInfo association;
+
         synchronized (mLock) {
-            association = mIdMap.remove(id);
+            association = mIdToAssociationMap.remove(id);
 
             if (association == null) {
-                Slog.w(TAG, "Can't remove association. It does not exist.");
+                Slog.w(TAG, "Can't remove association id=[" + id + "]. It does not exist.");
                 return;
             }
 
-            final MacAddress macAddress = association.getDeviceMacAddress();
-            if (macAddress != null) {
-                mAddressMap.get(macAddress).remove(id);
-            }
-
-            invalidateCacheForUserLocked(association.getUserId());
+            writeCacheToDisk(association.getUserId());
 
             Slog.i(TAG, "Done removing association.");
         }
 
-        broadcastChange(CHANGE_TYPE_REMOVED, association);
+        logRemoveAssociation(association.getDeviceProfile());
+
+        if (association.isActive()) {
+            broadcastChange(CHANGE_TYPE_REMOVED, association);
+        }
+    }
+
+    private void writeCacheToDisk(@UserIdInt int userId) {
+        mExecutor.execute(() -> {
+            Associations associations = new Associations();
+            synchronized (mLock) {
+                associations.setMaxId(mUserToMaxId.getOrDefault(userId, 0));
+                associations.setAssociations(
+                        CollectionUtils.filter(mIdToAssociationMap.values().stream().toList(),
+                                a -> a.getUserId() == userId));
+            }
+            mDiskStore.writeAssociationsForUser(userId, associations);
+        });
     }
 
     /**
-     * @return a "snapshot" of the current state of the existing associations.
+     * Get a copy of all associations including pending and revoked ones.
+     * Modifying the copy won't modify the actual associations.
+     *
+     * If a cache miss happens, read from disk.
      */
-    public @NonNull Collection<AssociationInfo> getAssociations() {
+    @NonNull
+    public List<AssociationInfo> getAssociations() {
         synchronized (mLock) {
-            // IMPORTANT: make and return a COPY of the mIdMap.values(), NOT a "direct" reference.
-            // The HashMap.values() returns a collection which is backed by the HashMap, so changes
-            // to the HashMap are reflected in this collection.
-            // For us this means that if mIdMap is modified while the iteration over mIdMap.values()
-            // is in progress it may lead to "undefined results" (according to the HashMap's
-            // documentation) or cause ConcurrentModificationExceptions in the iterator (according
-            // to the bugreports...).
-            return List.copyOf(mIdMap.values());
+            if (!mPersisted) {
+                refreshCache();
+            }
+            return List.copyOf(mIdToAssociationMap.values());
         }
     }
 
     /**
-     * Get associations for the user.
+     * Get a copy of active associations.
      */
-    public @NonNull List<AssociationInfo> getAssociationsForUser(@UserIdInt int userId) {
+    @NonNull
+    public List<AssociationInfo> getActiveAssociations() {
         synchronized (mLock) {
-            return getAssociationsForUserLocked(userId);
+            return CollectionUtils.filter(getAssociations(), AssociationInfo::isActive);
         }
     }
 
     /**
-     * Get associations for the package
+     * Get a copy of all associations by user.
      */
-    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);
+    @NonNull
+    public List<AssociationInfo> getAssociationsByUser(@UserIdInt int userId) {
+        synchronized (mLock) {
+            return CollectionUtils.filter(getAssociations(), a -> a.getUserId() == userId);
+        }
     }
 
     /**
-     * Get associations by mac address for the package.
+     * Get a copy of active associations by user.
      */
-    public @Nullable AssociationInfo getAssociationsForPackageWithAddress(
+    @NonNull
+    public List<AssociationInfo> getActiveAssociationsByUser(@UserIdInt int userId) {
+        synchronized (mLock) {
+            return CollectionUtils.filter(getActiveAssociations(), a -> a.getUserId() == userId);
+        }
+    }
+
+    /**
+     * Get a copy of all associations by package.
+     */
+    @NonNull
+    public List<AssociationInfo> getAssociationsByPackage(@UserIdInt int userId,
+            @NonNull String packageName) {
+        synchronized (mLock) {
+            return CollectionUtils.filter(getAssociationsByUser(userId),
+                    a -> a.getPackageName().equals(packageName));
+        }
+    }
+
+    /**
+     * Get a copy of active associations by package.
+     */
+    @NonNull
+    public List<AssociationInfo> getActiveAssociationsByPackage(@UserIdInt int userId,
+            @NonNull String packageName) {
+        synchronized (mLock) {
+            return CollectionUtils.filter(getActiveAssociationsByUser(userId),
+                    a -> a.getPackageName().equals(packageName));
+        }
+    }
+
+    /**
+     * Get the first active association with the mac address.
+     */
+    @Nullable
+    public AssociationInfo getFirstAssociationByAddress(
             @UserIdInt int userId, @NonNull String packageName, @NonNull String macAddress) {
-        final List<AssociationInfo> associations = getAssociationsByAddress(macAddress);
-        return CollectionUtils.find(associations,
-                it -> it.belongsToPackage(userId, packageName));
-    }
-
-    /**
-     * Get association by id.
-     */
-    public @Nullable AssociationInfo getAssociationById(int id) {
         synchronized (mLock) {
-            return mIdMap.get(id);
+            return CollectionUtils.find(getActiveAssociationsByPackage(userId, packageName),
+                    a -> a.getDeviceMacAddress() != null && a.getDeviceMacAddress()
+                            .equals(MacAddress.fromString(macAddress)));
         }
     }
 
     /**
-     * Get associations by mac address.
+     * Get the association by id.
+     */
+    @Nullable
+    public AssociationInfo getAssociationById(int id) {
+        synchronized (mLock) {
+            return mIdToAssociationMap.get(id);
+        }
+    }
+
+    /**
+     * Get a copy of active associations by mac address.
      */
     @NonNull
-    public List<AssociationInfo> getAssociationsByAddress(@NonNull String macAddress) {
-        final MacAddress address = MacAddress.fromString(macAddress);
-
+    public List<AssociationInfo> getActiveAssociationsByAddress(@NonNull String macAddress) {
         synchronized (mLock) {
-            final Set<Integer> ids = mAddressMap.get(address);
-            if (ids == null) return Collections.emptyList();
-
-            final List<AssociationInfo> associations = new ArrayList<>(ids.size());
-            for (Integer id : ids) {
-                associations.add(mIdMap.get(id));
-            }
-
-            return Collections.unmodifiableList(associations);
+            return CollectionUtils.filter(getActiveAssociations(),
+                    a -> a.getDeviceMacAddress() != null && a.getDeviceMacAddress()
+                            .equals(MacAddress.fromString(macAddress)));
         }
     }
 
-    @GuardedBy("mLock")
+    /**
+     * Get a copy of revoked associations.
+     */
     @NonNull
-    private 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);
-    }
-
-    /**
-     * Register a listener for association changes.
-     */
-    public void registerListener(@NonNull OnChangeListener listener) {
-        synchronized (mListeners) {
-            mListeners.add(listener);
+    public List<AssociationInfo> getRevokedAssociations() {
+        synchronized (mLock) {
+            return CollectionUtils.filter(getAssociations(), AssociationInfo::isRevoked);
         }
     }
 
     /**
-     * Unregister a listener previously registered for association changes.
+     * Get a copy of revoked associations for the package.
      */
-    public void unregisterListener(@NonNull OnChangeListener listener) {
-        synchronized (mListeners) {
-            mListeners.remove(listener);
+    @NonNull
+    public List<AssociationInfo> getRevokedAssociations(@UserIdInt int userId,
+            @NonNull String packageName) {
+        synchronized (mLock) {
+            return CollectionUtils.filter(getAssociations(),
+                    a -> packageName.equals(a.getPackageName()) && a.getUserId() == userId
+                            && a.isRevoked());
+        }
+    }
+
+    /**
+     * Get a copy of active associations.
+     */
+    @NonNull
+    public List<AssociationInfo> getPendingAssociations(@UserIdInt int userId,
+            @NonNull String packageName) {
+        synchronized (mLock) {
+            return CollectionUtils.filter(getAssociations(),
+                    a -> packageName.equals(a.getPackageName()) && a.getUserId() == userId
+                            && a.isPending());
+        }
+    }
+
+    /**
+     * Register a local listener for association changes.
+     */
+    public void registerLocalListener(@NonNull OnChangeListener listener) {
+        synchronized (mLocalListeners) {
+            mLocalListeners.add(listener);
+        }
+    }
+
+    /**
+     * Unregister a local listener previously registered for association changes.
+     */
+    public void unregisterLocalListener(@NonNull OnChangeListener listener) {
+        synchronized (mLocalListeners) {
+            mLocalListeners.remove(listener);
+        }
+    }
+
+    /**
+     * Register a remote listener for association changes.
+     */
+    public void registerRemoteListener(@NonNull IOnAssociationsChangedListener listener,
+            int userId) {
+        synchronized (mRemoteListeners) {
+            mRemoteListeners.register(listener, userId);
+        }
+    }
+
+    /**
+     * Unregister a remote listener previously registered for association changes.
+     */
+    public void unregisterRemoteListener(@NonNull IOnAssociationsChangedListener listener) {
+        synchronized (mRemoteListeners) {
+            mRemoteListeners.unregister(listener);
         }
     }
 
@@ -350,52 +490,41 @@
      */
     public void dump(@NonNull PrintWriter out) {
         out.append("Companion Device Associations: ");
-        if (getAssociations().isEmpty()) {
+        if (getActiveAssociations().isEmpty()) {
             out.append("<empty>\n");
         } else {
             out.append("\n");
-            for (AssociationInfo a : getAssociations()) {
+            for (AssociationInfo a : getActiveAssociations()) {
                 out.append("  ").append(a.toString()).append('\n');
             }
         }
     }
 
     private void broadcastChange(@ChangeType int changeType, AssociationInfo association) {
-        synchronized (mListeners) {
-            for (OnChangeListener listener : mListeners) {
+        Slog.i(TAG, "Broadcasting association changes - changeType=[" + changeType + "]...");
+
+        synchronized (mLocalListeners) {
+            for (OnChangeListener listener : mLocalListeners) {
                 listener.onAssociationChanged(changeType, association);
             }
         }
-    }
-
-    /**
-     * Set associations to cache. It will clear the existing cache.
-     */
-    public void setAssociationsToCache(Collection<AssociationInfo> associations) {
-        // Validity check first.
-        associations.forEach(AssociationStore::checkNotRevoked);
-
-        synchronized (mLock) {
-            mIdMap.clear();
-            mAddressMap.clear();
-            mCachedPerUser.clear();
-
-            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);
-                }
+        synchronized (mRemoteListeners) {
+            final int userId = association.getUserId();
+            final List<AssociationInfo> updatedAssociations = getActiveAssociationsByUser(userId);
+            // 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) {
+                mRemoteListeners.broadcast((listener, callbackUserId) -> {
+                    int listenerUserId = (int) callbackUserId;
+                    if (listenerUserId == userId || listenerUserId == UserHandle.USER_ALL) {
+                        try {
+                            listener.onAssociationsChanged(updatedAssociations);
+                        } catch (RemoteException ignored) {
+                        }
+                    }
+                });
             }
         }
     }
-
-    private static void checkNotRevoked(@NonNull AssociationInfo association) {
-        if (association.isRevoked()) {
-            throw new IllegalArgumentException(
-                    "Revoked (removed) associations MUST NOT appear in the AssociationStore");
-        }
-    }
 }
diff --git a/services/companion/java/com/android/server/companion/association/Associations.java b/services/companion/java/com/android/server/companion/association/Associations.java
new file mode 100644
index 0000000..7da3699
--- /dev/null
+++ b/services/companion/java/com/android/server/companion/association/Associations.java
@@ -0,0 +1,68 @@
+/*
+ * Copyright (C) 2024 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.association;
+
+import android.companion.AssociationInfo;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Represents associations per user. Should be only used by Association stores.
+ */
+public class Associations {
+
+    private int mVersion = 0;
+
+    private List<AssociationInfo> mAssociations = new ArrayList<>();
+
+    private int mMaxId = 0;
+
+    public Associations() {
+    }
+
+    public void setVersion(int version) {
+        mVersion = version;
+    }
+
+    /**
+     * Add an association.
+     */
+    public void addAssociation(AssociationInfo association) {
+        mAssociations.add(association);
+    }
+
+    public void setMaxId(int maxId) {
+        mMaxId = maxId;
+    }
+
+    public void setAssociations(List<AssociationInfo> associations) {
+        mAssociations = List.copyOf(associations);
+    }
+
+    public int getVersion() {
+        return mVersion;
+    }
+
+    public int getMaxId() {
+        return mMaxId;
+    }
+
+    public List<AssociationInfo> getAssociations() {
+        return mAssociations;
+    }
+}
diff --git a/services/companion/java/com/android/server/companion/association/DisassociationProcessor.java b/services/companion/java/com/android/server/companion/association/DisassociationProcessor.java
new file mode 100644
index 0000000..ec897791
--- /dev/null
+++ b/services/companion/java/com/android/server/companion/association/DisassociationProcessor.java
@@ -0,0 +1,218 @@
+/*
+ * Copyright (C) 2024 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.association;
+
+import static android.app.ActivityManager.RunningAppProcessInfo.IMPORTANCE_VISIBLE;
+import static android.companion.AssociationRequest.DEVICE_PROFILE_AUTOMOTIVE_PROJECTION;
+
+import static com.android.internal.util.CollectionUtils.any;
+import static com.android.server.companion.utils.RolesUtils.removeRoleHolderForAssociation;
+
+import android.annotation.NonNull;
+import android.annotation.SuppressLint;
+import android.annotation.UserIdInt;
+import android.app.ActivityManager;
+import android.companion.AssociationInfo;
+import android.content.Context;
+import android.content.pm.PackageManagerInternal;
+import android.os.Binder;
+import android.os.UserHandle;
+import android.util.Slog;
+
+import com.android.server.companion.CompanionApplicationController;
+import com.android.server.companion.datatransfer.SystemDataTransferRequestStore;
+import com.android.server.companion.presence.CompanionDevicePresenceMonitor;
+import com.android.server.companion.transport.CompanionTransportManager;
+
+/**
+ * A class response for Association removal.
+ */
+@SuppressLint("LongLogTag")
+public class DisassociationProcessor {
+
+    private static final String TAG = "CDM_DisassociationProcessor";
+    @NonNull
+    private final Context mContext;
+    @NonNull
+    private final AssociationStore mAssociationStore;
+    @NonNull
+    private final PackageManagerInternal mPackageManagerInternal;
+    @NonNull
+    private final CompanionDevicePresenceMonitor mDevicePresenceMonitor;
+    @NonNull
+    private final SystemDataTransferRequestStore mSystemDataTransferRequestStore;
+    @NonNull
+    private final CompanionApplicationController mCompanionAppController;
+    @NonNull
+    private final CompanionTransportManager mTransportManager;
+    private final OnPackageVisibilityChangeListener mOnPackageVisibilityChangeListener;
+    private final ActivityManager mActivityManager;
+
+    public DisassociationProcessor(@NonNull Context context,
+            @NonNull ActivityManager activityManager,
+            @NonNull AssociationStore associationStore,
+            @NonNull PackageManagerInternal packageManager,
+            @NonNull CompanionDevicePresenceMonitor devicePresenceMonitor,
+            @NonNull CompanionApplicationController applicationController,
+            @NonNull SystemDataTransferRequestStore systemDataTransferRequestStore,
+            @NonNull CompanionTransportManager companionTransportManager) {
+        mContext = context;
+        mActivityManager = activityManager;
+        mAssociationStore = associationStore;
+        mPackageManagerInternal = packageManager;
+        mOnPackageVisibilityChangeListener =
+                new OnPackageVisibilityChangeListener();
+        mDevicePresenceMonitor = devicePresenceMonitor;
+        mCompanionAppController = applicationController;
+        mSystemDataTransferRequestStore = systemDataTransferRequestStore;
+        mTransportManager = companionTransportManager;
+    }
+
+    /**
+     * Disassociate an association by id.
+     */
+    // TODO: also revoke notification access
+    public void disassociate(int id) {
+        Slog.i(TAG, "Disassociating id=[" + id + "]...");
+
+        final AssociationInfo association = mAssociationStore.getAssociationById(id);
+        if (association == null) {
+            Slog.e(TAG, "Can't disassociate id=[" + id + "]. It doesn't exist.");
+            return;
+        }
+
+        final int userId = association.getUserId();
+        final String packageName = association.getPackageName();
+        final String deviceProfile = association.getDeviceProfile();
+
+        final boolean isRoleInUseByOtherAssociations = deviceProfile != null
+                && any(mAssociationStore.getActiveAssociationsByPackage(userId, packageName),
+                    it -> deviceProfile.equals(it.getDeviceProfile()) && id != it.getId());
+
+        final int packageProcessImportance = getPackageProcessImportance(userId, packageName);
+        if (packageProcessImportance <= IMPORTANCE_VISIBLE && deviceProfile != null
+                && !isRoleInUseByOtherAssociations) {
+            // Need to remove the app from the list of role holders, but the process is visible
+            // to the user at the moment, so we'll need to do it later.
+            Slog.i(TAG, "Cannot disassociate id=[" + id + "] now - process is visible. "
+                    + "Start listening to package importance...");
+
+            AssociationInfo revokedAssociation = (new AssociationInfo.Builder(
+                    association)).setRevoked(true).build();
+            mAssociationStore.updateAssociation(revokedAssociation);
+            startListening();
+            return;
+        }
+
+        // Association cleanup.
+        mAssociationStore.removeAssociation(association.getId());
+        mSystemDataTransferRequestStore.removeRequestsByAssociationId(userId, id);
+
+        // Detach transport if exists
+        mTransportManager.detachSystemDataTransport(packageName, userId, id);
+
+        // If role is not in use by other associations, revoke the role.
+        // Do not need to remove the system role since it was pre-granted by the system.
+        if (!isRoleInUseByOtherAssociations && deviceProfile != null && !deviceProfile.equals(
+                DEVICE_PROFILE_AUTOMOTIVE_PROJECTION)) {
+            removeRoleHolderForAssociation(mContext, association.getUserId(),
+                    association.getPackageName(), association.getDeviceProfile());
+        }
+
+        // Unbind the app if needed.
+        final boolean wasPresent = mDevicePresenceMonitor.isDevicePresent(id);
+        if (!wasPresent || !association.isNotifyOnDeviceNearby()) {
+            return;
+        }
+        final boolean shouldStayBound = any(
+                mAssociationStore.getActiveAssociationsByPackage(userId, packageName),
+                it -> it.isNotifyOnDeviceNearby()
+                        && mDevicePresenceMonitor.isDevicePresent(it.getId()));
+        if (!shouldStayBound) {
+            mCompanionAppController.unbindCompanionApplication(userId, packageName);
+        }
+    }
+
+    @SuppressLint("MissingPermission")
+    private int getPackageProcessImportance(@UserIdInt int userId, @NonNull String packageName) {
+        return Binder.withCleanCallingIdentity(() -> {
+            final int uid =
+                    mPackageManagerInternal.getPackageUid(packageName, /* flags */0, userId);
+            return mActivityManager.getUidImportance(uid);
+        });
+    }
+
+    private void startListening() {
+        Slog.i(TAG, "Start listening to uid importance changes...");
+        try {
+            Binder.withCleanCallingIdentity(
+                    () -> mActivityManager.addOnUidImportanceListener(
+                            mOnPackageVisibilityChangeListener,
+                            ActivityManager.RunningAppProcessInfo.IMPORTANCE_VISIBLE));
+        }  catch (IllegalArgumentException e) {
+            Slog.e(TAG, "Failed to start listening to uid importance changes.");
+        }
+    }
+
+    private void stopListening() {
+        Slog.i(TAG, "Stop listening to uid importance changes.");
+        try {
+            Binder.withCleanCallingIdentity(() -> mActivityManager.removeOnUidImportanceListener(
+                    mOnPackageVisibilityChangeListener));
+        } catch (IllegalArgumentException e) {
+            Slog.e(TAG, "Failed to stop listening to uid importance changes.");
+        }
+    }
+
+    /**
+     * An OnUidImportanceListener class which watches the importance of the packages.
+     * In this class, we ONLY interested in the importance of the running process is greater than
+     * {@link ActivityManager.RunningAppProcessInfo#IMPORTANCE_VISIBLE}.
+     *
+     * Lastly remove the role holder for the revoked associations for the same packages.
+     *
+     * @see #disassociate(int)
+     */
+    private class OnPackageVisibilityChangeListener implements
+            ActivityManager.OnUidImportanceListener {
+
+        @Override
+        public void onUidImportance(int uid, int importance) {
+            if (importance <= ActivityManager.RunningAppProcessInfo.IMPORTANCE_VISIBLE) {
+                // The lower the importance value the more "important" the process is.
+                // We are only interested when the process ceases to be visible.
+                return;
+            }
+
+            final String packageName = mPackageManagerInternal.getNameForUid(uid);
+            if (packageName == null) {
+                // Not interested in this uid.
+                return;
+            }
+
+            int userId = UserHandle.getUserId(uid);
+            for (AssociationInfo association : mAssociationStore.getRevokedAssociations(userId,
+                    packageName)) {
+                disassociate(association.getId());
+            }
+
+            if (mAssociationStore.getRevokedAssociations().isEmpty()) {
+                stopListening();
+            }
+        }
+    }
+}
diff --git a/services/companion/java/com/android/server/companion/association/InactiveAssociationsRemovalService.java b/services/companion/java/com/android/server/companion/association/InactiveAssociationsRemovalService.java
index 894c49a..f287315 100644
--- a/services/companion/java/com/android/server/companion/association/InactiveAssociationsRemovalService.java
+++ b/services/companion/java/com/android/server/companion/association/InactiveAssociationsRemovalService.java
@@ -33,7 +33,7 @@
  * A Job Service responsible for clean up idle self-managed associations.
  *
  * The job will be executed only if the device is charging and in idle mode due to the application
- * will be killed if association/role are revoked. See {@link AssociationRevokeProcessor}
+ * will be killed if association/role are revoked. See {@link DisassociationProcessor}
  */
 public class InactiveAssociationsRemovalService extends JobService {
 
diff --git a/services/companion/java/com/android/server/companion/datatransfer/SystemDataTransferProcessor.java b/services/companion/java/com/android/server/companion/datatransfer/SystemDataTransferProcessor.java
index a08e0da..c5ca0bf 100644
--- a/services/companion/java/com/android/server/companion/datatransfer/SystemDataTransferProcessor.java
+++ b/services/companion/java/com/android/server/companion/datatransfer/SystemDataTransferProcessor.java
@@ -186,18 +186,14 @@
         intent.putExtras(extras);
 
         // Create a PendingIntent
-        final long token = Binder.clearCallingIdentity();
-        try {
-            return PendingIntent.getActivityAsUser(mContext, /*requestCode */ associationId, intent,
-                    FLAG_ONE_SHOT | FLAG_CANCEL_CURRENT | FLAG_IMMUTABLE,
-                    ActivityOptions.makeBasic()
-                            .setPendingIntentCreatorBackgroundActivityStartMode(
-                                    ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED)
-                            .toBundle(),
-                    UserHandle.CURRENT);
-        } finally {
-            Binder.restoreCallingIdentity(token);
-        }
+        return Binder.withCleanCallingIdentity(() ->
+                PendingIntent.getActivityAsUser(mContext, /*requestCode */ associationId,
+                        intent, FLAG_ONE_SHOT | FLAG_CANCEL_CURRENT | FLAG_IMMUTABLE,
+                        ActivityOptions.makeBasic()
+                                .setPendingIntentCreatorBackgroundActivityStartMode(
+                                        ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED)
+                                .toBundle(),
+                        UserHandle.CURRENT));
     }
 
     /**
@@ -228,8 +224,7 @@
         }
 
         // Start permission sync
-        final long callingIdentityToken = Binder.clearCallingIdentity();
-        try {
+        Binder.withCleanCallingIdentity(() -> {
             // TODO: refactor to work with streams of data
             mPermissionControllerManager.getRuntimePermissionBackup(UserHandle.of(userId),
                     mExecutor, backup -> {
@@ -237,39 +232,31 @@
                                 .requestPermissionRestore(associationId, backup);
                         translateFutureToCallback(future, callback);
                     });
-        } finally {
-            Binder.restoreCallingIdentity(callingIdentityToken);
-        }
+        });
     }
 
     /**
      * Enable perm sync for the association
      */
     public void enablePermissionsSync(int associationId) {
-        final long callingIdentityToken = Binder.clearCallingIdentity();
-        try {
+        Binder.withCleanCallingIdentity(() -> {
             int userId = mAssociationStore.getAssociationById(associationId).getUserId();
             PermissionSyncRequest request = new PermissionSyncRequest(associationId);
             request.setUserConsented(true);
             mSystemDataTransferRequestStore.writeRequest(userId, request);
-        } finally {
-            Binder.restoreCallingIdentity(callingIdentityToken);
-        }
+        });
     }
 
     /**
      * Disable perm sync for the association
      */
     public void disablePermissionsSync(int associationId) {
-        final long callingIdentityToken = Binder.clearCallingIdentity();
-        try {
+        Binder.withCleanCallingIdentity(() -> {
             int userId = mAssociationStore.getAssociationById(associationId).getUserId();
             PermissionSyncRequest request = new PermissionSyncRequest(associationId);
             request.setUserConsented(false);
             mSystemDataTransferRequestStore.writeRequest(userId, request);
-        } finally {
-            Binder.restoreCallingIdentity(callingIdentityToken);
-        }
+        });
     }
 
     /**
@@ -277,8 +264,7 @@
      */
     @Nullable
     public PermissionSyncRequest getPermissionSyncRequest(int associationId) {
-        final long callingIdentityToken = Binder.clearCallingIdentity();
-        try {
+        return Binder.withCleanCallingIdentity(() -> {
             int userId = mAssociationStore.getAssociationById(associationId).getUserId();
             List<SystemDataTransferRequest> requests =
                     mSystemDataTransferRequestStore.readRequestsByAssociationId(userId,
@@ -289,22 +275,17 @@
                 }
             }
             return null;
-        } finally {
-            Binder.restoreCallingIdentity(callingIdentityToken);
-        }
+        });
     }
 
     /**
      * Remove perm sync request for the association.
      */
     public void removePermissionSyncRequest(int associationId) {
-        final long callingIdentityToken = Binder.clearCallingIdentity();
-        try {
+        Binder.withCleanCallingIdentity(() -> {
             int userId = mAssociationStore.getAssociationById(associationId).getUserId();
             mSystemDataTransferRequestStore.removeRequestsByAssociationId(userId, associationId);
-        } finally {
-            Binder.restoreCallingIdentity(callingIdentityToken);
-        }
+        });
     }
 
     private void onReceivePermissionRestore(byte[] message) {
@@ -318,14 +299,12 @@
         Slog.i(LOG_TAG, "Applying permissions.");
         // Start applying permissions
         UserHandle user = mContext.getUser();
-        final long callingIdentityToken = Binder.clearCallingIdentity();
-        try {
+
+        Binder.withCleanCallingIdentity(() -> {
             // TODO: refactor to work with streams of data
             mPermissionControllerManager.stageAndApplyRuntimePermissionsBackup(
                     message, user);
-        } finally {
-            Binder.restoreCallingIdentity(callingIdentityToken);
-        }
+        });
     }
 
     private static void translateFutureToCallback(@NonNull Future<?> future,
diff --git a/services/companion/java/com/android/server/companion/presence/BleCompanionDeviceScanner.java b/services/companion/java/com/android/server/companion/presence/BleCompanionDeviceScanner.java
index 99466a9..c89ce11 100644
--- a/services/companion/java/com/android/server/companion/presence/BleCompanionDeviceScanner.java
+++ b/services/companion/java/com/android/server/companion/presence/BleCompanionDeviceScanner.java
@@ -106,7 +106,7 @@
         checkBleState();
         registerBluetoothStateBroadcastReceiver(context);
 
-        mAssociationStore.registerListener(this);
+        mAssociationStore.registerLocalListener(this);
     }
 
     @MainThread
@@ -183,7 +183,7 @@
 
         // Collect MAC addresses from all associations.
         final Set<String> macAddresses = new HashSet<>();
-        for (AssociationInfo association : mAssociationStore.getAssociations()) {
+        for (AssociationInfo association : mAssociationStore.getActiveAssociations()) {
             if (!association.isNotifyOnDeviceNearby()) continue;
 
             // Beware that BT stack does not consider low-case MAC addresses valid, while
@@ -255,7 +255,7 @@
         if (DEBUG) Log.i(TAG, "notifyDevice_Found()" + btDeviceToString(device));
 
         final List<AssociationInfo> associations =
-                mAssociationStore.getAssociationsByAddress(device.getAddress());
+                mAssociationStore.getActiveAssociationsByAddress(device.getAddress());
         if (DEBUG) Log.d(TAG, "  > associations=" + Arrays.toString(associations.toArray()));
 
         for (AssociationInfo association : associations) {
@@ -268,7 +268,7 @@
         if (DEBUG) Log.i(TAG, "notifyDevice_Lost()" + btDeviceToString(device));
 
         final List<AssociationInfo> associations =
-                mAssociationStore.getAssociationsByAddress(device.getAddress());
+                mAssociationStore.getActiveAssociationsByAddress(device.getAddress());
         if (DEBUG) Log.d(TAG, "  > associations=" + Arrays.toString(associations.toArray()));
 
         for (AssociationInfo association : associations) {
@@ -319,7 +319,7 @@
                 Log.v(TAG, "  > scanResult=" + result);
 
                 final List<AssociationInfo> associations =
-                        mAssociationStore.getAssociationsByAddress(device.getAddress());
+                        mAssociationStore.getActiveAssociationsByAddress(device.getAddress());
                 Log.v(TAG, "  > associations=" + Arrays.toString(associations.toArray()));
             }
 
diff --git a/services/companion/java/com/android/server/companion/presence/BluetoothCompanionDeviceConnectionListener.java b/services/companion/java/com/android/server/companion/presence/BluetoothCompanionDeviceConnectionListener.java
index 4da3f9b..cb363a7 100644
--- a/services/companion/java/com/android/server/companion/presence/BluetoothCompanionDeviceConnectionListener.java
+++ b/services/companion/java/com/android/server/companion/presence/BluetoothCompanionDeviceConnectionListener.java
@@ -93,7 +93,7 @@
 
         btAdapter.registerBluetoothConnectionCallback(
                 new HandlerExecutor(Handler.getMain()), /* callback */this);
-        mAssociationStore.registerListener(this);
+        mAssociationStore.registerLocalListener(this);
     }
 
     /**
@@ -168,7 +168,7 @@
     private void onDeviceConnectivityChanged(@NonNull BluetoothDevice device, boolean connected) {
         int userId = UserHandle.myUserId();
         final List<AssociationInfo> associations =
-                mAssociationStore.getAssociationsByAddress(device.getAddress());
+                mAssociationStore.getActiveAssociationsByAddress(device.getAddress());
         final List<ObservableUuid> observableUuids =
                 mObservableUuidStore.getObservableUuidsForUser(userId);
         final ParcelUuid[] bluetoothDeviceUuids = device.getUuids();
diff --git a/services/companion/java/com/android/server/companion/presence/CompanionDevicePresenceMonitor.java b/services/companion/java/com/android/server/companion/presence/CompanionDevicePresenceMonitor.java
index 37bbb93..7a1a83f 100644
--- a/services/companion/java/com/android/server/companion/presence/CompanionDevicePresenceMonitor.java
+++ b/services/companion/java/com/android/server/companion/presence/CompanionDevicePresenceMonitor.java
@@ -145,7 +145,7 @@
             Log.w(TAG, "BluetoothAdapter is NOT available.");
         }
 
-        mAssociationStore.registerListener(this);
+        mAssociationStore.registerLocalListener(this);
     }
 
     /**
@@ -481,7 +481,7 @@
      * BT connected and BLE presence and are not pending to report BLE lost.
      */
     private boolean canStopBleScan() {
-        for (AssociationInfo ai : mAssociationStore.getAssociations()) {
+        for (AssociationInfo ai : mAssociationStore.getActiveAssociations()) {
             int id = ai.getId();
             synchronized (mBtDisconnectedDevices) {
                 if (ai.isNotifyOnDeviceNearby() && !(isBtConnected(id)
diff --git a/services/companion/java/com/android/server/companion/presence/ObservableUuidStore.java b/services/companion/java/com/android/server/companion/presence/ObservableUuidStore.java
index ee8b106..db15da29 100644
--- a/services/companion/java/com/android/server/companion/presence/ObservableUuidStore.java
+++ b/services/companion/java/com/android/server/companion/presence/ObservableUuidStore.java
@@ -90,6 +90,8 @@
      * Remove the observable uuid.
      */
     public void removeObservableUuid(@UserIdInt int userId, ParcelUuid uuid, String packageName) {
+        Slog.i(TAG, "Removing uuid=[" + uuid.getUuid() + "] from store...");
+
         List<ObservableUuid> cachedObservableUuids;
 
         synchronized (mLock) {
@@ -108,7 +110,7 @@
      * Write the observable uuid.
      */
     public void writeObservableUuid(@UserIdInt int userId, ObservableUuid uuid) {
-        Slog.i(TAG, "Writing uuid=" + uuid.getUuid() + " to store.");
+        Slog.i(TAG, "Writing uuid=[" + uuid.getUuid() + "] to store...");
 
         List<ObservableUuid> cachedObservableUuids;
         synchronized (mLock) {
diff --git a/services/companion/java/com/android/server/companion/utils/DataStoreUtils.java b/services/companion/java/com/android/server/companion/utils/DataStoreUtils.java
index c75b1a5..369a925 100644
--- a/services/companion/java/com/android/server/companion/utils/DataStoreUtils.java
+++ b/services/companion/java/com/android/server/companion/utils/DataStoreUtils.java
@@ -64,8 +64,8 @@
      * IMPORTANT: the method will ALWAYS return the same {@link AtomicFile} object, which makes it
      * possible to synchronize reads and writes to the file using the returned object.
      *
-     * @param userId              the userId to retrieve the storage file
-     * @param fileName         the storage file name
+     * @param userId the userId to retrieve the storage file
+     * @param fileName the storage file name
      * @return an AtomicFile for the user
      */
     @NonNull
diff --git a/services/companion/java/com/android/server/companion/utils/RolesUtils.java b/services/companion/java/com/android/server/companion/utils/RolesUtils.java
index f798e21..dd12e04 100644
--- a/services/companion/java/com/android/server/companion/utils/RolesUtils.java
+++ b/services/companion/java/com/android/server/companion/utils/RolesUtils.java
@@ -93,8 +93,8 @@
 
         Slog.i(TAG, "Removing CDM role=" + deviceProfile
                 + " for userId=" + userId + ", packageName=" + packageName);
-        final long identity = Binder.clearCallingIdentity();
-        try {
+
+        Binder.withCleanCallingIdentity(() ->
             roleManager.removeRoleHolderAsUser(deviceProfile, packageName,
                     MANAGE_HOLDERS_FLAG_DONT_KILL_APP, userHandle, context.getMainExecutor(),
                     success -> {
@@ -103,11 +103,9 @@
                                     + packageName + " from the list of " + deviceProfile
                                     + " holders.");
                         }
-                    });
-        } finally {
-            Binder.restoreCallingIdentity(identity);
-        }
+                    })
+        );
     }
 
-    private RolesUtils() {};
+    private RolesUtils() {}
 }