Update gauge collection based on app update state

Author: tejasd

firebase-perf/src/main/java/com/google/firebase/perf/session/SessionManager.java

@@ -165,7 +165,7 @@ private void startOrStopCollectingGauges(ApplicationProcessState appState) {
         perfSession, "Session is not ready while trying to startOrStopCollectingGauges");
 
     if (perfSession.isGaugeAndEventCollectionEnabled()) {
-      gaugeManager.startCollectingGauges(perfSession, appState);
+      gaugeManager.startCollectingGauges(perfSession);
     } else {
       gaugeManager.stopCollectingGauges();
     }

firebase-perf/src/main/java/com/google/firebase/perf/session/gauges/GaugeManager.java

@@ -20,6 +20,7 @@
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import com.google.firebase.components.Lazy;
+import com.google.firebase.perf.application.AppStateUpdateHandler;
 import com.google.firebase.perf.config.ConfigResolver;
 import com.google.firebase.perf.logging.AndroidLogger;
 import com.google.firebase.perf.session.PerfSession;
@@ -31,7 +32,6 @@
 import com.google.firebase.perf.v1.GaugeMetadata;
 import com.google.firebase.perf.v1.GaugeMetric;
 import java.util.concurrent.Executors;
-import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
@@ -41,7 +41,7 @@
  * information and logging it to the Transport.
  */
 @Keep // Needed because of b/117526359.
