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
Merged-In: I1cdaf0027b69281cb7cb6f3c8ca923d03829b4dd

2 files changed, 61 insertions, 8 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) {