Optimize start info proto memory use

The proto reading and writing can allocate a lot of objects, re-use when possible.

NO_IFTTT=Not picking up labels correctly, will investigate separately.

Test: atest ApplicationStartInfoTest (testApplicationStartInfo covers persist and restore which leverages proto read/write)
Bug: 384539178
Flag: EXEMPT - small bugfix
Change-Id: I89b840b629a0b28deecf0e7078999a96e0c9cf9e

2 files changed, 80 insertions, 49 deletions

diff --git a/core/java/android/app/ApplicationStartInfo.java b/core/java/android/app/ApplicationStartInfo.java
index 3214bd8f01fc..2e8031dd22fe 100644
--- a/core/java/android/app/ApplicationStartInfo.java
+++ b/core/java/android/app/ApplicationStartInfo.java
@@ -840,7 +840,9 @@ public final class ApplicationStartInfo implements Parcelable {
      * @hide
      */
     // LINT.IfChange(write_proto)
-    public void writeToProto(ProtoOutputStream proto, long fieldId) throws IOException {
+    public void writeToProto(ProtoOutputStream proto, long fieldId,
+            ByteArrayOutputStream byteArrayOutputStream, ObjectOutputStream objectOutputStream,
+            TypedXmlSerializer typedXmlSerializer) throws IOException {
         final long token = proto.start(fieldId);
         proto.write(ApplicationStartInfoProto.PID, mPid);
         proto.write(ApplicationStartInfoProto.REAL_UID, mRealUid);
@@ -850,38 +852,38 @@ public final class ApplicationStartInfo implements Parcelable {
         proto.write(ApplicationStartInfoProto.STARTUP_STATE, mStartupState);
         proto.write(ApplicationStartInfoProto.REASON, mReason);
         if (mStartupTimestampsNs != null && mStartupTimestampsNs.size() > 0) {
-            ByteArrayOutputStream timestampsBytes = new ByteArrayOutputStream();
-            ObjectOutputStream timestampsOut = new ObjectOutputStream(timestampsBytes);
-            TypedXmlSerializer serializer = Xml.resolveSerializer(timestampsOut);
-            serializer.startDocument(null, true);
-            serializer.startTag(null, PROTO_SERIALIZER_ATTRIBUTE_TIMESTAMPS);
+            byteArrayOutputStream = new ByteArrayOutputStream();
+            objectOutputStream = new ObjectOutputStream(byteArrayOutputStream);
+            typedXmlSerializer = Xml.resolveSerializer(objectOutputStream);
+            typedXmlSerializer.startDocument(null, true);
+            typedXmlSerializer.startTag(null, PROTO_SERIALIZER_ATTRIBUTE_TIMESTAMPS);
             for (int i = 0; i < mStartupTimestampsNs.size(); i++) {
-                serializer.startTag(null, PROTO_SERIALIZER_ATTRIBUTE_TIMESTAMP);
-                serializer.attributeInt(null, PROTO_SERIALIZER_ATTRIBUTE_KEY,
+                typedXmlSerializer.startTag(null, PROTO_SERIALIZER_ATTRIBUTE_TIMESTAMP);
+                typedXmlSerializer.attributeInt(null, PROTO_SERIALIZER_ATTRIBUTE_KEY,
                         mStartupTimestampsNs.keyAt(i));
-                serializer.attributeLong(null, PROTO_SERIALIZER_ATTRIBUTE_TS,
+                typedXmlSerializer.attributeLong(null, PROTO_SERIALIZER_ATTRIBUTE_TS,
                         mStartupTimestampsNs.valueAt(i));
-                serializer.endTag(null, PROTO_SERIALIZER_ATTRIBUTE_TIMESTAMP);
+                typedXmlSerializer.endTag(null, PROTO_SERIALIZER_ATTRIBUTE_TIMESTAMP);
             }
-            serializer.endTag(null, PROTO_SERIALIZER_ATTRIBUTE_TIMESTAMPS);
-            serializer.endDocument();
+            typedXmlSerializer.endTag(null, PROTO_SERIALIZER_ATTRIBUTE_TIMESTAMPS);
+            typedXmlSerializer.endDocument();
             proto.write(ApplicationStartInfoProto.STARTUP_TIMESTAMPS,
-                    timestampsBytes.toByteArray());
-            timestampsOut.close();
+                    byteArrayOutputStream.toByteArray());
+            objectOutputStream.close();
         }
         proto.write(ApplicationStartInfoProto.START_TYPE, mStartType);
         if (mStartIntent != null) {
-            ByteArrayOutputStream intentBytes = new ByteArrayOutputStream();
-            ObjectOutputStream intentOut = new ObjectOutputStream(intentBytes);
-            TypedXmlSerializer serializer = Xml.resolveSerializer(intentOut);
-            serializer.startDocument(null, true);
-            serializer.startTag(null, PROTO_SERIALIZER_ATTRIBUTE_INTENT);
-            mStartIntent.saveToXml(serializer);
-            serializer.endTag(null, PROTO_SERIALIZER_ATTRIBUTE_INTENT);
-            serializer.endDocument();
+            byteArrayOutputStream = new ByteArrayOutputStream();
+            objectOutputStream = new ObjectOutputStream(byteArrayOutputStream);
+            typedXmlSerializer = Xml.resolveSerializer(objectOutputStream);
+            typedXmlSerializer.startDocument(null, true);
+            typedXmlSerializer.startTag(null, PROTO_SERIALIZER_ATTRIBUTE_INTENT);
+            mStartIntent.saveToXml(typedXmlSerializer);
+            typedXmlSerializer.endTag(null, PROTO_SERIALIZER_ATTRIBUTE_INTENT);
+            typedXmlSerializer.endDocument();
             proto.write(ApplicationStartInfoProto.START_INTENT,
-                    intentBytes.toByteArray());
-            intentOut.close();
+                    byteArrayOutputStream.toByteArray());
+            objectOutputStream.close();
         }
         proto.write(ApplicationStartInfoProto.LAUNCH_MODE, mLaunchMode);
         proto.write(ApplicationStartInfoProto.WAS_FORCE_STOPPED, mWasForceStopped);
@@ -900,7 +902,9 @@ public final class ApplicationStartInfo implements Parcelable {
      * @hide
      */
     // LINT.IfChange(read_proto)
-    public void readFromProto(ProtoInputStream proto, long fieldId)
+    public void readFromProto(ProtoInputStream proto, long fieldId,
+            ByteArrayInputStream byteArrayInputStream, ObjectInputStream objectInputStream,
+            TypedXmlPullParser typedXmlPullParser)
             throws IOException, WireTypeMismatchException, ClassNotFoundException {
         final long token = proto.start(fieldId);
         while (proto.nextField() != ProtoInputStream.NO_MORE_FIELDS) {
@@ -927,19 +931,21 @@ public final class ApplicationStartInfo implements Parcelable {
                     mReason = proto.readInt(ApplicationStartInfoProto.REASON);
                     break;
                 case (int) ApplicationStartInfoProto.STARTUP_TIMESTAMPS:
-                    ByteArrayInputStream timestampsBytes = new ByteArrayInputStream(proto.readBytes(
+                    byteArrayInputStream = new ByteArrayInputStream(proto.readBytes(
                             ApplicationStartInfoProto.STARTUP_TIMESTAMPS));
-                    ObjectInputStream timestampsIn = new ObjectInputStream(timestampsBytes);
+                    objectInputStream = new ObjectInputStream(byteArrayInputStream);
                     mStartupTimestampsNs = new ArrayMap<Integer, Long>();
                     try {
-                        TypedXmlPullParser parser = Xml.resolvePullParser(timestampsIn);
-                        XmlUtils.beginDocument(parser, PROTO_SERIALIZER_ATTRIBUTE_TIMESTAMPS);
-                        int depth = parser.getDepth();
-                        while (XmlUtils.nextElementWithin(parser, depth)) {
-                            if (PROTO_SERIALIZER_ATTRIBUTE_TIMESTAMP.equals(parser.getName())) {
-                                int key = parser.getAttributeInt(null,
+                        typedXmlPullParser = Xml.resolvePullParser(objectInputStream);
+                        XmlUtils.beginDocument(typedXmlPullParser,
+                                PROTO_SERIALIZER_ATTRIBUTE_TIMESTAMPS);
+                        int depth = typedXmlPullParser.getDepth();
+                        while (XmlUtils.nextElementWithin(typedXmlPullParser, depth)) {
+                            if (PROTO_SERIALIZER_ATTRIBUTE_TIMESTAMP.equals(
+                                    typedXmlPullParser.getName())) {
+                                int key = typedXmlPullParser.getAttributeInt(null,
                                         PROTO_SERIALIZER_ATTRIBUTE_KEY);
-                                long ts = parser.getAttributeLong(null,
+                                long ts = typedXmlPullParser.getAttributeLong(null,
                                         PROTO_SERIALIZER_ATTRIBUTE_TS);
                                 mStartupTimestampsNs.put(key, ts);
                             }
@@ -947,23 +953,24 @@ public final class ApplicationStartInfo implements Parcelable {
                     } catch (XmlPullParserException e) {
                         // Timestamps lost
                     }
-                    timestampsIn.close();
+                    objectInputStream.close();
                     break;
                 case (int) ApplicationStartInfoProto.START_TYPE:
                     mStartType = proto.readInt(ApplicationStartInfoProto.START_TYPE);
                     break;
                 case (int) ApplicationStartInfoProto.START_INTENT:
-                    ByteArrayInputStream intentBytes = new ByteArrayInputStream(proto.readBytes(
+                    byteArrayInputStream = new ByteArrayInputStream(proto.readBytes(
                             ApplicationStartInfoProto.START_INTENT));
-                    ObjectInputStream intentIn = new ObjectInputStream(intentBytes);
+                    objectInputStream = new ObjectInputStream(byteArrayInputStream);
                     try {
-                        TypedXmlPullParser parser = Xml.resolvePullParser(intentIn);
-                        XmlUtils.beginDocument(parser, PROTO_SERIALIZER_ATTRIBUTE_INTENT);
-                        mStartIntent = Intent.restoreFromXml(parser);
+                        typedXmlPullParser = Xml.resolvePullParser(objectInputStream);
+                        XmlUtils.beginDocument(typedXmlPullParser,
+                                PROTO_SERIALIZER_ATTRIBUTE_INTENT);
+                        mStartIntent = Intent.restoreFromXml(typedXmlPullParser);
                     } catch (XmlPullParserException e) {
                         // Intent lost
                     }
-                    intentIn.close();
+                    objectInputStream.close();
                     break;
                 case (int) ApplicationStartInfoProto.LAUNCH_MODE:
                     mLaunchMode = proto.readInt(ApplicationStartInfoProto.LAUNCH_MODE);
diff --git a/services/core/java/com/android/server/am/AppStartInfoTracker.java b/services/core/java/com/android/server/am/AppStartInfoTracker.java
index 961022b7231b..517279bd7527 100644
--- a/services/core/java/com/android/server/am/AppStartInfoTracker.java
+++ b/services/core/java/com/android/server/am/AppStartInfoTracker.java
@@ -54,15 +54,21 @@ import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.app.ProcessMap;
 import com.android.internal.os.Clock;
 import com.android.internal.os.MonotonicClock;
+import com.android.modules.utils.TypedXmlPullParser;
+import com.android.modules.utils.TypedXmlSerializer;
 import com.android.server.IoThread;
 import com.android.server.ServiceThread;
 import com.android.server.SystemServiceManager;
 import com.android.server.wm.WindowProcessController;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileOutputStream;
 import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -1006,6 +1012,12 @@ public final class AppStartInfoTracker {
             throws IOException, WireTypeMismatchException, ClassNotFoundException {
         long token = proto.start(fieldId);
         String pkgName = "";
+
+        // Create objects for reuse.
+        ByteArrayInputStream byteArrayInputStream = null;
+        ObjectInputStream objectInputStream = null;
+        TypedXmlPullParser typedXmlPullParser = null;
+
         for (int next = proto.nextField();
                 next != ProtoInputStream.NO_MORE_FIELDS;
                 next = proto.nextField()) {
@@ -1017,7 +1029,7 @@ public final class AppStartInfoTracker {
                     AppStartInfoContainer container =
                             new AppStartInfoContainer(mAppStartInfoHistoryListSize);
                     int uid = container.readFromProto(proto, AppsStartInfoProto.Package.USERS,
-                            pkgName);
+                            pkgName, byteArrayInputStream, objectInputStream, typedXmlPullParser);
 
                     // If the isolated process flag is enabled and the uid is that of an isolated
                     // process, then break early so that the container will not be added to mData.
@@ -1052,6 +1064,12 @@ public final class AppStartInfoTracker {
             out = af.startWrite();
             ProtoOutputStream proto = new ProtoOutputStream(out);
             proto.write(AppsStartInfoProto.LAST_UPDATE_TIMESTAMP, now);
+
+            // Create objects for reuse.
+            ByteArrayOutputStream byteArrayOutputStream = null;
+            ObjectOutputStream objectOutputStream = null;
+            TypedXmlSerializer typedXmlSerializer = null;
+
             synchronized (mLock) {
                 succeeded = forEachPackageLocked(
                         (packageName, records) -> {
@@ -1060,8 +1078,9 @@ public final class AppStartInfoTracker {
                             int uidArraySize = records.size();
                             for (int j = 0; j < uidArraySize; j++) {
                                 try {
-                                    records.valueAt(j)
-                                            .writeToProto(proto, AppsStartInfoProto.Package.USERS);
+                                    records.valueAt(j).writeToProto(proto,
+                                            AppsStartInfoProto.Package.USERS, byteArrayOutputStream,
+                                            objectOutputStream, typedXmlSerializer);
                                 } catch (IOException e) {
                                     Slog.w(TAG, "Unable to write app start info into persistent"
                                             + "storage: " + e);
@@ -1414,19 +1433,23 @@ public final class AppStartInfoTracker {
         }
 
         @GuardedBy("mLock")
-        void writeToProto(ProtoOutputStream proto, long fieldId) throws IOException {
+        void writeToProto(ProtoOutputStream proto, long fieldId,
+                ByteArrayOutputStream byteArrayOutputStream, ObjectOutputStream objectOutputStream,
+                TypedXmlSerializer typedXmlSerializer) throws IOException {
             long token = proto.start(fieldId);
             proto.write(AppsStartInfoProto.Package.User.UID, mUid);
             int size = mInfos.size();
             for (int i = 0; i < size; i++) {
-                mInfos.get(i)
-                        .writeToProto(proto, AppsStartInfoProto.Package.User.APP_START_INFO);
+                mInfos.get(i).writeToProto(proto, AppsStartInfoProto.Package.User.APP_START_INFO,
+                        byteArrayOutputStream, objectOutputStream, typedXmlSerializer);
             }
             proto.write(AppsStartInfoProto.Package.User.MONITORING_ENABLED, mMonitoringModeEnabled);
             proto.end(token);
         }
 
-        int readFromProto(ProtoInputStream proto, long fieldId, String packageName)
+        int readFromProto(ProtoInputStream proto, long fieldId, String packageName,
+                ByteArrayInputStream byteArrayInputStream, ObjectInputStream objectInputStream,
+                TypedXmlPullParser typedXmlPullParser)
                 throws IOException, WireTypeMismatchException, ClassNotFoundException {
             long token = proto.start(fieldId);
             for (int next = proto.nextField();
@@ -1440,7 +1463,8 @@ public final class AppStartInfoTracker {
                         // Create record with monotonic time 0 in case the persisted record does not
                         // have a create time.
                         ApplicationStartInfo info = new ApplicationStartInfo(0);
-                        info.readFromProto(proto, AppsStartInfoProto.Package.User.APP_START_INFO);
+                        info.readFromProto(proto, AppsStartInfoProto.Package.User.APP_START_INFO,
+                                byteArrayInputStream, objectInputStream, typedXmlPullParser);
                         info.setPackageName(packageName);
                         mInfos.add(info);
                         break;