-public class GaugeManager {
+public class GaugeManager extends AppStateUpdateHandler {
 
   private static final AndroidLogger logger = AndroidLogger.getInstance();
   private static final GaugeManager instance = new GaugeManager();
@@ -59,8 +59,8 @@ public class GaugeManager {
   private final TransportManager transportManager;
 
   @Nullable private GaugeMetadataManager gaugeMetadataManager;
-  @Nullable private ScheduledFuture gaugeManagerDataCollectionJob = null;
-  @Nullable private String sessionId = null;
+  @Nullable private ScheduledFuture<?> gaugeManagerDataCollectionJob = null;
+  @Nullable private PerfSession session = null;
   private ApplicationProcessState applicationProcessState =
       ApplicationProcessState.APPLICATION_PROCESS_STATE_UNKNOWN;
 
@@ -91,13 +91,27 @@ private GaugeManager() {
     this.gaugeMetadataManager = gaugeMetadataManager;
     this.cpuGaugeCollector = cpuGaugeCollector;
     this.memoryGaugeCollector = memoryGaugeCollector;
+    registerForAppState();
   }
 
   /** Initializes GaugeMetadataManager which requires application context. */
   public void initializeGaugeMetadataManager(Context appContext) {
     this.gaugeMetadataManager = new GaugeMetadataManager(appContext);
   }
 
+  @Override
+  public void onUpdateAppState(ApplicationProcessState applicationProcessState) {
+    this.applicationProcessState = applicationProcessState;
+
+    if (session == null) {
+      // If the session is null, it means there's no gauges being collected.
+      return;
+    }
+
+    stopCollectingGauges();
+    startCollectingGauges(this.applicationProcessState, session.getTimer());
+  }
+
   /** Returns the singleton instance of this class. */
   public static synchronized GaugeManager getInstance() {
     return instance;
@@ -112,47 +126,22 @@ public static synchronized GaugeManager getInstance() {
    * will then be associated with the same or new sessionId and applicationProcessState.
    *
    * @param session The {@link PerfSession} to which the collected gauges will be associated with.
-   * @param applicationProcessState The {@link ApplicationProcessState} the collected GaugeMetrics
-   *     will be associated with.
    * @note: This method is NOT thread safe - {@link this.startCollectingGauges()} and {@link
-   *     this.stopCollectingGauges()} should always be called from the same thread.
+   * this.stopCollectingGauges()} should always be called from the same thread.
    */
-  public void startCollectingGauges(
-      PerfSession session, ApplicationProcessState applicationProcessState) {
-    if (this.sessionId != null) {
+  public void startCollectingGauges(PerfSession session) {
+    if (this.session != null) {
       stopCollectingGauges();
     }
 
-    long collectionFrequency = startCollectingGauges(applicationProcessState, session.getTimer());
+    long collectionFrequency =
+        startCollectingGauges(this.applicationProcessState, session.getTimer());
     if (collectionFrequency == INVALID_GAUGE_COLLECTION_FREQUENCY) {
       logger.warn("Invalid gauge collection frequency. Unable to start collecting Gauges.");
       return;
     }
 
-    this.sessionId = session.sessionId();
-    this.applicationProcessState = applicationProcessState;
-
-    // This is needed, otherwise the Runnable might use a stale value.
-    final String sessionIdForScheduledTask = sessionId;
-    final ApplicationProcessState applicationProcessStateForScheduledTask = applicationProcessState;
-
-    // TODO(b/394127311): Switch to using AQS.
-    try {
-      gaugeManagerDataCollectionJob =
-          gaugeManagerExecutor
-              .get()
-              .scheduleWithFixedDelay(
-                  () -> {
-                    syncFlush(sessionIdForScheduledTask, applicationProcessStateForScheduledTask);
-                  },
-                  /* initialDelay= */ collectionFrequency
-                      * APPROX_NUMBER_OF_DATA_POINTS_PER_GAUGE_METRIC,
-                  /* period= */ collectionFrequency * APPROX_NUMBER_OF_DATA_POINTS_PER_GAUGE_METRIC,
-                  TimeUnit.MILLISECONDS);
-
-    } catch (RejectedExecutionException e) {
-      logger.warn("Unable to start collecting Gauges: " + e.getMessage());
-    }
+    this.session = session;
   }
 
   /**
@@ -190,12 +179,12 @@ private long startCollectingGauges(ApplicationProcessState appState, Timer refer
    *     this.stopCollectingGauges()} should always be called from the same thread.
    */
   public void stopCollectingGauges() {
-    if (this.sessionId == null) {
+    if (session == null) {
       return;
     }
 
     // This is needed, otherwise the Runnable might use a stale value.
-    final String sessionIdForScheduledTask = sessionId;
+    final String sessionIdForScheduledTask = session.sessionId();
     final ApplicationProcessState applicationProcessStateForScheduledTask = applicationProcessState;
 
     cpuGaugeCollector.get().stopCollecting();
@@ -205,10 +194,9 @@ public void stopCollectingGauges() {
       gaugeManagerDataCollectionJob.cancel(false);
     }
 
-    // TODO(b/394127311): Switch to using AQS.
     // Flush any data that was collected for this session one last time.
     @SuppressWarnings("FutureReturnValueIgnored")
-    ScheduledFuture unusedFuture =
+    ScheduledFuture<?> unusedFuture =
         gaugeManagerExecutor
             .get()
             .schedule(
@@ -218,7 +206,7 @@ public void stopCollectingGauges() {
                 TIME_TO_WAIT_BEFORE_FLUSHING_GAUGES_QUEUE_MS,
                 TimeUnit.MILLISECONDS);
 
-    this.sessionId = null;
+    this.session = null;
     this.applicationProcessState = ApplicationProcessState.APPLICATION_PROCESS_STATE_UNKNOWN;
   }
 
@@ -227,7 +215,7 @@ public void stopCollectingGauges() {
    * proto and logs it to transport.
    *
    * @param sessionId The sessionId to which the collected GaugeMetrics should be associated with.
-   * @param appState The app state for which these gauges are collected.
+   * @param appState The app state for which these gauges are attributed to.
    */
   private void syncFlush(String sessionId, ApplicationProcessState appState) {
     GaugeMetric.Builder gaugeMetricBuilder = GaugeMetric.newBuilder();
@@ -244,7 +232,6 @@ private void syncFlush(String sessionId, ApplicationProcessState appState) {
     }
 
     // Adding Session ID info.
-    // TODO(b/394127311): Switch to using AQS.
     gaugeMetricBuilder.setSessionId(sessionId);
 
     transportManager.log(gaugeMetricBuilder.build(), appState);
@@ -253,16 +240,16 @@ private void syncFlush(String sessionId, ApplicationProcessState appState) {
   /**
    * Log the Gauge Metadata information to the transport.
    *
-   * @param aqsSessionId The {@link PerfSession#aqsSessionId()} ()} to which the collected Gauge Metrics
+   * @param sessionId The {@link PerfSession#sessionId()} ()} to which the collected Gauge Metrics
    *     should be associated with.
    * @param appState The {@link ApplicationProcessState} for which these gauges are collected.
    * @return true if GaugeMetadata was logged, false otherwise.
    */
-  public boolean logGaugeMetadata(String aqsSessionId, ApplicationProcessState appState) {
+  public boolean logGaugeMetadata(String sessionId, ApplicationProcessState appState) {
     if (gaugeMetadataManager != null) {
       GaugeMetric gaugeMetric =
           GaugeMetric.newBuilder()
-              .setSessionId(aqsSessionId)
+              .setSessionId(sessionId)
               .setGaugeMetadata(getGaugeMetadata())
               .build();
       transportManager.log(gaugeMetric, appState);

firebase-perf/src/test/java/com/google/firebase/perf/session/SessionManagerTest.java

@@ -144,8 +144,7 @@ public void testUpdatePerfSessionStartsCollectingGaugesIfSessionIsVerbose() {
     testSessionManager.setApplicationContext(mockApplicationContext);
 
     verify(mockGaugeManager, times(1)).initializeGaugeMetadataManager(mockApplicationContext);
-    verify(mockGaugeManager, times(1))
-        .startCollectingGauges(newSession, ApplicationProcessState.FOREGROUND);
+    verify(mockGaugeManager, times(1)).startCollectingGauges(newSession);
   }
 
   @Test