Flush the partial bucket when startd shuts down or config updated.

Test: statsd test

BUG: b/77556036
Change-Id: Ie4a04ace55e07c4529cdff5906ba874f8815f620
diff --git a/cmds/statsd/src/StatsLogProcessor.cpp b/cmds/statsd/src/StatsLogProcessor.cpp
index 41527f6..90ce735 100644
--- a/cmds/statsd/src/StatsLogProcessor.cpp
+++ b/cmds/statsd/src/StatsLogProcessor.cpp
@@ -212,7 +212,10 @@
     sp<MetricsManager> newMetricsManager =
         new MetricsManager(key, config, mTimeBaseSec, (timestampNs - 1) / NS_PER_SEC + 1, mUidMap,
                            mAnomalyAlarmMonitor, mPeriodicAlarmMonitor);
-
+    auto it = mMetricsManagers.find(key);
+    if (it != mMetricsManagers.end()) {
+        WriteDataToDiskLocked(it->first);
+    }
     if (newMetricsManager->isConfigValid()) {
         mUidMap->OnConfigUpdated(key);
         if (newMetricsManager->shouldAddUidMapListener()) {
@@ -251,19 +254,10 @@
  * onDumpReport dumps serialized ConfigMetricsReportList into outData.
  */
 void StatsLogProcessor::onDumpReport(const ConfigKey& key, const int64_t dumpTimeStampNs,
+                                     const bool include_current_partial_bucket,
                                      vector<uint8_t>* outData) {
     std::lock_guard<std::mutex> lock(mMetricsMutex);
 
-    auto it = mMetricsManagers.find(key);
-    if (it == mMetricsManagers.end()) {
-        ALOGW("Config source %s does not exist", key.ToString().c_str());
-        return;
-    }
-
-    // This allows another broadcast to be sent within the rate-limit period if we get close to
-    // filling the buffer again soon.
-    mLastBroadcastTimes.erase(key);
-
     ProtoOutputStream proto;
 
     // Start of ConfigKey.
@@ -273,18 +267,26 @@
     proto.end(configKeyToken);
     // End of ConfigKey.
 
-    // Start of ConfigMetricsReport (reports).
-    uint64_t reportsToken =
-            proto.start(FIELD_TYPE_MESSAGE | FIELD_COUNT_REPEATED | FIELD_ID_REPORTS);
-    onConfigMetricsReportLocked(key, dumpTimeStampNs, &proto);
-    proto.end(reportsToken);
-    // End of ConfigMetricsReport (reports).
-
-
     // Then, check stats-data directory to see there's any file containing
     // ConfigMetricsReport from previous shutdowns to concatenate to reports.
     StorageManager::appendConfigMetricsReport(key, &proto);
 
+    auto it = mMetricsManagers.find(key);
+    if (it != mMetricsManagers.end()) {
+        // This allows another broadcast to be sent within the rate-limit period if we get close to
+        // filling the buffer again soon.
+        mLastBroadcastTimes.erase(key);
+
+        // Start of ConfigMetricsReport (reports).
+        uint64_t reportsToken =
+                proto.start(FIELD_TYPE_MESSAGE | FIELD_COUNT_REPEATED | FIELD_ID_REPORTS);
+        onConfigMetricsReportLocked(key, dumpTimeStampNs, include_current_partial_bucket, &proto);
+        proto.end(reportsToken);
+        // End of ConfigMetricsReport (reports).
+    } else {
+        ALOGW("Config source %s does not exist", key.ToString().c_str());
+    }
+
     if (outData != nullptr) {
         outData->clear();
         outData->resize(proto.size());
@@ -298,7 +300,7 @@
         }
     }
 
-    StatsdStats::getInstance().noteMetricsReportSent(key);
+    StatsdStats::getInstance().noteMetricsReportSent(key, proto.size());
 }
 
 /*
@@ -306,16 +308,20 @@
  */
 void StatsLogProcessor::onConfigMetricsReportLocked(const ConfigKey& key,
                                                     const int64_t dumpTimeStampNs,
+                                                    const bool include_current_partial_bucket,
                                                     ProtoOutputStream* proto) {
     // We already checked whether key exists in mMetricsManagers in
     // WriteDataToDisk.
     auto it = mMetricsManagers.find(key);
+    if (it == mMetricsManagers.end()) {
+        return;
+    }
     int64_t lastReportTimeNs = it->second->getLastReportTimeNs();
     int64_t lastReportWallClockNs = it->second->getLastReportWallClockNs();
 
     // First, fill in ConfigMetricsReport using current data on memory, which
     // starts from filling in StatsLogReport's.
-    it->second->onDumpReport(dumpTimeStampNs, proto);
+    it->second->onDumpReport(dumpTimeStampNs, include_current_partial_bucket, proto);
 
     // Fill in UidMap.
     uint64_t uidMapToken = proto->start(FIELD_TYPE_MESSAGE | FIELD_ID_UID_MAP);
@@ -331,7 +337,6 @@
                 (long long)lastReportWallClockNs);
     proto->write(FIELD_TYPE_INT64 | FIELD_ID_CURRENT_REPORT_WALL_CLOCK_NANOS,
                 (long long)getWallClockNs());
-
 }
 
 void StatsLogProcessor::resetIfConfigTtlExpiredLocked(const int64_t timestampNs) {
@@ -361,6 +366,7 @@
     std::lock_guard<std::mutex> lock(mMetricsMutex);
     auto it = mMetricsManagers.find(key);
     if (it != mMetricsManagers.end()) {
+        WriteDataToDiskLocked(key);
         mMetricsManagers.erase(it);
         mUidMap->OnConfigRemoved(key);
     }
@@ -406,22 +412,30 @@
     }
 }
 
+void StatsLogProcessor::WriteDataToDiskLocked(const ConfigKey& key) {
+    ProtoOutputStream proto;
+    onConfigMetricsReportLocked(key, getElapsedRealtimeNs(),
+                                true /* include_current_partial_bucket*/, &proto);
+    string file_name = StringPrintf("%s/%ld_%d_%lld", STATS_DATA_DIR,
+         (long)getWallClockSec(), key.GetUid(), (long long)key.GetId());
+    android::base::unique_fd fd(open(file_name.c_str(),
+                                O_WRONLY | O_CREAT | O_CLOEXEC, S_IRUSR | S_IWUSR));
+    if (fd == -1) {
+        ALOGE("Attempt to write %s but failed", file_name.c_str());
+        return;
+    }
+    proto.flush(fd.get());
+}
+
+void StatsLogProcessor::WriteDataToDiskLocked() {
+    for (auto& pair : mMetricsManagers) {
+        WriteDataToDiskLocked(pair.first);
+    }
+}
+
 void StatsLogProcessor::WriteDataToDisk() {
     std::lock_guard<std::mutex> lock(mMetricsMutex);
-    for (auto& pair : mMetricsManagers) {
-        const ConfigKey& key = pair.first;
-        ProtoOutputStream proto;
-        onConfigMetricsReportLocked(key, getElapsedRealtimeNs(), &proto);
-        string file_name = StringPrintf("%s/%ld_%d_%lld", STATS_DATA_DIR,
-             (long)getWallClockSec(), key.GetUid(), (long long)key.GetId());
-        android::base::unique_fd fd(open(file_name.c_str(),
-                                    O_WRONLY | O_CREAT | O_CLOEXEC, S_IRUSR | S_IWUSR));
-        if (fd == -1) {
-            VLOG("Attempt to write %s but failed", file_name.c_str());
-            return;
-        }
-        proto.flush(fd.get());
-    }
+    WriteDataToDiskLocked();
 }
 
 }  // namespace statsd