Merge "libsnapshot: Add RecoveryCreateSnapshotDevices(device)" into rvc-dev
diff --git a/init/Android.bp b/init/Android.bp
index 72a7bfe..d2bdf98 100644
--- a/init/Android.bp
+++ b/init/Android.bp
@@ -240,6 +240,7 @@
         "firmware_handler_test.cpp",
         "init_test.cpp",
         "keychords_test.cpp",
+        "oneshot_on_test.cpp",
         "persistent_properties_test.cpp",
         "property_service_test.cpp",
         "property_type_test.cpp",
diff --git a/init/README.md b/init/README.md
index 4f0a7ec..13f1bac 100644
--- a/init/README.md
+++ b/init/README.md
@@ -720,23 +720,35 @@
   characteristics in a device agnostic manner.
 
 Init responds to properties that begin with `ctl.`.  These properties take the format of
-`ctl.<command>` and the _value_ of the system property is used as a parameter, for example:
-`SetProperty("ctl.start", "logd")` will run the `start` command on `logd`.  Note that these
+`ctl.[<target>_]<command>` and the _value_ of the system property is used as a parameter.  The
+_target_ is optional and specifies the service option that _value_ is meant to match with.  There is
+only one option for _target_, `interface` which indicates that _value_ will refer to an interface
+that a service provides and not the service name itself.
+
+For example:
+
+`SetProperty("ctl.start", "logd")` will run the `start` command on `logd`.
+
+`SetProperty("ctl.interface_start", "aidl/aidl_lazy_test_1")` will run the `start` command on the
+service that exposes the `aidl aidl_lazy_test_1` interface.
+
+Note that these
 properties are only settable; they will have no value when read.
 
-`ctl.start` \
-`ctl.restart` \
-`ctl.stop`
-> These are equivalent to using the `start`, `restart`, and `stop` commands on the service specified
+The _commands_ are listed below.
+
+`start` \
+`restart` \
+`stop` \
+These are equivalent to using the `start`, `restart`, and `stop` commands on the service specified
 by the _value_ of the property.
 
-`ctl.interface_start` \
-`ctl.interface_restart` \
-`ctl.interface_stop`
-> These are equivalent to using the `interface_start`, `interface_restart`, and `interface_stop`
-commands on the interface specified by the _value_ of the property.
+`oneshot_one` and `oneshot_off` will turn on or off the _oneshot_
+flag for the service specified by the _value_ of the property.  This is
+particularly intended for services that are conditionally lazy HALs.  When
+they are lazy HALs, oneshot must be on, otherwise oneshot should be off.
 
-`ctl.sigstop_on` and `ctl.sigstop_off` will turn on or off the _sigstop_ feature for the service
+`sigstop_on` and `sigstop_off` will turn on or off the _sigstop_ feature for the service
 specified by the _value_ of the property.  See the _Debugging init_ section below for more details
 about this feature.
 
diff --git a/init/init.cpp b/init/init.cpp
index 4289dcf..5444a32 100644
--- a/init/init.cpp
+++ b/init/init.cpp
@@ -81,6 +81,7 @@
 using namespace std::string_literals;
 
 using android::base::boot_clock;
+using android::base::ConsumePrefix;
 using android::base::GetProperty;
 using android::base::ReadFileToString;
 using android::base::SetProperty;
@@ -367,40 +368,27 @@
     INTERFACE,  // action gets called for every service that holds this interface
 };
 
-struct ControlMessageFunction {
-    ControlTarget target;
-    std::function<Result<void>(Service*)> action;
-};
+using ControlMessageFunction = std::function<Result<void>(Service*)>;
 
