Store up to 5 network log batches if needed.

In the normal mode when the DO fetches the logs ASAP, there will still be
no more than one last full batch in memory at once. If the DO is too slow,
or the broadcast queue is too crowded we will store up to 5 of them,
discarding older ones when there are more than 5.

Also the batch gets discarded 5 minutes after it has been retrieved or
another more recent batch has been retrieved. Previously the last batch
would stay in memory until the next one is ready. But it seems
unreasonable for the DO to rely on it since there are no guarantees.
This would probably even save some memory under normal conditions on
average.

Bug: 35753013
Test: cts-tradefed run cts -m CtsDevicePolicyManagerTestCases -t com.android.cts.devicepolicy.DeviceOwnerTest#testNetworkLoggingWithSingleUser
Change-Id: Ib8e91a98103d804375cb0d7423f93175b4b9bcb6
(cherry picked from commit 48733074d7ba80755e40432b7ff02b66e27d3edb)
Merged-in: Ib8e91a98103d804375cb0d7423f93175b4b9bcb6
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/NetworkLoggingHandler.java b/services/devicepolicy/java/com/android/server/devicepolicy/NetworkLoggingHandler.java
index 9b4de043..70c7e58 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/NetworkLoggingHandler.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/NetworkLoggingHandler.java
@@ -26,6 +26,7 @@
 import android.os.Message;
 import android.os.SystemClock;
 import android.util.Log;
+import android.util.LongSparseArray;
 
 import com.android.internal.annotations.GuardedBy;
 
@@ -44,12 +45,21 @@
 
     // If this value changes, update DevicePolicyManager#retrieveNetworkLogs() javadoc
     private static final int MAX_EVENTS_PER_BATCH = 1200;
-    private static final long BATCH_FINALIZATION_TIMEOUT_MS = TimeUnit.MINUTES.toMillis(90);
-    private static final long BATCH_FINALIZATION_TIMEOUT_ALARM_INTERVAL_MS =
-            TimeUnit.MINUTES.toMillis(30);
+
+    /**
+     * Maximum number of batches to store in memory. If more batches are generated and the DO
+     * doesn't fetch them, we will discard the oldest one.
+     */
+    private static final int MAX_BATCHES = 5;
+
+    private static final long BATCH_FINALIZATION_TIMEOUT_MS = 90 * 60 * 1000; // 1.5h
+    private static final long BATCH_FINALIZATION_TIMEOUT_ALARM_INTERVAL_MS = 30 * 60 * 1000; // 30m
 
     private static final String NETWORK_LOGGING_TIMEOUT_ALARM_TAG = "NetworkLogging.batchTimeout";
 
+    /** Delay after which older batches get discarded after a retrieval. */
+    private static final long RETRIEVED_BATCH_DISCARD_DELAY_MS = 5 * 60 * 1000; // 5m
+
     private final DevicePolicyManagerService mDpm;
     private final AlarmManager mAlarmManager;
 
@@ -66,22 +76,27 @@
 
     static final int LOG_NETWORK_EVENT_MSG = 1;
 
-    // threadsafe as it's Handler's thread confined
+    /** Network events accumulated so far to be finalized into a batch at some point. */
     @GuardedBy("this")
-    private ArrayList<NetworkEvent> mNetworkEvents = new ArrayList<NetworkEvent>();
+    private ArrayList<NetworkEvent> mNetworkEvents = new ArrayList<>();
 
+    /**
+     * Up to {@code MAX_BATCHES} finalized batches of logs ready to be retrieved by the DO. Already
+     * retrieved batches are discarded after {@code RETRIEVED_BATCH_DISCARD_DELAY_MS}.
+     */
     @GuardedBy("this")
-    private ArrayList<NetworkEvent> mFullBatch;
+    private final LongSparseArray<ArrayList<NetworkEvent>> mBatches =
+            new LongSparseArray<>(MAX_BATCHES);
 
     @GuardedBy("this")
     private boolean mPaused = false;
 
     // each full batch is represented by its token, which the DPC has to provide back to retrieve it
     @GuardedBy("this")
-    private long mCurrentFullBatchToken;
+    private long mCurrentBatchToken;
 
     @GuardedBy("this")
