Fix StatsEvent memory usage for pulled events

Add usePooledBuffer flag to the Builder which determines whether to reuse
the Buffer's byte array in StatsEvent or use a copy.
The build() function also calls release() on the Buffer
if a copy of the Buffer's byte array is passed to StatsEvent.

Also, for pushed events, release the StatsEvent object and consequently,
the Buffer in StatsLog.write(StatsEvent)

Fixes: 145026572
Fixes: 144126444
Test: bit FrameworksCoreTests:android.util.StatsEventTest
Change-Id: I1cdaf0027b69281cb7cb6f3c8ca923d03829b4dd

4 files changed, 68 insertions, 9 deletions

diff --git a/core/java/android/util/StatsEvent.java b/core/java/android/util/StatsEvent.java
index 7e7164042781..c7659457bdf9 100644
--- a/core/java/android/util/StatsEvent.java
+++ b/core/java/android/util/StatsEvent.java
@@ -31,14 +31,23 @@ import com.android.internal.annotations.VisibleForTesting;
  *
  * <p>Usage:</p>
  * <pre>
+ *      // Pushed event
  *      StatsEvent statsEvent = StatsEvent.newBuilder()
  *          .setAtomId(atomId)
  *          .writeBoolean(false)
  *          .writeString("annotated String field")
  *          .addBooleanAnnotation(annotationId, true)
+ *          .usePooledBuffer()
  *          .build();
- *
  *      StatsLog.write(statsEvent);
+ *
+ *      // Pulled event
+ *      StatsEvent statsEvent = StatsEvent.newBuilder()
+ *          .setAtomId(atomId)
+ *          .writeBoolean(false)
+ *          .writeString("annotated String field")
+ *          .addBooleanAnnotation(annotationId, true)
+ *          .build();
  * </pre>
  * @hide
  **/