-static const std::map<std::string, ControlMessageFunction>& get_control_message_map() {
+static const std::map<std::string, ControlMessageFunction, std::less<>>& GetControlMessageMap() {
     // clang-format off
-    static const std::map<std::string, ControlMessageFunction> control_message_functions = {
-        {"sigstop_on",        {ControlTarget::SERVICE,
-                               [](auto* service) { service->set_sigstop(true); return Result<void>{}; }}},
-        {"sigstop_off",       {ControlTarget::SERVICE,
-                               [](auto* service) { service->set_sigstop(false); return Result<void>{}; }}},
-        {"start",             {ControlTarget::SERVICE,   DoControlStart}},
-        {"stop",              {ControlTarget::SERVICE,   DoControlStop}},
-        {"restart",           {ControlTarget::SERVICE,   DoControlRestart}},
-        {"interface_start",   {ControlTarget::INTERFACE, DoControlStart}},
-        {"interface_stop",    {ControlTarget::INTERFACE, DoControlStop}},
-        {"interface_restart", {ControlTarget::INTERFACE, DoControlRestart}},
+    static const std::map<std::string, ControlMessageFunction, std::less<>> control_message_functions = {
+        {"sigstop_on",        [](auto* service) { service->set_sigstop(true); return Result<void>{}; }},
+        {"sigstop_off",       [](auto* service) { service->set_sigstop(false); return Result<void>{}; }},
+        {"oneshot_on",        [](auto* service) { service->set_oneshot(true); return Result<void>{}; }},
+        {"oneshot_off",       [](auto* service) { service->set_oneshot(false); return Result<void>{}; }},
+        {"start",             DoControlStart},
+        {"stop",              DoControlStop},
+        {"restart",           DoControlRestart},
     };
     // clang-format on
 
     return control_message_functions;
 }
 
-bool HandleControlMessage(const std::string& msg, const std::string& name, pid_t pid) {
-    const auto& map = get_control_message_map();
-    const auto it = map.find(msg);
-
-    if (it == map.end()) {
-        LOG(ERROR) << "Unknown control msg '" << msg << "'";
-        return false;
-    }
-
-    std::string cmdline_path = StringPrintf("proc/%d/cmdline", pid);
+static bool HandleControlMessage(std::string_view message, const std::string& name,
+                                 pid_t from_pid) {
+    std::string cmdline_path = StringPrintf("proc/%d/cmdline", from_pid);
     std::string process_cmdline;
     if (ReadFileToString(cmdline_path, &process_cmdline)) {
         std::replace(process_cmdline.begin(), process_cmdline.end(), '\0', ' ');
@@ -409,37 +397,37 @@
         process_cmdline = "unknown process";
     }
 
-    const ControlMessageFunction& function = it->second;
-
-    Service* svc = nullptr;
-
-    switch (function.target) {
-        case ControlTarget::SERVICE:
-            svc = ServiceList::GetInstance().FindService(name);
-            break;
-        case ControlTarget::INTERFACE:
-            svc = ServiceList::GetInstance().FindInterface(name);
-            break;
-        default:
-            LOG(ERROR) << "Invalid function target from static map key ctl." << msg << ": "
-                       << static_cast<std::underlying_type<ControlTarget>::type>(function.target);
-            return false;
+    Service* service = nullptr;
+    auto action = message;
+    if (ConsumePrefix(&action, "interface_")) {
+        service = ServiceList::GetInstance().FindInterface(name);
+    } else {
+        service = ServiceList::GetInstance().FindService(name);
     }
 
-    if (svc == nullptr) {
-        LOG(ERROR) << "Control message: Could not find '" << name << "' for ctl." << msg
-                   << " from pid: " << pid << " (" << process_cmdline << ")";
+    if (service == nullptr) {
+        LOG(ERROR) << "Control message: Could not find '" << name << "' for ctl." << message
+                   << " from pid: " << from_pid << " (" << process_cmdline << ")";
         return false;
     }
 
-    if (auto result = function.action(svc); !result.ok()) {
-        LOG(ERROR) << "Control message: Could not ctl." << msg << " for '" << name
-                   << "' from pid: " << pid << " (" << process_cmdline << "): " << result.error();
+    const auto& map = GetControlMessageMap();
+    const auto it = map.find(action);
+    if (it == map.end()) {
+        LOG(ERROR) << "Unknown control msg '" << message << "'";
+        return false;
+    }
+    const auto& function = it->second;
+
+    if (auto result = function(service); !result.ok()) {
+        LOG(ERROR) << "Control message: Could not ctl." << message << " for '" << name
+                   << "' from pid: " << from_pid << " (" << process_cmdline
+                   << "): " << result.error();
         return false;
     }
 
-    LOG(INFO) << "Control message: Processed ctl." << msg << " for '" << name
-              << "' from pid: " << pid << " (" << process_cmdline << ")";
+    LOG(INFO) << "Control message: Processed ctl." << message << " for '" << name
+              << "' from pid: " << from_pid << " (" << process_cmdline << ")";
     return true;
 }
 
diff --git a/init/oneshot_on_test.cpp b/init/oneshot_on_test.cpp
new file mode 100644
index 0000000..7e7cc36
--- /dev/null
+++ b/init/oneshot_on_test.cpp
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2020 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 <gtest/gtest.h>
+
+#include <chrono>
+
+#include <android-base/properties.h>
+
+using android::base::GetProperty;
+using android::base::SetProperty;
+using android::base::WaitForProperty;
+using namespace std::literals;
+
+TEST(init, oneshot_on) {
+    // Bootanim shouldn't be running once the device has booted.
+    ASSERT_EQ("stopped", GetProperty("init.svc.bootanim", ""));
+
+    SetProperty("ctl.oneshot_off", "bootanim");
+    SetProperty("ctl.start", "bootanim");
+
+    // Bootanim exits quickly when the device is fully booted, so check that it goes back to the
+    // 'restarting' state that non-oneshot services enter once they've restarted.
+    EXPECT_TRUE(WaitForProperty("init.svc.bootanim", "restarting", 10s));
+
+    SetProperty("ctl.oneshot_on", "bootanim");
+    SetProperty("ctl.start", "bootanim");
+
+    // Now that oneshot is enabled again, bootanim should transition into the 'stopped' state.
+    EXPECT_TRUE(WaitForProperty("init.svc.bootanim", "stopped", 10s));
+}
diff --git a/init/service.h b/init/service.h
index cf3f0c2..9f1d697 100644
--- a/init/service.h
+++ b/init/service.h
@@ -130,6 +130,13 @@
     bool is_updatable() const { return updatable_; }
     bool is_post_data() const { return post_data_; }
     bool is_from_apex() const { return from_apex_; }
+    void set_oneshot(bool value) {
+        if (value) {
+            flags_ |= SVC_ONESHOT;
+        } else {
+            flags_ &= ~SVC_ONESHOT;
+        }
+    }
 
   private:
     void NotifyStateChange(const std::string& new_state) const;
diff --git a/libstats/pull/Android.bp b/libstats/pull/Android.bp
index 1a9cb92..0fb8f1b 100644
--- a/libstats/pull/Android.bp
+++ b/libstats/pull/Android.bp
@@ -65,3 +65,25 @@
         "//frameworks/base/apex/statsd/tests/libstatspull",
     ],
 }
+
+// Note: These unit tests only test PullAtomMetadata.
+// For full E2E tests of libstatspull, use LibStatsPullTests
+cc_test {
+    name: "libstatspull_test",
+    srcs: [
+        "tests/pull_atom_metadata_test.cpp",
+    ],
+    shared_libs: [
+        "libstatspull",
+        "libstatssocket",
+    ],
+    test_suites: ["general-tests"],
+    cflags: [
+        "-Wall",
+        "-Werror",
+        "-Wno-missing-field-initializers",
+        "-Wno-unused-variable",
+        "-Wno-unused-function",
+        "-Wno-unused-parameter",
+    ],
+}
\ No newline at end of file
diff --git a/libstats/pull/TEST_MAPPING b/libstats/pull/TEST_MAPPING
new file mode 100644
index 0000000..76f4f02
--- /dev/null
+++ b/libstats/pull/TEST_MAPPING
@@ -0,0 +1,7 @@
+{
+  "presubmit" : [
+    {
+      "name" : "libstatspull_test"
+    }
+  ]
+}
\ No newline at end of file
diff --git a/libstats/pull/include/stats_pull_atom_callback.h b/libstats/pull/include/stats_pull_atom_callback.h
index 0b0df2b..c976c68 100644
--- a/libstats/pull/include/stats_pull_atom_callback.h
+++ b/libstats/pull/include/stats_pull_atom_callback.h
@@ -45,17 +45,27 @@
 void AStatsManager_PullAtomMetadata_release(AStatsManager_PullAtomMetadata* metadata);
 
 /**
- * Set the cool down time of the pull in nanoseconds. If two successive pulls are issued
+ * Set the cool down time of the pull in milliseconds. If two successive pulls are issued
  * within the cool down, a cached version of the first will be used for the second.
  */
-void AStatsManager_PullAtomMetadata_setCoolDownNs(AStatsManager_PullAtomMetadata* metadata,
-                                                  int64_t cool_down_ns);
+void AStatsManager_PullAtomMetadata_setCoolDownMillis(AStatsManager_PullAtomMetadata* metadata,
+                                                      int64_t cool_down_millis);
 
 /**
- * Set the maximum time the pull can take in nanoseconds.
+ * Get the cool down time of the pull in milliseconds.
  */
-void AStatsManager_PullAtomMetadata_setTimeoutNs(AStatsManager_PullAtomMetadata* metadata,
-                                                 int64_t timeout_ns);
+int64_t AStatsManager_PullAtomMetadata_getCoolDownMillis(AStatsManager_PullAtomMetadata* metadata);
+
+/**
+ * Set the maximum time the pull can take in milliseconds.
+ */
+void AStatsManager_PullAtomMetadata_setTimeoutMillis(AStatsManager_PullAtomMetadata* metadata,
+                                                     int64_t timeout_millis);
+
+/**
+ * Get the maximum time the pull can take in milliseconds.
+ */
+int64_t AStatsManager_PullAtomMetadata_getTimeoutMillis(AStatsManager_PullAtomMetadata* metadata);
 
 /**
  * Set the additive fields of this pulled atom.
@@ -65,7 +75,25 @@
  * will be combined when the non-additive fields are the same.
  */
 void AStatsManager_PullAtomMetadata_setAdditiveFields(AStatsManager_PullAtomMetadata* metadata,
-                                                      int* additive_fields, int num_fields);
+                                                      int32_t* additive_fields, int32_t num_fields);
+
+/**
+ * Get the number of additive fields for this pulled atom. This is intended to be called before
+ * AStatsManager_PullAtomMetadata_getAdditiveFields to determine the size of the array.
+ */
+int32_t AStatsManager_PullAtomMetadata_getNumAdditiveFields(
+        AStatsManager_PullAtomMetadata* metadata);
+
+/**
+ * Get the additive fields of this pulled atom.
+ *
+ * \param fields an output parameter containing the additive fields for this PullAtomMetadata.
+ *               Fields is an array and it is assumed that it is at least as large as the number of
+ *               additive fields, which can be obtained by calling
+ *               AStatsManager_PullAtomMetadata_getNumAdditiveFields.
+ */
+void AStatsManager_PullAtomMetadata_getAdditiveFields(AStatsManager_PullAtomMetadata* metadata,
+                                                      int32_t* fields);
 
 /**
  * Return codes for the result of a pull.
@@ -108,7 +136,7 @@
 typedef AStatsManager_PullAtomCallbackReturn (*AStatsManager_PullAtomCallback)(
         int32_t atom_tag, AStatsEventList* data, void* cookie);
 /**
- * Registers a callback for an atom when that atom is to be pulled. The stats service will
+ * Sets a callback for an atom when that atom is to be pulled. The stats service will
  * invoke the callback when the stats service determines that this atom needs to be
  * pulled.
  *
@@ -122,19 +150,18 @@
  * \param cookie            A pointer that will be passed back to the callback.
  *                          It has no meaning to statsd.
  */
-void AStatsManager_registerPullAtomCallback(int32_t atom_tag,
-                                            AStatsManager_PullAtomCallback callback,
-                                            AStatsManager_PullAtomMetadata* metadata, void* cookie);
+void AStatsManager_setPullAtomCallback(int32_t atom_tag, AStatsManager_PullAtomMetadata* metadata,
+                                       AStatsManager_PullAtomCallback callback, void* cookie);
 
 /**
- * Unregisters a callback for an atom when that atom is to be pulled. Note that any ongoing
+ * Clears a callback for an atom when that atom is to be pulled. Note that any ongoing
  * pulls will still occur.
  *
  * Requires the REGISTER_STATS_PULL_ATOM permission.
  *
  * \param atomTag           The tag of the atom of which to unregister
  */
-void AStatsManager_unregisterPullAtomCallback(int32_t atom_tag);
+void AStatsManager_clearPullAtomCallback(int32_t atom_tag);
 
 #ifdef __cplusplus
 }
diff --git a/libstats/pull/libstatspull.map.txt b/libstats/pull/libstatspull.map.txt
index dc3fd8b..e0e851a 100644
--- a/libstats/pull/libstatspull.map.txt
+++ b/libstats/pull/libstatspull.map.txt
@@ -2,12 +2,16 @@
     global:
         AStatsManager_PullAtomMetadata_obtain; # apex # introduced=30
         AStatsManager_PullAtomMetadata_release; # apex # introduced=30
-        AStatsManager_PullAtomMetadata_setCoolDownNs; # apex # introduced=30
-        AStatsManager_PullAtomMetadata_setTimeoutNs; # apex # introduced=30
+        AStatsManager_PullAtomMetadata_setCoolDownMillis; # apex # introduced=30
+        AStatsManager_PullAtomMetadata_getCoolDownMillis; # apex # introduced=30
+        AStatsManager_PullAtomMetadata_setTimeoutMillis; # apex # introduced=30
+        AStatsManager_PullAtomMetadata_getTimeoutMillis; # apex # introduced=30
         AStatsManager_PullAtomMetadata_setAdditiveFields; # apex # introduced=30
+        AStatsManager_PullAtomMetadata_getNumAdditiveFields; # apex # introduced=30
+        AStatsManager_PullAtomMetadata_getAdditiveFields; # apex # introduced=30
         AStatsEventList_addStatsEvent; # apex # introduced=30
-        AStatsManager_registerPullAtomCallback; # apex # introduced=30
-        AStatsManager_unregisterPullAtomCallback; # apex # introduced=30
+        AStatsManager_setPullAtomCallback; # apex # introduced=30
+        AStatsManager_clearPullAtomCallback; # apex # introduced=30
     local:
         *;
 };
diff --git a/libstats/pull/stats_pull_atom_callback.cpp b/libstats/pull/stats_pull_atom_callback.cpp
index 27e9d29..2d184bd 100644
--- a/libstats/pull/stats_pull_atom_callback.cpp
+++ b/libstats/pull/stats_pull_atom_callback.cpp
@@ -46,19 +46,19 @@
     return event;
 }
 
