register callback to pull NetworkBpfMapInfo atom

Test: statsd_testdrive 10161, atest BpfNetMaps
Bug: 217624062
Bug: 245228691

Change-Id: I630a86cb5a34aaa1810eb75b374588dba570f0ff
diff --git a/service/Android.bp b/service/Android.bp
index b68d389..68c1722 100644
--- a/service/Android.bp
+++ b/service/Android.bp
@@ -143,6 +143,7 @@
         "src/**/*.java",
         ":framework-connectivity-shared-srcs",
         ":services-connectivity-shared-srcs",
+        ":statslog-connectivity-java-gen",
     ],
     libs: [
         "framework-annotations-lib",
@@ -152,6 +153,7 @@
         "framework-wifi.stubs.module_lib",
         "unsupportedappusage",
         "ServiceConnectivityResources",
+        "framework-statsd.stubs.module_lib",
     ],
     static_libs: [
         // Do not add libs here if they are already included
@@ -332,3 +334,10 @@
         "--output $(out)",
     visibility: ["//visibility:private"],
 }
+
+genrule {
+  name: "statslog-connectivity-java-gen",
+  tools: ["stats-log-api-gen"],
+  cmd: "$(location stats-log-api-gen) --java $(out) --module connectivity --javaPackage com.android.server --javaClass ConnectivityStatsLog",
+  out: ["com/android/server/ConnectivityStatsLog.java"],
+}
diff --git a/service/src/com/android/server/BpfNetMaps.java b/service/src/com/android/server/BpfNetMaps.java
index 4eeef5c..69b6486 100644
--- a/service/src/com/android/server/BpfNetMaps.java
+++ b/service/src/com/android/server/BpfNetMaps.java
@@ -33,6 +33,9 @@
 import static android.system.OsConstants.ENOENT;
 import static android.system.OsConstants.EOPNOTSUPP;
 
+import static com.android.server.ConnectivityStatsLog.NETWORK_BPF_MAP_INFO;
+
+import android.app.StatsManager;
 import android.content.Context;
 import android.net.INetd;
 import android.os.RemoteException;
@@ -42,12 +45,15 @@
 import android.system.Os;
 import android.util.ArraySet;
 import android.util.Log;
+import android.util.StatsEvent;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.modules.utils.BackgroundThread;
 import com.android.modules.utils.build.SdkLevel;
 import com.android.net.module.util.BpfMap;
 import com.android.net.module.util.DeviceConfigUtils;
 import com.android.net.module.util.IBpfMap;
+import com.android.net.module.util.Struct;
 import com.android.net.module.util.Struct.U32;
 import com.android.net.module.util.Struct.U8;
 import com.android.net.module.util.bpf.CookieTagMapKey;
@@ -55,6 +61,7 @@
 
 import java.io.FileDescriptor;
 import java.io.IOException;
+import java.util.List;
 import java.util.Set;
 
 /**
@@ -274,6 +281,15 @@
         public int synchronizeKernelRCU() {
             return native_synchronizeKernelRCU();
         }
+
+        /**
+         * Build Stats Event for NETWORK_BPF_MAP_INFO atom
+         */
+        public StatsEvent buildStatsEvent(final int cookieTagMapSize, final int uidOwnerMapSize,
+                final int uidPermissionMapSize) {
+            return ConnectivityStatsLog.buildStatsEvent(NETWORK_BPF_MAP_INFO, cookieTagMapSize,
+                    uidOwnerMapSize, uidPermissionMapSize);
+        }
     }
 
     /** Constructor used after T that doesn't need to use netd anymore. */
@@ -849,6 +865,43 @@
         }
     }
 
