Java API for pulled atoms

This creates a java API for registering pullers. Will implement the
statsd side in a follow up CL.

Test: builds, boots
Change-Id: Ib6735984297ce3148839a6370a3c15b2a585baf5

Android.bp

@@ -405,6 +405,8 @@ filegroup {
 filegroup {
     name: "statsd_aidl",
     srcs: [
+        "core/java/android/os/IPullAtomCallback.aidl",
+        "core/java/android/os/IPullAtomResultReceiver.aidl",
         "core/java/android/os/IStatsCompanionService.aidl",
         "core/java/android/os/IStatsManager.aidl",
         "core/java/android/os/IStatsPullerCallback.aidl",

apex/statsd/service/java/com/android/server/stats/StatsCompanionService.java

@@ -76,6 +76,7 @@ import android.os.FileUtils;
 import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.IBinder;
+import android.os.IPullAtomCallback;
 import android.os.IStatsCompanionService;
 import android.os.IStatsManager;
 import android.os.IStoraged;
@@ -165,6 +166,7 @@ import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Objects;
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
@@ -272,6 +274,72 @@ public class StatsCompanionService extends IStatsCompanionService.Stub {
     private final BroadcastReceiver mAppUpdateReceiver;
     private final BroadcastReceiver mUserUpdateReceiver;
     private final ShutdownEventReceiver mShutdownEventReceiver;
+
+    private static final class PullerKey {
+        private final int mUid;
+        private final int mAtomTag;
+
+        PullerKey(int uid, int atom) {
+            mUid = uid;
+            mAtomTag = atom;
+        }
+
+        public int getUid() {
+            return mUid;
+        }
+
+        public int getAtom() {
+            return mAtomTag;
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(mUid, mAtomTag);
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            if (obj instanceof PullerKey) {
+                PullerKey other = (PullerKey) obj;
+                return this.mUid == other.getUid() && this.mAtomTag == other.getAtom();
+            }
+            return false;
+        }
+    }
+
+    private static final class PullerValue {
+        private final long mCoolDownNs;
+        private final long mTimeoutNs;
+        private int[] mAdditiveFields;
+        private IPullAtomCallback mCallback;
+
+        PullerValue(long coolDownNs, long timeoutNs, int[] additiveFields,
+                IPullAtomCallback callback) {
+            mCoolDownNs = coolDownNs;
+            mTimeoutNs = timeoutNs;
+            mAdditiveFields = additiveFields;
+            mCallback = callback;
+        }
+
+        public long getCoolDownNs() {
+            return mCoolDownNs;
+        }
+
+        public long getTimeoutNs() {
+            return mTimeoutNs;
+        }
+
+        public int[] getAdditiveFields() {
+            return mAdditiveFields;
+        }
+
+        public IPullAtomCallback getCallback() {
+            return mCallback;
+        }
+    }
+
+    private final HashMap<PullerKey, PullerValue> mPullers = new HashMap<>();
+
     private final KernelWakelockReader mKernelWakelockReader = new KernelWakelockReader();
     private final KernelWakelockStats mTmpWakelockStats = new KernelWakelockStats();
     private IWifiManager mWifiManager = null;
@@ -323,7 +391,6 @@ public class StatsCompanionService extends IStatsCompanionService.Stub {
             @Override
             public void onReceive(Context context, Intent intent) {
                 synchronized (sStatsdLock) {
-                    sStatsd = fetchStatsdService();
                     if (sStatsd == null) {
                         Slog.w(TAG, "Could not access statsd for UserUpdateReceiver");
                         return;
@@ -2553,10 +2620,40 @@ public class StatsCompanionService extends IStatsCompanionService.Stub {
         mContext.enforceCallingPermission(android.Manifest.permission.STATSCOMPANION, null);
     }
 
+    @Override
+    public void registerPullAtomCallback(int atomTag, long coolDownNs, long timeoutNs,
+            int[] additiveFields, IPullAtomCallback pullerCallback) {
+        synchronized (sStatsdLock) {
+            // Always cache the puller in SCS.
+            // If statsd is down, we will register it when it comes back up.
+            int callingUid = Binder.getCallingUid();
+            final long token = Binder.clearCallingIdentity();
+            PullerKey key = new PullerKey(callingUid, atomTag);
+            PullerValue val = new PullerValue(
+                    coolDownNs, timeoutNs, additiveFields, pullerCallback);
+            mPullers.put(key, val);
+
+            if (sStatsd == null) {
+                Slog.w(TAG, "Could not access statsd for registering puller for atom " + atomTag);
+                return;
+            }
+            try {
+                sStatsd.registerPullAtomCallback(
+                        callingUid, atomTag, coolDownNs, timeoutNs, additiveFields, pullerCallback);
+            } catch (RemoteException e) {
+                Slog.e(TAG, "Failed to access statsd to register puller for atom " + atomTag);
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
+        }
+    }
+
     // Lifecycle and related code
 
     /**
-     * Fetches the statsd IBinder service
+     * Fetches the statsd IBinder service.
+     * Note: This should only be called from sayHiToStatsd. All other clients should use the cached
+     * sStatsd with a null check.
      */
     private static IStatsManager fetchStatsdService() {
         return IStatsManager.Stub.asInterface(ServiceManager.getService("stats"));
@@ -2654,6 +2751,8 @@ public class StatsCompanionService extends IStatsCompanionService.Stub {
                     // Pull the latest state of UID->app name, version mapping when
                     // statsd starts.
                     informAllUidsLocked(mContext);
+                    // Register all pullers. If SCS has just started, this should be empty.
+                    registerAllPullersLocked();
                 } finally {
                     restoreCallingIdentity(token);
                 }
@@ -2665,10 +2764,21 @@ public class StatsCompanionService extends IStatsCompanionService.Stub {
         }
     }
 
+    @GuardedBy("sStatsdLock")
+    private void registerAllPullersLocked() throws RemoteException {
+        // TODO: pass in one call, using a file descriptor (similar to uidmap).
+        for (Map.Entry<PullerKey, PullerValue> entry : mPullers.entrySet()) {
+            PullerKey key = entry.getKey();
+            PullerValue val = entry.getValue();
+            sStatsd.registerPullAtomCallback(key.getUid(), key.getAtom(), val.getCoolDownNs(),
+                    val.getTimeoutNs(), val.getAdditiveFields(), val.getCallback());
+        }
+    }
+
     private class StatsdDeathRecipient implements IBinder.DeathRecipient {
         @Override
         public void binderDied() {
-            Slog.i(TAG, "Statsd is dead - erase all my knowledge.");
+            Slog.i(TAG, "Statsd is dead - erase all my knowledge, except pullers");
             synchronized (sStatsdLock) {
                 long now = SystemClock.elapsedRealtime();
                 for (Long timeMillis : mDeathTimeMillis) {

cmds/statsd/src/StatsService.cpp

@@ -1289,6 +1289,13 @@ Status StatsService::registerPullerCallback(int32_t atomTag,
     return Status::ok();
 }
 
+Status StatsService::registerPullAtomCallback(int32_t uid, int32_t atomTag, int64_t coolDownNs,
+                                    int64_t timeoutNs, const std::vector<int32_t>& additiveFields,
+                                    const sp<android::os::IPullAtomCallback>& pullerCallback) {
+    VLOG("StatsService::registerPuller called.");
+    return Status::ok();
+}
+
 Status StatsService::unregisterPullerCallback(int32_t atomTag, const String16& packageName) {
     ENFORCE_DUMP_AND_USAGE_STATS(packageName);
 

cmds/statsd/src/StatsService.h

@@ -179,6 +179,13 @@ public:
         const sp<android::os::IStatsPullerCallback>& pullerCallback,
         const String16& packageName) override;
 
+    /**
+     * Binder call to register a callback function for a pulled atom.
+     */
+    virtual Status registerPullAtomCallback(int32_t uid, int32_t atomTag, int64_t coolDownNs,
+            int64_t timeoutNs, const std::vector<int32_t>& additiveFields,
+            const sp<android::os::IPullAtomCallback>& pullerCallback) override;
+
     /**
      * Binder call to unregister any existing callback function for a vendor pulled atom.
      */

core/java/android/app/StatsManager.java

@@ -24,12 +24,22 @@ import android.annotation.RequiresPermission;
 import android.annotation.SystemApi;
 import android.content.Context;
 import android.os.IBinder;
+import android.os.IPullAtomCallback;
+import android.os.IPullAtomResultReceiver;
+import android.os.IStatsCompanionService;
 import android.os.IStatsManager;
 import android.os.IStatsPullerCallback;
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.util.AndroidException;
 import android.util.Slog;
+import android.util.StatsEvent;
+
+import com.android.internal.annotations.GuardedBy;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.Executor;
 
 /**
  * API for statsd clients to send configurations and retrieve data.
@@ -43,8 +53,12 @@ public final class StatsManager {
 
     private final Context mContext;
 
+    @GuardedBy("this")
     private IStatsManager mService;
 
+    @GuardedBy("this")
+    private IStatsCompanionService mStatsCompanion;
+
     /**
      * Long extra of uid that added the relevant stats config.
      */
@@ -449,7 +463,9 @@ public final class StatsManager {
      * @throws StatsUnavailableException if unsuccessful due to failing to connect to stats service
      *
      * @hide
+     * @deprecated Please use registerPullAtomCallback
      */
+    @Deprecated
     @RequiresPermission(allOf = { DUMP, PACKAGE_USAGE_STATS })
     public void setPullerCallback(int atomTag, IStatsPullerCallback callback)
             throws StatsUnavailableException {
@@ -472,6 +488,75 @@ public final class StatsManager {
         }
     }
 
+
+    /**
+     * Registers a callback for an atom when that atom is to be pulled. The stats service will
+     * invoke pullData in the callback when the stats service determines that this atom needs to be
+     * pulled.
+     *
+     * @param atomTag           The tag of the atom for this puller callback.
+     * @param coolDownNs        The minimum time between successive pulls. A cache of the previous
+     *                          pull will be used if the time between pulls is less than coolDownNs.
+     * @param timeoutNs         The maximum time a pull should take. Statsd will wait timeoutNs for
+     *                          the pull to complete before timing out and marking the pull as
+     *                          failed.
+     * @param additiveFields    Fields that are added when mapping isolated uids to host uids.
+     * @param callback          The callback to be invoked when the stats service pulls the atom.
+     * @throws RemoteException  if unsuccessful due to failing to connect to system server.
+     *
+     * @hide
+     */
+    public void registerPullAtomCallback(int atomTag, long coolDownNs, long timeoutNs,
+            int[] additiveFields, @NonNull StatsPullAtomCallback callback,
+            @NonNull Executor executor) throws RemoteException, SecurityException {
+        synchronized (this) {
+            IStatsCompanionService service = getIStatsCompanionServiceLocked();
+            PullAtomCallbackInternal rec =
+                    new PullAtomCallbackInternal(atomTag, callback, executor);
+            service.registerPullAtomCallback(atomTag, coolDownNs, timeoutNs, additiveFields, rec);
+        }
+    }
+
+    private static class  PullAtomCallbackInternal extends IPullAtomCallback.Stub {
+        public final int mAtomId;
+        public final StatsPullAtomCallback mCallback;
+        public final Executor mExecutor;
+
+        PullAtomCallbackInternal(int atomId, StatsPullAtomCallback callback, Executor executor) {
+            mAtomId = atomId;
+            mCallback = callback;
+            mExecutor = executor;
+        }
+
+        @Override
+        public void onPullAtom(int atomTag, IPullAtomResultReceiver resultReceiver) {
+            mExecutor.execute(() -> {
+                List<StatsEvent> data = new ArrayList<>();
+                boolean success = mCallback.onPullAtom(atomTag, data);
+                StatsEvent[] arr = new StatsEvent[data.size()];
+                arr = data.toArray(arr);
+                try {
+                    resultReceiver.pullFinished(atomTag, success, arr);
+                } catch (RemoteException e) {
+                    Slog.w(TAG, "StatsPullResultReceiver failed for tag " + mAtomId);
+                }
+            });
+        }
+    }
+
+    /**
+     * Callback interface for pulling atoms requested by the stats service.
+     *
+     * @hide
+     */
+    public interface StatsPullAtomCallback {
+        /**
+         * Pull data for the specified atom tag, filling in the provided list of StatsEvent data.
+         * @return if the pull was successful
+         */
+        boolean onPullAtom(int atomTag, List<StatsEvent> data);
+    }
+
     private class StatsdDeathRecipient implements IBinder.DeathRecipient {
         @Override
         public void binderDied() {
@@ -481,6 +566,7 @@ public final class StatsManager {
         }
     }
 
+    @GuardedBy("this")
     private IStatsManager getIStatsManagerLocked() throws StatsUnavailableException {
         if (mService != null) {
             return mService;
@@ -497,6 +583,16 @@ public final class StatsManager {
         return mService;
     }
 
+    @GuardedBy("this")
+    private IStatsCompanionService getIStatsCompanionServiceLocked() {
+        if (mStatsCompanion != null) {
+            return mStatsCompanion;
+        }
+        mStatsCompanion = IStatsCompanionService.Stub.asInterface(
+                ServiceManager.getService("statscompanion"));
+        return mStatsCompanion;
+    }
+
     /**
      * Exception thrown when communication with the stats service fails (eg if it is not available).
      * This might be thrown early during boot before the stats service has started or if it crashed.

core/java/android/os/IPullAtomCallback.aidl

@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2019 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 android.os;
+
+import android.os.IPullAtomResultReceiver;
+
+/**
+  * Binder interface to pull atoms for the stats service.
+  * {@hide}
+  */
+interface IPullAtomCallback {
+    /**
+     * Initiate a request for a pull for an atom.
+     */
+     void onPullAtom(int atomTag, IPullAtomResultReceiver resultReceiver);
+
+}

core/java/android/os/IPullAtomResultReceiver.aidl

@@ -0,0 +1,32 @@
+/*
+ * Copyright (C) 2019 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 android.os;
+
+import android.util.StatsEvent;
+
+/**
+  * Binder interface to pull atoms for the stats service.
+  * {@hide}
+  */
+interface IPullAtomResultReceiver {
+
+    /**
+     * Indicate that a pull request for an atom is complete.
+     */
+     oneway void pullFinished(int atomTag, boolean success, in StatsEvent[] output);
+
+}

core/java/android/os/IStatsCompanionService.aidl

@@ -16,6 +16,7 @@
 
 package android.os;
 
+import android.os.IPullAtomCallback;
 import android.os.StatsDimensionsValue;
 import android.os.StatsLogEventWrapper;
 
@@ -85,4 +86,8 @@ interface IStatsCompanionService {
 
     /** Tells StatsCompaionService to grab the uid map snapshot and send it to statsd. */
     oneway void triggerUidSnapshot();
+
+    /** Tells StatsCompanionService to tell statsd to register a puller for the given atom id */
+    oneway void registerPullAtomCallback(int atomTag, long coolDownNs, long timeoutNs,
+            in int[] additiveFields, IPullAtomCallback pullerCallback);
 }

core/java/android/os/IStatsManager.aidl

@@ -17,6 +17,7 @@
 package android.os;
 
 import android.os.IStatsPullerCallback;
+import android.os.IPullAtomCallback;
 import android.os.ParcelFileDescriptor;
 
 /**
@@ -188,10 +189,18 @@ interface IStatsManager {
      * for the specified vendor atom tag.
      *
      * Requires Manifest.permission.DUMP and Manifest.permission.PACKAGE_USAGE_STATS
+     * @deprecated please use registerPullAtomCallback.
      */
     oneway void registerPullerCallback(int atomTag, IStatsPullerCallback pullerCallback,
                                        String packageName);
 
+   /**
+    * Registers a puller callback function that, when invoked, pulls the data
+    * for the specified atom tag.
+    */
+    oneway void registerPullAtomCallback(int uid, int atomTag, long coolDownNs, long timeoutNs,
+                           in int[] additiveFields, IPullAtomCallback pullerCallback);
+
    /**
     * Unregisters a puller callback function for the given vendor atom.
     *

core/java/android/os/IStatsPullerCallback.aidl

@@ -19,6 +19,7 @@ package android.os;
 import android.os.StatsLogEventWrapper;
 
 /**
+  * DEPRECATED
   * Binder interface to pull atoms for the stats service.
   * {@hide}
   */

core/java/android/util/StatsEvent.aidl

@@ -0,0 +1,19 @@
+/**
+ * Copyright (c) 2019, 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 android.util;
+
+parcelable StatsEvent cpp_header "android/util/StatsEvent.h";

core/java/android/util/StatsEvent.java

@@ -20,6 +20,8 @@ import static java.nio.charset.StandardCharsets.UTF_8;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.os.Parcel;
+import android.os.Parcelable;
 import android.os.SystemClock;
 
 import com.android.internal.annotations.GuardedBy;
@@ -42,7 +44,7 @@ import com.android.internal.annotations.VisibleForTesting;
  * </pre>
  * @hide
  **/
-public final class StatsEvent {
+public final class StatsEvent implements Parcelable {
     private static final int LOGGER_ENTRY_MAX_PAYLOAD = 4068;
 
     // Max payload size is 4 bytes less as 4 bytes are reserved for statsEventTag.
@@ -631,4 +633,39 @@ public final class StatsEvent {
             return 0;
         }
     }
+
+    /**
+     * Boilerplate for Parcel.
+     *
+     * @hide
+     */
+    public static final @NonNull Parcelable.Creator<StatsEvent> CREATOR =
+            new Parcelable.Creator<StatsEvent>() {
+                public StatsEvent createFromParcel(Parcel in) {
+                    // Purposefully leaving this method not implemented.
+                    throw new RuntimeException("Not implemented");
+                }
+
+                public StatsEvent[] newArray(int size) {
+                    // Purposefully leaving this method not implemented.
+                    throw new RuntimeException("Not implemented");
+                }
+            };
+
+    /**
+     * @hide
+     */
+    public void writeToParcel(Parcel out, int flags) {
+        out.writeInt(mAtomId);
+        out.writeInt(getNumBytes());
+        out.writeByteArray(getBytes());
+    }
+
+    /**
+     * Boilerplate for Parcel.
+     */
+    public int describeContents() {
+        return 0;
+    }
+
 }

libs/services/Android.bp

@@ -22,6 +22,7 @@ cc_library_shared {
         "src/os/DropBoxManager.cpp",
         "src/os/StatsDimensionsValue.cpp",
         "src/os/StatsLogEventWrapper.cpp",
+        "src/util/StatsEvent.cpp",
     ],
 
     shared_libs: [

libs/services/include/android/util/StatsEvent.h

@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2019 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.
+ */
+#ifndef STATS_EVENT_H
+#define STATS_EVENT_H
+
+#include <binder/Parcel.h>
+#include <binder/Parcelable.h>
+#include <binder/Status.h>
+#include <vector>
+
+namespace android {
+namespace util {
+class StatsEvent : public android::Parcelable {
+    public:
+        StatsEvent();
+
+        StatsEvent(StatsEvent&& in) = default;
+
+        android::status_t writeToParcel(android::Parcel* out) const;
+
+        android::status_t readFromParcel(const android::Parcel* in);
+
+    private:
+        int mAtomTag;
+        std::vector<uint8_t> mBuffer;
+};
+} // Namespace util
+} // Namespace android
+
+#endif  // STATS_ EVENT_H

libs/services/src/util/StatsEvent.cpp

@@ -0,0 +1,58 @@
+/*
+ * Copyright (C) 2019 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 <android/util/StatsEvent.h>
+
+#include <binder/Parcel.h>
+#include <binder/Parcelable.h>
+#include <binder/Status.h>
+#include <vector>
+
+using android::Parcel;
+using android::Parcelable;
+using android::status_t;
+using std::vector;
+
+namespace android {
+namespace util {
+
+StatsEvent::StatsEvent(){};
+
+status_t StatsEvent::writeToParcel(Parcel* out) const {
+    // Implement me if desired. We don't currently use this.
+    ALOGE("Cannot do c++ StatsEvent.writeToParcel(); it is not implemented.");
+    (void)out;  // To prevent compile error of unused parameter 'out'
+    return UNKNOWN_ERROR;
+};
+
+status_t StatsEvent::readFromParcel(const Parcel* in) {
+    status_t res = OK;
+    if (in == NULL) {
+        ALOGE("statsd received parcel argument was NULL.");
+        return BAD_VALUE;
+    }
+    if ((res = in->readInt32(&mAtomTag)) != OK) {
+        ALOGE("statsd could not read atom tag from parcel");
+        return res;
+    }
+    if ((res = in->readByteVector(&mBuffer)) != OK) {
+        ALOGE("statsd could not read buffer from parcel");
+        return res;
+    }
+    return NO_ERROR;
+};
+
+} // Namespace util
+} // Namespace android