-static const int64_t DEFAULT_COOL_DOWN_NS = 1000000000LL;  // 1 second.
-static const int64_t DEFAULT_TIMEOUT_NS = 10000000000LL;   // 10 seconds.
+static const int64_t DEFAULT_COOL_DOWN_MILLIS = 1000LL;  // 1 second.
+static const int64_t DEFAULT_TIMEOUT_MILLIS = 10000LL;   // 10 seconds.
 
 struct AStatsManager_PullAtomMetadata {
-    int64_t cool_down_ns;
-    int64_t timeout_ns;
+    int64_t cool_down_millis;
+    int64_t timeout_millis;
     std::vector<int32_t> additive_fields;
 };
 
 AStatsManager_PullAtomMetadata* AStatsManager_PullAtomMetadata_obtain() {
     AStatsManager_PullAtomMetadata* metadata = new AStatsManager_PullAtomMetadata();
-    metadata->cool_down_ns = DEFAULT_COOL_DOWN_NS;
-    metadata->timeout_ns = DEFAULT_TIMEOUT_NS;
+    metadata->cool_down_millis = DEFAULT_COOL_DOWN_MILLIS;
+    metadata->timeout_millis = DEFAULT_TIMEOUT_MILLIS;
     metadata->additive_fields = std::vector<int32_t>();
     return metadata;
 }
