Add wipeOnError flag to NetworkStatsRecorder

If reading data happens exception while doing data migration,
the file will be deleted by legacy recorders. This would cause
legacy persistent data being lost and cannot be retrieved by
any method. To avoid the files being deleted, add a wipeOnError
flag to recorder which indicates this recorder will wipe on
error or not . If the flag is set to true then deletes all files
when it throws, otherwise keeps all files.

Ignore-AOSP-First: urgent fix and will cherry-pick immediately after.
Bug: 233828210
Test: FrameworksNetTests:NetworkStatsRecorderTest
Change-Id: Id7a3d8bebf8a00d814f9e84bf4c10d927e6ff749

4 files changed, 120 insertions, 17 deletions

diff --git a/service-t/src/com/android/server/net/NetworkStatsRecorder.java b/service-t/src/com/android/server/net/NetworkStatsRecorder.java
index d73e3428e2..768f3ebf14 100644
--- a/service-t/src/com/android/server/net/NetworkStatsRecorder.java
+++ b/service-t/src/com/android/server/net/NetworkStatsRecorder.java
@@ -79,6 +79,7 @@ public class NetworkStatsRecorder {
 
     private final long mBucketDuration;
     private final boolean mOnlyTags;
+    private final boolean mWipeOnError;
 
     private long mPersistThresholdBytes = 2 * MB_IN_BYTES;
     private NetworkStats mLastSnapshot;
@@ -103,6 +104,7 @@ public class NetworkStatsRecorder {
         // slack to avoid overflow
         mBucketDuration = YEAR_IN_MILLIS;
         mOnlyTags = false;
+        mWipeOnError = true;
 
         mPending = null;
         mSinceBoot = new NetworkStatsCollection(mBucketDuration);
@@ -114,7 +116,8 @@ public class NetworkStatsRecorder {
      * Persisted recorder.
      */
     public NetworkStatsRecorder(FileRotator rotator, NonMonotonicObserver<String> observer,
-            DropBoxManager dropBox, String cookie, long bucketDuration, boolean onlyTags) {
+            DropBoxManager dropBox, String cookie, long bucketDuration, boolean onlyTags,
+            boolean wipeOnError) {
         mRotator = Objects.requireNonNull(rotator, "missing FileRotator");
         mObserver = Objects.requireNonNull(observer, "missing NonMonotonicObserver");
         mDropBox = Objects.requireNonNull(dropBox, "missing DropBoxManager");
@@ -122,6 +125,7 @@ public class NetworkStatsRecorder {
 
         mBucketDuration = bucketDuration;
         mOnlyTags = onlyTags;
+        mWipeOnError = wipeOnError;
 
         mPending = new NetworkStatsCollection(bucketDuration);
         mSinceBoot = new NetworkStatsCollection(bucketDuration);
@@ -593,7 +597,9 @@ public class NetworkStatsRecorder {
             }
             mDropBox.addData(TAG_NETSTATS_DUMP, os.toByteArray(), 0);
         }
-
-        mRotator.deleteAll();
+        // Delete all files if this recorder is set wipe on error.
+        if (mWipeOnError) {
+            mRotator.deleteAll();
+        }
     }
 }
diff --git a/service-t/src/com/android/server/net/NetworkStatsService.java b/service-t/src/com/android/server/net/NetworkStatsService.java
index b37f93d271..b955db9f02 100644
--- a/service-t/src/com/android/server/net/NetworkStatsService.java
+++ b/service-t/src/com/android/server/net/NetworkStatsService.java
@@ -800,11 +800,14 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
             mSystemReady = true;
 
             // create data recorders along with historical rotators
-            mDevRecorder = buildRecorder(PREFIX_DEV, mSettings.getDevConfig(), false, mStatsDir);
-            mXtRecorder = buildRecorder(PREFIX_XT, mSettings.getXtConfig(), false, mStatsDir);
-            mUidRecorder = buildRecorder(PREFIX_UID, mSettings.getUidConfig(), false, mStatsDir);
+            mDevRecorder = buildRecorder(PREFIX_DEV, mSettings.getDevConfig(), false, mStatsDir,
+                    true /* wipeOnError */);
+            mXtRecorder = buildRecorder(PREFIX_XT, mSettings.getXtConfig(), false, mStatsDir,
+                    true /* wipeOnError */);
+            mUidRecorder = buildRecorder(PREFIX_UID, mSettings.getUidConfig(), false, mStatsDir,
+                    true /* wipeOnError */);
             mUidTagRecorder = buildRecorder(PREFIX_UID_TAG, mSettings.getUidTagConfig(), true,
-                    mStatsDir);
+                    mStatsDir, true /* wipeOnError */);
 
             updatePersistThresholdsLocked();
 
@@ -869,12 +872,13 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
 
     private NetworkStatsRecorder buildRecorder(
             String prefix, NetworkStatsSettings.Config config, boolean includeTags,
-            File baseDir) {
+            File baseDir, boolean wipeOnError) {
         final DropBoxManager dropBox = (DropBoxManager) mContext.getSystemService(
                 Context.DROPBOX_SERVICE);
         return new NetworkStatsRecorder(new FileRotator(
                 baseDir, prefix, config.rotateAgeMillis, config.deleteAgeMillis),
-                mNonMonotonicObserver, dropBox, prefix, config.bucketDuration, includeTags);
+                mNonMonotonicObserver, dropBox, prefix, config.bucketDuration, includeTags,
+                wipeOnError);
     }
 
     @GuardedBy("mStatsLock")
