Fix PullUidProvider unregistering on config update

Previously, MetricsManagers would unregister themselves as a
PullUidProvider for a given ConfigKey in the destructor. This caused all
pulls to fail after a config update because the new MetricsManager would
register itself before the old MetricsManager was destructed and
unregistered. This resulted in the old MetricsManager removing the new
config since they shared the same config key. The fix is for the
PullerManager to check that the PullUidProviders are equal in the
unregister function before actually erasing it.

Test: bit statsd_test:* (wrote a failing test that now passes).
Test: statsd_localdrive to manually update a config, ensured pulls still
worked.
Bug: 154544328

Change-Id: Id7af3b3b407e24bee74fc34bd1c2b9e0575e9c9e
diff --git a/cmds/statsd/src/external/StatsPullerManager.cpp b/cmds/statsd/src/external/StatsPullerManager.cpp
index ebe9610..cfd5d14 100644
--- a/cmds/statsd/src/external/StatsPullerManager.cpp
+++ b/cmds/statsd/src/external/StatsPullerManager.cpp
@@ -252,9 +252,13 @@
     mPullUidProviders[configKey] = provider;
 }
 
-void StatsPullerManager::UnregisterPullUidProvider(const ConfigKey& configKey) {
+void StatsPullerManager::UnregisterPullUidProvider(const ConfigKey& configKey,
+                                                   wp<PullUidProvider> provider) {
     std::lock_guard<std::mutex> _l(mLock);
-    mPullUidProviders.erase(configKey);
+    const auto& it = mPullUidProviders.find(configKey);
+    if (it != mPullUidProviders.end() && it->second == provider) {
+        mPullUidProviders.erase(it);
+    }
 }
 
 void StatsPullerManager::OnAlarmFired(int64_t elapsedTimeNs) {
diff --git a/cmds/statsd/src/external/StatsPullerManager.h b/cmds/statsd/src/external/StatsPullerManager.h
index ab0ccee..5e18aaa 100644
--- a/cmds/statsd/src/external/StatsPullerManager.h
+++ b/cmds/statsd/src/external/StatsPullerManager.h
@@ -78,11 +78,12 @@
                                     wp<PullDataReceiver> receiver);
 
     // Registers a pull uid provider for the config key. When pulling atoms, it will be used to
-    // determine which atoms to pull from.
+    // determine which uids to pull from.
     virtual void RegisterPullUidProvider(const ConfigKey& configKey, wp<PullUidProvider> provider);
 
     // Unregister a pull uid provider.
-    virtual void UnregisterPullUidProvider(const ConfigKey& configKey);
+    virtual void UnregisterPullUidProvider(const ConfigKey& configKey,
+                                           wp<PullUidProvider> provider);
 
     // Verify if we know how to pull for this matcher
     bool PullerForMatcherExists(int tagId) const;
@@ -180,6 +181,8 @@
     FRIEND_TEST(ValueMetricE2eTest, TestPulledEvents);
     FRIEND_TEST(ValueMetricE2eTest, TestPulledEvents_LateAlarm);
     FRIEND_TEST(ValueMetricE2eTest, TestPulledEvents_WithActivation);
+
+    FRIEND_TEST(StatsLogProcessorTest, TestPullUidProviderSetOnConfigUpdate);
 };
 
 }  // namespace statsd
diff --git a/cmds/statsd/src/metrics/MetricsManager.cpp b/cmds/statsd/src/metrics/MetricsManager.cpp
index bcca1fd9..e03de0b 100644
--- a/cmds/statsd/src/metrics/MetricsManager.cpp
+++ b/cmds/statsd/src/metrics/MetricsManager.cpp
@@ -189,7 +189,7 @@
             StateManager::getInstance().unregisterListener(atomId, it);
         }
     }
-    mPullerManager->UnregisterPullUidProvider(mConfigKey);
+    mPullerManager->UnregisterPullUidProvider(mConfigKey, this);
 
     VLOG("~MetricsManager()");
 }
diff --git a/cmds/statsd/tests/MetricsManager_test.cpp b/cmds/statsd/tests/MetricsManager_test.cpp
index 1075fe4..a44541d 100644
--- a/cmds/statsd/tests/MetricsManager_test.cpp
+++ b/cmds/statsd/tests/MetricsManager_test.cpp
@@ -528,7 +528,7 @@
             }));
     sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();
     EXPECT_CALL(*pullerManager, RegisterPullUidProvider(kConfigKey, _)).Times(1);
-    EXPECT_CALL(*pullerManager, UnregisterPullUidProvider(kConfigKey)).Times(1);
+    EXPECT_CALL(*pullerManager, UnregisterPullUidProvider(kConfigKey, _)).Times(1);
 
     sp<AlarmMonitor> anomalyAlarmMonitor;
     sp<AlarmMonitor> periodicAlarmMonitor;
diff --git a/cmds/statsd/tests/StatsLogProcessor_test.cpp b/cmds/statsd/tests/StatsLogProcessor_test.cpp
index d29394b..b809286 100644
--- a/cmds/statsd/tests/StatsLogProcessor_test.cpp
+++ b/cmds/statsd/tests/StatsLogProcessor_test.cpp
@@ -301,6 +301,29 @@
     EXPECT_TRUE(noData);
 }
 
+TEST(StatsLogProcessorTest, TestPullUidProviderSetOnConfigUpdate) {
+    // Setup simple config key corresponding to empty config.
+    sp<UidMap> m = new UidMap();
+    sp<StatsPullerManager> pullerManager = new StatsPullerManager();
+    sp<AlarmMonitor> anomalyAlarmMonitor;
+    sp<AlarmMonitor> subscriberAlarmMonitor;
+    StatsLogProcessor p(
+            m, pullerManager, anomalyAlarmMonitor, subscriberAlarmMonitor, 0,
+            [](const ConfigKey& key) { return true; },
+            [](const int&, const vector<int64_t>&) { return true; });
+    ConfigKey key(3, 4);
+    StatsdConfig config = MakeConfig(false);
+    p.OnConfigUpdated(0, key, config);
+    EXPECT_NE(pullerManager->mPullUidProviders.find(key), pullerManager->mPullUidProviders.end());
+
+    config.add_default_pull_packages("AID_STATSD");
+    p.OnConfigUpdated(5, key, config);
+    EXPECT_NE(pullerManager->mPullUidProviders.find(key), pullerManager->mPullUidProviders.end());
+
+    p.OnConfigRemoved(key);
+    EXPECT_EQ(pullerManager->mPullUidProviders.find(key), pullerManager->mPullUidProviders.end());
+}
+
 TEST(StatsLogProcessorTest, TestActiveConfigMetricDiskWriteRead) {
     int uid = 1111;
 
diff --git a/cmds/statsd/tests/metrics/metrics_test_helper.h b/cmds/statsd/tests/metrics/metrics_test_helper.h
index 69f7e3f..be410b1 100644
--- a/cmds/statsd/tests/metrics/metrics_test_helper.h
+++ b/cmds/statsd/tests/metrics/metrics_test_helper.h
@@ -44,7 +44,8 @@
                             vector<std::shared_ptr<LogEvent>>* data, bool useUids));
     MOCK_METHOD2(RegisterPullUidProvider,
                  void(const ConfigKey& configKey, wp<PullUidProvider> provider));
-    MOCK_METHOD1(UnregisterPullUidProvider, void(const ConfigKey& configKey));
+    MOCK_METHOD2(UnregisterPullUidProvider,
+                 void(const ConfigKey& configKey, wp<PullUidProvider> provider));
 };
 
 class MockUidMap : public UidMap {