-    private long mLastRetrievedFullBatchToken;
+    private long mLastRetrievedBatchToken;
 
     NetworkLoggingHandler(Looper looper, DevicePolicyManagerService dpm) {
         super(looper);
@@ -93,7 +108,7 @@
     public void handleMessage(Message msg) {
         switch (msg.what) {
             case LOG_NETWORK_EVENT_MSG: {
-                NetworkEvent networkEvent = msg.getData().getParcelable(NETWORK_EVENT_KEY);
+                final NetworkEvent networkEvent = msg.getData().getParcelable(NETWORK_EVENT_KEY);
                 if (networkEvent != null) {
                     synchronized (NetworkLoggingHandler.this) {
                         mNetworkEvents.add(networkEvent);
@@ -113,6 +128,8 @@
 
     void scheduleBatchFinalization() {
         final long when = SystemClock.elapsedRealtime() + BATCH_FINALIZATION_TIMEOUT_MS;
+        // We use alarm manager and not just postDelayed here to ensure the batch gets finalized
+        // even if the device goes to sleep.
         mAlarmManager.setWindow(AlarmManager.ELAPSED_REALTIME_WAKEUP, when,
                 BATCH_FINALIZATION_TIMEOUT_ALARM_INTERVAL_MS, NETWORK_LOGGING_TIMEOUT_ALARM_TAG,
                 mBatchTimeoutAlarmListener, this);
@@ -131,62 +148,80 @@
             return;
         }
 
-        Log.d(TAG, "Resumed network logging. Current batch="
-                + mCurrentFullBatchToken + ", LastRetrievedBatch=" + mLastRetrievedFullBatchToken);
+        Log.d(TAG, "Resumed network logging. Current batch=" + mCurrentBatchToken
+                + ", LastRetrievedBatch=" + mLastRetrievedBatchToken);
         mPaused = false;
 
-        // If there is a full batch ready that the device owner hasn't been notified about, do it
-        // now.
-        if (mFullBatch != null && mFullBatch.size() > 0
-                && mLastRetrievedFullBatchToken != mCurrentFullBatchToken) {
+        // If there is a batch ready that the device owner hasn't been notified about, do it now.
+        if (mBatches.size() > 0 && mLastRetrievedBatchToken != mCurrentBatchToken) {
             scheduleBatchFinalization();
             notifyDeviceOwnerLocked();
         }
     }
 
     synchronized void discardLogs() {
-        mFullBatch = null;
-        mNetworkEvents = new ArrayList<NetworkEvent>();
+        mBatches.clear();
+        mNetworkEvents = new ArrayList<>();
         Log.d(TAG, "Discarded all network logs");
     }
 
     @GuardedBy("this")
     private void finalizeBatchAndNotifyDeviceOwnerLocked() {
         if (mNetworkEvents.size() > 0) {
-            // finalize the batch and start a new one from scratch
-            mFullBatch = mNetworkEvents;
-            mCurrentFullBatchToken++;
-            mNetworkEvents = new ArrayList<NetworkEvent>();
+            // Finalize the batch and start a new one from scratch.
+            if (mBatches.size() >= MAX_BATCHES) {
+                // Remove the oldest batch if we hit the limit.
+                mBatches.removeAt(0);
+            }
+            mCurrentBatchToken++;
+            mBatches.append(mCurrentBatchToken, mNetworkEvents);
+            mNetworkEvents = new ArrayList<>();
             if (!mPaused) {
                 notifyDeviceOwnerLocked();
             }
         } else {
-            // don't notify the DO, since there are no events; DPC can still retrieve
+            // Don't notify the DO, since there are no events; DPC can still retrieve
             // the last full batch if not paused.
             Log.d(TAG, "Was about to finalize the batch, but there were no events to send to"
-                    + " the DPC, the batchToken of last available batch: "
-                    + mCurrentFullBatchToken);
+                    + " the DPC, the batchToken of last available batch: " + mCurrentBatchToken);
         }
-        // regardless of whether the batch was non-empty schedule a new finalization after timeout
+        // Regardless of whether the batch was non-empty schedule a new finalization after timeout.
         scheduleBatchFinalization();
     }
 
+    /** Sends a notification to the DO. Should only be called when there is a batch available. */
     @GuardedBy("this")
     private void notifyDeviceOwnerLocked() {
-        Bundle extras = new Bundle();
-        extras.putLong(DeviceAdminReceiver.EXTRA_NETWORK_LOGS_TOKEN, mCurrentFullBatchToken);
-        extras.putInt(DeviceAdminReceiver.EXTRA_NETWORK_LOGS_COUNT, mFullBatch.size());
+        final Bundle extras = new Bundle();
+        final int lastBatchSize = mBatches.valueAt(mBatches.size() - 1).size();
+        extras.putLong(DeviceAdminReceiver.EXTRA_NETWORK_LOGS_TOKEN, mCurrentBatchToken);
+        extras.putInt(DeviceAdminReceiver.EXTRA_NETWORK_LOGS_COUNT, lastBatchSize);
         Log.d(TAG, "Sending network logging batch broadcast to device owner, batchToken: "
-                + mCurrentFullBatchToken);
+                + mCurrentBatchToken);
         mDpm.sendDeviceOwnerCommand(DeviceAdminReceiver.ACTION_NETWORK_LOGS_AVAILABLE, extras);
     }
 
-    synchronized List<NetworkEvent> retrieveFullLogBatch(long batchToken) {
-        if (batchToken != mCurrentFullBatchToken) {
+    synchronized List<NetworkEvent> retrieveFullLogBatch(final long batchToken) {
+        final int index = mBatches.indexOfKey(batchToken);
+        if (index < 0) {
+            // Invalid token or batch has already been discarded.
             return null;
         }
-        mLastRetrievedFullBatchToken = mCurrentFullBatchToken;
-        return mFullBatch;
+
+        // Schedule this and older batches to be discarded after a delay to lessen memory load
+        // without interfering with the admin's ability to collect logs out-of-order.
+        // It isn't critical and we allow it to be delayed further if the phone sleeps, so we don't
+        // use the alarm manager here.
+        postDelayed(() -> {
+            synchronized(this) {
+                while (mBatches.size() > 0 && mBatches.keyAt(0) <= batchToken) {
+                    mBatches.removeAt(0);
+                }
+            }
+        }, RETRIEVED_BATCH_DISCARD_DELAY_MS);
+
+        mLastRetrievedBatchToken = batchToken;
+        return mBatches.valueAt(index);
     }
 }