@@ -67,30 +67,49 @@
     delete metadata;
 }
 
-void AStatsManager_PullAtomMetadata_setCoolDownNs(AStatsManager_PullAtomMetadata* metadata,
-                                                  int64_t cool_down_ns) {
-    metadata->cool_down_ns = cool_down_ns;
+void AStatsManager_PullAtomMetadata_setCoolDownMillis(AStatsManager_PullAtomMetadata* metadata,
+                                                      int64_t cool_down_millis) {
+    metadata->cool_down_millis = cool_down_millis;
 }
 
-void AStatsManager_PullAtomMetadata_setTimeoutNs(AStatsManager_PullAtomMetadata* metadata,
-                                                 int64_t timeout_ns) {
-    metadata->timeout_ns = timeout_ns;
+int64_t AStatsManager_PullAtomMetadata_getCoolDownMillis(AStatsManager_PullAtomMetadata* metadata) {
+    return metadata->cool_down_millis;
+}
+
+void AStatsManager_PullAtomMetadata_setTimeoutMillis(AStatsManager_PullAtomMetadata* metadata,
+                                                     int64_t timeout_millis) {
+    metadata->timeout_millis = timeout_millis;
+}
+
+int64_t AStatsManager_PullAtomMetadata_getTimeoutMillis(AStatsManager_PullAtomMetadata* metadata) {
+    return metadata->timeout_millis;
 }
 
 void AStatsManager_PullAtomMetadata_setAdditiveFields(AStatsManager_PullAtomMetadata* metadata,
-                                                      int* additive_fields, int num_fields) {
+                                                      int32_t* additive_fields,
+                                                      int32_t num_fields) {
     metadata->additive_fields.assign(additive_fields, additive_fields + num_fields);
 }
 
+int32_t AStatsManager_PullAtomMetadata_getNumAdditiveFields(
+        AStatsManager_PullAtomMetadata* metadata) {
+    return metadata->additive_fields.size();
+}
+
+void AStatsManager_PullAtomMetadata_getAdditiveFields(AStatsManager_PullAtomMetadata* metadata,
+                                                      int32_t* fields) {
+    std::copy(metadata->additive_fields.begin(), metadata->additive_fields.end(), fields);
+}
+
 class StatsPullAtomCallbackInternal : public BnPullAtomCallback {
   public:
     StatsPullAtomCallbackInternal(const AStatsManager_PullAtomCallback callback, void* cookie,
-                                  const int64_t coolDownNs, const int64_t timeoutNs,
+                                  const int64_t coolDownMillis, const int64_t timeoutMillis,
                                   const std::vector<int32_t> additiveFields)
         : mCallback(callback),
           mCookie(cookie),
-          mCoolDownNs(coolDownNs),
-          mTimeoutNs(timeoutNs),
+          mCoolDownMillis(coolDownMillis),
+          mTimeoutMillis(timeoutMillis),
           mAdditiveFields(additiveFields) {}
 
     Status onPullAtom(int32_t atomTag,
@@ -119,15 +138,15 @@
         return Status::ok();
     }
 
-    const int64_t& getCoolDownNs() const { return mCoolDownNs; }
-    const int64_t& getTimeoutNs() const { return mTimeoutNs; }
+    int64_t getCoolDownMillis() const { return mCoolDownMillis; }
+    int64_t getTimeoutMillis() const { return mTimeoutMillis; }
     const std::vector<int32_t>& getAdditiveFields() const { return mAdditiveFields; }
 
   private:
     const AStatsManager_PullAtomCallback mCallback;
     void* mCookie;
-    const int64_t mCoolDownNs;
-    const int64_t mTimeoutNs;
+    const int64_t mCoolDownMillis;
+    const int64_t mTimeoutMillis;
     const std::vector<int32_t> mAdditiveFields;
 };
 
@@ -156,8 +175,8 @@
         pullersCopy = mPullers;
     }
     for (const auto& it : pullersCopy) {
-        statsService->registerNativePullAtomCallback(it.first, it.second->getCoolDownNs(),
-                                                     it.second->getTimeoutNs(),
+        statsService->registerNativePullAtomCallback(it.first, it.second->getCoolDownMillis(),
+                                                     it.second->getTimeoutMillis(),
                                                      it.second->getAdditiveFields(), it.second);
     }
 }
@@ -186,8 +205,8 @@
         return;
     }
 