@@ -971,12 +975,17 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
         final NetworkStatsRecorder[] legacyRecorders;
         if (runComparison) {
             final File legacyBaseDir = mDeps.getLegacyStatsDir();
+            // Set wipeOnError flag false so the recorder won't damage persistent data if reads
+            // failed and calling deleteAll.
             legacyRecorders = new NetworkStatsRecorder[]{
-                    buildRecorder(PREFIX_DEV, mSettings.getDevConfig(), false, legacyBaseDir),
-                    buildRecorder(PREFIX_XT, mSettings.getXtConfig(), false, legacyBaseDir),
-                    buildRecorder(PREFIX_UID, mSettings.getUidConfig(), false, legacyBaseDir),
-                    buildRecorder(PREFIX_UID_TAG, mSettings.getUidTagConfig(), true, legacyBaseDir)
-            };
+                buildRecorder(PREFIX_DEV, mSettings.getDevConfig(), false, legacyBaseDir,
+                        false /* wipeOnError */),
+                buildRecorder(PREFIX_XT, mSettings.getXtConfig(), false, legacyBaseDir,
+                        false /* wipeOnError */),
+                buildRecorder(PREFIX_UID, mSettings.getUidConfig(), false, legacyBaseDir,
+                        false /* wipeOnError */),
+                buildRecorder(PREFIX_UID_TAG, mSettings.getUidTagConfig(), true, legacyBaseDir,
+                        false /* wipeOnError */)};
         } else {
             legacyRecorders = null;
         }
diff --git a/tests/unit/java/android/net/NetworkStatsRecorderTest.java b/tests/unit/java/android/net/NetworkStatsRecorderTest.java
new file mode 100644
index 0000000000..fad11a3ff7
--- /dev/null
+++ b/tests/unit/java/android/net/NetworkStatsRecorderTest.java
@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *i
+ * 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.android.server.net;
+
+import static android.text.format.DateUtils.HOUR_IN_MILLIS;
+
+import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
+
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.anyLong;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import android.net.NetworkStats;
+import android.os.DropBoxManager;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.internal.util.FileRotator;
+import com.android.testutils.DevSdkIgnoreRule;
+import com.android.testutils.DevSdkIgnoreRunner;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.io.IOException;
+
+@RunWith(DevSdkIgnoreRunner.class)
+@SmallTest
+@DevSdkIgnoreRule.IgnoreUpTo(SC_V2)
+public final class NetworkStatsRecorderTest {
+    private static final String TAG = NetworkStatsRecorderTest.class.getSimpleName();
+
+    private static final String TEST_PREFIX = "test";
+
+    @Mock private DropBoxManager mDropBox;
+    @Mock private NetworkStats.NonMonotonicObserver mObserver;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+    }
+
+    private NetworkStatsRecorder buildRecorder(FileRotator rotator, boolean wipeOnError) {
+        return new NetworkStatsRecorder(rotator, mObserver, mDropBox, TEST_PREFIX,
+                    HOUR_IN_MILLIS, false /* includeTags */, wipeOnError);
+    }
+
+    @Test
+    public void testWipeOnError() throws Exception {
+        final FileRotator rotator = mock(FileRotator.class);
+        final NetworkStatsRecorder wipeOnErrorRecorder = buildRecorder(rotator, true);
+
+        // Assuming that the rotator gets an exception happened when read data.
+        doThrow(new IOException()).when(rotator).readMatching(any(), anyLong(), anyLong());
+        wipeOnErrorRecorder.getOrLoadPartialLocked(Long.MIN_VALUE, Long.MAX_VALUE);
+        // Verify that the files will be deleted.
+        verify(rotator, times(1)).deleteAll();
+        reset(rotator);
+
+        final NetworkStatsRecorder noWipeOnErrorRecorder = buildRecorder(rotator, false);
+        doThrow(new IOException()).when(rotator).readMatching(any(), anyLong(), anyLong());
+        noWipeOnErrorRecorder.getOrLoadPartialLocked(Long.MIN_VALUE, Long.MAX_VALUE);
+        // Verify that the rotator won't delete files.
+        verify(rotator, never()).deleteAll();
+    }
+}
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index dbc1e2e90c..a7c5877f31 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -2010,18 +2010,18 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
     }
 
     private NetworkStatsRecorder makeTestRecorder(File directory, String prefix, Config config,
-            boolean includeTags) {
+            boolean includeTags, boolean wipeOnError) {
         final NetworkStats.NonMonotonicObserver observer =
                 mock(NetworkStats.NonMonotonicObserver.class);
         final DropBoxManager dropBox = mock(DropBoxManager.class);
         return new NetworkStatsRecorder(new FileRotator(
                 directory, prefix, config.rotateAgeMillis, config.deleteAgeMillis),
-                observer, dropBox, prefix, config.bucketDuration, includeTags);
+                observer, dropBox, prefix, config.bucketDuration, includeTags, wipeOnError);
     }
 
     private NetworkStatsCollection getLegacyCollection(String prefix, boolean includeTags) {
         final NetworkStatsRecorder recorder = makeTestRecorder(mLegacyStatsDir, prefix,
-                mSettings.getDevConfig(), includeTags);
+                mSettings.getDevConfig(), includeTags, false);
         return recorder.getOrLoadCompleteLocked();
     }