diff --git a/core/java/com/android/internal/os/LongMultiStateCounter.java b/core/java/com/android/internal/os/LongMultiStateCounter.java
new file mode 100644
index 0000000000000..518d9ab42d21c
--- /dev/null
+++ b/core/java/com/android/internal/os/LongMultiStateCounter.java
@@ -0,0 +1,210 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * 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.internal.os;
+
+import android.os.Parcel;
+import android.os.Parcelable;
+
+import com.android.internal.util.Preconditions;
+
+import dalvik.annotation.optimization.CriticalNative;
+import dalvik.annotation.optimization.FastNative;
+
+import libcore.util.NativeAllocationRegistry;
+
+/**
+ * Performs per-state counting of long integers over time. The tracked "value" is expected
+ * to increase monotonously. The counter keeps track of the current state. When the
+ * updateValue method is called, the delta from the previous invocation of this method
+ * and the new value is added to the counter corresponding to the current state. If the
+ * state changed in the interim, the delta is distributed proptionally.
+ *
+ * The class's behavior is illustrated by this example:
+ *
+ * // At 0 ms, the state of the tracked object is 0 and the initial tracked value is 100
+ * counter.setState(0, 0);
+ * counter.updateValue(100, 0);
+ *
+ * // At 1000 ms, the state changes to 1
+ * counter.setState(1, 1000);
+ *
+ * // At 3000 ms, the tracked value is updated to 130
+ * counter.updateValue(130, 3000);
+ *
+ * // The delta (130 - 100 = 30) is distributed between states 0 and 1 according to the time
+ * // spent in those respective states; in this specific case, 1000 and 2000 ms.
+ * long countForState0 == counter.getCount(0); // 10
+ * long countForState1 == counter.getCount(1); // 20
+ *
+ *
+ * The tracked values are expected to increase monotonically.
+ *
+ * @hide
+ */
+public final class LongMultiStateCounter implements Parcelable {
+
+ private static final NativeAllocationRegistry sRegistry =
+ NativeAllocationRegistry.createMalloced(
+ LongMultiStateCounter.class.getClassLoader(), native_getReleaseFunc());
+
+ private final int mStateCount;
+
+ // Visible to other objects in this package so that it can be passed to @CriticalNative
+ // methods.
+ final long mNativeObject;
+
+ public LongMultiStateCounter(int stateCount) {
+ Preconditions.checkArgumentPositive(stateCount, "stateCount must be greater than 0");
+ mStateCount = stateCount;
+ mNativeObject = native_init(stateCount);
+ sRegistry.registerNativeAllocation(this, mNativeObject);
+ }
+
+ private LongMultiStateCounter(Parcel in) {
+ mNativeObject = native_initFromParcel(in);
+ sRegistry.registerNativeAllocation(this, mNativeObject);
+
+ mStateCount = native_getStateCount(mNativeObject);
+ }
+
+ public int getStateCount() {
+ return mStateCount;
+ }
+
+ /**
+ * Enables or disables the counter. When the counter is disabled, it does not
+ * accumulate counts supplied by the {@link #updateValue} method.
+ */
+ public void setEnabled(boolean enabled, long timestampMs) {
+ native_setEnabled(mNativeObject, enabled, timestampMs);
+ }
+
+ /**
+ * Sets the current state to the supplied value.
+ *
+ * @param state The new state
+ * @param timestampMs The time when the state change occurred, e.g.
+ * SystemClock.elapsedRealtime()
+ */
+ public void setState(int state, long timestampMs) {
+ if (state < 0 || state >= mStateCount) {
+ throw new IllegalArgumentException(
+ "State: " + state + ", outside the range: [0-" + (mStateCount - 1) + "]");
+ }
+ native_setState(mNativeObject, state, timestampMs);
+ }
+
+ /**
+ * Sets the new values. The delta between the previously set values and these values
+ * is distributed among the state according to the time the object spent in those states
+ * since the previous call to updateValues.
+ */
+ public void updateValue(long value, long timestampMs) {
+ native_updateValue(mNativeObject, value, timestampMs);
+ }
+
+ /**
+ * Adds the supplied values to the current accumulated values in the counter.
+ */
+ public void addCount(long count) {
+ native_addCount(mNativeObject, count);
+ }
+
+ /**
+ * Resets the accumulated counts to 0.
+ */
+ public void reset() {
+ native_reset(mNativeObject);
+ }
+
+ /**
+ * Returns the accumulated count for the specified state.
+ */
+ public long getCount(int state) {
+ if (state < 0 || state >= mStateCount) {
+ throw new IllegalArgumentException(
+ "State: " + state + ", outside the range: [0-" + mStateCount + "]");
+ }
+ return native_getCount(mNativeObject, state);
+ }
+
+ @Override
+ public String toString() {
+ return native_toString(mNativeObject);
+ }
+
+ @Override
+ public void writeToParcel(Parcel dest, int flags) {
+ native_writeToParcel(mNativeObject, dest, flags);
+ }
+
+ @Override
+ public int describeContents() {
+ return 0;
+ }
+
+ public static final Creator CREATOR =
+ new Creator() {
+ @Override
+ public LongMultiStateCounter createFromParcel(Parcel in) {
+ return new LongMultiStateCounter(in);
+ }
+
+ @Override
+ public LongMultiStateCounter[] newArray(int size) {
+ return new LongMultiStateCounter[size];
+ }
+ };
+
+
+ @CriticalNative
+ private static native long native_init(int stateCount);
+
+ @CriticalNative
+ private static native long native_getReleaseFunc();
+
+ @CriticalNative
+ private static native void native_setEnabled(long nativeObject, boolean enabled,
+ long timestampMs);
+
+ @CriticalNative
+ private static native void native_setState(long nativeObject, int state, long timestampMs);
+
+ @CriticalNative
+ private static native void native_updateValue(long nativeObject, long value, long timestampMs);
+
+ @CriticalNative
+ private static native void native_addCount(long nativeObject, long count);
+
+ @CriticalNative
+ private static native void native_reset(long nativeObject);
+
+ @CriticalNative
+ private static native long native_getCount(long nativeObject, int state);
+
+ @FastNative
+ private native String native_toString(long nativeObject);
+
+ @FastNative
+ private native void native_writeToParcel(long nativeObject, Parcel dest, int flags);
+
+ @FastNative
+ private static native long native_initFromParcel(Parcel parcel);
+
+ @CriticalNative
+ private static native int native_getStateCount(long nativeObject);
+}
diff --git a/core/jni/Android.bp b/core/jni/Android.bp
index 945adb7cd3ddb..1a1a8badf82b1 100644
--- a/core/jni/Android.bp
+++ b/core/jni/Android.bp
@@ -218,6 +218,7 @@ cc_library_shared {
"com_android_internal_os_KernelSingleProcessCpuThreadReader.cpp",
"com_android_internal_os_KernelSingleUidTimeReader.cpp",
"com_android_internal_os_LongArrayMultiStateCounter.cpp",
+ "com_android_internal_os_LongMultiStateCounter.cpp",
"com_android_internal_os_Zygote.cpp",
"com_android_internal_os_ZygoteCommandBuffer.cpp",
"com_android_internal_os_ZygoteInit.cpp",
diff --git a/core/jni/AndroidRuntime.cpp b/core/jni/AndroidRuntime.cpp
index 91179c2770120..5c9b6df0ddfa3 100644
--- a/core/jni/AndroidRuntime.cpp
+++ b/core/jni/AndroidRuntime.cpp
@@ -205,6 +205,7 @@ extern int register_com_android_internal_os_KernelCpuUidBpfMapReader(JNIEnv *env
extern int register_com_android_internal_os_KernelSingleProcessCpuThreadReader(JNIEnv* env);
extern int register_com_android_internal_os_KernelSingleUidTimeReader(JNIEnv *env);
extern int register_com_android_internal_os_LongArrayMultiStateCounter(JNIEnv* env);
+extern int register_com_android_internal_os_LongMultiStateCounter(JNIEnv* env);
extern int register_com_android_internal_os_Zygote(JNIEnv *env);
extern int register_com_android_internal_os_ZygoteCommandBuffer(JNIEnv *env);
extern int register_com_android_internal_os_ZygoteInit(JNIEnv *env);
@@ -1588,6 +1589,7 @@ static const RegJNIRec gRegJNI[] = {
REG_JNI(register_com_android_internal_net_NetworkUtilsInternal),
REG_JNI(register_com_android_internal_os_ClassLoaderFactory),
REG_JNI(register_com_android_internal_os_LongArrayMultiStateCounter),
+ REG_JNI(register_com_android_internal_os_LongMultiStateCounter),
REG_JNI(register_com_android_internal_os_Zygote),
REG_JNI(register_com_android_internal_os_ZygoteCommandBuffer),
REG_JNI(register_com_android_internal_os_ZygoteInit),
diff --git a/core/jni/com_android_internal_os_LongMultiStateCounter.cpp b/core/jni/com_android_internal_os_LongMultiStateCounter.cpp
new file mode 100644
index 0000000000000..9ffc7574ddb0e
--- /dev/null
+++ b/core/jni/com_android_internal_os_LongMultiStateCounter.cpp
@@ -0,0 +1,195 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * 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.
+ */
+
+#include
+#include
+#include
+#include
+
+#include
+
+#include "MultiStateCounter.h"
+#include "core_jni_helpers.h"
+
+namespace android {
+
+namespace battery {
+
+typedef battery::MultiStateCounter LongMultiStateCounter;
+
+template <>
+bool LongMultiStateCounter::delta(const int64_t &previousValue, const int64_t &newValue,
+ int64_t *outValue) const {
+ *outValue = newValue - previousValue;
+ return *outValue >= 0;
+}
+
+template <>
+void LongMultiStateCounter::add(int64_t *value1, const int64_t &value2, const uint64_t numerator,
+ const uint64_t denominator) const {
+ if (numerator != denominator) {
+ // The caller ensures that denominator != 0
+ *value1 += value2 * numerator / denominator;
+ } else {
+ *value1 += value2;
+ }
+}
+
+template <>
+std::string LongMultiStateCounter::valueToString(const int64_t &v) const {
+ return std::to_string(v);
+}
+
+} // namespace battery
+
+static inline battery::LongMultiStateCounter *asLongMultiStateCounter(const jlong nativePtr) {
+ return reinterpret_cast(nativePtr);
+}
+
+static jlong native_init(jint stateCount) {
+ battery::LongMultiStateCounter *counter = new battery::LongMultiStateCounter(stateCount, 0);
+ return reinterpret_cast(counter);
+}
+
+static void native_dispose(void *nativePtr) {
+ delete reinterpret_cast(nativePtr);
+}
+
+static jlong native_getReleaseFunc() {
+ return reinterpret_cast(native_dispose);
+}
+
+static void native_setEnabled(jlong nativePtr, jboolean enabled, jlong timestamp) {
+ asLongMultiStateCounter(nativePtr)->setEnabled(enabled, timestamp);
+}
+
+static void native_setState(jlong nativePtr, jint state, jlong timestamp) {
+ asLongMultiStateCounter(nativePtr)->setState(state, timestamp);
+}
+
+static void native_updateValue(jlong nativePtr, jlong value, jlong timestamp) {
+ asLongMultiStateCounter(nativePtr)->updateValue((int64_t)value, timestamp);
+}
+
+static void native_addCount(jlong nativePtr, jlong count) {
+ asLongMultiStateCounter(nativePtr)->addValue(count);
+}
+
+static void native_reset(jlong nativePtr) {
+ asLongMultiStateCounter(nativePtr)->reset();
+}
+
+static jlong native_getCount(jlong nativePtr, jint state) {
+ return asLongMultiStateCounter(nativePtr)->getCount(state);
+}
+
+static jobject native_toString(JNIEnv *env, jobject self, jlong nativePtr) {
+ return env->NewStringUTF(asLongMultiStateCounter(nativePtr)->toString().c_str());
+}
+
+static void throwWriteRE(JNIEnv *env, binder_status_t status) {
+ ALOGE("Could not write LongMultiStateCounter to Parcel, status = %d", status);
+ jniThrowRuntimeException(env, "Could not write LongMultiStateCounter to Parcel");
+}
+
+#define THROW_ON_WRITE_ERROR(expr) \
+ { \
+ binder_status_t status = expr; \
+ if (status != STATUS_OK) { \
+ throwWriteRE(env, status); \
+ } \
+ }
+
+static void native_writeToParcel(JNIEnv *env, jobject self, jlong nativePtr, jobject jParcel,
+ jint flags) {
+ battery::LongMultiStateCounter *counter = asLongMultiStateCounter(nativePtr);
+ AParcel *parcel = AParcel_fromJavaParcel(env, jParcel);
+
+ uint16_t stateCount = counter->getStateCount();
+ THROW_ON_WRITE_ERROR(AParcel_writeInt32(parcel, stateCount));
+
+ for (battery::state_t state = 0; state < stateCount; state++) {
+ THROW_ON_WRITE_ERROR(AParcel_writeInt64(parcel, counter->getCount(state)));
+ }
+}
+
+static void throwReadRE(JNIEnv *env, binder_status_t status) {
+ ALOGE("Could not read LongMultiStateCounter from Parcel, status = %d", status);
+ jniThrowRuntimeException(env, "Could not read LongMultiStateCounter from Parcel");
+}
+
+#define THROW_ON_READ_ERROR(expr) \
+ { \
+ binder_status_t status = expr; \
+ if (status != STATUS_OK) { \
+ throwReadRE(env, status); \
+ } \
+ }
+
+static jlong native_initFromParcel(JNIEnv *env, jclass theClass, jobject jParcel) {
+ AParcel *parcel = AParcel_fromJavaParcel(env, jParcel);
+
+ int32_t stateCount;
+ THROW_ON_READ_ERROR(AParcel_readInt32(parcel, &stateCount));
+
+ battery::LongMultiStateCounter *counter = new battery::LongMultiStateCounter(stateCount, 0);
+
+ for (battery::state_t state = 0; state < stateCount; state++) {
+ int64_t value;
+ THROW_ON_READ_ERROR(AParcel_readInt64(parcel, &value));
+ counter->setValue(state, value);
+ }
+
+ return reinterpret_cast(counter);
+}
+
+static jint native_getStateCount(jlong nativePtr) {
+ return asLongMultiStateCounter(nativePtr)->getStateCount();
+}
+
+static const JNINativeMethod g_methods[] = {
+ // @CriticalNative
+ {"native_init", "(I)J", (void *)native_init},
+ // @CriticalNative
+ {"native_getReleaseFunc", "()J", (void *)native_getReleaseFunc},
+ // @CriticalNative
+ {"native_setEnabled", "(JZJ)V", (void *)native_setEnabled},
+ // @CriticalNative
+ {"native_setState", "(JIJ)V", (void *)native_setState},
+ // @CriticalNative
+ {"native_updateValue", "(JJJ)V", (void *)native_updateValue},
+ // @CriticalNative
+ {"native_addCount", "(JJ)V", (void *)native_addCount},
+ // @CriticalNative
+ {"native_reset", "(J)V", (void *)native_reset},
+ // @CriticalNative
+ {"native_getCount", "(JI)J", (void *)native_getCount},
+ // @FastNative
+ {"native_toString", "(J)Ljava/lang/String;", (void *)native_toString},
+ // @FastNative
+ {"native_writeToParcel", "(JLandroid/os/Parcel;I)V", (void *)native_writeToParcel},
+ // @FastNative
+ {"native_initFromParcel", "(Landroid/os/Parcel;)J", (void *)native_initFromParcel},
+ // @CriticalNative
+ {"native_getStateCount", "(J)I", (void *)native_getStateCount},
+};
+
+int register_com_android_internal_os_LongMultiStateCounter(JNIEnv *env) {
+ return RegisterMethodsOrDie(env, "com/android/internal/os/LongMultiStateCounter", g_methods,
+ NELEM(g_methods));
+}
+
+} // namespace android
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryStatsTests.java b/core/tests/coretests/src/com/android/internal/os/BatteryStatsTests.java
index 90a9572b55608..92c2d43ea2a35 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryStatsTests.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryStatsTests.java
@@ -64,6 +64,8 @@ import org.junit.runners.Suite;
KernelSingleProcessCpuThreadReaderTest.class,
KernelSingleUidTimeReaderTest.class,
KernelWakelockReaderTest.class,
+ LongArrayMultiStateCounterTest.class,
+ LongMultiStateCounterTest.class,
LongSamplingCounterTest.class,
LongSamplingCounterArrayTest.class,
MobileRadioPowerCalculatorTest.class,
diff --git a/core/tests/coretests/src/com/android/internal/os/LongMultiStateCounterTest.java b/core/tests/coretests/src/com/android/internal/os/LongMultiStateCounterTest.java
new file mode 100644
index 0000000000000..8540adb96171b
--- /dev/null
+++ b/core/tests/coretests/src/com/android/internal/os/LongMultiStateCounterTest.java
@@ -0,0 +1,134 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * 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.internal.os;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.os.Parcel;
+
+import androidx.test.filters.SmallTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+public class LongMultiStateCounterTest {
+
+ @Test
+ public void setStateAndUpdateValue() {
+ LongMultiStateCounter counter = new LongMultiStateCounter(2);
+
+ counter.updateValue(0, 1000);
+ counter.setState(0, 1000);
+ counter.setState(1, 2000);
+ counter.setState(0, 4000);
+ counter.updateValue(100, 9000);
+
+ assertThat(counter.getCount(0)).isEqualTo(75);
+ assertThat(counter.getCount(1)).isEqualTo(25);
+
+ assertThat(counter.toString()).isEqualTo("[0: 75, 1: 25] updated: 9000 currentState: 0");
+ }
+
+ @Test
+ public void setEnabled() {
+ LongMultiStateCounter counter = new LongMultiStateCounter(2);
+ counter.setState(0, 1000);
+ counter.updateValue(0, 1000);
+ counter.updateValue(100, 2000);
+
+ assertThat(counter.getCount(0)).isEqualTo(100);
+
+ counter.setEnabled(false, 3000);
+
+ // Partially included, because the counter is disabled after the previous update
+ counter.updateValue(200, 4000);
+
+ // Count only 50%, because the counter was disabled for 50% of the time
+ assertThat(counter.getCount(0)).isEqualTo(150);
+
+ // Not counted because the counter is disabled
+ counter.updateValue(250, 5000);
+
+ counter.setEnabled(true, 6000);
+
+ counter.updateValue(300, 7000);
+
+ // Again, take 50% of the delta
+ assertThat(counter.getCount(0)).isEqualTo(175);
+ }
+
+ @Test
+ public void reset() {
+ LongMultiStateCounter counter = new LongMultiStateCounter(2);
+ counter.setState(0, 1000);
+ counter.updateValue(0, 1000);
+ counter.updateValue(100, 2000);
+
+ assertThat(counter.getCount(0)).isEqualTo(100);
+
+ counter.reset();
+
+ assertThat(counter.getCount(0)).isEqualTo(0);
+
+ counter.updateValue(200, 3000);
+ counter.updateValue(300, 4000);
+
+ assertThat(counter.getCount(0)).isEqualTo(100);
+ }
+
+ @Test
+ public void parceling() {
+ LongMultiStateCounter counter = new LongMultiStateCounter(2);
+ counter.updateValue(0, 1000);
+ counter.setState(0, 1000);
+ counter.updateValue(100, 2000);
+ counter.setState(1, 2000);
+ counter.updateValue(101, 3000);
+
+ assertThat(counter.getCount(0)).isEqualTo(100);
+ assertThat(counter.getCount(1)).isEqualTo(1);
+
+ Parcel parcel = Parcel.obtain();
+ counter.writeToParcel(parcel, 0);
+ byte[] bytes = parcel.marshall();
+ parcel.recycle();
+
+ parcel = Parcel.obtain();
+ parcel.unmarshall(bytes, 0, bytes.length);
+ parcel.setDataPosition(0);
+
+ LongMultiStateCounter newCounter = LongMultiStateCounter.CREATOR.createFromParcel(parcel);
+ parcel.recycle();
+
+ assertThat(newCounter.getCount(0)).isEqualTo(100);
+ assertThat(newCounter.getCount(1)).isEqualTo(1);
+
+ // ==== Verify that the counter keeps accumulating after unparceling.
+
+ // State, last update timestamp and current counts are undefined at this point.
+ newCounter.setState(0, 100);
+ newCounter.updateValue(300, 100);
+
+ // A new base state and counters are established; we can continue accumulating deltas
+ newCounter.updateValue(316, 200);
+
+ assertThat(newCounter.getCount(0)).isEqualTo(116);
+ }
+}