Merge "Allow annotations to be added to Atom Id" into rvc-dev

2 files changed, 219 insertions, 13 deletions

diff --git a/apex/statsd/framework/java/android/util/StatsEvent.java b/apex/statsd/framework/java/android/util/StatsEvent.java
index 1a45c4a5b7f6..8bd36a516b12 100644
--- a/apex/statsd/framework/java/android/util/StatsEvent.java
+++ b/apex/statsd/framework/java/android/util/StatsEvent.java
@@ -188,6 +188,12 @@ public final class StatsEvent {
     @VisibleForTesting
     public static final int ERROR_ATTRIBUTION_UIDS_TAGS_SIZES_NOT_EQUAL = 0x1000;
 
+    /**
+     * @hide
+     **/
+    @VisibleForTesting
+    public static final int ERROR_ATOM_ID_INVALID_POSITION = 0x2000;
+
     // Size limits.
 
     /**
@@ -350,19 +356,32 @@ public final class StatsEvent {
             mPos = 0;
             writeTypeId(TYPE_OBJECT);
 
-            // Set mPos to after atom id's location in the buffer.
-            // First 2 elements in the buffer are event timestamp followed by the atom id.
-            mPos = POS_ATOM_ID + Byte.BYTES + Integer.BYTES;
-            mPosLastField = 0;
-            mLastType = 0;
+            // Write timestamp.
+            mPos = POS_TIMESTAMP_NS;
+            writeLong(mTimestampNs);
         }
 
         /**
          * Sets the atom id for this StatsEvent.
+         *
+         * This should be called immediately after StatsEvent.newBuilder()
+         * and should only be called once.
+         * Not calling setAtomId will result in ERROR_NO_ATOM_ID.
+         * Calling setAtomId out of order will result in ERROR_ATOM_ID_INVALID_POSITION.
          **/
         @NonNull
         public Builder setAtomId(final int atomId) {
-            mAtomId = atomId;
+            if (0 == mAtomId) {
+                mAtomId = atomId;
+
+                if (1 == mNumElements) { // Only timestamp is written so far.
+                    writeInt(atomId);
+                } else {
+                    // setAtomId called out of order.
+                    mErrorMask |= ERROR_ATOM_ID_INVALID_POSITION;
+                }
+            }
+
             return this;
         }
 
@@ -557,7 +576,7 @@ public final class StatsEvent {
         public Builder addBooleanAnnotation(
                 final byte annotationId, final boolean value) {
             // Ensure there's a field written to annotate.
-            if (0 == mPosLastField) {
+            if (mNumElements < 2) {
                 mErrorMask |= ERROR_ANNOTATION_DOES_NOT_FOLLOW_FIELD;
             } else if (mCurrentAnnotationCount >= MAX_ANNOTATION_COUNT) {
                 mErrorMask |= ERROR_TOO_MANY_ANNOTATIONS;
@@ -568,6 +587,7 @@ public final class StatsEvent {
                 mCurrentAnnotationCount++;
                 writeAnnotationCount();
             }
+
             return this;
         }
 
@@ -576,7 +596,7 @@ public final class StatsEvent {
          **/
         @NonNull
         public Builder addIntAnnotation(final byte annotationId, final int value) {
-            if (0 == mPosLastField) {
+            if (mNumElements < 2) {
                 mErrorMask |= ERROR_ANNOTATION_DOES_NOT_FOLLOW_FIELD;
             } else if (mCurrentAnnotationCount >= MAX_ANNOTATION_COUNT) {
                 mErrorMask |= ERROR_TOO_MANY_ANNOTATIONS;
@@ -587,6 +607,7 @@ public final class StatsEvent {
                 mCurrentAnnotationCount++;
                 writeAnnotationCount();
             }
+
             return this;
         }
 
@@ -619,19 +640,20 @@ public final class StatsEvent {
                 mErrorMask |= ERROR_TOO_MANY_FIELDS;
             }
 
-            int size = mPos;
-            mPos = POS_TIMESTAMP_NS;
-            writeLong(mTimestampNs);
-            writeInt(mAtomId);
             if (0 == mErrorMask) {
                 mBuffer.putByte(POS_NUM_ELEMENTS, (byte) mNumElements);
             } else {
+                // Write atom id and error mask. Overwrite any annotations for atom Id.
+                mPos = POS_ATOM_ID;
+                mPos += mBuffer.putByte(mPos, TYPE_INT);
+                mPos += mBuffer.putInt(mPos, mAtomId);
                 mPos += mBuffer.putByte(mPos, TYPE_ERRORS);
                 mPos += mBuffer.putInt(mPos, mErrorMask);
                 mBuffer.putByte(POS_NUM_ELEMENTS, (byte) 3);
-                size = mPos;
             }
 
+            final int size = mPos;
+
             if (mUsePooledBuffer) {
                 return new StatsEvent(mAtomId, mBuffer, mBuffer.getBytes(), size);
             } else {
diff --git a/apex/statsd/framework/test/src/android/util/StatsEventTest.java b/apex/statsd/framework/test/src/android/util/StatsEventTest.java
index ac25e2734ac9..7b511553a26f 100644
--- a/apex/statsd/framework/test/src/android/util/StatsEventTest.java
+++ b/apex/statsd/framework/test/src/android/util/StatsEventTest.java
@@ -85,6 +85,45 @@ public class StatsEventTest {
     }
 
     @Test
+    public void testOnlyAtomId() {
+        final int expectedAtomId = 109;
+
+        final long minTimestamp = SystemClock.elapsedRealtimeNanos();
+        final StatsEvent statsEvent = StatsEvent.newBuilder()
+                .setAtomId(expectedAtomId)
+                .usePooledBuffer()
+                .build();
+        final long maxTimestamp = SystemClock.elapsedRealtimeNanos();
+
+        assertThat(statsEvent.getAtomId()).isEqualTo(expectedAtomId);
+
+        final ByteBuffer buffer =
+                ByteBuffer.wrap(statsEvent.getBytes()).order(ByteOrder.LITTLE_ENDIAN);
+
+        assertWithMessage("Root element in buffer is not TYPE_OBJECT")
+                .that(buffer.get()).isEqualTo(StatsEvent.TYPE_OBJECT);
+
+        assertWithMessage("Incorrect number of elements in root object")
+                .that(buffer.get()).isEqualTo(2);
+
+        assertWithMessage("First element is not timestamp")
+                .that(buffer.get()).isEqualTo(StatsEvent.TYPE_LONG);
+
+        assertWithMessage("Incorrect timestamp")
+                .that(buffer.getLong()).isIn(Range.closed(minTimestamp, maxTimestamp));
+
+        assertWithMessage("Second element is not atom id")
+                .that(buffer.get()).isEqualTo(StatsEvent.TYPE_INT);
+
+        assertWithMessage("Incorrect atom id")
+                .that(buffer.getInt()).isEqualTo(expectedAtomId);
+
+        assertThat(statsEvent.getNumBytes()).isEqualTo(buffer.position());
+
+        statsEvent.release();
+    }
+
+    @Test
     public void testIntBooleanIntInt() {
         final int expectedAtomId = 109;
         final int field1 = 1;
@@ -460,6 +499,151 @@ public class StatsEventTest {
         statsEvent.release();
     }
 
+    @Test
+    public void testAtomIdAnnotations() {
+        final int expectedAtomId = 109;
+        final byte atomAnnotationId = 84;
+        final int atomAnnotationValue = 9;
+        final int field1 = 1;
+        final byte field1AnnotationId = 45;
+        final boolean field1AnnotationValue = false;
+        final boolean field2 = true;
+        final byte field2AnnotationId = 1;
+        final int field2AnnotationValue = 23;
+
+        final long minTimestamp = SystemClock.elapsedRealtimeNanos();
+        final StatsEvent statsEvent = StatsEvent.newBuilder()
+                .setAtomId(expectedAtomId)
+                .addIntAnnotation(atomAnnotationId, atomAnnotationValue)
+                .writeInt(field1)
+                .addBooleanAnnotation(field1AnnotationId, field1AnnotationValue)
+                .writeBoolean(field2)
+                .addIntAnnotation(field2AnnotationId, field2AnnotationValue)
+                .usePooledBuffer()
+                .build();
+        final long maxTimestamp = SystemClock.elapsedRealtimeNanos();
+
+        assertThat(statsEvent.getAtomId()).isEqualTo(expectedAtomId);
+
+        final ByteBuffer buffer =
+                ByteBuffer.wrap(statsEvent.getBytes()).order(ByteOrder.LITTLE_ENDIAN);
+
+        assertWithMessage("Root element in buffer is not TYPE_OBJECT")
+                .that(buffer.get()).isEqualTo(StatsEvent.TYPE_OBJECT);
+
+        assertWithMessage("Incorrect number of elements in root object")
+                .that(buffer.get()).isEqualTo(4);
+
+        assertWithMessage("First element is not timestamp")
+                .that(buffer.get()).isEqualTo(StatsEvent.TYPE_LONG);
+
+        assertWithMessage("Incorrect timestamp")
+                .that(buffer.getLong()).isIn(Range.closed(minTimestamp, maxTimestamp));
+
+        final byte atomIdHeader = buffer.get();
+        final int atomIdAnnotationValueCount = atomIdHeader >> 4;
+        final byte atomIdValueType = (byte) (atomIdHeader & 0x0F);
+        assertWithMessage("Second element is not atom id")
+                .that(atomIdValueType).isEqualTo(StatsEvent.TYPE_INT);
+        assertWithMessage("Atom id annotation count is wrong")
+                .that(atomIdAnnotationValueCount).isEqualTo(1);
+        assertWithMessage("Incorrect atom id")
+                .that(buffer.getInt()).isEqualTo(expectedAtomId);
+        assertWithMessage("Atom id's annotation id is wrong")
+                .that(buffer.get()).isEqualTo(atomAnnotationId);
+        assertWithMessage("Atom id's annotation type is wrong")
+                .that(buffer.get()).isEqualTo(StatsEvent.TYPE_INT);
+        assertWithMessage("Atom id's annotation value is wrong")
+                .that(buffer.getInt()).isEqualTo(atomAnnotationValue);
+
+        final byte field1Header = buffer.get();
+        final int field1AnnotationValueCount = field1Header >> 4;
+        final byte field1Type = (byte) (field1Header & 0x0F);
+        assertWithMessage("First field is not Int")
+                .that(field1Type).isEqualTo(StatsEvent.TYPE_INT);
+        assertWithMessage("First field annotation count is wrong")
+                .that(field1AnnotationValueCount).isEqualTo(1);
+        assertWithMessage("Incorrect field 1")
+                .that(buffer.getInt()).isEqualTo(field1);
+        assertWithMessage("First field's annotation id is wrong")
+                .that(buffer.get()).isEqualTo(field1AnnotationId);
+        assertWithMessage("First field's annotation type is wrong")
+                .that(buffer.get()).isEqualTo(StatsEvent.TYPE_BOOLEAN);
+        assertWithMessage("First field's annotation value is wrong")
+                .that(buffer.get()).isEqualTo(field1AnnotationValue ? 1 : 0);
+
+        final byte field2Header = buffer.get();
+        final int field2AnnotationValueCount = field2Header >> 4;
+        final byte field2Type = (byte) (field2Header & 0x0F);
+        assertWithMessage("Second field is not boolean")
+                .that(field2Type).isEqualTo(StatsEvent.TYPE_BOOLEAN);
+        assertWithMessage("Second field annotation count is wrong")
+                .that(field2AnnotationValueCount).isEqualTo(1);
+        assertWithMessage("Incorrect field 2")
+                .that(buffer.get()).isEqualTo(field2 ? 1 : 0);
+        assertWithMessage("Second field's annotation id is wrong")
+                .that(buffer.get()).isEqualTo(field2AnnotationId);
+        assertWithMessage("Second field's annotation type is wrong")
+                .that(buffer.get()).isEqualTo(StatsEvent.TYPE_INT);
+        assertWithMessage("Second field's annotation value is wrong")
+                .that(buffer.getInt()).isEqualTo(field2AnnotationValue);
+
+        assertThat(statsEvent.getNumBytes()).isEqualTo(buffer.position());
+
+        statsEvent.release();
+    }
+
+    @Test
+    public void testSetAtomIdNotCalledImmediately() {
+        final int expectedAtomId = 109;
+        final int field1 = 25;
+        final boolean field2 = true;
+
+        final long minTimestamp = SystemClock.elapsedRealtimeNanos();
+        final StatsEvent statsEvent = StatsEvent.newBuilder()
+                .writeInt(field1)
+                .setAtomId(expectedAtomId)
+                .writeBoolean(field2)
+                .usePooledBuffer()
+                .build();
+        final long maxTimestamp = SystemClock.elapsedRealtimeNanos();
+
+        assertThat(statsEvent.getAtomId()).isEqualTo(expectedAtomId);
+
+        final ByteBuffer buffer =
+                ByteBuffer.wrap(statsEvent.getBytes()).order(ByteOrder.LITTLE_ENDIAN);
+
+        assertWithMessage("Root element in buffer is not TYPE_OBJECT")
+                .that(buffer.get()).isEqualTo(StatsEvent.TYPE_OBJECT);
+
+        assertWithMessage("Incorrect number of elements in root object")
+                .that(buffer.get()).isEqualTo(3);
+
+        assertWithMessage("First element is not timestamp")
+                .that(buffer.get()).isEqualTo(StatsEvent.TYPE_LONG);
+
+        assertWithMessage("Incorrect timestamp")
+                .that(buffer.getLong()).isIn(Range.closed(minTimestamp, maxTimestamp));
+
+        assertWithMessage("Second element is not atom id")
+                .that(buffer.get()).isEqualTo(StatsEvent.TYPE_INT);
+
+        assertWithMessage("Incorrect atom id")
+                .that(buffer.getInt()).isEqualTo(expectedAtomId);
+
+        assertWithMessage("Third element is not errors type")
+                .that(buffer.get()).isEqualTo(StatsEvent.TYPE_ERRORS);
+
+        final int errorMask = buffer.getInt();
+
+        assertWithMessage("ERROR_ATOM_ID_INVALID_POSITION should be the only error in the mask")
+                .that(errorMask).isEqualTo(StatsEvent.ERROR_ATOM_ID_INVALID_POSITION);
+
+        assertThat(statsEvent.getNumBytes()).isEqualTo(buffer.position());
+
+        statsEvent.release();
+    }
+
     private static byte[] getByteArrayFromByteBuffer(final ByteBuffer buffer) {
         final int numBytes = buffer.getInt();
         byte[] bytes = new byte[numBytes];