Finish message broadcast after sending Intent to bound Service.

firebase-messaging/src/main/java/com/google/firebase/messaging/FcmBroadcastProcessor.java

@@ -118,11 +118,14 @@ private static Task<Integer> bindToMessagingService(Context context, Intent inte
     if (Log.isLoggable(TAG, Log.DEBUG)) {
       Log.d(TAG, "Binding to service");
     }
+    if (ServiceStarter.getInstance().hasWakeLockPermission(context)) {
+      WakeLockHolder.sendWakefulServiceIntent(
+          context, getServiceConnection(context, ServiceStarter.ACTION_MESSAGING_EVENT), intent);
+    } else {
+      getServiceConnection(context, ServiceStarter.ACTION_MESSAGING_EVENT).sendIntent(intent);
+    }
 
-    return getServiceConnection(context, ServiceStarter.ACTION_MESSAGING_EVENT)
-        .sendIntent(intent)
-        // ok to use direct executor because we're just immediately returning an int
-        .continueWith(Runnable::run, t -> ServiceStarter.SUCCESS);
+    return Tasks.forResult(ServiceStarter.SUCCESS);
   }
 
   /** Connect to a service via bind. This is used to process intents in Android O+ */

firebase-messaging/src/main/java/com/google/firebase/messaging/WakeLockHolder.java

@@ -38,7 +38,7 @@ final class WakeLockHolder {
   private static final String EXTRA_WAKEFUL_INTENT =
       "com.google.firebase.iid.WakeLockHolder.wakefulintent";
   /** Release wakelocks after 60s, because we don't expect operations to take longer than that. */
-  private static final long WAKE_LOCK_ACQUIRE_TIMEOUT_MILLIS = TimeUnit.MINUTES.toMillis(1);
+  static final long WAKE_LOCK_ACQUIRE_TIMEOUT_MILLIS = TimeUnit.MINUTES.toMillis(1);
   // Object to sync threads
   private static final Object syncObject = new Object();
 
@@ -86,6 +86,34 @@ static ComponentName startWakefulService(@NonNull Context context, @NonNull Inte
     }
   }
 
+  /**
+   * Sends an Intent to a Service, binding to it, if necessary. Acquires a WakeLock based on the
+   * Intent and holds until the Service has finished processing the Intent or after a certain amount
+   * of time.
+   *
+   * @param context Application context.
+   * @param connection ServiceConnection to send the Intent to.
+   * @param intent Intent for starting the service.
+   */
+  static void sendWakefulServiceIntent(
+      Context context, WithinAppServiceConnection connection, Intent intent) {
+    synchronized (syncObject) {
+      checkAndInitWakeLock(context);
+
+      boolean isWakeLockAlreadyAcquired = isWakefulIntent(intent);
+
+      setAsWakefulIntent(intent, true);
+
+      connection
+          .sendIntent(intent)
+          .addOnCompleteListener(Runnable::run, t -> completeWakefulIntent(intent));
+
+      if (!isWakeLockAlreadyAcquired) {
+        wakeLock.acquire(WAKE_LOCK_ACQUIRE_TIMEOUT_MILLIS);
+      }
+    }
+  }
+
   private static void setAsWakefulIntent(@NonNull Intent intent, boolean isWakeful) {
     intent.putExtra(EXTRA_WAKEFUL_INTENT, isWakeful);
   }

firebase-messaging/src/main/java/com/google/firebase/messaging/WithinAppServiceConnection.java

@@ -14,6 +14,7 @@
 package com.google.firebase.messaging;
 
 import static com.google.firebase.messaging.FirebaseMessaging.TAG;
+import static com.google.firebase.messaging.WakeLockHolder.WAKE_LOCK_ACQUIRE_TIMEOUT_MILLIS;
 
 import android.content.ComponentName;
 import android.content.Context;
