Fix wakelock not stopping early when notifications are disabled (#6424)

If notifications for a device are disabled when there is no connection with the HS, the push registration will still exist, so the device can still receive push notifications.

In that cases, we were running into an issue where the wakelock for push notifications was started immediately after receiving a push but was never stopped and it ran for 3 minutes until its timeout, keeping the device awake for no reason.

This patch changes `DefaultPushHandler` so if we don't need the wakelock it returns `false` and we can stop the wakelock early.

libraries/pushproviders/api/src/main/kotlin/io/element/android/libraries/pushproviders/api/PushHandler.kt

@@ -9,11 +9,21 @@
 package io.element.android.libraries.pushproviders.api
 
 interface PushHandler {
+    /**
+     * Handle a push received from the provider.
+     *
+     * @param pushData the data of the push, containing the client secret and the push content.
+     * @param providerInfo an identifier of the provider that sent the push, for logging and debugging purposes.
+     * @return `true` if the push was handled successfully and is now enqueued for processing, false otherwise.
+     */
     suspend fun handle(
         pushData: PushData,
         providerInfo: String,
-    )
+    ): Boolean
 
+    /**
+     * Handle an invalid push received from the provider.
+     */
     suspend fun handleInvalid(
         providerInfo: String,
         data: String,

libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/VectorFirebaseMessagingService.kt

@@ -58,11 +58,17 @@ class VectorFirebaseMessagingService : FirebaseMessagingService() {
                         "$it: ${message.data[it]}"
                     },
                 )
+                pushHandlingWakeLock.unlock()
             } else {
-                pushHandler.handle(
+                val handled = pushHandler.handle(
                     pushData = pushData,
                     providerInfo = FirebaseConfig.NAME,
                 )
+
+                // If we failed to handle the push, we should release the wakelock early to avoid keeping the device awake for too long.
+                if (!handled) {
+                    pushHandlingWakeLock.unlock()
+                }
             }
         }
     }

libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/VectorFirebaseMessagingServiceTest.kt

