Merge "Read per-process GPU memory when pulling ProcessMemorySnapshot"
diff --git a/core/java/com/android/internal/os/DmabufInfoReader.java b/core/java/com/android/internal/os/KernelAllocationStats.java
similarity index 72%
rename from core/java/com/android/internal/os/DmabufInfoReader.java
rename to core/java/com/android/internal/os/KernelAllocationStats.java
index 786a6ee..1c3f8b0 100644
--- a/core/java/com/android/internal/os/DmabufInfoReader.java
+++ b/core/java/com/android/internal/os/KernelAllocationStats.java
@@ -18,9 +18,9 @@
 
 import android.annotation.Nullable;
 
-/** Wrapper around libdmabufinfo. */
-public final class DmabufInfoReader {
-    private DmabufInfoReader() {}
+/** JNI wrapper around libmeminfo for kernel memory allocation stats (dmabufs, gpu driver). */
+public final class KernelAllocationStats {
+    private KernelAllocationStats() {}
 
     /** Process dma-buf stats. */
     public static final class ProcessDmabuf {
@@ -47,5 +47,19 @@
      * stats could not be read.
      */
     @Nullable
-    public static native ProcessDmabuf getProcessStats(int pid);
+    public static native ProcessDmabuf getDmabufAllocations(int pid);
+
+    /** Pid to gpu memory size. */
+    public static final class ProcessGpuMem {
+        public final int pid;
+        public final int gpuMemoryKb;
+
+        ProcessGpuMem(int pid, int gpuMemoryKb) {
+            this.pid = pid;
+            this.gpuMemoryKb = gpuMemoryKb;
+        }
+    }
+
+    /** Return list of pid to gpu memory size. */
+    public static native ProcessGpuMem[] getGpuAllocations();
 }
diff --git a/core/jni/Android.bp b/core/jni/Android.bp
index 1a1a8ba..da62863 100644
--- a/core/jni/Android.bp
+++ b/core/jni/Android.bp
@@ -210,8 +210,8 @@
                 "com_android_internal_content_om_OverlayConfig.cpp",
                 "com_android_internal_net_NetworkUtilsInternal.cpp",
                 "com_android_internal_os_ClassLoaderFactory.cpp",
-                "com_android_internal_os_DmabufInfoReader.cpp",
                 "com_android_internal_os_FuseAppLoop.cpp",
+                "com_android_internal_os_KernelAllocationStats.cpp",
                 "com_android_internal_os_KernelCpuBpfTracking.cpp",
                 "com_android_internal_os_KernelCpuTotalBpfMapReader.cpp",
                 "com_android_internal_os_KernelCpuUidBpfMapReader.cpp",
diff --git a/core/jni/AndroidRuntime.cpp b/core/jni/AndroidRuntime.cpp
index 18e85b6..04fafb4 100644
--- a/core/jni/AndroidRuntime.cpp
+++ b/core/jni/AndroidRuntime.cpp
@@ -197,8 +197,8 @@
 extern int register_com_android_internal_content_om_OverlayConfig(JNIEnv *env);
 extern int register_com_android_internal_net_NetworkUtilsInternal(JNIEnv* env);
 extern int register_com_android_internal_os_ClassLoaderFactory(JNIEnv* env);
-extern int register_com_android_internal_os_DmabufInfoReader(JNIEnv* env);
 extern int register_com_android_internal_os_FuseAppLoop(JNIEnv* env);
+extern int register_com_android_internal_os_KernelAllocationStats(JNIEnv* env);
 extern int register_com_android_internal_os_KernelCpuBpfTracking(JNIEnv* env);
 extern int register_com_android_internal_os_KernelCpuTotalBpfMapReader(JNIEnv* env);
 extern int register_com_android_internal_os_KernelCpuUidBpfMapReader(JNIEnv *env);
@@ -1655,8 +1655,8 @@
         REG_JNI(register_android_security_Scrypt),
         REG_JNI(register_com_android_internal_content_F2fsUtils),
         REG_JNI(register_com_android_internal_content_NativeLibraryHelper),
-        REG_JNI(register_com_android_internal_os_DmabufInfoReader),
         REG_JNI(register_com_android_internal_os_FuseAppLoop),
+        REG_JNI(register_com_android_internal_os_KernelAllocationStats),
         REG_JNI(register_com_android_internal_os_KernelCpuBpfTracking),
         REG_JNI(register_com_android_internal_os_KernelCpuTotalBpfMapReader),
         REG_JNI(register_com_android_internal_os_KernelCpuUidBpfMapReader),
diff --git a/core/jni/com_android_internal_os_DmabufInfoReader.cpp b/core/jni/com_android_internal_os_DmabufInfoReader.cpp
deleted file mode 100644
index 4b0a6ac..0000000
--- a/core/jni/com_android_internal_os_DmabufInfoReader.cpp
+++ /dev/null
@@ -1,60 +0,0 @@
-/*
- * 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 <dmabufinfo/dmabufinfo.h>
-#include "core_jni_helpers.h"
-
-namespace android {
-
-static jobject DmabufInfoReader_getProcessStats(JNIEnv *env, jobject, jint pid) {
-    std::vector<dmabufinfo::DmaBuffer> buffers;
-    if (!dmabufinfo::ReadDmaBufMapRefs(pid, &buffers)) {
-        return nullptr;
-    }
-    jint mappedSize = 0;
-    jint mappedCount = buffers.size();
-    for (const auto &buffer : buffers) {
-        mappedSize += buffer.size();
-    }
-    mappedSize /= 1024;
-
-    jint retainedSize = -1;
-    jint retainedCount = -1;
-    if (dmabufinfo::ReadDmaBufFdRefs(pid, &buffers)) {
-        retainedCount = buffers.size();
-        retainedSize = 0;
-        for (const auto &buffer : buffers) {
-            retainedSize += buffer.size();
-        }
-        retainedSize /= 1024;
-    }
-
-    jclass clazz = FindClassOrDie(env, "com/android/internal/os/DmabufInfoReader$ProcessDmabuf");
-    jmethodID constructID = GetMethodIDOrDie(env, clazz, "<init>", "(IIII)V");
-    return env->NewObject(clazz, constructID, retainedSize, retainedCount, mappedSize, mappedCount);
-}
-
-static const JNINativeMethod methods[] = {
-        {"getProcessStats", "(I)Lcom/android/internal/os/DmabufInfoReader$ProcessDmabuf;",
-         (void *)DmabufInfoReader_getProcessStats},
-};
-
-int register_com_android_internal_os_DmabufInfoReader(JNIEnv *env) {
-    return RegisterMethodsOrDie(env, "com/android/internal/os/DmabufInfoReader", methods,
-                                NELEM(methods));
-}
-
-} // namespace android
diff --git a/core/jni/com_android_internal_os_KernelAllocationStats.cpp b/core/jni/com_android_internal_os_KernelAllocationStats.cpp
new file mode 100644
index 0000000..e0a24430
--- /dev/null
+++ b/core/jni/com_android_internal_os_KernelAllocationStats.cpp
@@ -0,0 +1,97 @@
+/*
+ * 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 <dmabufinfo/dmabufinfo.h>
+#include <jni.h>
+#include <meminfo/sysmeminfo.h>
+
+#include "core_jni_helpers.h"
+
+namespace {
+static jclass gProcessDmabufClazz;
+static jmethodID gProcessDmabufCtor;
+static jclass gProcessGpuMemClazz;
+static jmethodID gProcessGpuMemCtor;
+} // namespace
+
+namespace android {
+
+static jobject KernelAllocationStats_getDmabufAllocations(JNIEnv *env, jobject, jint pid) {
+    std::vector<dmabufinfo::DmaBuffer> buffers;
+    if (!dmabufinfo::ReadDmaBufMapRefs(pid, &buffers)) {
+        return nullptr;
+    }
+    jint mappedSize = 0;
+    jint mappedCount = buffers.size();
+    for (const auto &buffer : buffers) {
+        mappedSize += buffer.size();
+    }
+    mappedSize /= 1024;
+
+    jint retainedSize = -1;
+    jint retainedCount = -1;
+    if (dmabufinfo::ReadDmaBufFdRefs(pid, &buffers)) {
+        retainedCount = buffers.size();
+        retainedSize = 0;
+        for (const auto &buffer : buffers) {
+            retainedSize += buffer.size();
+        }
+        retainedSize /= 1024;
+    }
+    return env->NewObject(gProcessDmabufClazz, gProcessDmabufCtor, retainedSize, retainedCount,
+                          mappedSize, mappedCount);
+}
+
+static jobject KernelAllocationStats_getGpuAllocations(JNIEnv *env) {
+    std::unordered_map<uint32_t, uint64_t> out;
+    meminfo::ReadPerProcessGpuMem(&out);
+    jobjectArray result = env->NewObjectArray(out.size(), gProcessGpuMemClazz, nullptr);
+    if (result == NULL) {
+        jniThrowRuntimeException(env, "Cannot create result array");
+        return nullptr;
+    }
+    int idx = 0;
+    for (const auto &entry : out) {
+        jobject pidStats =
+                env->NewObject(gProcessGpuMemClazz, gProcessGpuMemCtor, entry.first, entry.second);
+        env->SetObjectArrayElement(result, idx, pidStats);
+        env->DeleteLocalRef(pidStats);
+        ++idx;
+    }
+    return result;
+}
+
+static const JNINativeMethod methods[] = {
+        {"getDmabufAllocations", "(I)Lcom/android/internal/os/KernelAllocationStats$ProcessDmabuf;",
+         (void *)KernelAllocationStats_getDmabufAllocations},
+        {"getGpuAllocations", "()[Lcom/android/internal/os/KernelAllocationStats$ProcessGpuMem;",
+         (void *)KernelAllocationStats_getGpuAllocations},
+};
+
+int register_com_android_internal_os_KernelAllocationStats(JNIEnv *env) {
+    int res = RegisterMethodsOrDie(env, "com/android/internal/os/KernelAllocationStats", methods,
+                                   NELEM(methods));
+    jclass clazz =
+            FindClassOrDie(env, "com/android/internal/os/KernelAllocationStats$ProcessDmabuf");
+    gProcessDmabufClazz = MakeGlobalRefOrDie(env, clazz);
+    gProcessDmabufCtor = GetMethodIDOrDie(env, gProcessDmabufClazz, "<init>", "(IIII)V");
+
+    clazz = FindClassOrDie(env, "com/android/internal/os/KernelAllocationStats$ProcessGpuMem");
+    gProcessGpuMemClazz = MakeGlobalRefOrDie(env, clazz);
+    gProcessGpuMemCtor = GetMethodIDOrDie(env, gProcessGpuMemClazz, "<init>", "(II)V");
+    return res;
+}
+
+} // namespace android
diff --git a/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java b/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
index 45f1b16..b1cc517 100644
--- a/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
+++ b/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
@@ -172,6 +172,7 @@
 import android.util.Pair;
 import android.util.Slog;
 import android.util.SparseArray;
+import android.util.SparseIntArray;
 import android.util.StatsEvent;
 import android.util.proto.ProtoOutputStream;
 import android.view.Display;
@@ -183,7 +184,7 @@
 import com.android.internal.os.BatterySipper;
 import com.android.internal.os.BatteryStatsHelper;
 import com.android.internal.os.BinderCallsStats.ExportedCallStat;
-import com.android.internal.os.DmabufInfoReader;
+import com.android.internal.os.KernelAllocationStats;
 import com.android.internal.os.KernelCpuBpfTracking;
 import com.android.internal.os.KernelCpuThreadReader;
 import com.android.internal.os.KernelCpuThreadReaderDiff;
@@ -425,7 +426,6 @@
     private final Object mSystemUptimeLock = new Object();
     private final Object mProcessMemoryStateLock = new Object();
     private final Object mProcessMemoryHighWaterMarkLock = new Object();
-    private final Object mProcessMemorySnapshotLock = new Object();
     private final Object mSystemIonHeapSizeLock = new Object();
     private final Object mIonHeapSizeLock = new Object();
     private final Object mProcessSystemIonHeapSizeLock = new Object();
@@ -563,9 +563,7 @@
                             return pullProcessMemoryHighWaterMarkLocked(atomTag, data);
                         }
                     case FrameworkStatsLog.PROCESS_MEMORY_SNAPSHOT:
-                        synchronized (mProcessMemorySnapshotLock) {
-                            return pullProcessMemorySnapshotLocked(atomTag, data);
-                        }
+                        return pullProcessMemorySnapshot(atomTag, data);
                     case FrameworkStatsLog.SYSTEM_ION_HEAP_SIZE:
                         synchronized (mSystemIonHeapSizeLock) {
                             return pullSystemIonHeapSizeLocked(atomTag, data);
@@ -2218,10 +2216,16 @@
         );
     }
 
-    int pullProcessMemorySnapshotLocked(int atomTag, List<StatsEvent> pulledData) {
+    int pullProcessMemorySnapshot(int atomTag, List<StatsEvent> pulledData) {
         List<ProcessMemoryState> managedProcessList =
                 LocalServices.getService(ActivityManagerInternal.class)
                         .getMemoryStateForProcesses();
+        KernelAllocationStats.ProcessGpuMem[] gpuAllocations =
+                KernelAllocationStats.getGpuAllocations();
+        SparseIntArray gpuMemPerPid = new SparseIntArray(gpuAllocations.length);
+        for (KernelAllocationStats.ProcessGpuMem processGpuMem : gpuAllocations) {
+            gpuMemPerPid.put(processGpuMem.pid, processGpuMem.gpuMemoryKb);
+        }
         for (ProcessMemoryState managedProcess : managedProcessList) {
             final MemorySnapshot snapshot = readMemorySnapshotFromProcfs(managedProcess.pid);
             if (snapshot == null) {
@@ -2230,7 +2234,8 @@
             pulledData.add(FrameworkStatsLog.buildStatsEvent(atomTag, managedProcess.uid,
                     managedProcess.processName, managedProcess.pid, managedProcess.oomScore,
                     snapshot.rssInKilobytes, snapshot.anonRssInKilobytes, snapshot.swapInKilobytes,
-                    snapshot.anonRssInKilobytes + snapshot.swapInKilobytes));
+                    snapshot.anonRssInKilobytes + snapshot.swapInKilobytes,
+                    gpuMemPerPid.get(managedProcess.pid)));
         }
         // Complement the data with native system processes. Given these measurements can be taken
         // in response to LMKs happening, we want to first collect the managed app stats (to
@@ -2248,7 +2253,8 @@
                     processCmdlines.valueAt(i), pid,
                     -1001 /*Placeholder for native processes, OOM_SCORE_ADJ_MIN - 1.*/,
                     snapshot.rssInKilobytes, snapshot.anonRssInKilobytes, snapshot.swapInKilobytes,
-                    snapshot.anonRssInKilobytes + snapshot.swapInKilobytes));
+                    snapshot.anonRssInKilobytes + snapshot.swapInKilobytes,
+                    gpuMemPerPid.get(pid)));
         }
         return StatsManager.PULL_SUCCESS;
     }
@@ -2328,7 +2334,8 @@
             if (process.uid == Process.SYSTEM_UID) {
                 continue;
             }
-            DmabufInfoReader.ProcessDmabuf proc = DmabufInfoReader.getProcessStats(process.pid);
+            KernelAllocationStats.ProcessDmabuf proc =
+                    KernelAllocationStats.getDmabufAllocations(process.pid);
             if (proc == null || (proc.retainedBuffersCount <= 0 && proc.mappedBuffersCount <= 0)) {
                 continue;
             }