Add a foreground service with a wakelock for fetching push notifications (#6321)

* Create `PushHandlingWakeLock` to start a foreground service:

When receiving a push and scheduling the notification fetching, several problems can happen:

1. Some async operation is waiting for a timeout and it takes way longer than that to finish (i.e. timeout of 10s but it took 30s to advance).
2. The same, but when starting new coroutines. I've seen the time between scheduling a coroutine and it running sometimes take up to 1 minute.
3. Notification fetching can be scheduled immediately, but it can take a while to actually run because the OS understands the app is now in Doze.

Having a wakelock that runs as soon as the push handling starts fixes these: it continues the previous wakelock held by either Firebase or the UnifiedPush distributor.

* Acquire the wakelock as soon as we received the pushes in both receivers

* Also release the wakelock ahead of time if possible

libraries/push/impl/src/main/AndroidManifest.xml

@@ -11,6 +11,11 @@
     <uses-permission android:name="android.permission.REQUEST_IGNORE_BATTERY_OPTIMIZATIONS" />
 
     <application>
+        <service android:name=".push.FetchPushForegroundService"
+            android:foregroundServiceType="shortService"
+            android:exported="false"
+            android:enabled="true" />
+
         <receiver
             android:name=".notifications.TestNotificationReceiver"
             android:exported="false" />

libraries/push/impl/src/main/kotlin/io/element/android/libraries/push/impl/di/PushBindings.kt

@@ -0,0 +1,17 @@
+/*
+ * Copyright (c) 2026 Element Creations Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial.
+ * Please see LICENSE files in the repository root for full details.
+ */
+
+package io.element.android.libraries.push.impl.di
+
+import dev.zacsweers.metro.AppScope
+import dev.zacsweers.metro.ContributesTo
+import io.element.android.libraries.push.impl.push.FetchPushForegroundService
+
+@ContributesTo(AppScope::class)
+interface PushBindings {
+    fun inject(fetchPushForegroundService: FetchPushForegroundService)
+}

libraries/push/impl/src/main/kotlin/io/element/android/libraries/push/impl/notifications/channels/NotificationChannels.kt

@@ -58,6 +58,8 @@ interface NotificationChannels {
      * Get the channel for test notifications.
      */
     fun getChannelIdForTest(): String
+
+    fun getSilentChannelId(): String = SILENT_NOTIFICATION_CHANNEL_ID
 }
 
 @ChecksSdkIntAtLeast(api = Build.VERSION_CODES.O)

libraries/push/impl/src/main/kotlin/io/element/android/libraries/push/impl/push/DefaultPushHandler.kt

@@ -31,7 +31,10 @@ import io.element.android.libraries.workmanager.api.WorkManagerScheduler
 import io.element.android.services.analytics.api.AnalyticsLongRunningTransaction
 import io.element.android.services.analytics.api.AnalyticsService
 import io.element.android.services.toolbox.api.systemclock.SystemClock
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.launch
 import timber.log.Timber
 
 private val loggerTag = LoggerTag("PushHandler", LoggerTag.PushLoggerTag)
@@ -71,7 +74,12 @@ class DefaultPushHandler(
         if (buildMeta.lowPrivacyLoggingEnabled) {
             Timber.tag(loggerTag.value).d("## pushData: $pushData")
         }
-        incrementPushDataStore.incrementPushCounter()
+
+        // Update the push counter without blocking the coroutine execution, as it is not critical to be updated before handling the push
+        CoroutineScope(currentCoroutineContext()).launch {
+            incrementPushDataStore.incrementPushCounter()
+        }
+
         // Diagnostic Push
         if (pushData.eventId == DefaultTestPush.TEST_EVENT_ID) {
             pushHistoryService.onDiagnosticPush(providerInfo)
@@ -130,6 +138,7 @@ class DefaultPushHandler(
 
             Timber.d("Queueing notification: $pushRequest")
             pushHistoryService.insertOrUpdatePushRequest(pushRequest)
+            Timber.d("Queueing notification finished")
 
             if (!workManagerScheduler.hasPendingWork(userId, WorkManagerRequestType.NOTIFICATION_SYNC)) {
                 Timber.d("No pending worker for push notifications found")

libraries/push/impl/src/main/kotlin/io/element/android/libraries/push/impl/push/DefaultPushHandlingWakeLock.kt

@@ -0,0 +1,45 @@
+/*
+ * Copyright (c) 2026 Element Creations Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial.
+ * Please see LICENSE files in the repository root for full details.
+ */
+
+package io.element.android.libraries.push.impl.push
+
+import android.content.Context
+import dev.zacsweers.metro.AppScope
+import dev.zacsweers.metro.ContributesBinding
+import dev.zacsweers.metro.SingleIn
+import io.element.android.libraries.di.annotations.ApplicationContext
+import io.element.android.libraries.push.api.push.PushHandlingWakeLock
+import timber.log.Timber
+import java.util.concurrent.atomic.AtomicInteger
+import kotlin.time.Duration
+
+@ContributesBinding(AppScope::class)
+@SingleIn(AppScope::class)
+class DefaultPushHandlingWakeLock(
+    @ApplicationContext private val context: Context,
+) : PushHandlingWakeLock {
+    private val count = AtomicInteger(0)
+
+    override fun lock(time: Duration) {
+        Timber.d("Acquiring wakelock for push handling, starting service.")
+        FetchPushForegroundService.startIfNeeded(context)
+
+        count.incrementAndGet()
+    }
+
+    override fun unlock() {
+        Timber.d("Releasing wakelock used for push handling.")
+        FetchPushForegroundService.stop(context)
+        if (count.decrementAndGet() <= 0) {
+            Timber.d("No more wakelock needed for push handling, stopping service.")
+            count.set(0)
+        } else {
+            Timber.d("Wakelock still needed for push handling, restarting service | count: ${count.get()}.")
+            FetchPushForegroundService.startIfNeeded(context)
+        }
+    }
+}

libraries/push/impl/src/main/kotlin/io/element/android/libraries/push/impl/push/FetchPushForegroundService.kt

@@ -0,0 +1,115 @@
+/*
+ * Copyright (c) 2026 Element Creations Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial.
+ * Please see LICENSE files in the repository root for full details.
+ */
+
+package io.element.android.libraries.push.impl.push
+
+import android.app.Service
+import android.content.Context
+import android.content.Intent
+import android.os.Build
+import android.os.IBinder
+import android.os.PowerManager
+import androidx.core.app.NotificationCompat
+import dev.zacsweers.metro.Inject
+import io.element.android.libraries.architecture.bindings
+import io.element.android.libraries.designsystem.utils.CommonDrawables
+import io.element.android.libraries.di.annotations.AppCoroutineScope
+import io.element.android.libraries.push.api.push.PushHandlingWakeLock
+import io.element.android.libraries.push.impl.di.PushBindings
+import io.element.android.libraries.push.impl.notifications.channels.NotificationChannels
+import io.element.android.libraries.ui.strings.CommonStrings
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.launch
+import kotlin.time.Duration.Companion.minutes
+
+private const val NOTIFICATION_ID = 1001
+
+// This kind of foreground service can only last up to 3 minutes before onTimeout is called
+private val wakelockTimeout = 3.minutes.inWholeMilliseconds
+
+/**
+ * Foreground service used to ensure the device stays awake while we handle the pushes and schedule and run the work to fetch the notification content.
+ */
+class FetchPushForegroundService : Service() {
+    override fun onBind(intent: Intent?): IBinder? {
+        return null
+    }
+
+    @Inject lateinit var notificationChannels: NotificationChannels
+    @Inject lateinit var pushHandlingWakeLock: PushHandlingWakeLock
+    @Inject @AppCoroutineScope lateinit var coroutineScope: CoroutineScope
+
+    private val wakelock: PowerManager.WakeLock by lazy {
+        val powerManager = getSystemService(POWER_SERVICE) as PowerManager
+        powerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, "FetchPushService:WakeLock").apply {
+            setReferenceCounted(false)
+        }
+    }
+
+    override fun onStartCommand(intent: Intent?, flags: Int, startId: Int): Int {
+        bindings<PushBindings>().inject(this)
+
+        wakelock.acquire(wakelockTimeout)
+
+        val notificationCompat = NotificationCompat.Builder(this, notificationChannels.getSilentChannelId())
+            .setSmallIcon(CommonDrawables.ic_notification)
+            .setContentTitle(getString(CommonStrings.common_android_fetching_notifications_title))
+            .setProgress(0, 0, true)
+            .setVibrate(longArrayOf(0))
+            .setSound(null)
+            .build()
+        startForeground(NOTIFICATION_ID, notificationCompat)
+
+        // The timeout is not automatic before Android 15, so we need to schedule it ourselves
+        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.VANILLA_ICE_CREAM) {
+            coroutineScope.launch {
+                delay(wakelockTimeout)
+                onTimeout(startId)
+            }
+        }
+
+        return START_NOT_STICKY
+    }
+
+    override fun stopService(intent: Intent?): Boolean {
+        wakelock.release()
+
+        stopForeground(STOP_FOREGROUND_REMOVE)
+        return super.stopService(intent)
+    }
+
+    override fun onTimeout(startId: Int) {
+        super.onTimeout(startId)
+
+        pushHandlingWakeLock.unlock()
+    }
+
+    companion object {
+        fun startIfNeeded(context: Context) {
+            // Don't start the foreground service if the device is already awake
+            val powerManager = context.getSystemService(POWER_SERVICE) as PowerManager
+            if (powerManager.isInteractive) return
+
+            start(context)
+        }
+
+        fun start(context: Context) {
+            val intent = Intent(context, FetchPushForegroundService::class.java)
+            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.O) {
+                context.startForegroundService(intent)
+            } else {
+                context.startService(intent)
+            }
+        }
+
+        fun stop(context: Context) {
+            val intent = Intent(context, FetchPushForegroundService::class.java)
+            context.stopService(intent)
+        }
+    }
+}

libraries/push/impl/src/main/kotlin/io/element/android/libraries/push/impl/workmanager/FetchPendingNotificationsWorker.kt

@@ -25,6 +25,7 @@ import io.element.android.libraries.matrix.api.auth.SessionRestorationException
 import io.element.android.libraries.matrix.api.core.SessionId
 import io.element.android.libraries.matrix.api.exception.ClientException
 import io.element.android.libraries.matrix.api.exception.isNetworkError
+import io.element.android.libraries.push.api.push.PushHandlingWakeLock
 import io.element.android.libraries.push.impl.db.PushRequest
 import io.element.android.libraries.push.impl.history.PushHistoryService
 import io.element.android.libraries.push.impl.notifications.NotifiableEventResolver
@@ -57,6 +58,7 @@ class FetchPendingNotificationsWorker(
     private val resultProcessor: NotificationResultProcessor,
     private val analyticsService: AnalyticsService,
     private val systemClock: SystemClock,
+    private val pushHandlingWakeLock: PushHandlingWakeLock,
 ) : CoroutineWorker(context, params) {
     override suspend fun doWork(): Result {
         Timber.d("FetchNotificationsWorker started")
@@ -65,6 +67,8 @@ class FetchPendingNotificationsWorker(
             inputData.getString(SyncPendingNotificationsRequestBuilder.SESSION_ID)?.let(::SessionId)
         }.getOrNull() ?: return Result.failure()
 
+        pushHandlingWakeLock.unlock()
+
         // Fetch pending requests in the last 24 hours
         val fetchSince = Instant.fromEpochMilliseconds(systemClock.epochMillis()).minus(1.days)
         val requests = pushHistoryService.getPendingPushRequests(sessionId, fetchSince).getOrNull() ?: return Result.failure()
@@ -101,9 +105,9 @@ class FetchPendingNotificationsWorker(
 
                     results
                 },
-                onFailure = {
+                onFailure = { throwable ->
                     // This is a failure at the fetch notification setup, not a failure for a single fetch notification operation
-                    return handleSetupError(sessionId, requests, pendingAnalyticTransactions, it)
+                    return handleSetupError(sessionId, requests, pendingAnalyticTransactions, throwable)
                 }
             )
 

libraries/push/impl/src/test/kotlin/io/element/android/libraries/push/impl/push/DefaultPushHandlerTest.kt

@@ -42,6 +42,7 @@ import io.element.android.tests.testutils.lambda.lambdaRecorder
 import io.element.android.tests.testutils.lambda.value
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.advanceTimeBy
+import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import kotlin.time.Duration.Companion.milliseconds
@@ -173,6 +174,10 @@ class DefaultPushHandlerTest {
                 workManagerScheduler = workManagerScheduler,
             )
             defaultPushHandler.handle(aPushData, A_PUSHER_INFO)
+
+            // Give it enough time to increment the push counter
+            runCurrent()
+
             submitWorkLambda.assertions()
                 .isNeverCalled()
             incrementPushCounterResult.assertions()

libraries/push/impl/src/test/kotlin/io/element/android/libraries/push/impl/workmanager/DefaultSyncPendingNotificationsRequestBuilderTest.kt

@@ -24,7 +24,9 @@ import io.element.android.services.toolbox.test.sdk.FakeBuildVersionSdkIntProvid
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.robolectric.annotation.Config
 
+@Config(sdk = [33])
 @RunWith(AndroidJUnit4::class)
 class DefaultSyncPendingNotificationsRequestBuilderTest {
     @Test

libraries/push/impl/src/test/kotlin/io/element/android/libraries/push/impl/workmanager/FetchPendingNotificationWorkerTest.kt

@@ -28,6 +28,7 @@ import io.element.android.libraries.push.impl.notifications.FakeNotificationResu
 import io.element.android.libraries.push.impl.notifications.fixtures.aPushRequest
 import io.element.android.libraries.push.impl.notifications.model.ResolvedPushEvent
 import io.element.android.libraries.push.impl.push.SyncOnNotifiableEvent
+import io.element.android.libraries.push.test.push.FakePushHandlingWakeLock
 import io.element.android.libraries.workmanager.api.WorkManagerRequestBuilder
 import io.element.android.libraries.workmanager.api.di.MetroWorkerFactory
 import io.element.android.services.analytics.test.FakeAnalyticsService
@@ -238,6 +239,7 @@ class FetchPendingNotificationWorkerTest {
         pushHistoryService: FakePushHistoryService = FakePushHistoryService(),
         resultProcessor: FakeNotificationResultProcessor = FakeNotificationResultProcessor(),
         systemClock: FakeSystemClock = FakeSystemClock(),
+        pushHandlingWakeLock: FakePushHandlingWakeLock = FakePushHandlingWakeLock(),
     ) = FetchPendingNotificationsWorker(
         params = createWorkerParams(workDataOf("session_id" to input)),
         context = InstrumentationRegistry.getInstrumentation().context,
@@ -248,6 +250,7 @@ class FetchPendingNotificationWorkerTest {
         pushHistoryService = pushHistoryService,
         resultProcessor = resultProcessor,
         systemClock = systemClock,
+        pushHandlingWakeLock = pushHandlingWakeLock,
     )
 
     private fun TestScope.createWorkerParams(