+    /** Register callback for statsd to pull atom. */
+    public void setPullAtomCallback(final Context context) {
+        throwIfPreT("setPullAtomCallback is not available on pre-T devices");
+
+        final StatsManager statsManager = context.getSystemService(StatsManager.class);
+        statsManager.setPullAtomCallback(NETWORK_BPF_MAP_INFO, null /* metadata */,
+                BackgroundThread.getExecutor(), this::pullBpfMapInfoAtom);
+    }
+
+    private <K extends Struct, V extends Struct> int getMapSize(IBpfMap<K, V> map)
+            throws ErrnoException {
+        // forEach could restart iteration from the beginning if there is a concurrent entry
+        // deletion. netd and skDestroyListener could delete CookieTagMap entry concurrently.
+        // So using Set to count the number of entry in the map.
+        Set<K> keySet = new ArraySet<>();
+        map.forEach((k, v) -> keySet.add(k));
+        return keySet.size();
+    }
+
+    /** Callback for StatsManager#setPullAtomCallback */
+    @VisibleForTesting
+    public int pullBpfMapInfoAtom(final int atomTag, final List<StatsEvent> data) {
+        if (atomTag != NETWORK_BPF_MAP_INFO) {
+            Log.e(TAG, "Unexpected atom tag: " + atomTag);
+            return StatsManager.PULL_SKIP;
+        }
+
+        try {
+            data.add(mDeps.buildStatsEvent(getMapSize(sCookieTagMap), getMapSize(sUidOwnerMap),
+                    getMapSize(sUidPermissionMap)));
+        } catch (ErrnoException e) {
+            Log.e(TAG, "Failed to pull NETWORK_BPF_MAP_INFO atom: " + e);
+            return StatsManager.PULL_SKIP;
+        }
+        return StatsManager.PULL_SUCCESS;
+    }
+
     /**
      * Dump BPF maps
      *
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index afb3ca2..36c81dd 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -3040,6 +3040,11 @@
         if (!ConnectivitySettingsManager.getMobileDataPreferredUids(mContext).isEmpty()) {
             updateMobileDataPreferredUids();
         }
+
+        // On T+ devices, register callback for statsd to pull NETWORK_BPF_MAP_INFO atom
+        if (SdkLevel.isAtLeastT()) {
+            mBpfNetMaps.setPullAtomCallback(mContext);
+        }
     }
 
     /**
diff --git a/tests/unit/java/com/android/server/BpfNetMapsTest.java b/tests/unit/java/com/android/server/BpfNetMapsTest.java
index b7d80ac..eb5d2ef 100644
--- a/tests/unit/java/com/android/server/BpfNetMapsTest.java
+++ b/tests/unit/java/com/android/server/BpfNetMapsTest.java
@@ -30,6 +30,7 @@
 import static android.net.INetd.PERMISSION_NONE;
 import static android.net.INetd.PERMISSION_UNINSTALLED;
 import static android.net.INetd.PERMISSION_UPDATE_DEVICE_STATS;
+import static android.system.OsConstants.EINVAL;
 import static android.system.OsConstants.EPERM;
 
 import static com.android.server.BpfNetMaps.DOZABLE_MATCH;
@@ -40,6 +41,7 @@
 import static com.android.server.BpfNetMaps.PENALTY_BOX_MATCH;
 import static com.android.server.BpfNetMaps.POWERSAVE_MATCH;
 import static com.android.server.BpfNetMaps.RESTRICTED_MATCH;
+import static com.android.server.ConnectivityStatsLog.NETWORK_BPF_MAP_INFO;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -47,13 +49,18 @@
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assume.assumeFalse;
+import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 
+import android.app.StatsManager;
 import android.content.Context;
 import android.net.INetd;
 import android.os.Build;
 import android.os.ServiceSpecificException;
+import android.system.ErrnoException;
 
 import androidx.test.filters.SmallTest;
 
@@ -76,6 +83,7 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.util.ArrayList;
 import java.util.List;
 
 @RunWith(DevSdkIgnoreRunner.class)
@@ -120,7 +128,7 @@
             new TestBpfMap<>(U32.class, UidOwnerValue.class);
     private final IBpfMap<U32, U8> mUidPermissionMap = new TestBpfMap<>(U32.class, U8.class);
     private final IBpfMap<CookieTagMapKey, CookieTagMapValue> mCookieTagMap =
-            new TestBpfMap<>(CookieTagMapKey.class, CookieTagMapValue.class);
+            spy(new TestBpfMap<>(CookieTagMapKey.class, CookieTagMapValue.class));
 
     @Before
     public void setUp() throws Exception {
@@ -882,4 +890,40 @@
 
         assertThrows(ServiceSpecificException.class, () -> mBpfNetMaps.swapActiveStatsMap());
     }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testPullBpfMapInfo() throws Exception {
+        // mCookieTagMap has 1 entry
+        mCookieTagMap.updateEntry(new CookieTagMapKey(0), new CookieTagMapValue(0, 0));
+
+        // mUidOwnerMap has 2 entries
+        mUidOwnerMap.updateEntry(new U32(0), new UidOwnerValue(0, 0));
+        mUidOwnerMap.updateEntry(new U32(1), new UidOwnerValue(0, 0));
+
+        // mUidPermissionMap has 3 entries
+        mUidPermissionMap.updateEntry(new U32(0), new U8((short) 0));
+        mUidPermissionMap.updateEntry(new U32(1), new U8((short) 0));
+        mUidPermissionMap.updateEntry(new U32(2), new U8((short) 0));
+
+        final int ret = mBpfNetMaps.pullBpfMapInfoAtom(NETWORK_BPF_MAP_INFO, new ArrayList<>());
+        assertEquals(StatsManager.PULL_SUCCESS, ret);
+        verify(mDeps).buildStatsEvent(
+                1 /* cookieTagMapSize */, 2 /* uidOwnerMapSize */, 3 /* uidPermissionMapSize */);
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testPullBpfMapInfoGetMapSizeFailure() throws Exception {
+        doThrow(new ErrnoException("", EINVAL)).when(mCookieTagMap).forEach(any());
+        final int ret = mBpfNetMaps.pullBpfMapInfoAtom(NETWORK_BPF_MAP_INFO, new ArrayList<>());
+        assertEquals(StatsManager.PULL_SKIP, ret);
+    }
+
+    @Test
+    @IgnoreUpTo(Build.VERSION_CODES.S_V2)
+    public void testPullBpfMapInfoUnexpectedAtomTag() {
+        final int ret = mBpfNetMaps.pullBpfMapInfoAtom(-1 /* atomTag */, new ArrayList<>());
+        assertEquals(StatsManager.PULL_SKIP, ret);
+    }
 }