Merge changes I7dbd511f,I87dc63bf into main

* changes:
  Register death recipient for discovery callbacks
  Add wakelocks for endpoint discovery callbacks

4 files changed, 127 insertions, 22 deletions

diff --git a/core/java/android/hardware/location/ContextHubManager.java b/core/java/android/hardware/location/ContextHubManager.java
index 9181bd0cb2ed..953ee08800cf 100644
--- a/core/java/android/hardware/location/ContextHubManager.java
+++ b/core/java/android/hardware/location/ContextHubManager.java
@@ -771,6 +771,7 @@ public final class ContextHubManager {
      */
     @FlaggedApi(Flags.FLAG_OFFLOAD_API)
     private IContextHubEndpointDiscoveryCallback createDiscoveryCallback(
+            IContextHubService service,
             Executor executor,
             HubEndpointDiscoveryCallback callback,
             @Nullable String serviceDescriptor) {
@@ -779,6 +780,7 @@ public final class ContextHubManager {
             public void onEndpointsStarted(HubEndpointInfo[] hubEndpointInfoList) {
                 if (hubEndpointInfoList.length == 0) {
                     Log.w(TAG, "onEndpointsStarted: received empty discovery list");
+                    invokeCallbackFinished(service);
                     return;
                 }
                 executor.execute(
@@ -791,6 +793,7 @@ public final class ContextHubManager {
                             } else {
                                 callback.onEndpointsStarted(discoveryList);
                             }
+                            invokeCallbackFinished(service);
                         });
             }
 
@@ -798,6 +801,7 @@ public final class ContextHubManager {
             public void onEndpointsStopped(HubEndpointInfo[] hubEndpointInfoList, int reason) {
                 if (hubEndpointInfoList.length == 0) {
                     Log.w(TAG, "onEndpointsStopped: received empty discovery list");
+                    invokeCallbackFinished(service);
                     return;
                 }
                 executor.execute(
@@ -810,8 +814,17 @@ public final class ContextHubManager {
                             } else {
                                 callback.onEndpointsStopped(discoveryList, reason);
                             }
+                            invokeCallbackFinished(service);
                         });
             }
+
+            private void invokeCallbackFinished(IContextHubService service) {
+                try {
+                    service.onDiscoveryCallbackFinished();
+                } catch (RemoteException e) {
+                    e.rethrowFromSystemServer();
+                }
+            }
         };
     }
 