@@ -41,8 +42,8 @@
  */
 class WithinAppServiceConnection implements ServiceConnection {
 
-  // Time out requests by finishing the pending result and thus unblocking the receiver after
-  // 9s to avoid causing the receiver to time out after 10s.
+  // Time out WakeLock after 9s to match previous behavior of forcing the broadcast to finish after
+  // that much time.
   private static final int REQUEST_TIMEOUT_MS = 9000;
 
   static class BindRequest {
@@ -54,17 +55,21 @@ static class BindRequest {
     }
 
     void arrangeTimeout(ScheduledExecutorService executor) {
+      // Allow the same maximum WakeLock duration for high priority messages as when using
+      // startService(), otherwise allow up to 9 seconds to match the time that the service was
+      // allowed to run before finishing the BroadcastReceiver.
+      boolean isHighPriority = (intent.getFlags() & Intent.FLAG_RECEIVER_FOREGROUND) != 0;
       ScheduledFuture<?> timeoutFuture =
           executor.schedule(
               () -> {
                 Log.w(
                     TAG,
                     "Service took too long to process intent: "
                         + intent.getAction()
-                        + " App may get closed.");
+                        + " Releasing WakeLock.");
                 finish();
               },
-              REQUEST_TIMEOUT_MS,
+              isHighPriority ? WAKE_LOCK_ACQUIRE_TIMEOUT_MILLIS : REQUEST_TIMEOUT_MS,
               TimeUnit.MILLISECONDS);
 
       getTask()

firebase-messaging/src/test/java/com/google/firebase/messaging/FcmBroadcastProcessorRoboTest.java

@@ -0,0 +1,106 @@
+// Copyright 2022 Google LLC
+//
+// 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.google.firebase.messaging;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.robolectric.Shadows.shadowOf;
+
+import android.app.Application;
+import android.content.Context;
+import android.content.Intent;
+import android.os.Build.VERSION_CODES;
+import androidx.test.core.app.ApplicationProvider;
+import com.google.android.gms.tasks.Task;
+import com.google.firebase.messaging.testing.FakeScheduledExecutorService;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.annotation.Config;
+import org.robolectric.shadows.ShadowPowerManager;
+
+/** Robolectric test for FcmBroadcastProcessor. */
+@RunWith(RobolectricTestRunner.class)
+public class FcmBroadcastProcessorRoboTest {
+
+  private static final String ACTION_FCM_MESSAGE = "com.google.android.c2dm.intent.RECEIVE";
+
+  @Rule public final MockitoRule mocks = MockitoJUnit.rule();
+
+  private Application context;
+  private FcmBroadcastProcessor processor;
+  private FakeScheduledExecutorService fakeExecutorService;
+  @Mock private ServiceStarter serviceStarter;
+
+  @Before
+  public void setUp() {
+    context = ApplicationProvider.getApplicationContext();
+    ServiceStarter.setForTesting(serviceStarter);
+    fakeExecutorService = new FakeScheduledExecutorService();
+    processor = new FcmBroadcastProcessor(context, fakeExecutorService);
+  }
+
+  @After
+  public void resetStaticState() {
+    ServiceStarter.setForTesting(null);
+    FcmBroadcastProcessor.reset();
+    WakeLockHolder.reset();
+    ShadowPowerManager.clearWakeLocks();
+  }
+
+  @Test
+  @Config(sdk = VERSION_CODES.O)
+  public void testStartMessagingService_NormalPriorityBackgroundCheck() {
+    // Subject to background check when run on Android O and targetSdkVersion set to O.
+    context.getApplicationInfo().targetSdkVersion = VERSION_CODES.O;
+    when(serviceStarter.hasWakeLockPermission(any(Context.class))).thenReturn(true);
+
+    Task<Integer> startServiceTask =
+        processor.startMessagingService(context, new Intent(ACTION_FCM_MESSAGE));
+
+    // Should return immediately with SUCCESS, bind to the Service, and acquire a WakeLock.
+    assertThat(startServiceTask.getResult()).isEqualTo(ServiceStarter.SUCCESS);
+    verify(serviceStarter, never()).startMessagingService(any(), any());
+    assertThat(shadowOf(context).getBoundServiceConnections()).hasSize(1);
+    assertThat(ShadowPowerManager.getLatestWakeLock()).isNotNull();
+    assertThat(ShadowPowerManager.getLatestWakeLock().isHeld()).isTrue();
+  }
+
+  @Test
+  @Config(sdk = VERSION_CODES.O)
+  public void testStartMessagingService_bindNoWakeLockPermission() {
+    // Subject to background check when run on Android O and targetSdkVersion set to O.
+    context.getApplicationInfo().targetSdkVersion = VERSION_CODES.O;
+    when(serviceStarter.hasWakeLockPermission(any(Context.class))).thenReturn(false);
+
+    Task<Integer> startServiceTask =
+        processor.startMessagingService(context, new Intent(ACTION_FCM_MESSAGE));
+
+    // Should return immediately with SUCCESS and bind to the Service, but not try to acquire a
+    // WakeLock since it doesn't hold the permission.
+    assertThat(startServiceTask.getResult()).isEqualTo(ServiceStarter.SUCCESS);
+    verify(serviceStarter, never()).startMessagingService(any(), any());
+    assertThat(shadowOf(context).getBoundServiceConnections()).hasSize(1);
+    assertThat(ShadowPowerManager.getLatestWakeLock()).isNull();
+  }
+}

firebase-messaging/src/test/java/com/google/firebase/messaging/WakeLockHolderRoboTest.java

@@ -14,11 +14,15 @@
 package com.google.firebase.messaging;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import android.app.Application;
 import android.content.Intent;
 import android.os.PowerManager.WakeLock;
 import androidx.test.core.app.ApplicationProvider;
+import com.google.android.gms.tasks.TaskCompletionSource;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -51,6 +55,31 @@ public void testStartWakefulService_AcquiresWakeLock() throws Exception {
     assertThat(ShadowPowerManager.getLatestWakeLock().isHeld()).isTrue();
   }
 
+  @Test
+  public void testSendWakefulServiceIntent_AcquiresWakeLock() {
+    TaskCompletionSource<Void> taskCompletionSource = new TaskCompletionSource<>();
+    WithinAppServiceConnection mockConnection = mock(WithinAppServiceConnection.class);
+    when(mockConnection.sendIntent(any(Intent.class))).thenReturn(taskCompletionSource.getTask());
+    WakeLockHolder.sendWakefulServiceIntent(context, mockConnection, new Intent());
+
+    assertThat(ShadowPowerManager.getLatestWakeLock()).isNotNull();
+    assertThat(ShadowPowerManager.getLatestWakeLock().isHeld()).isTrue();
+  }
+
+  @Test
+  public void testSendWakefulServiceIntent_ReleasesWakeLock() {
+    TaskCompletionSource<Void> taskCompletionSource = new TaskCompletionSource<>();
+    WithinAppServiceConnection mockConnection = mock(WithinAppServiceConnection.class);
+    when(mockConnection.sendIntent(any(Intent.class))).thenReturn(taskCompletionSource.getTask());
+    WakeLockHolder.sendWakefulServiceIntent(context, mockConnection, new Intent());
+
+    // Verify that the WakeLock is released once the Intent has been handled by the Service.
+    WakeLock wakeLock = ShadowPowerManager.getLatestWakeLock();
+    taskCompletionSource.setResult(null);
+
+    assertThat(wakeLock.isHeld()).isFalse();
+  }
+
   @Test
   public void testCompleteWakefulIntent_ReleasesWakeLockIfPresent() throws Exception {
     Intent intent = newIntent("Any_Action");

firebase-messaging/src/test/java/com/google/firebase/messaging/WithinAppServiceConnectionRoboTest.java

@@ -46,6 +46,9 @@ public class WithinAppServiceConnectionRoboTest {
   // The amount of time that a broadcast receiver takes to time out
   private static final long RECEIVER_TIMEOUT_S = 10;
 
+  // The amount of time to allow a foreground broadcast's service to run.
+  private static final long FOREGROUND_RECEIVER_TIMEOUT_S = 60;
+
   private Application context;
   private FakeScheduledExecutorService fakeExecutor;
 
@@ -66,14 +69,34 @@ public void testReceiverTimesOut() {
     Task<Void> pendingResult = connection.sendIntent(new Intent());
     assertThat(pendingResult.isComplete()).isFalse();
 
-    // Check the runnable doesn't run early, and that after it shuld have run the pending
+    // Check the runnable doesn't run early, and that after it should have run the pending
     // result is finished.
     fakeExecutor.simulateNormalOperationFor(RECEIVER_TIMEOUT_S - 2, TimeUnit.SECONDS);
     assertThat(pendingResult.isComplete()).isFalse();
     fakeExecutor.simulateNormalOperationFor(1, TimeUnit.SECONDS);
     assertThat(pendingResult.isComplete()).isTrue();
   }
 
+  @Test
+  public void testReceiverTimesOut_ForegroundBroadcast() {
+    WithinAppServiceConnection connection =
+        new WithinAppServiceConnection(context, TEST_BIND_ACTION, fakeExecutor);
+    setMockBinder(TEST_BIND_ACTION);
+
+    // Send a foreground broadcst intent, verify the pending result isn't finished
+    Intent foregroundBroadcastIntent = new Intent();
+    foregroundBroadcastIntent.addFlags(Intent.FLAG_RECEIVER_FOREGROUND);
+    Task<Void> pendingResult = connection.sendIntent(foregroundBroadcastIntent);
+    assertThat(pendingResult.isComplete()).isFalse();
+
+    // Check the runnable doesn't run early, and that after it should have run the pending
+    // result is finished.
+    fakeExecutor.simulateNormalOperationFor(FOREGROUND_RECEIVER_TIMEOUT_S - 1, TimeUnit.SECONDS);
+    assertThat(pendingResult.isComplete()).isFalse();
+    fakeExecutor.simulateNormalOperationFor(1, TimeUnit.SECONDS);
+    assertThat(pendingResult.isComplete()).isTrue();
+  }
+
   @Test
   public void testTimeoutGetsCancelled() {
     WithinAppServiceConnection connection =