Add support for createHintSessionWithConfig

This patch adds support for the new "createHintSessionWithConfig" call
which supports session tagging and session IDs, and updates the
performance_hint NDK implementation to use it.

Bug: 330553312
Bug: 315894228
Test: atest PerformanceHintNativeTestCases
Test: atest HintManagerServiceTest
Change-Id: Idc1c16be6bb53983cc3ec102596b9613935f7301
diff --git a/core/java/android/os/IHintManager.aidl b/core/java/android/os/IHintManager.aidl
index d97ea54..e057a85 100644
--- a/core/java/android/os/IHintManager.aidl
+++ b/core/java/android/os/IHintManager.aidl
@@ -18,16 +18,21 @@
 package android.os;
 
 import android.os.IHintSession;
+import android.hardware.power.SessionConfig;
+import android.hardware.power.SessionTag;
 
 /** {@hide} */
 interface IHintManager {
     /**
      * Creates a {@link Session} for the given set of threads and associates to a binder token.
+     * Returns a config if creation is not supported, and HMS had to use the
+     * legacy creation method.
      */
-    IHintSession createHintSession(in IBinder token, in int[] tids, long durationNanos);
+    IHintSession createHintSessionWithConfig(in IBinder token, in int[] threadIds,
+            in long durationNanos, in SessionTag tag, out @nullable SessionConfig config);
 
     /**
-     * Get preferred rate limit in nano second.
+     * Get preferred rate limit in nanoseconds.
      */
     long getHintSessionPreferredRate();
 
diff --git a/native/android/performance_hint.cpp b/native/android/performance_hint.cpp
index 882afca..fbb35e2 100644
--- a/native/android/performance_hint.cpp
+++ b/native/android/performance_hint.cpp
@@ -59,7 +59,8 @@
     ~APerformanceHintManager() = default;
 
     APerformanceHintSession* createSession(const int32_t* threadIds, size_t size,
-                                           int64_t initialTargetWorkDurationNanos);
+                                           int64_t initialTargetWorkDurationNanos,
+                                           hal::SessionTag tag = hal::SessionTag::OTHER);
     int64_t getPreferredRateNanos() const;
 
 private:
@@ -84,7 +85,8 @@
 public:
     APerformanceHintSession(std::shared_ptr<IHintManager> hintManager,
                             std::shared_ptr<IHintSession> session, int64_t preferredRateNanos,
-                            int64_t targetDurationNanos);
+                            int64_t targetDurationNanos,
+                            std::optional<hal::SessionConfig> sessionConfig);
     APerformanceHintSession() = delete;
     ~APerformanceHintSession();
 
@@ -116,9 +118,10 @@
     // Cached samples
     std::vector<hal::WorkDuration> mActualWorkDurations;
     std::string mSessionName;
-    static int32_t sIDCounter;
+    static int64_t sIDCounter;
     // The most recent set of thread IDs
     std::vector<int32_t> mLastThreadIDs;
+    std::optional<hal::SessionConfig> mSessionConfig;
     // Tracing helpers
     void traceThreads(std::vector<int32_t>& tids);
     void tracePowerEfficient(bool powerEfficient);
@@ -129,7 +132,8 @@
 
 static std::shared_ptr<IHintManager>* gIHintManagerForTesting = nullptr;
 static APerformanceHintManager* gHintManagerForTesting = nullptr;
-int32_t APerformanceHintSession::sIDCounter = 0;
+// Start above the int32 range so we don't collide with config sessions
+int64_t APerformanceHintSession::sIDCounter = INT32_MAX;
 
 // ===================================== APerformanceHintManager implementation
 APerformanceHintManager::APerformanceHintManager(std::shared_ptr<IHintManager> manager,
@@ -174,16 +178,20 @@
 }
 
 APerformanceHintSession* APerformanceHintManager::createSession(
-        const int32_t* threadIds, size_t size, int64_t initialTargetWorkDurationNanos) {
+        const int32_t* threadIds, size_t size, int64_t initialTargetWorkDurationNanos,
+        hal::SessionTag tag) {
     std::vector<int32_t> tids(threadIds, threadIds + size);
     std::shared_ptr<IHintSession> session;
-    ndk::ScopedAStatus ret =
-            mHintManager->createHintSession(mToken, tids, initialTargetWorkDurationNanos, &session);
+    ndk::ScopedAStatus ret;
+    std::optional<hal::SessionConfig> sessionConfig;
+    ret = mHintManager->createHintSessionWithConfig(mToken, tids, initialTargetWorkDurationNanos,
+                                                    tag, &sessionConfig, &session);
+
     if (!ret.isOk() || !session) {
         return nullptr;
     }
     auto out = new APerformanceHintSession(mHintManager, std::move(session), mPreferredRateNanos,
-                                           initialTargetWorkDurationNanos);
+                                           initialTargetWorkDurationNanos, sessionConfig);
     out->traceThreads(tids);
     out->traceTargetDuration(initialTargetWorkDurationNanos);
     out->tracePowerEfficient(false);