@@ -210,12 +219,15 @@ public final class StatsEvent {
     private static final int MAX_PAYLOAD_SIZE = LOGGER_ENTRY_MAX_PAYLOAD - 4;
 
     private final int mAtomId;
-    private final Buffer mBuffer;
+    private final byte[] mPayload;
+    private Buffer mBuffer;
     private final int mNumBytes;
 
-    private StatsEvent(final int atomId, @NonNull final Buffer buffer, final int numBytes) {
+    private StatsEvent(final int atomId, @Nullable final Buffer buffer,
+            @NonNull final byte[] payload, final int numBytes) {
         mAtomId = atomId;
         mBuffer = buffer;
+        mPayload = payload;
         mNumBytes = numBytes;
     }
 
@@ -243,7 +255,7 @@ public final class StatsEvent {
      **/
     @NonNull
     public byte[] getBytes() {
-        return mBuffer.getBytes();
+        return mPayload;
     }
 
     /**
@@ -256,10 +268,14 @@ public final class StatsEvent {
     }
 
     /**
-     * Recycle this StatsEvent object.
+     * Recycle resources used by this StatsEvent object.
+     * No actions should be taken on this StatsEvent after release() is called.
      **/
     public void release() {
-        mBuffer.release();
+        if (mBuffer != null) {
+            mBuffer.release();
+            mBuffer = null;
+        }
     }
 
     /**
@@ -280,7 +296,18 @@ public final class StatsEvent {
      *         optional string field3 = 3 [(annotation1) = true];
      *     }
      *
-     *     // StatsEvent construction.
+     *     // StatsEvent construction for pushed event.
+     *     StatsEvent.newBuilder()
+     *     StatsEvent statsEvent = StatsEvent.newBuilder()
+     *         .setAtomId(atomId)
+     *         .writeInt(3) // field1
+     *         .writeLong(8L) // field2
+     *         .writeString("foo") // field 3
+     *         .addBooleanAnnotation(annotation1Id, true)
+     *         .usePooledBuffer()
+     *         .build();
+     *
+     *     // StatsEvent construction for pulled event.
      *     StatsEvent.newBuilder()
      *     StatsEvent statsEvent = StatsEvent.newBuilder()
      *         .setAtomId(atomId)
@@ -306,6 +333,7 @@ public final class StatsEvent {
         private byte mLastType;
         private int mNumElements;
         private int mErrorMask;
+        private boolean mUsePooledBuffer = false;
 
         private Builder(final Buffer buffer) {
             mBuffer = buffer;
@@ -569,6 +597,17 @@ public final class StatsEvent {
         }
 
         /**
+         * Indicates to reuse Buffer's byte array as the underlying payload in StatsEvent.
+         * This should be called for pushed events to reduce memory allocations and garbage
+         * collections.
+         **/
+        @NonNull
+        public Builder usePooledBuffer() {
+            mUsePooledBuffer = true;
+            return this;
+        }
+
+        /**
          * Builds a StatsEvent object with values entered in this Builder.
          **/
         @NonNull
@@ -599,7 +638,18 @@ public final class StatsEvent {
                 size = mPos;
             }
 
-            return new StatsEvent(mAtomId, mBuffer, size);
+            if (mUsePooledBuffer) {
+                return new StatsEvent(mAtomId, mBuffer, mBuffer.getBytes(), size);
+            } else {
+                // Create a copy of the buffer with the required number of bytes.
+                final byte[] payload = new byte[size];
+                System.arraycopy(mBuffer.getBytes(), 0, payload, 0, size);
+
+                // Return Buffer instance to the pool.
+                mBuffer.release();
+
+                return new StatsEvent(mAtomId, null, payload, size);
+            }
         }
 
         private void writeTypeId(final byte typeId) {
diff --git a/core/java/android/util/StatsLog.java b/core/java/android/util/StatsLog.java
index 64e15cfb7948..23fd4f2f61d3 100644
--- a/core/java/android/util/StatsLog.java
+++ b/core/java/android/util/StatsLog.java
@@ -246,12 +246,15 @@ public final class StatsLog extends StatsLogInternal {
 
     /**
      * Write an event to stats log using the raw format encapsulated in StatsEvent.
+     * After writing to stats log, release() is called on the StatsEvent object.
+     * No further action should be taken on the StatsEvent object following this call.
      *
      * @param statsEvent    The StatsEvent object containing the encoded buffer of data to write.
      * @hide
      */
     public static void write(@NonNull final StatsEvent statsEvent) {
         writeImpl(statsEvent.getBytes(), statsEvent.getNumBytes(), statsEvent.getAtomId());
+        statsEvent.release();
     }
 
     private static void enforceDumpCallingPermission(Context context) {
diff --git a/core/tests/coretests/src/android/util/StatsEventTest.java b/core/tests/coretests/src/android/util/StatsEventTest.java
index 097badadcea9..ac25e2734ac9 100644
--- a/core/tests/coretests/src/android/util/StatsEventTest.java
+++ b/core/tests/coretests/src/android/util/StatsEventTest.java
@@ -44,7 +44,7 @@ public class StatsEventTest {
     @Test
     public void testNoFields() {
         final long minTimestamp = SystemClock.elapsedRealtimeNanos();
-        final StatsEvent statsEvent = StatsEvent.newBuilder().build();
+        final StatsEvent statsEvent = StatsEvent.newBuilder().usePooledBuffer().build();
         final long maxTimestamp = SystemClock.elapsedRealtimeNanos();
 
         final int expectedAtomId = 0;
@@ -99,6 +99,7 @@ public class StatsEventTest {
                 .writeBoolean(field2)
                 .writeInt(field3)
                 .writeInt(field4)
+                .usePooledBuffer()
                 .build();
         final long maxTimestamp = SystemClock.elapsedRealtimeNanos();
 
@@ -167,6 +168,7 @@ public class StatsEventTest {
                 .writeString(field1)
                 .writeFloat(field2)
                 .writeByteArray(field3)
+                .usePooledBuffer()
                 .build();
         final long maxTimestamp = SystemClock.elapsedRealtimeNanos();
 
@@ -230,6 +232,7 @@ public class StatsEventTest {
                 .setAtomId(expectedAtomId)
                 .writeAttributionChain(uids, tags)
                 .writeLong(field2)
+                .usePooledBuffer()
                 .build();
         final long maxTimestamp = SystemClock.elapsedRealtimeNanos();
 
@@ -299,6 +302,7 @@ public class StatsEventTest {
         final StatsEvent statsEvent = StatsEvent.newBuilder()
                 .setAtomId(expectedAtomId)
                 .writeKeyValuePairs(intMap, longMap, stringMap, floatMap)
+                .usePooledBuffer()
                 .build();
         final long maxTimestamp = SystemClock.elapsedRealtimeNanos();
 
@@ -392,6 +396,7 @@ public class StatsEventTest {
                 .addBooleanAnnotation(field1AnnotationId, field1AnnotationValue)
                 .writeBoolean(field2)
                 .addIntAnnotation(field2AnnotationId, field2AnnotationValue)
+                .usePooledBuffer()
                 .build();
         final long maxTimestamp = SystemClock.elapsedRealtimeNanos();
 
diff --git a/tools/stats_log_api_gen/java_writer.cpp b/tools/stats_log_api_gen/java_writer.cpp
index 4899fbd3b669..3814f691c5f2 100644
--- a/tools/stats_log_api_gen/java_writer.cpp
+++ b/tools/stats_log_api_gen/java_writer.cpp
@@ -186,6 +186,7 @@ static int write_java_methods(
         }
 
         fprintf(out, "\n");
+        fprintf(out, "%s        builder.usePooledBuffer();\n", indent.c_str());
         fprintf(out, "%s        StatsLog.write(builder.build());\n", indent.c_str());
 
         // Add support for writing using Q schema if this is not the default module.