-    statsService->registerNativePullAtomCallback(atomTag, cb->getCoolDownNs(), cb->getTimeoutNs(),
-                                                 cb->getAdditiveFields(), cb);
+    statsService->registerNativePullAtomCallback(
+            atomTag, cb->getCoolDownMillis(), cb->getTimeoutMillis(), cb->getAdditiveFields(), cb);
 }
 
 void unregisterStatsPullAtomCallbackBlocking(int32_t atomTag) {
@@ -200,12 +219,11 @@
     statsService->unregisterNativePullAtomCallback(atomTag);
 }
 
-void AStatsManager_registerPullAtomCallback(int32_t atom_tag,
-                                            AStatsManager_PullAtomCallback callback,
-                                            AStatsManager_PullAtomMetadata* metadata,
-                                            void* cookie) {
-    int64_t coolDownNs = metadata == nullptr ? DEFAULT_COOL_DOWN_NS : metadata->cool_down_ns;
-    int64_t timeoutNs = metadata == nullptr ? DEFAULT_TIMEOUT_NS : metadata->timeout_ns;
+void AStatsManager_setPullAtomCallback(int32_t atom_tag, AStatsManager_PullAtomMetadata* metadata,
+                                       AStatsManager_PullAtomCallback callback, void* cookie) {
+    int64_t coolDownMillis =
+            metadata == nullptr ? DEFAULT_COOL_DOWN_MILLIS : metadata->cool_down_millis;
+    int64_t timeoutMillis = metadata == nullptr ? DEFAULT_TIMEOUT_MILLIS : metadata->timeout_millis;
 
     std::vector<int32_t> additiveFields;
     if (metadata != nullptr) {
@@ -213,8 +231,8 @@
     }
 
     std::shared_ptr<StatsPullAtomCallbackInternal> callbackBinder =
-            SharedRefBase::make<StatsPullAtomCallbackInternal>(callback, cookie, coolDownNs,
-                                                               timeoutNs, additiveFields);
+            SharedRefBase::make<StatsPullAtomCallbackInternal>(callback, cookie, coolDownMillis,
+                                                               timeoutMillis, additiveFields);
 
     {
         std::lock_guard<std::mutex> lg(pullAtomMutex);
@@ -226,7 +244,7 @@
     registerThread.detach();
 }
 
-void AStatsManager_unregisterPullAtomCallback(int32_t atom_tag) {
+void AStatsManager_clearPullAtomCallback(int32_t atom_tag) {
     {
         std::lock_guard<std::mutex> lg(pullAtomMutex);
         // Always remove the puller from our map.
diff --git a/libstats/pull/tests/pull_atom_metadata_test.cpp b/libstats/pull/tests/pull_atom_metadata_test.cpp
new file mode 100644
index 0000000..cf19303
--- /dev/null
+++ b/libstats/pull/tests/pull_atom_metadata_test.cpp
@@ -0,0 +1,92 @@
+/*
+ * Copyright (C) 2020, 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 <gtest/gtest.h>
+
+#include <stats_pull_atom_callback.h>
+
+namespace {
+
+static const int64_t DEFAULT_COOL_DOWN_MILLIS = 1000LL;  // 1 second.
+static const int64_t DEFAULT_TIMEOUT_MILLIS = 10000LL;   // 10 seconds.
+
+}  // anonymous namespace
+
+TEST(AStatsManager_PullAtomMetadataTest, TestEmpty) {
+    AStatsManager_PullAtomMetadata* metadata = AStatsManager_PullAtomMetadata_obtain();
+    EXPECT_EQ(AStatsManager_PullAtomMetadata_getCoolDownMillis(metadata), DEFAULT_COOL_DOWN_MILLIS);
+    EXPECT_EQ(AStatsManager_PullAtomMetadata_getTimeoutMillis(metadata), DEFAULT_TIMEOUT_MILLIS);
+    EXPECT_EQ(AStatsManager_PullAtomMetadata_getNumAdditiveFields(metadata), 0);
+    AStatsManager_PullAtomMetadata_release(metadata);
+}
+
+TEST(AStatsManager_PullAtomMetadataTest, TestSetTimeoutMillis) {
+    int64_t timeoutMillis = 500;
+    AStatsManager_PullAtomMetadata* metadata = AStatsManager_PullAtomMetadata_obtain();
+    AStatsManager_PullAtomMetadata_setTimeoutMillis(metadata, timeoutMillis);
+    EXPECT_EQ(AStatsManager_PullAtomMetadata_getCoolDownMillis(metadata), DEFAULT_COOL_DOWN_MILLIS);
+    EXPECT_EQ(AStatsManager_PullAtomMetadata_getTimeoutMillis(metadata), timeoutMillis);
+    EXPECT_EQ(AStatsManager_PullAtomMetadata_getNumAdditiveFields(metadata), 0);
+    AStatsManager_PullAtomMetadata_release(metadata);
+}
+
+TEST(AStatsManager_PullAtomMetadataTest, TestSetCoolDownMillis) {
+    int64_t coolDownMillis = 10000;
+    AStatsManager_PullAtomMetadata* metadata = AStatsManager_PullAtomMetadata_obtain();
+    AStatsManager_PullAtomMetadata_setCoolDownMillis(metadata, coolDownMillis);
+    EXPECT_EQ(AStatsManager_PullAtomMetadata_getCoolDownMillis(metadata), coolDownMillis);
+    EXPECT_EQ(AStatsManager_PullAtomMetadata_getTimeoutMillis(metadata), DEFAULT_TIMEOUT_MILLIS);
+    EXPECT_EQ(AStatsManager_PullAtomMetadata_getNumAdditiveFields(metadata), 0);
+    AStatsManager_PullAtomMetadata_release(metadata);
+}
+
+TEST(AStatsManager_PullAtomMetadataTest, TestSetAdditiveFields) {
+    const int numFields = 3;
+    int inputFields[numFields] = {2, 4, 6};
+    AStatsManager_PullAtomMetadata* metadata = AStatsManager_PullAtomMetadata_obtain();
+    AStatsManager_PullAtomMetadata_setAdditiveFields(metadata, inputFields, numFields);
+    EXPECT_EQ(AStatsManager_PullAtomMetadata_getCoolDownMillis(metadata), DEFAULT_COOL_DOWN_MILLIS);
+    EXPECT_EQ(AStatsManager_PullAtomMetadata_getTimeoutMillis(metadata), DEFAULT_TIMEOUT_MILLIS);
+    EXPECT_EQ(AStatsManager_PullAtomMetadata_getNumAdditiveFields(metadata), numFields);
+    int outputFields[numFields];
+    AStatsManager_PullAtomMetadata_getAdditiveFields(metadata, outputFields);
+    for (int i = 0; i < numFields; i++) {
+        EXPECT_EQ(inputFields[i], outputFields[i]);
+    }
+    AStatsManager_PullAtomMetadata_release(metadata);
+}
+
+TEST(AStatsManager_PullAtomMetadataTest, TestSetAllElements) {
+    int64_t timeoutMillis = 500;
+    int64_t coolDownMillis = 10000;
+    const int numFields = 3;
+    int inputFields[numFields] = {2, 4, 6};
+
+    AStatsManager_PullAtomMetadata* metadata = AStatsManager_PullAtomMetadata_obtain();
+    AStatsManager_PullAtomMetadata_setTimeoutMillis(metadata, timeoutMillis);
+    AStatsManager_PullAtomMetadata_setCoolDownMillis(metadata, coolDownMillis);
+    AStatsManager_PullAtomMetadata_setAdditiveFields(metadata, inputFields, numFields);
+
+    EXPECT_EQ(AStatsManager_PullAtomMetadata_getCoolDownMillis(metadata), coolDownMillis);
+    EXPECT_EQ(AStatsManager_PullAtomMetadata_getTimeoutMillis(metadata), timeoutMillis);
+    EXPECT_EQ(AStatsManager_PullAtomMetadata_getNumAdditiveFields(metadata), numFields);
+    int outputFields[numFields];
+    AStatsManager_PullAtomMetadata_getAdditiveFields(metadata, outputFields);
+    for (int i = 0; i < numFields; i++) {
+        EXPECT_EQ(inputFields[i], outputFields[i]);
+    }
+    AStatsManager_PullAtomMetadata_release(metadata);
+}