@@ -199,19 +207,23 @@
 APerformanceHintSession::APerformanceHintSession(std::shared_ptr<IHintManager> hintManager,
                                                  std::shared_ptr<IHintSession> session,
                                                  int64_t preferredRateNanos,
-                                                 int64_t targetDurationNanos)
+                                                 int64_t targetDurationNanos,
+                                                 std::optional<hal::SessionConfig> sessionConfig)
       : mHintManager(hintManager),
         mHintSession(std::move(session)),
         mPreferredRateNanos(preferredRateNanos),
         mTargetDurationNanos(targetDurationNanos),
         mFirstTargetMetTimestamp(0),
-        mLastTargetMetTimestamp(0) {
-    const std::vector<hal::SessionHint> sessionHintRange{ndk::enum_range<hal::SessionHint>()
-                                                                 .begin(),
-                                                         ndk::enum_range<hal::SessionHint>().end()};
-
-    mLastHintSentTimestamp = std::vector<int64_t>(sessionHintRange.size(), 0);
-    mSessionName = android::base::StringPrintf("ADPF Session %" PRId32, ++sIDCounter);
+        mLastTargetMetTimestamp(0),
+        mSessionConfig(sessionConfig) {
+    if (sessionConfig->id > INT32_MAX) {
+        ALOGE("Session ID too large, must fit 32-bit integer");
+    }
+    constexpr int numEnums =
+            ndk::enum_range<hal::SessionHint>().end() - ndk::enum_range<hal::SessionHint>().begin();
+    mLastHintSentTimestamp = std::vector<int64_t>(numEnums, 0);
+    int64_t traceId = sessionConfig.has_value() ? sessionConfig->id : ++sIDCounter;
+    mSessionName = android::base::StringPrintf("ADPF Session %" PRId64, traceId);
 }
 
 APerformanceHintSession::~APerformanceHintSession() {
diff --git a/native/android/tests/performance_hint/PerformanceHintNativeTest.cpp b/native/android/tests/performance_hint/PerformanceHintNativeTest.cpp
index bfbe34e..974e6e6 100644
--- a/native/android/tests/performance_hint/PerformanceHintNativeTest.cpp
+++ b/native/android/tests/performance_hint/PerformanceHintNativeTest.cpp
@@ -16,6 +16,8 @@
 
 #define LOG_TAG "PerformanceHintNativeTest"
 
+#include <aidl/android/hardware/power/SessionConfig.h>
+#include <aidl/android/hardware/power/SessionTag.h>
 #include <aidl/android/hardware/power/WorkDuration.h>
 #include <aidl/android/os/IHintManager.h>
 #include <android/binder_manager.h>
@@ -28,6 +30,8 @@
 #include <memory>
 #include <vector>
 
+using aidl::android::hardware::power::SessionConfig;
+using aidl::android::hardware::power::SessionTag;
 using aidl::android::hardware::power::WorkDuration;
 using aidl::android::os::IHintManager;
 using aidl::android::os::IHintSession;
@@ -39,8 +43,9 @@
 
 class MockIHintManager : public IHintManager {
 public:
-    MOCK_METHOD(ScopedAStatus, createHintSession,
+    MOCK_METHOD(ScopedAStatus, createHintSessionWithConfig,
                 (const SpAIBinder& token, const ::std::vector<int32_t>& tids, int64_t durationNanos,
+                 SessionTag tag, std::optional<SessionConfig>* config,
                  std::shared_ptr<IHintSession>* _aidl_return),
                 (override));
     MOCK_METHOD(ScopedAStatus, getHintSessionPreferredRate, (int64_t * _aidl_return), (override));
@@ -92,14 +97,18 @@
     APerformanceHintSession* createSession(APerformanceHintManager* manager,
                                            int64_t targetDuration = 56789L) {
         mMockSession = ndk::SharedRefBase::make<NiceMock<MockIHintSession>>();
-
+        int64_t sessionId = 123;
         std::vector<int32_t> tids;
         tids.push_back(1);
         tids.push_back(2);
 
-        ON_CALL(*mMockIHintManager, createHintSession(_, Eq(tids), Eq(targetDuration), _))
-                .WillByDefault(DoAll(SetArgPointee<3>(std::shared_ptr<IHintSession>(mMockSession)),
+        ON_CALL(*mMockIHintManager,
+                createHintSessionWithConfig(_, Eq(tids), Eq(targetDuration), _, _, _))
+                .WillByDefault(DoAll(SetArgPointee<4>(
+                                             std::make_optional<SessionConfig>({.id = sessionId})),
+                                     SetArgPointee<5>(std::shared_ptr<IHintSession>(mMockSession)),
                                      [] { return ScopedAStatus::ok(); }));
+
         ON_CALL(*mMockIHintManager, setHintSessionThreads(_, _)).WillByDefault([] {
             return ScopedAStatus::ok();
         });
@@ -115,7 +124,6 @@
         ON_CALL(*mMockSession, reportActualWorkDuration2(_)).WillByDefault([] {
             return ScopedAStatus::ok();
         });
-
         return APerformanceHint_createSession(manager, tids.data(), tids.size(), targetDuration);
     }
 
@@ -178,6 +186,14 @@
     APerformanceHint_closeSession(session);
 }
 
+TEST_F(PerformanceHintTest, TestUpdatedSessionCreation) {
+    EXPECT_CALL(*mMockIHintManager, createHintSessionWithConfig(_, _, _, _, _, _)).Times(1);
+    APerformanceHintManager* manager = createManager();
+    APerformanceHintSession* session = createSession(manager);
+    ASSERT_TRUE(session);
+    APerformanceHint_closeSession(session);
+}
+
 TEST_F(PerformanceHintTest, SetThreads) {
     APerformanceHintManager* manager = createManager();
 
diff --git a/services/core/java/com/android/server/power/hint/HintManagerService.java b/services/core/java/com/android/server/power/hint/HintManagerService.java
index 101983e..6e17fd3 100644
--- a/services/core/java/com/android/server/power/hint/HintManagerService.java
+++ b/services/core/java/com/android/server/power/hint/HintManagerService.java
@@ -20,11 +20,14 @@
 import static com.android.server.power.hint.Flags.powerhintThreadCleanup;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.app.ActivityManager;
 import android.app.ActivityManagerInternal;
 import android.app.StatsManager;
 import android.app.UidObserver;
 import android.content.Context;
+import android.hardware.power.SessionConfig;
+import android.hardware.power.SessionTag;
 import android.hardware.power.WorkDuration;
 import android.os.Binder;
 import android.os.Handler;
@@ -67,6 +70,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /** An hint service implementation that runs in System Server process. */
 public final class HintManagerService extends SystemService {
@@ -87,7 +91,7 @@
     @GuardedBy("mLock")
     private final ArrayMap<Integer, ArrayMap<IBinder, ArraySet<AppHintSession>>> mActiveSessions;
 
-    /** Lock to protect HAL handles and listen list. */
+    /** Lock to protect mActiveSessions. */
     private final Object mLock = new Object();
 
     @GuardedBy("mNonIsolatedTidsLock")
@@ -104,6 +108,8 @@
 
     private final Context mContext;
 
+    private AtomicBoolean mConfigCreationSupport = new AtomicBoolean(true);
+
     private static final String PROPERTY_SF_ENABLE_CPU_HINT = "debug.sf.enable_adpf_cpu_hint";
     private static final String PROPERTY_HWUI_ENABLE_HINT_MANAGER = "debug.hwui.use_hint_manager";
 
@@ -217,6 +223,9 @@
         private static native long nativeCreateHintSession(int tgid, int uid, int[] tids,
                 long durationNanos);
 
+        private static native long nativeCreateHintSessionWithConfig(int tgid, int uid, int[] tids,
+                long durationNanos, int tag, SessionConfig config);
+
         private static native void nativePauseHintSession(long halPtr);
 
         private static native void nativeResumeHintSession(long halPtr);
@@ -253,6 +262,12 @@
             return nativeCreateHintSession(tgid, uid, tids, durationNanos);
         }
 
+        /** Wrapper for HintManager.nativeCreateHintSessionWithConfig */
+        public long halCreateHintSessionWithConfig(
+                int tgid, int uid, int[] tids, long durationNanos, int tag, SessionConfig config) {
+            return nativeCreateHintSessionWithConfig(tgid, uid, tids, durationNanos, tag, config);
+        }
+
         /** Wrapper for HintManager.nativePauseHintSession */
         public void halPauseHintSession(long halPtr) {
             nativePauseHintSession(halPtr);
@@ -612,8 +627,12 @@
     @VisibleForTesting
     final class BinderService extends IHintManager.Stub {
         @Override
-        public IHintSession createHintSession(IBinder token, int[] tids, long durationNanos) {
-            if (!isHalSupported()) return null;
+        public IHintSession createHintSessionWithConfig(@NonNull IBinder token,
+                @NonNull int[] tids, long durationNanos, @SessionTag int tag,
+                @Nullable SessionConfig config) {
+            if (!isHalSupported()) {
+                throw new UnsupportedOperationException("PowerHAL is not supported!");
+            }
 
             java.util.Objects.requireNonNull(token);
             java.util.Objects.requireNonNull(tids);
@@ -634,8 +653,35 @@
                     throw new SecurityException(errMsg);
                 }
 
-                long halSessionPtr = mNativeWrapper.halCreateHintSession(callingTgid, callingUid,
-                        tids, durationNanos);
+                Long halSessionPtr = null;
+                if (mConfigCreationSupport.get()) {
+                    try {
+                        halSessionPtr = mNativeWrapper.halCreateHintSessionWithConfig(
+                                callingTgid, callingUid, tids, durationNanos, tag, config);
+                    } catch (UnsupportedOperationException e) {
+                        mConfigCreationSupport.set(false);
+                    } catch (IllegalStateException e) {
+                        Slog.e("createHintSessionWithConfig failed: ", e.getMessage());
+                        throw new IllegalStateException(
+                            "createHintSessionWithConfig failed: " + e.getMessage());
+                    }
+                }
+
+                if (halSessionPtr == null) {
+                    try {
+                        halSessionPtr = mNativeWrapper.halCreateHintSession(callingTgid,
+                                callingUid, tids, durationNanos);
+                    } catch (UnsupportedOperationException e) {
+                        Slog.w("createHintSession unsupported: ", e.getMessage());
+                        throw new UnsupportedOperationException(
+                            "createHintSession unsupported: " + e.getMessage());
+                    } catch (IllegalStateException e) {
+                        Slog.e("createHintSession failed: ", e.getMessage());
+                        throw new IllegalStateException(
+                            "createHintSession failed: " + e.getMessage());
+                    }
+                }
+
                 if (powerhintThreadCleanup()) {
                     synchronized (mNonIsolatedTidsLock) {
                         for (int i = nonIsolated.size() - 1; i >= 0; i--) {
@@ -644,9 +690,6 @@
                         }
                     }
                 }
-                if (halSessionPtr == 0) {
-                    return null;
-                }
 
                 AppHintSession hs = new AppHintSession(callingUid, callingTgid, tids, token,
                         halSessionPtr, durationNanos);
diff --git a/services/core/jni/com_android_server_hint_HintManagerService.cpp b/services/core/jni/com_android_server_hint_HintManagerService.cpp
index 5b8ef19..be18835 100644
--- a/services/core/jni/com_android_server_hint_HintManagerService.cpp
+++ b/services/core/jni/com_android_server_hint_HintManagerService.cpp
@@ -34,13 +34,13 @@
 
 #include "jni.h"
 
+using aidl::android::hardware::power::SessionConfig;
 using aidl::android::hardware::power::SessionHint;
 using aidl::android::hardware::power::SessionMode;
+using aidl::android::hardware::power::SessionTag;
 using aidl::android::hardware::power::WorkDuration;
 using android::power::PowerHintSessionWrapper;
 
-using android::base::StringPrintf;
-
 namespace android {
 
 static struct {
@@ -66,6 +66,15 @@
     return rate;
 }
 
+void throwUnsupported(JNIEnv* env, const char* msg) {
+    env->ThrowNew(env->FindClass("java/lang/UnsupportedOperationException"), msg);
+}
+
+void throwFailed(JNIEnv* env, const char* msg) {
+    // We throw IllegalStateException for all errors other than the "unsupported" ones
+    env->ThrowNew(env->FindClass("java/lang/IllegalStateException"), msg);
+}
+
 static jlong createHintSession(JNIEnv* env, int32_t tgid, int32_t uid,
                                std::vector<int32_t>& threadIds, int64_t durationNanos) {
     auto result = gPowerHalController.createHintSession(tgid, uid, threadIds, durationNanos);
@@ -76,10 +85,38 @@
         return res.second ? session_ptr : 0;
     } else if (result.isFailed()) {
         ALOGW("createHintSession failed with message: %s", result.errorMessage());
+        throwFailed(env, result.errorMessage());
+    } else if (result.isUnsupported()) {
+        throwUnsupported(env, result.errorMessage());
+        return -1;
     }
     return 0;
 }
 
+static jlong createHintSessionWithConfig(JNIEnv* env, int32_t tgid, int32_t uid,
+                                         std::vector<int32_t> threadIds, int64_t durationNanos,
+                                         int32_t sessionTag, SessionConfig& config) {
+    auto result =
+            gPowerHalController.createHintSessionWithConfig(tgid, uid, threadIds, durationNanos,
+                                                            static_cast<SessionTag>(sessionTag),
+                                                            &config);
+    if (result.isOk()) {
+        jlong session_ptr = reinterpret_cast<jlong>(result.value().get());
+        std::scoped_lock sessionLock(gSessionMapLock);
+        auto res = gSessionMap.insert({session_ptr, result.value()});
+        if (!res.second) {
+            throwFailed(env, "PowerHAL provided an invalid session");
+            return 0;
+        }
+        return session_ptr;
+    } else if (result.isUnsupported()) {
+        throwUnsupported(env, result.errorMessage());
+        return -1;
+    }
+    throwFailed(env, result.errorMessage());
+    return 0;
+}
+
 static void pauseHintSession(JNIEnv* env, int64_t session_ptr) {
     auto appSession = reinterpret_cast<PowerHintSessionWrapper*>(session_ptr);
     appSession->pause();
@@ -136,13 +173,34 @@
                                      jintArray tids, jlong durationNanos) {
     ScopedIntArrayRO tidArray(env, tids);
     if (nullptr == tidArray.get() || tidArray.size() == 0) {
-        ALOGW("GetIntArrayElements returns nullptr.");
+        ALOGW("nativeCreateHintSession: GetIntArrayElements returns nullptr.");
         return 0;
     }
     std::vector<int32_t> threadIds(tidArray.get(), tidArray.get() + tidArray.size());
     return createHintSession(env, tgid, uid, threadIds, durationNanos);
 }
 
+static jlong nativeCreateHintSessionWithConfig(JNIEnv* env, jclass /* clazz */, jint tgid, jint uid,
+                                               jintArray tids, jlong durationNanos, jint sessionTag,
+                                               jobject sessionConfig) {
+    ScopedIntArrayRO tidArray(env, tids);
+    if (nullptr == tidArray.get() || tidArray.size() == 0) {
+        ALOGW("nativeCreateHintSessionWithConfig: GetIntArrayElements returns nullptr.");
+        return 0;
+    }
+    std::vector<int32_t> threadIds(tidArray.get(), tidArray.get() + tidArray.size());
+    SessionConfig config;
+    jlong out = createHintSessionWithConfig(env, tgid, uid, std::move(threadIds), durationNanos,
+                                            sessionTag, config);
+    if (out <= 0) {
+        return out;
+    }
+    static jclass configClass = env->FindClass("android/hardware/power/SessionConfig");
+    static jfieldID fid = env->GetFieldID(configClass, "id", "J");
+    env->SetLongField(sessionConfig, fid, config.id);
+    return out;
+}
+
 static void nativePauseHintSession(JNIEnv* env, jclass /* clazz */, jlong session_ptr) {
     pauseHintSession(env, session_ptr);
 }
@@ -215,6 +273,8 @@
         {"nativeInit", "()V", (void*)nativeInit},
         {"nativeGetHintSessionPreferredRate", "()J", (void*)nativeGetHintSessionPreferredRate},
         {"nativeCreateHintSession", "(II[IJ)J", (void*)nativeCreateHintSession},
+        {"nativeCreateHintSessionWithConfig", "(II[IJILandroid/hardware/power/SessionConfig;)J",
+         (void*)nativeCreateHintSessionWithConfig},
         {"nativePauseHintSession", "(J)V", (void*)nativePauseHintSession},
         {"nativeResumeHintSession", "(J)V", (void*)nativeResumeHintSession},
         {"nativeCloseHintSession", "(J)V", (void*)nativeCloseHintSession},
diff --git a/services/tests/servicestests/src/com/android/server/power/hint/HintManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/power/hint/HintManagerServiceTest.java
index 510e7c4..5902caa 100644
--- a/services/tests/servicestests/src/com/android/server/power/hint/HintManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/power/hint/HintManagerServiceTest.java
@@ -22,6 +22,7 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
@@ -41,6 +42,8 @@
 import android.app.ActivityManager;
 import android.app.ActivityManagerInternal;
 import android.content.Context;
+import android.hardware.power.SessionConfig;
+import android.hardware.power.SessionTag;
 import android.hardware.power.WorkDuration;
 import android.os.Binder;
 import android.os.IBinder;
@@ -63,6 +66,8 @@
 import org.junit.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -97,6 +102,7 @@
 
     private static final long DEFAULT_HINT_PREFERRED_RATE = 16666666L;
     private static final long DEFAULT_TARGET_DURATION = 16666666L;
+    private static final long DOUBLED_TARGET_DURATION = 33333333L;
     private static final long CONCURRENCY_TEST_DURATION_SEC = 10;
     private static final int UID = Process.myUid();
     private static final int TID = Process.myPid();
@@ -106,6 +112,8 @@
     private static final int[] SESSION_TIDS_C = new int[] {TID};
     private static final long[] DURATIONS_THREE = new long[] {1L, 100L, 1000L};
     private static final long[] TIMESTAMPS_THREE = new long[] {1L, 2L, 3L};
+    private static final long[] SESSION_PTRS = new long[] {11L, 22L, 33L};
+    private static final long[] SESSION_IDS = new long[] {1L, 11L, 111L};
     private static final long[] DURATIONS_ZERO = new long[] {};
     private static final long[] TIMESTAMPS_ZERO = new long[] {};
     private static final long[] TIMESTAMPS_TWO = new long[] {1L, 2L};
@@ -129,21 +137,61 @@
 
     private HintManagerService mService;
 
+    private static Answer<Long> fakeCreateWithConfig(Long ptr, Long sessionId) {
+        return new Answer<Long>() {
+            public Long answer(InvocationOnMock invocation) {
+                ((SessionConfig) invocation.getArguments()[5]).id = sessionId;
+                return ptr;
+            }
+        };
+    }
+
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
         when(mNativeWrapperMock.halGetHintSessionPreferredRate())
                 .thenReturn(DEFAULT_HINT_PREFERRED_RATE);
         when(mNativeWrapperMock.halCreateHintSession(eq(TGID), eq(UID), eq(SESSION_TIDS_A),
-                eq(DEFAULT_TARGET_DURATION))).thenReturn(1L);
+                eq(DEFAULT_TARGET_DURATION))).thenReturn(SESSION_PTRS[0]);
         when(mNativeWrapperMock.halCreateHintSession(eq(TGID), eq(UID), eq(SESSION_TIDS_B),
-                eq(DEFAULT_TARGET_DURATION))).thenReturn(2L);
+                eq(DOUBLED_TARGET_DURATION))).thenReturn(SESSION_PTRS[1]);
         when(mNativeWrapperMock.halCreateHintSession(eq(TGID), eq(UID), eq(SESSION_TIDS_C),
-                eq(0L))).thenReturn(1L);
+                eq(0L))).thenReturn(SESSION_PTRS[2]);
+        when(mNativeWrapperMock.halCreateHintSessionWithConfig(eq(TGID), eq(UID),
+                eq(SESSION_TIDS_A), eq(DEFAULT_TARGET_DURATION), anyInt(),
+                any(SessionConfig.class))).thenAnswer(fakeCreateWithConfig(SESSION_PTRS[0],
+                    SESSION_IDS[0]));
+        when(mNativeWrapperMock.halCreateHintSessionWithConfig(eq(TGID), eq(UID),
+                eq(SESSION_TIDS_B), eq(DOUBLED_TARGET_DURATION), anyInt(),
+                any(SessionConfig.class))).thenAnswer(fakeCreateWithConfig(SESSION_PTRS[1],
+                    SESSION_IDS[1]));
+        when(mNativeWrapperMock.halCreateHintSessionWithConfig(eq(TGID), eq(UID),
+                eq(SESSION_TIDS_C), eq(0L), anyInt(),
+                any(SessionConfig.class))).thenAnswer(fakeCreateWithConfig(SESSION_PTRS[2],
+                    SESSION_IDS[2]));
+
         LocalServices.removeServiceForTest(ActivityManagerInternal.class);
         LocalServices.addService(ActivityManagerInternal.class, mAmInternalMock);
     }
 
+    /**
+     * Mocks the creation calls, but without support for new createHintSessionWithConfig method
+     */
+    public void makeConfigCreationUnsupported() {
+        reset(mNativeWrapperMock);
+        when(mNativeWrapperMock.halGetHintSessionPreferredRate())
+                .thenReturn(DEFAULT_HINT_PREFERRED_RATE);
+        when(mNativeWrapperMock.halCreateHintSession(eq(TGID), eq(UID), eq(SESSION_TIDS_A),
+                eq(DEFAULT_TARGET_DURATION))).thenReturn(SESSION_PTRS[0]);
+        when(mNativeWrapperMock.halCreateHintSession(eq(TGID), eq(UID), eq(SESSION_TIDS_B),
+                eq(DOUBLED_TARGET_DURATION))).thenReturn(SESSION_PTRS[1]);
+        when(mNativeWrapperMock.halCreateHintSession(eq(TGID), eq(UID), eq(SESSION_TIDS_C),
+                eq(0L))).thenReturn(SESSION_PTRS[2]);
+        when(mNativeWrapperMock.halCreateHintSessionWithConfig(anyInt(), anyInt(),
+            any(int[].class), anyLong(), anyInt(),
+            any(SessionConfig.class))).thenThrow(new UnsupportedOperationException());
+    }
+
     static class NativeWrapperFake extends NativeWrapper {
         @Override
         public void halInit() {
@@ -160,6 +208,12 @@
         }
 
         @Override
+        public long halCreateHintSessionWithConfig(int tgid, int uid, int[] tids,
+                long durationNanos, int tag, SessionConfig config) {
+            return 1;
+        }
+
+        @Override
         public void halPauseHintSession(long halPtr) {
         }
 
@@ -224,27 +278,57 @@
         IBinder token = new Binder();
         // Make sure we throw exception when adding a TID doesn't belong to the processes
         // In this case, we add `init` PID into the list.
+        SessionConfig config = new SessionConfig();
         assertThrows(SecurityException.class,
-                () -> service.getBinderServiceInstance().createHintSession(token,
-                        new int[]{TID, 1}, DEFAULT_TARGET_DURATION));
+                () -> service.getBinderServiceInstance().createHintSessionWithConfig(token,
+                        new int[]{TID, 1}, DEFAULT_TARGET_DURATION, SessionTag.OTHER, config));
     }
 
     @Test
-    public void testCreateHintSession() throws Exception {
+    public void testCreateHintSessionFallback() throws Exception {
+        HintManagerService service = createService();
+        IBinder token = new Binder();
+        makeConfigCreationUnsupported();
+
+        IHintSession a = service.getBinderServiceInstance().createHintSessionWithConfig(token,
+                SESSION_TIDS_A, DEFAULT_TARGET_DURATION, SessionTag.OTHER, new SessionConfig());
+        assertNotNull(a);
+
+        IHintSession b = service.getBinderServiceInstance().createHintSessionWithConfig(token,
+                SESSION_TIDS_B, DOUBLED_TARGET_DURATION, SessionTag.OTHER, new SessionConfig());
+        assertNotEquals(a, b);
+
+        IHintSession c = service.getBinderServiceInstance().createHintSessionWithConfig(token,
+                SESSION_TIDS_C, 0L, SessionTag.OTHER, new SessionConfig());
+        assertNotNull(c);
+        verify(mNativeWrapperMock, times(3)).halCreateHintSession(anyInt(), anyInt(),
+                                                                  any(int[].class), anyLong());
+    }
+
+    @Test
+    public void testCreateHintSessionWithConfig() throws Exception {
         HintManagerService service = createService();
         IBinder token = new Binder();
 
-        IHintSession a = service.getBinderServiceInstance().createHintSession(token,
-                SESSION_TIDS_A, DEFAULT_TARGET_DURATION);
+        SessionConfig config = new SessionConfig();
+        IHintSession a = service.getBinderServiceInstance().createHintSessionWithConfig(token,
+                SESSION_TIDS_A, DEFAULT_TARGET_DURATION, SessionTag.OTHER, config);
         assertNotNull(a);
+        assertEquals(SESSION_IDS[0], config.id);
 
-        IHintSession b = service.getBinderServiceInstance().createHintSession(token,
-                SESSION_TIDS_B, DEFAULT_TARGET_DURATION);
+        SessionConfig config2 = new SessionConfig();
+        IHintSession b = service.getBinderServiceInstance().createHintSessionWithConfig(token,
+                SESSION_TIDS_B, DOUBLED_TARGET_DURATION, SessionTag.APP, config2);
         assertNotEquals(a, b);
+        assertEquals(SESSION_IDS[1], config2.id);
 
-        IHintSession c = service.getBinderServiceInstance().createHintSession(token,
-                SESSION_TIDS_C, 0L);
+        SessionConfig config3 = new SessionConfig();
+        IHintSession c = service.getBinderServiceInstance().createHintSessionWithConfig(token,
+                SESSION_TIDS_C, 0L, SessionTag.GAME, config3);
         assertNotNull(c);
+        assertEquals(SESSION_IDS[2], config3.id);
+        verify(mNativeWrapperMock, times(3)).halCreateHintSessionWithConfig(anyInt(), anyInt(),
+                any(int[].class), anyLong(), anyInt(), any(SessionConfig.class));
     }
 
     @Test
@@ -253,7 +337,8 @@
         IBinder token = new Binder();
 
         AppHintSession a = (AppHintSession) service.getBinderServiceInstance()
-                .createHintSession(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION);
+                .createHintSessionWithConfig(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION,
+                        SessionTag.OTHER, new SessionConfig());
 
         // Set session to background and calling updateHintAllowed() would invoke pause();
         service.mUidObserver.onUidStateChanged(
@@ -288,8 +373,8 @@
         HintManagerService service = createService();
         IBinder token = new Binder();
 
-        IHintSession a = service.getBinderServiceInstance().createHintSession(token,
-                SESSION_TIDS_A, DEFAULT_TARGET_DURATION);
+        IHintSession a = service.getBinderServiceInstance().createHintSessionWithConfig(token,
+                SESSION_TIDS_A, DEFAULT_TARGET_DURATION, SessionTag.OTHER, new SessionConfig());
 
         a.close();
         verify(mNativeWrapperMock, times(1)).halCloseHintSession(anyLong());
@@ -300,8 +385,8 @@
         HintManagerService service = createService();
         IBinder token = new Binder();
 
-        IHintSession a = service.getBinderServiceInstance().createHintSession(token,
-                SESSION_TIDS_A, DEFAULT_TARGET_DURATION);
+        IHintSession a = service.getBinderServiceInstance().createHintSessionWithConfig(token,
+                SESSION_TIDS_A, DEFAULT_TARGET_DURATION, SessionTag.OTHER, new SessionConfig());
 
         assertThrows(IllegalArgumentException.class, () -> {
             a.updateTargetWorkDuration(-1L);
@@ -321,7 +406,8 @@
         IBinder token = new Binder();
 
         AppHintSession a = (AppHintSession) service.getBinderServiceInstance()
-                .createHintSession(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION);
+                .createHintSessionWithConfig(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION,
+                        SessionTag.OTHER, new SessionConfig());
 
         a.updateTargetWorkDuration(100L);
         a.reportActualWorkDuration(DURATIONS_THREE, TIMESTAMPS_THREE);
@@ -363,7 +449,8 @@
         IBinder token = new Binder();
 
         AppHintSession a = (AppHintSession) service.getBinderServiceInstance()
-                .createHintSession(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION);
+                .createHintSessionWithConfig(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION,
+                    SessionTag.OTHER, new SessionConfig());
 
         a.sendHint(PerformanceHintManager.Session.CPU_LOAD_RESET);
         verify(mNativeWrapperMock, times(1)).halSendHint(anyLong(),
@@ -389,7 +476,8 @@
         IBinder token = new Binder();
 
         AppHintSession a = (AppHintSession) service.getBinderServiceInstance()
-                .createHintSession(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION);
+                .createHintSessionWithConfig(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION,
+                    SessionTag.OTHER, new SessionConfig());
 
         service.mUidObserver.onUidStateChanged(
                 a.mUid, ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND, 0, 0);
@@ -410,7 +498,8 @@
         IBinder token = new Binder();
 
         AppHintSession a = (AppHintSession) service.getBinderServiceInstance()
-                .createHintSession(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION);
+                .createHintSessionWithConfig(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION,
+                        SessionTag.OTHER, new SessionConfig());
 
         service.mUidObserver.onUidStateChanged(
                 a.mUid, ActivityManager.PROCESS_STATE_IMPORTANT_FOREGROUND, 0, 0);
@@ -423,7 +512,8 @@
         IBinder token = new Binder();
 
         AppHintSession a = (AppHintSession) service.getBinderServiceInstance()
-                .createHintSession(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION);
+                .createHintSessionWithConfig(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION,
+                        SessionTag.OTHER, new SessionConfig());
 
         a.updateTargetWorkDuration(100L);
 
@@ -454,10 +544,12 @@
         int threadCount = 3;
         int[] tids1 = createThreads(threadCount, stopLatch1);
         long sessionPtr1 = 111;
-        when(mNativeWrapperMock.halCreateHintSession(eq(TGID), eq(UID), eq(tids1),
-                eq(DEFAULT_TARGET_DURATION))).thenReturn(sessionPtr1);
+        when(mNativeWrapperMock.halCreateHintSessionWithConfig(eq(TGID), eq(UID), eq(tids1),
+                eq(DEFAULT_TARGET_DURATION), anyInt(), any(SessionConfig.class)))
+                .thenReturn(sessionPtr1);
         AppHintSession session1 = (AppHintSession) service.getBinderServiceInstance()
-                .createHintSession(token, tids1, DEFAULT_TARGET_DURATION);
+                .createHintSessionWithConfig(token, tids1, DEFAULT_TARGET_DURATION,
+                        SessionTag.OTHER, new SessionConfig());
         assertNotNull(session1);
 
         // for test only to avoid conflicting with any real thread that exists on device
@@ -473,10 +565,12 @@
         tids2WithIsolated[threadCount] = isoProc1;
         tids2WithIsolated[threadCount + 1] = isoProc2;
         long sessionPtr2 = 222;
-        when(mNativeWrapperMock.halCreateHintSession(eq(TGID), eq(UID), eq(tids2WithIsolated),
-                eq(DEFAULT_TARGET_DURATION))).thenReturn(sessionPtr2);
+        when(mNativeWrapperMock.halCreateHintSessionWithConfig(eq(TGID), eq(UID),
+                eq(tids2WithIsolated), eq(DEFAULT_TARGET_DURATION), anyInt(),
+                any(SessionConfig.class))).thenReturn(sessionPtr2);
         AppHintSession session2 = (AppHintSession) service.getBinderServiceInstance()
-                .createHintSession(token, tids2WithIsolated, DEFAULT_TARGET_DURATION);
+                .createHintSessionWithConfig(token, tids2WithIsolated,
+                        DEFAULT_TARGET_DURATION, SessionTag.OTHER, new SessionConfig());
         assertNotNull(session2);
 
         // trigger clean up through UID state change by making the process background
@@ -608,7 +702,8 @@
         IBinder token = new Binder();
 
         AppHintSession a = (AppHintSession) service.getBinderServiceInstance()
-                .createHintSession(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION);
+                .createHintSessionWithConfig(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION,
+                        SessionTag.OTHER, new SessionConfig());
 
         a.setMode(0, true);
         verify(mNativeWrapperMock, times(1)).halSetMode(anyLong(),
@@ -726,7 +821,8 @@
             AtomicReference<Boolean> shouldRun) throws Exception {
         IBinder token = new Binder();
         AppHintSession a = (AppHintSession) service.getBinderServiceInstance()
-                .createHintSession(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION);
+                .createHintSessionWithConfig(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION,
+                        SessionTag.OTHER, new SessionConfig());
         // we will start some threads and get their valid TIDs to update
         int threadCount = 3;
         // the list of TIDs
@@ -793,7 +889,8 @@
         IBinder token = new Binder();
 
         AppHintSession a = (AppHintSession) service.getBinderServiceInstance()
-                .createHintSession(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION);
+                .createHintSessionWithConfig(token, SESSION_TIDS_A, DEFAULT_TARGET_DURATION,
+                        SessionTag.OTHER, new SessionConfig());
 
         a.updateTargetWorkDuration(100L);
         a.reportActualWorkDuration2(WORK_DURATIONS_FIVE);