@@ -29,6 +29,7 @@ import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.robolectric.RobolectricTestRunner
+import kotlin.time.Duration
 
 @RunWith(RobolectricTestRunner::class)
 class VectorFirebaseMessagingServiceTest {
@@ -56,7 +57,7 @@ class VectorFirebaseMessagingServiceTest {
 
     @Test
     fun `test receiving valid data`() = runTest {
-        val lambda = lambdaRecorder<PushData, String, Unit> { _, _ -> }
+        val lambda = lambdaRecorder<PushData, String, Boolean> { _, _ -> true }
         val vectorFirebaseMessagingService = createVectorFirebaseMessagingService(
             pushHandler = FakePushHandler(handleResult = lambda)
         )
@@ -78,6 +79,68 @@ class VectorFirebaseMessagingServiceTest {
             )
     }
 
+    @Test
+    fun `test pushHandler returning true locks and does not unlock the wakelock so it continues running`() = runTest {
+        val lockLambda = lambdaRecorder<Duration, Unit> { _ -> }
+        val unlockLambda = lambdaRecorder<Unit> { }
+        val vectorFirebaseMessagingService = createVectorFirebaseMessagingService(
+            pushHandler = FakePushHandler(handleResult = { _, _ -> true }),
+            pushHandlingWakeLock = FakePushHandlingWakeLock(
+                lock = lockLambda,
+                unlock = unlockLambda
+            )
+        )
+        vectorFirebaseMessagingService.onMessageReceived(
+            message = RemoteMessage(
+                Bundle().apply {
+                    putString("event_id", AN_EVENT_ID.value)
+                    putString("room_id", A_ROOM_ID.value)
+                    putString("cs", A_SECRET)
+                },
+            )
+        )
+
+        // The wakelock should be locked but not unlocked
+        lockLambda.assertions().isCalledOnce()
+        unlockLambda.assertions().isNeverCalled()
+
+        advanceUntilIdle()
+
+        // After handling the push, the wakelock should still not be unlocked
+        unlockLambda.assertions().isNeverCalled()
+    }
+
+    @Test
+    fun `test pushHandler returning false locks and unlocks the wakelock early`() = runTest {
+        val lockLambda = lambdaRecorder<Duration, Unit> { _ -> }
+        val unlockLambda = lambdaRecorder<Unit> { }
+        val vectorFirebaseMessagingService = createVectorFirebaseMessagingService(
+            pushHandler = FakePushHandler(handleResult = { _, _ -> false }),
+            pushHandlingWakeLock = FakePushHandlingWakeLock(
+                lock = lockLambda,
+                unlock = unlockLambda
+            )
+        )
+        vectorFirebaseMessagingService.onMessageReceived(
+            message = RemoteMessage(
+                Bundle().apply {
+                    putString("event_id", AN_EVENT_ID.value)
+                    putString("room_id", A_ROOM_ID.value)
+                    putString("cs", A_SECRET)
+                },
+            )
+        )
+
+        // The wakelock should be locked but not unlocked
+        lockLambda.assertions().isCalledOnce()
+        unlockLambda.assertions().isNeverCalled()
+
+        advanceUntilIdle()
+
+        // After handling the push, the wakelock should be unlocked
+        unlockLambda.assertions().isCalledOnce()
+    }
+
     @Test
     fun `test new token is forwarded to the handler`() = runTest {
         val lambda = lambdaRecorder<String, Unit> { }

libraries/pushproviders/unifiedpush/src/main/kotlin/io/element/android/libraries/pushproviders/unifiedpush/VectorUnifiedPushMessagingReceiver.kt

@@ -71,11 +71,17 @@ class VectorUnifiedPushMessagingReceiver : MessagingReceiver() {
                     providerInfo = "${UnifiedPushConfig.NAME} - $instance",
                     data = String(message.content),
                 )
+                pushHandlingWakeLock.unlock()
             } else {
-                pushHandler.handle(
+                val handled = pushHandler.handle(
                     pushData = pushData,
                     providerInfo = "${UnifiedPushConfig.NAME} - $instance",
                 )
+
+                // If we failed to handle the push, we should release the wakelock early to avoid keeping the device awake for too long.
+                if (!handled) {
+                    pushHandlingWakeLock.unlock()
+                }
             }
         }
     }

libraries/pushproviders/unifiedpush/src/test/kotlin/io/element/android/libraries/pushproviders/unifiedpush/VectorUnifiedPushMessagingReceiverTest.kt

@@ -39,6 +39,7 @@ import org.unifiedpush.android.connector.FailedReason
 import org.unifiedpush.android.connector.data.PublicKeySet
 import org.unifiedpush.android.connector.data.PushEndpoint
 import org.unifiedpush.android.connector.data.PushMessage
+import kotlin.time.Duration
 
 @RunWith(RobolectricTestRunner::class)
 class VectorUnifiedPushMessagingReceiverTest {
@@ -76,7 +77,7 @@ class VectorUnifiedPushMessagingReceiverTest {
     @Test
     fun `onMessage valid invokes the push handler`() = runTest {
         val context = InstrumentationRegistry.getInstrumentation().context
-        val pushHandlerResult = lambdaRecorder<PushData, String, Unit> { _, _ -> }
+        val pushHandlerResult = lambdaRecorder<PushData, String, Boolean> { _, _ -> true }
         val vectorUnifiedPushMessagingReceiver = createVectorUnifiedPushMessagingReceiver(
             pushHandler = FakePushHandler(
                 handleResult = pushHandlerResult
@@ -101,6 +102,60 @@ class VectorUnifiedPushMessagingReceiverTest {
             )
     }
 
+    @Test
+    fun `pushHandler returning true locks the wake lock but does not unlock it so it continues to run`() = runTest {
+        val context = InstrumentationRegistry.getInstrumentation().context
+        val pushHandlerResult = lambdaRecorder<PushData, String, Boolean> { _, _ -> true }
+        val lockLambda = lambdaRecorder<Duration, Unit> { _ -> }
+        val unlockLambda = lambdaRecorder<Unit> { }
+        val vectorUnifiedPushMessagingReceiver = createVectorUnifiedPushMessagingReceiver(
+            pushHandler = FakePushHandler(
+                handleResult = pushHandlerResult
+            ),
+            pushHandlingWakeLock = FakePushHandlingWakeLock(
+                lock = lockLambda,
+                unlock = unlockLambda,
+            ),
+        )
+        vectorUnifiedPushMessagingReceiver.onMessage(context, aPushMessage(), A_SECRET)
+
+        // The wakelock should be locked but not unlocked, so it should continue to run
+        lockLambda.assertions().isCalledOnce()
+        unlockLambda.assertions().isNeverCalled()
+
+        advanceUntilIdle()
+
+        // After waiting for any possible timeout, the lock is still locked
+        unlockLambda.assertions().isNeverCalled()
+    }
+
+    @Test
+    fun `pushHandler returning false locks and unlocks the wakelock early`() = runTest {
+        val context = InstrumentationRegistry.getInstrumentation().context
+        val pushHandlerResult = lambdaRecorder<PushData, String, Boolean> { _, _ -> false }
+        val lockLambda = lambdaRecorder<Duration, Unit> { _ -> }
+        val unlockLambda = lambdaRecorder<Unit> { }
+        val vectorUnifiedPushMessagingReceiver = createVectorUnifiedPushMessagingReceiver(
+            pushHandler = FakePushHandler(
+                handleResult = pushHandlerResult
+            ),
+            pushHandlingWakeLock = FakePushHandlingWakeLock(
+                lock = lockLambda,
+                unlock = unlockLambda,
+            ),
+        )
+        vectorUnifiedPushMessagingReceiver.onMessage(context, aPushMessage(), A_SECRET)
+
+        // The wakelock should be locked but not unlocked, so it should continue to run
+        lockLambda.assertions().isCalledOnce()
+        unlockLambda.assertions().isNeverCalled()
+
+        advanceUntilIdle()
+
+        // After waiting for a bit, the lock should have been unlocked since the handler returned false
+        unlockLambda.assertions().isCalledOnce()
+    }
+
     @Test
     fun `onMessage invalid invokes the push handler invalid method`() = runTest {
         val context = InstrumentationRegistry.getInstrumentation().context