@@ -873,7 +886,7 @@ public final class ContextHubManager {
         Objects.requireNonNull(executor, "executor cannot be null");
         Objects.requireNonNull(callback, "callback cannot be null");
         IContextHubEndpointDiscoveryCallback iCallback =
-                createDiscoveryCallback(executor, callback, null);
+                createDiscoveryCallback(mService, executor, callback, null);
         try {
             mService.registerEndpointDiscoveryCallbackId(endpointId, iCallback);
         } catch (RemoteException e) {
@@ -919,7 +932,7 @@ public final class ContextHubManager {
         }
 
         IContextHubEndpointDiscoveryCallback iCallback =
-                createDiscoveryCallback(executor, callback, serviceDescriptor);
+                createDiscoveryCallback(mService, executor, callback, serviceDescriptor);
         try {
             mService.registerEndpointDiscoveryCallbackDescriptor(serviceDescriptor, iCallback);
         } catch (RemoteException e) {
diff --git a/core/java/android/hardware/location/IContextHubService.aidl b/core/java/android/hardware/location/IContextHubService.aidl
index d5b3fa251e82..bb5491d98cf9 100644
--- a/core/java/android/hardware/location/IContextHubService.aidl
+++ b/core/java/android/hardware/location/IContextHubService.aidl
@@ -150,4 +150,8 @@ interface IContextHubService {
     // Unregister an endpoint with the context hub
     @EnforcePermission("ACCESS_CONTEXT_HUB")
     void unregisterEndpointDiscoveryCallback(in IContextHubEndpointDiscoveryCallback callback);
+
+    // Called when a discovery callback is finished executing
+    @EnforcePermission("ACCESS_CONTEXT_HUB")
+    void onDiscoveryCallbackFinished();
 }
diff --git a/services/core/java/com/android/server/location/contexthub/ContextHubService.java b/services/core/java/com/android/server/location/contexthub/ContextHubService.java
index 2615a76ac279..2b0ca145372b 100644
--- a/services/core/java/com/android/server/location/contexthub/ContextHubService.java
+++ b/services/core/java/com/android/server/location/contexthub/ContextHubService.java
@@ -334,7 +334,7 @@ public class ContextHubService extends IContextHubService.Stub {
         if (Flags.offloadApi() && Flags.offloadImplementation()) {
             HubInfoRegistry registry;
             try {
-                registry = new HubInfoRegistry(mContextHubWrapper);
+                registry = new HubInfoRegistry(mContext, mContextHubWrapper);
                 mEndpointManager =
                         new ContextHubEndpointManager(
                                 mContext, mContextHubWrapper, registry, mTransactionManager);
@@ -821,6 +821,13 @@ public class ContextHubService extends IContextHubService.Stub {
         mHubInfoRegistry.unregisterEndpointDiscoveryCallback(callback);
     }
 
+    @android.annotation.EnforcePermission(android.Manifest.permission.ACCESS_CONTEXT_HUB)
+    @Override
+    public void onDiscoveryCallbackFinished() throws RemoteException {
+        super.onDiscoveryCallbackFinished_enforcePermission();
+        mHubInfoRegistry.onDiscoveryCallbackFinished();
+    }
+
     private void checkEndpointDiscoveryPreconditions() {
         if (mHubInfoRegistry == null) {
             Log.e(TAG, "Hub endpoint registry failed to initialize");
diff --git a/services/core/java/com/android/server/location/contexthub/HubInfoRegistry.java b/services/core/java/com/android/server/location/contexthub/HubInfoRegistry.java
index bf54fd720d42..711383bbca37 100644
--- a/services/core/java/com/android/server/location/contexthub/HubInfoRegistry.java
+++ b/services/core/java/com/android/server/location/contexthub/HubInfoRegistry.java
@@ -16,12 +16,18 @@
 
 package com.android.server.location.contexthub;
 
+import android.content.Context;
 import android.hardware.contexthub.HubEndpointInfo;
 import android.hardware.contexthub.HubServiceInfo;
 import android.hardware.contexthub.IContextHubEndpointDiscoveryCallback;
 import android.hardware.location.HubInfo;
-import android.os.DeadObjectException;
+import android.os.Binder;
+import android.os.IBinder;
+import android.os.PowerManager;
+import android.os.PowerManager.WakeLock;
+import android.os.Process;
 import android.os.RemoteException;
+import android.os.WorkSource;
 import android.util.ArrayMap;
 import android.util.IndentingPrintWriter;
 import android.util.Log;
@@ -34,10 +40,15 @@ import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.BiConsumer;
 
 class HubInfoRegistry implements ContextHubHalEndpointCallback.IEndpointLifecycleCallback {
     private static final String TAG = "HubInfoRegistry";
+
+    /** The duration of wakelocks acquired during discovery callbacks */
+    private static final long WAKELOCK_TIMEOUT_MILLIS = 5 * 1000;
+
     private final Object mLock = new Object();
 
     private final IContextHubWrapper mContextHubWrapper;
@@ -53,21 +64,37 @@ class HubInfoRegistry implements ContextHubHalEndpointCallback.IEndpointLifecycl
      * A wrapper class that is used to store arguments to
      * ContextHubManager.registerEndpointCallback.
      */
-    private static class DiscoveryCallback {
+    private static class DiscoveryCallback implements IBinder.DeathRecipient {
+        private final HubInfoRegistry mHubInfoRegistry;
         private final IContextHubEndpointDiscoveryCallback mCallback;
         private final Optional<Long> mEndpointId;
         private final Optional<String> mServiceDescriptor;
 
-        DiscoveryCallback(IContextHubEndpointDiscoveryCallback callback, long endpointId) {
+        // True if the binder death recipient fired
+        private final AtomicBoolean mBinderDied = new AtomicBoolean(false);
+
+        DiscoveryCallback(
+                HubInfoRegistry registry,
+                IContextHubEndpointDiscoveryCallback callback,
+                long endpointId)
+                throws RemoteException {
+            mHubInfoRegistry = registry;
             mCallback = callback;
             mEndpointId = Optional.of(endpointId);
             mServiceDescriptor = Optional.empty();
+            attachDeathRecipient();
         }
 
-        DiscoveryCallback(IContextHubEndpointDiscoveryCallback callback, String serviceDescriptor) {
+        DiscoveryCallback(
+                HubInfoRegistry registry,
+                IContextHubEndpointDiscoveryCallback callback,
+                String serviceDescriptor)
+                throws RemoteException {
+            mHubInfoRegistry = registry;
             mCallback = callback;
             mEndpointId = Optional.empty();
             mServiceDescriptor = Optional.of(serviceDescriptor);
+            attachDeathRecipient();
         }
 
         public IContextHubEndpointDiscoveryCallback getCallback() {
@@ -79,6 +106,10 @@ class HubInfoRegistry implements ContextHubHalEndpointCallback.IEndpointLifecycl
          * @return true if info matches
          */
         public boolean isMatch(HubEndpointInfo info) {
+            if (mBinderDied.get()) {
+                Log.w(TAG, "Callback died, isMatch returning false");
+                return false;
+            }
             if (mEndpointId.isPresent()) {
                 return mEndpointId.get() == info.getIdentifier().getEndpoint();
             }
@@ -91,6 +122,17 @@ class HubInfoRegistry implements ContextHubHalEndpointCallback.IEndpointLifecycl
             }
             return false;
         }
+
+        @Override
+        public void binderDied() {
+            Log.d(TAG, "Binder died for discovery callback");
+            mBinderDied.set(true);
+            mHubInfoRegistry.unregisterEndpointDiscoveryCallback(mCallback);
+        }
+
+        private void attachDeathRecipient() throws RemoteException {
+            mCallback.asBinder().linkToDeath(this, 0 /* flags */);
+        }
     }
 
     /* The list of discovery callbacks registered with the service */
@@ -99,7 +141,11 @@ class HubInfoRegistry implements ContextHubHalEndpointCallback.IEndpointLifecycl
 
     private final Object mCallbackLock = new Object();
 
-    HubInfoRegistry(IContextHubWrapper contextHubWrapper) throws InstantiationException {
+    /** Wakelock held while endpoint callbacks are being invoked */
+    private final WakeLock mWakeLock;
+
+    HubInfoRegistry(Context context, IContextHubWrapper contextHubWrapper)
+            throws InstantiationException {
         mContextHubWrapper = contextHubWrapper;
         try {
             refreshCachedHubs();
@@ -109,6 +155,16 @@ class HubInfoRegistry implements ContextHubHalEndpointCallback.IEndpointLifecycl
             Log.e(TAG, error, e);
             throw new InstantiationException(error);
         }
+
+        PowerManager powerManager = context.getSystemService(PowerManager.class);
+        if (powerManager == null) {
+            String error = "PowerManager was null";
+            Log.e(TAG, error);
+            throw new InstantiationError(error);
+        }
+        mWakeLock = powerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, TAG);
+        mWakeLock.setWorkSource(new WorkSource(Process.myUid(), context.getPackageName()));
+        mWakeLock.setReferenceCounted(true);
     }
 
     /** Retrieve the list of hubs available. */
@@ -178,12 +234,7 @@ class HubInfoRegistry implements ContextHubHalEndpointCallback.IEndpointLifecycl
                     try {
                         cb.onEndpointsStarted(infoList);
                     } catch (RemoteException e) {
-                        if (e instanceof DeadObjectException) {
-                            Log.w(TAG, "onEndpointStarted: callback died, unregistering");
-                            unregisterEndpointDiscoveryCallback(cb);
-                        } else {
-                            Log.e(TAG, "Exception while calling onEndpointsStarted", e);
-                        }
+                        Log.e(TAG, "Exception while calling onEndpointsStarted", e);
                     }
                 });
     }
@@ -208,12 +259,7 @@ class HubInfoRegistry implements ContextHubHalEndpointCallback.IEndpointLifecycl
                         cb.onEndpointsStopped(
                                 infoList, ContextHubServiceUtil.toAppHubEndpointReason(reason));
                     } catch (RemoteException e) {
-                        if (e instanceof DeadObjectException) {
-                            Log.w(TAG, "onEndpointStopped: callback died, unregistering");
-                            unregisterEndpointDiscoveryCallback(cb);
-                        } else {
-                            Log.e(TAG, "Exception while calling onEndpointsStopped", e);
-                        }
+                        Log.e(TAG, "Exception while calling onEndpointsStopped", e);
                     }
                 });
     }
@@ -254,7 +300,11 @@ class HubInfoRegistry implements ContextHubHalEndpointCallback.IEndpointLifecycl
         Objects.requireNonNull(callback, "callback cannot be null");
         synchronized (mCallbackLock) {
             checkCallbackAlreadyRegistered(callback);
-            mEndpointDiscoveryCallbacks.add(new DiscoveryCallback(callback, endpointId));
+            try {
+                mEndpointDiscoveryCallbacks.add(new DiscoveryCallback(this, callback, endpointId));
+            } catch (RemoteException e) {
+                Log.e(TAG, "RemoteException while adding discovery callback", e);
+            }
         }
     }
 
@@ -264,7 +314,12 @@ class HubInfoRegistry implements ContextHubHalEndpointCallback.IEndpointLifecycl
         Objects.requireNonNull(callback, "callback cannot be null");
         synchronized (mCallbackLock) {
             checkCallbackAlreadyRegistered(callback);
-            mEndpointDiscoveryCallbacks.add(new DiscoveryCallback(callback, serviceDescriptor));
+            try {
+                mEndpointDiscoveryCallbacks.add(
+                        new DiscoveryCallback(this, callback, serviceDescriptor));
+            } catch (RemoteException e) {
+                Log.e(TAG, "RemoteException while adding discovery callback", e);
+            }
         }
     }
 
@@ -282,6 +337,11 @@ class HubInfoRegistry implements ContextHubHalEndpointCallback.IEndpointLifecycl
         }
     }
 
+    /* package */
+    void onDiscoveryCallbackFinished() {
+        releaseWakeLock();
+    }
+
     private void checkCallbackAlreadyRegistered(
             IContextHubEndpointDiscoveryCallback callback) {
         synchronized (mCallbackLock) {
@@ -315,6 +375,7 @@ class HubInfoRegistry implements ContextHubHalEndpointCallback.IEndpointLifecycl
                     }
                 }
 
+                acquireWakeLock();
                 consumer.accept(
                         discoveryCallback.getCallback(),
                         infoList.toArray(new HubEndpointInfo[infoList.size()]));
@@ -322,6 +383,26 @@ class HubInfoRegistry implements ContextHubHalEndpointCallback.IEndpointLifecycl
         }
     }
 
+    private void acquireWakeLock() {
+        Binder.withCleanCallingIdentity(
+                () -> {
+                    mWakeLock.acquire(WAKELOCK_TIMEOUT_MILLIS);
+                });
+    }
+
+    private void releaseWakeLock() {
+        Binder.withCleanCallingIdentity(
+                () -> {
+                    if (mWakeLock.isHeld()) {
+                        try {
+                            mWakeLock.release();
+                        } catch (RuntimeException e) {
+                            Log.e(TAG, "Releasing the wakelock fails - ", e);
+                        }
+                    }
+                });
+    }
+
     void dump(IndentingPrintWriter ipw) {
         synchronized (mLock) {
             dumpLocked(ipw);