Add support for multi train logging

Test: atest GtsStatsdHostTestCases
Change-Id: I4ca7089347d6dfab8a33c07acd5ef9c252413ec9

13 files changed, 333 insertions, 243 deletions

diff --git a/apex/statsd/aidl/android/os/IStatsd.aidl b/apex/statsd/aidl/android/os/IStatsd.aidl
index 253b2c17423d..10b1e5b26d12 100644
--- a/apex/statsd/aidl/android/os/IStatsd.aidl
+++ b/apex/statsd/aidl/android/os/IStatsd.aidl
@@ -222,12 +222,6 @@ interface IStatsd {
     const int FLAG_REQUIRE_LOW_LATENCY_MONITOR = 0x04;
 
     /**
-     * Logs an event for watchdog rollbacks.
-     */
-     oneway void sendWatchdogRollbackOccurredAtom(in int rollbackType, in String packageName,
-         in long packageVersionCode, in int rollbackReason, in String failingPackageName);
-
-    /**
      * Returns the most recently registered experiment IDs.
      */
     long[] getRegisteredExperimentIds();
diff --git a/apex/statsd/framework/java/android/util/StatsLog.java b/apex/statsd/framework/java/android/util/StatsLog.java
index 791073794b0e..e7659d88dd81 100644
--- a/apex/statsd/framework/java/android/util/StatsLog.java
+++ b/apex/statsd/framework/java/android/util/StatsLog.java
@@ -171,52 +171,11 @@ public final class StatsLog {
                 state,
                 proto.getBytes(),
                 0,
-                0);
+                0,
+                false);
         return true;
     }
 
-    /**
-     * Logs an event for watchdog rollbacks.
-     *
-     * @param rollbackType          state of the rollback.
-     * @param packageName           package name being rolled back.
-     * @param packageVersionCode    version of the package being rolled back.
-     * @param rollbackReason        reason the package is being rolled back.
-     * @param failingPackageName    the package name causing the failure.
-     *
-     * @return True if the log request was sent to statsd.
-     *
-     * @hide
-     */
-    @RequiresPermission(allOf = {DUMP, PACKAGE_USAGE_STATS})
-    public static boolean logWatchdogRollbackOccurred(int rollbackType, String packageName,
-            long packageVersionCode, int rollbackReason, String failingPackageName) {
-        synchronized (sLogLock) {
-            try {
-                IStatsd service = getIStatsdLocked();
-                if (service == null) {
-                    if (DEBUG) {
-                        Slog.d(TAG, "Failed to find statsd when logging event");
-                    }
-                    return false;
-                }
-
-                service.sendWatchdogRollbackOccurredAtom(rollbackType, packageName,
-                        packageVersionCode, rollbackReason, failingPackageName);
-                return true;
-            } catch (RemoteException e) {
-                sService = null;
-                if (DEBUG) {
-                    Slog.d(TAG,
-                            "Failed to connect to StatsCompanionService when logging "
-                                    + "WatchdogRollbackOccurred");
-                }
-                return false;
-            }
-        }
-    }
-
-
     private static IStatsd getIStatsdLocked() throws RemoteException {
         if (sService != null) {
             return sService;
diff --git a/cmds/statsd/src/StatsLogProcessor.cpp b/cmds/statsd/src/StatsLogProcessor.cpp
index bde15a5cdaae..6e7f08135f1e 100644
--- a/cmds/statsd/src/StatsLogProcessor.cpp
+++ b/cmds/statsd/src/StatsLogProcessor.cpp
@@ -196,17 +196,27 @@ void StatsLogProcessor::onBinaryPushStateChangedEventLocked(LogEvent* event) {
         !checkPermissionForIds(kPermissionUsage, pid, uid)) {
         return;
     }
-    status_t err = NO_ERROR, err2 = NO_ERROR, err3 = NO_ERROR, err4 = NO_ERROR;
-    string trainName = string(event->GetString(1 /*train name field id*/, &err));
-    int64_t trainVersionCode = event->GetLong(2 /*train version field id*/, &err2);
-    int32_t state = int32_t(event->GetLong(6 /*state field id*/, &err3));
+    // The Get* functions don't modify the status on success, they only write in
+    // failure statuses, so we can use one status variable for all calls then
+    // check if it is no longer NO_ERROR.
+    status_t err = NO_ERROR;
+    InstallTrainInfo trainInfo;
+    trainInfo.trainName = string(event->GetString(1 /*train name field id*/, &err));
+    trainInfo.trainVersionCode = event->GetLong(2 /*train version field id*/, &err);
+    trainInfo.requiresStaging = event->GetBool(3 /*requires staging field id*/, &err);
+    trainInfo.rollbackEnabled = event->GetBool(4 /*rollback enabled field id*/, &err);
+    trainInfo.requiresLowLatencyMonitor =
+            event->GetBool(5 /*requires low latency monitor field id*/, &err);
+    trainInfo.status = int32_t(event->GetLong(6 /*state field id*/, &err));
 #ifdef NEW_ENCODING_SCHEME
     std::vector<uint8_t> trainExperimentIdBytes =
-        event->GetStorage(7 /*experiment ids field id*/, &err4);
+            event->GetStorage(7 /*experiment ids field id*/, &err);
 #else
-    string trainExperimentIdString = event->GetString(7 /*experiment ids field id*/, &err4);
+    string trainExperimentIdString = event->GetString(7 /*experiment ids field id*/, &err);
 #endif
-    if (err != NO_ERROR || err2 != NO_ERROR || err3 != NO_ERROR || err4 != NO_ERROR) {
+    bool is_rollback = event->GetBool(10 /*is rollback field id*/, &err);
+
+    if (err != NO_ERROR) {
         ALOGE("Failed to parse fields in binary push state changed log event");
         return;
     }
@@ -220,83 +230,154 @@ void StatsLogProcessor::onBinaryPushStateChangedEventLocked(LogEvent* event) {
         ALOGE("Failed to parse experimentids in binary push state changed.");
         return;
     }
-    vector<int64_t> experimentIdVector = {trainExperimentIds.experiment_id().begin(),
-                                          trainExperimentIds.experiment_id().end()};
+    trainInfo.experimentIds = {trainExperimentIds.experiment_id().begin(),
+                               trainExperimentIds.experiment_id().end()};
+
     // Update the train info on disk and get any data the logevent is missing.
-    getAndUpdateTrainInfoOnDisk(
-        state, &trainVersionCode, &trainName, &experimentIdVector);
+    getAndUpdateTrainInfoOnDisk(is_rollback, &trainInfo);
 
     std::vector<uint8_t> trainExperimentIdProto;
-    writeExperimentIdsToProto(experimentIdVector, &trainExperimentIdProto);
+    writeExperimentIdsToProto(trainInfo.experimentIds, &trainExperimentIdProto);
     int32_t userId = multiuser_get_user_id(uid);
 
-    event->updateValue(1 /*train name field id*/, trainName, STRING);
-    event->updateValue(2 /*train version field id*/, trainVersionCode, LONG);
+    event->updateValue(2 /*train version field id*/, trainInfo.trainVersionCode, LONG);
 #ifdef NEW_ENCODING_SCHEME
     event->updateValue(7 /*experiment ids field id*/, trainExperimentIdProto, STORAGE);
 #else
     event->updateValue(7 /*experiment ids field id*/, trainExperimentIdProto, STRING);
 #endif
     event->updateValue(8 /*user id field id*/, userId, INT);
+
+    if (is_rollback) {
+        int bit = trainInfo.requiresStaging ? 1 : 0;
+        event->updateValue(3 /*requires staging field id*/, bit, INT);
+        bit = trainInfo.rollbackEnabled ? 1 : 0;
+        event->updateValue(4 /*rollback enabled field id*/, bit, INT);
+        bit = trainInfo.requiresLowLatencyMonitor ? 1 : 0;
+        event->updateValue(5 /*requires low latency monitor field id*/, bit, INT);
+    }
 }
 
-void StatsLogProcessor::getAndUpdateTrainInfoOnDisk(int32_t state,
-                                         int64_t* trainVersionCode,
-                                         string* trainName,
-                                         std::vector<int64_t>* experimentIds) {
+void StatsLogProcessor::getAndUpdateTrainInfoOnDisk(bool is_rollback,
+                                                    InstallTrainInfo* trainInfo) {
+    // If the train name is empty, we don't know which train to attribute the
+    // event to, so return early.
+    if (trainInfo->trainName.empty()) {
+        return;
+    }
     bool readTrainInfoSuccess = false;
     InstallTrainInfo trainInfoOnDisk;
-    readTrainInfoSuccess = StorageManager::readTrainInfo(trainInfoOnDisk);
+    readTrainInfoSuccess = StorageManager::readTrainInfo(trainInfo->trainName, trainInfoOnDisk);
 
     bool resetExperimentIds = false;
     if (readTrainInfoSuccess) {
         // Keep the old train version if we received an empty version.
-        if (*trainVersionCode == -1) {
-            *trainVersionCode = trainInfoOnDisk.trainVersionCode;
-        } else if (*trainVersionCode != trainInfoOnDisk.trainVersionCode) {
-        // Reset experiment ids if we receive a new non-empty train version.
-            resetExperimentIds = true;
-        }
-
-        // Keep the old train name if we received an empty train name.
-        if (trainName->size() == 0) {
-            *trainName = trainInfoOnDisk.trainName;
-        } else if (*trainName != trainInfoOnDisk.trainName) {
-            // Reset experiment ids if we received a new valid train name.
+        if (trainInfo->trainVersionCode == -1) {
+            trainInfo->trainVersionCode = trainInfoOnDisk.trainVersionCode;
+        } else if (trainInfo->trainVersionCode != trainInfoOnDisk.trainVersionCode) {
+            // Reset experiment ids if we receive a new non-empty train version.
             resetExperimentIds = true;
         }
 
         // Reset if we received a different experiment id.
-        if (!experimentIds->empty() &&
-                (trainInfoOnDisk.experimentIds.empty() ||
-                 experimentIds->at(0) != trainInfoOnDisk.experimentIds[0])) {
+        if (!trainInfo->experimentIds.empty() &&
+            (trainInfoOnDisk.experimentIds.empty() ||
+             trainInfo->experimentIds.at(0) != trainInfoOnDisk.experimentIds[0])) {
             resetExperimentIds = true;
         }
     }
 
     // Find the right experiment IDs
-    if (!resetExperimentIds && readTrainInfoSuccess) {
-        *experimentIds = trainInfoOnDisk.experimentIds;
+    if ((!resetExperimentIds || is_rollback) && readTrainInfoSuccess) {
+        trainInfo->experimentIds = trainInfoOnDisk.experimentIds;
     }
 
-    if (!experimentIds->empty()) {
-        int64_t firstId = experimentIds->at(0);
-        switch (state) {
+    if (!trainInfo->experimentIds.empty()) {
+        int64_t firstId = trainInfo->experimentIds.at(0);
+        switch (trainInfo->status) {
             case android::util::BINARY_PUSH_STATE_CHANGED__STATE__INSTALL_SUCCESS:
-                experimentIds->push_back(firstId + 1);
+                trainInfo->experimentIds.push_back(firstId + 1);
                 break;
             case android::util::BINARY_PUSH_STATE_CHANGED__STATE__INSTALLER_ROLLBACK_INITIATED:
-                experimentIds->push_back(firstId + 2);
+                trainInfo->experimentIds.push_back(firstId + 2);
                 break;
             case android::util::BINARY_PUSH_STATE_CHANGED__STATE__INSTALLER_ROLLBACK_SUCCESS:
-                experimentIds->push_back(firstId + 3);
+                trainInfo->experimentIds.push_back(firstId + 3);
                 break;
         }
     }
 
-    StorageManager::writeTrainInfo(*trainVersionCode, *trainName, state, *experimentIds);
+    if (is_rollback) {
+        trainInfo->requiresStaging = trainInfoOnDisk.requiresStaging;
+        trainInfo->rollbackEnabled = trainInfoOnDisk.rollbackEnabled;
+        trainInfo->requiresLowLatencyMonitor = trainInfoOnDisk.requiresLowLatencyMonitor;
+    }
+
+    StorageManager::writeTrainInfo(*trainInfo);
 }
 
+void StatsLogProcessor::onWatchdogRollbackOccurredLocked(LogEvent* event) {
+    pid_t pid = event->GetPid();
+    uid_t uid = event->GetUid();
+    if (!checkPermissionForIds(kPermissionDump, pid, uid) ||
+        !checkPermissionForIds(kPermissionUsage, pid, uid)) {
+        return;
+    }
+    // The Get* functions don't modify the status on success, they only write in
+    // failure statuses, so we can use one status variable for all calls then
+    // check if it is no longer NO_ERROR.
+    status_t err = NO_ERROR;
+    int32_t rollbackType = int32_t(event->GetInt(1 /*rollback type field id*/, &err));
+    string packageName = string(event->GetString(2 /*package name field id*/, &err));
+
+    if (err != NO_ERROR) {
+        ALOGE("Failed to parse fields in watchdog rollback occurred log event");
+        return;
+    }
+
+    vector<int64_t> experimentIds =
+        processWatchdogRollbackOccurred(rollbackType, packageName);
+    vector<uint8_t> experimentIdProto;
+    writeExperimentIdsToProto(experimentIds, &experimentIdProto);
+
+#ifdef NEW_ENCODING_SCHEME
+    event->updateValue(6 /*experiment ids field id*/, experimentIdProto, STORAGE);
+#else
+    event->updateValue(6 /*experiment ids field id*/, experimentIdProto, STRING);
+#endif
+}
+
+vector<int64_t> StatsLogProcessor::processWatchdogRollbackOccurred(const int32_t rollbackTypeIn,
+                                                                    const string& packageNameIn) {
+    // If the package name is empty, we can't attribute it to any train, so
+    // return early.
+    if (packageNameIn.empty()) {
+      return vector<int64_t>();
+    }
+    bool readTrainInfoSuccess = false;
+    InstallTrainInfo trainInfoOnDisk;
+    readTrainInfoSuccess = StorageManager::readTrainInfo(packageNameIn, trainInfoOnDisk);
+
+    if (!readTrainInfoSuccess) {
+        return vector<int64_t>();
+    }
+
+    if (trainInfoOnDisk.experimentIds.empty()) {
+        return vector<int64_t>();
+    }
+    switch (rollbackTypeIn) {
+        case android::util::WATCHDOG_ROLLBACK_OCCURRED__ROLLBACK_TYPE__ROLLBACK_INITIATE:
+            trainInfoOnDisk.experimentIds.push_back(trainInfoOnDisk.experimentIds[0] + 4);
+            StorageManager::writeTrainInfo(trainInfoOnDisk);
+            break;
+        case android::util::WATCHDOG_ROLLBACK_OCCURRED__ROLLBACK_TYPE__ROLLBACK_SUCCESS:
+            trainInfoOnDisk.experimentIds.push_back(trainInfoOnDisk.experimentIds[0] + 5);
+            StorageManager::writeTrainInfo(trainInfoOnDisk);
+            break;
+    }
+
+    return trainInfoOnDisk.experimentIds;
+}
 
 void StatsLogProcessor::resetConfigs() {
     std::lock_guard<std::mutex> lock(mMetricsMutex);
@@ -321,7 +402,13 @@ void StatsLogProcessor::OnLogEvent(LogEvent* event, int64_t elapsedRealtimeNs) {
     // Hard-coded logic to update train info on disk and fill in any information
     // this log event may be missing.
     if (event->GetTagId() == android::util::BINARY_PUSH_STATE_CHANGED) {
-      onBinaryPushStateChangedEventLocked(event);
+        onBinaryPushStateChangedEventLocked(event);
+    }
+
+    // Hard-coded logic to update experiment ids on disk for certain rollback
+    // types and fill the rollback atom with experiment ids
+    if (event->GetTagId() == android::util::WATCHDOG_ROLLBACK_OCCURRED) {
+        onWatchdogRollbackOccurredLocked(event);
     }
 
 #ifdef VERY_VERBOSE_PRINTING
diff --git a/cmds/statsd/src/StatsLogProcessor.h b/cmds/statsd/src/StatsLogProcessor.h
index c49f2e0ec68d..42e56760fb89 100644
--- a/cmds/statsd/src/StatsLogProcessor.h
+++ b/cmds/statsd/src/StatsLogProcessor.h
@@ -18,6 +18,7 @@
 
 #include <gtest/gtest_prod.h>
 #include "config/ConfigListener.h"
+#include "logd/LogEvent.h"
 #include "metrics/MetricsManager.h"
 #include "packages/UidMap.h"
 #include "external/StatsPullerManager.h"
@@ -199,10 +200,18 @@ private:
     // Handler over the binary push state changed event.
     void onBinaryPushStateChangedEventLocked(LogEvent* event);
 
+    // Handler over the watchdog rollback occurred event.
+    void onWatchdogRollbackOccurredLocked(LogEvent* event);
+
     // Updates train info on disk based on binary push state changed info and
     // write disk info into parameters.
-    void getAndUpdateTrainInfoOnDisk(int32_t state, int64_t* trainVersionCode,
-                                     string* trainName, std::vector<int64_t>* experimentIds);
+    void getAndUpdateTrainInfoOnDisk(bool is_rollback, InstallTrainInfo* trainInfoIn);
+
+    // Gets experiment ids on disk for associated train and updates them
+    // depending on rollback type. Then writes them back to disk and returns
+    // them.
+    std::vector<int64_t> processWatchdogRollbackOccurred(const int32_t rollbackTypeIn,
+                                                          const string& packageName);
 
     // Reset all configs.
     void resetConfigsLocked(const int64_t timestampNs);
diff --git a/cmds/statsd/src/StatsService.cpp b/cmds/statsd/src/StatsService.cpp
index a06e59c8e409..168833fc8c8b 100644
--- a/cmds/statsd/src/StatsService.cpp
+++ b/cmds/statsd/src/StatsService.cpp
@@ -1293,81 +1293,24 @@ Status StatsService::unregisterNativePullAtomCallback(int32_t atomTag) {
     return Status::ok();
 }
 
-Status StatsService::sendWatchdogRollbackOccurredAtom(const int32_t rollbackTypeIn,
-                                                      const android::String16& packageNameIn,
-                                                      const int64_t packageVersionCodeIn,
-                                                      const int32_t rollbackReasonIn,
-                                                      const android::String16&
-                                                       failingPackageNameIn) {
-    // Note: We skip the usage stats op check here since we do not have a package name.
-    // This is ok since we are overloading the usage_stats permission.
-    // This method only sends data, it does not receive it.
-    pid_t pid = IPCThreadState::self()->getCallingPid();
-    uid_t uid = IPCThreadState::self()->getCallingUid();
-    // Root, system, and shell always have access
-    if (uid != AID_ROOT && uid != AID_SYSTEM && uid != AID_SHELL) {
-        // Caller must be granted these permissions
-        if (!checkPermission(kPermissionDump)) {
-            return exception(binder::Status::EX_SECURITY,
-                             StringPrintf("UID %d / PID %d lacks permission %s", uid, pid,
-                                          kPermissionDump));
-        }
-        if (!checkPermission(kPermissionUsage)) {
-            return exception(binder::Status::EX_SECURITY,
-                             StringPrintf("UID %d / PID %d lacks permission %s", uid, pid,
-                                          kPermissionUsage));
-        }
-    }
-
-    android::util::stats_write(android::util::WATCHDOG_ROLLBACK_OCCURRED,
-            rollbackTypeIn, String8(packageNameIn).string(), packageVersionCodeIn,
-            rollbackReasonIn, String8(failingPackageNameIn).string());
-
-    // Fast return to save disk read.
-    if (rollbackTypeIn != android::util::WATCHDOG_ROLLBACK_OCCURRED__ROLLBACK_TYPE__ROLLBACK_SUCCESS
-            && rollbackTypeIn !=
-                    android::util::WATCHDOG_ROLLBACK_OCCURRED__ROLLBACK_TYPE__ROLLBACK_INITIATE) {
-        return Status::ok();
-    }
-
-    bool readTrainInfoSuccess = false;
-    InstallTrainInfo trainInfoOnDisk;
-    readTrainInfoSuccess = StorageManager::readTrainInfo(trainInfoOnDisk);
-
-    if (!readTrainInfoSuccess) {
-        return Status::ok();
-    }
-    std::vector<int64_t> experimentIds = trainInfoOnDisk.experimentIds;
-    if (experimentIds.empty()) {
-        return Status::ok();
-    }
-    switch (rollbackTypeIn) {
-        case android::util::WATCHDOG_ROLLBACK_OCCURRED__ROLLBACK_TYPE__ROLLBACK_INITIATE:
-            experimentIds.push_back(experimentIds[0] + 4);
-            break;
-        case android::util::WATCHDOG_ROLLBACK_OCCURRED__ROLLBACK_TYPE__ROLLBACK_SUCCESS:
-            experimentIds.push_back(experimentIds[0] + 5);
-            break;
-    }
-    StorageManager::writeTrainInfo(trainInfoOnDisk.trainVersionCode, trainInfoOnDisk.trainName,
-            trainInfoOnDisk.status, experimentIds);
-    return Status::ok();
-}
-
 Status StatsService::getRegisteredExperimentIds(std::vector<int64_t>* experimentIdsOut) {
     ENFORCE_UID(AID_SYSTEM);
     // TODO: add verifier permission
 
+    experimentIdsOut->clear();
     // Read the latest train info
-    InstallTrainInfo trainInfo;
-    if (!StorageManager::readTrainInfo(trainInfo)) {
+    vector<InstallTrainInfo> trainInfoList = StorageManager::readAllTrainInfo();
+    if (trainInfoList.empty()) {
         // No train info means no experiment IDs, return an empty list
-        experimentIdsOut->clear();
         return Status::ok();
     }
 
     // Copy the experiment IDs to the out vector
-    experimentIdsOut->assign(trainInfo.experimentIds.begin(), trainInfo.experimentIds.end());
+    for (InstallTrainInfo& trainInfo : trainInfoList) {
+        experimentIdsOut->insert(experimentIdsOut->end(),
+                                 trainInfo.experimentIds.begin(),
+                                 trainInfo.experimentIds.end());
+    }
     return Status::ok();
 }
 
diff --git a/cmds/statsd/src/StatsService.h b/cmds/statsd/src/StatsService.h
index 82a5a5305df4..0527d435fdf4 100644
--- a/cmds/statsd/src/StatsService.h
+++ b/cmds/statsd/src/StatsService.h
@@ -187,16 +187,6 @@ public:
     virtual Status unregisterNativePullAtomCallback(int32_t atomTag) override;
 
     /**
-     * Binder call to log WatchdogRollbackOccurred atom.
-     */
-    virtual Status sendWatchdogRollbackOccurredAtom(
-            const int32_t rollbackTypeIn,
-            const android::String16& packageNameIn,
-            const int64_t packageVersionCodeIn,
-            const int32_t rollbackReasonIn,
-            const android::String16& failingPackageNameIn) override;
-
-    /**
      * Binder call to get registered experiment IDs.
      */
     virtual Status getRegisteredExperimentIds(std::vector<int64_t>* expIdsOut);
diff --git a/cmds/statsd/src/atoms.proto b/cmds/statsd/src/atoms.proto
index 71afc32686f1..5e2dbf37fb27 100644
--- a/cmds/statsd/src/atoms.proto
+++ b/cmds/statsd/src/atoms.proto
@@ -1874,6 +1874,8 @@ message WatchdogRollbackOccurred {
     // Set by RollbackPackageHealthObserver to be the package that is failing when a rollback
     // is initiated. Empty if the package is unknown.
     optional string failing_package_name = 5;
+
+    optional TrainExperimentIds experiment_ids = 6 [(log_mode) = MODE_BYTES];
 }
 
 /**
@@ -3770,6 +3772,8 @@ message BinaryPushStateChanged {
     // user id
     optional int32 user_id = 8;
     optional int32 reason = 9;
+    // Whether or not this is a rollback event
+    optional bool is_rollback = 10;
 }
 
 /* Test atom, is not logged anywhere */
diff --git a/cmds/statsd/src/external/TrainInfoPuller.cpp b/cmds/statsd/src/external/TrainInfoPuller.cpp
index 9d0924297912..a7d8d4ebebae 100644
--- a/cmds/statsd/src/external/TrainInfoPuller.cpp
+++ b/cmds/statsd/src/external/TrainInfoPuller.cpp
@@ -37,14 +37,16 @@ TrainInfoPuller::TrainInfoPuller() :
 }
 
 bool TrainInfoPuller::PullInternal(vector<shared_ptr<LogEvent>>* data) {
-    InstallTrainInfo trainInfo;
-    bool ret = StorageManager::readTrainInfo(trainInfo);
-    if (!ret) {
-        ALOGW("Failed to read train info.");
-        return false;
+    vector<InstallTrainInfo> trainInfoList =
+        StorageManager::readAllTrainInfo();
+    if (trainInfoList.empty()) {
+        ALOGW("Train info was empty.");
+        return true;
+    }
+    for (InstallTrainInfo& trainInfo : trainInfoList) {
+        auto event = make_shared<LogEvent>(getWallClockNs(), getElapsedRealtimeNs(), trainInfo);
+        data->push_back(event);
     }
-    auto event = make_shared<LogEvent>(getWallClockNs(), getElapsedRealtimeNs(), trainInfo);
-    data->push_back(event);
     return true;
 }
 
diff --git a/cmds/statsd/src/logd/LogEvent.h b/cmds/statsd/src/logd/LogEvent.h
index 5509c093a6f6..a67bef451be7 100644
--- a/cmds/statsd/src/logd/LogEvent.h
+++ b/cmds/statsd/src/logd/LogEvent.h
@@ -56,6 +56,9 @@ struct InstallTrainInfo {
     std::string trainName;
     int32_t status;
     std::vector<int64_t> experimentIds;
+    bool requiresStaging;
+    bool rollbackEnabled;
+    bool requiresLowLatencyMonitor;
 };
 
 /**
diff --git a/cmds/statsd/src/storage/StorageManager.cpp b/cmds/statsd/src/storage/StorageManager.cpp
index 507297c6c401..1bac19ed2c5d 100644
--- a/cmds/statsd/src/storage/StorageManager.cpp
+++ b/cmds/statsd/src/storage/StorageManager.cpp
@@ -44,7 +44,7 @@ using std::map;
 #define TRAIN_INFO_PATH "/data/misc/train-info/train-info.bin"
 
 // Magic word at the start of the train info file, change this if changing the file format
-const uint32_t TRAIN_INFO_FILE_MAGIC = 0xff7447ff;
+const uint32_t TRAIN_INFO_FILE_MAGIC = 0xfb7447bf;
 
 // for ConfigMetricsReportList
 const int FIELD_ID_REPORTS = 2;
@@ -75,6 +75,29 @@ string StorageManager::getDataHistoryFileName(long wallClockSec, int uid, int64_
                         (long long)id);
 }
 
+static const char* findTrainInfoFileNameLocked(const string& trainName) {
+    unique_ptr<DIR, decltype(&closedir)> dir(opendir(TRAIN_INFO_DIR), closedir);
+    if (dir == NULL) {
+        VLOG("Path %s does not exist", TRAIN_INFO_DIR);
+        return nullptr;
+    }
+    dirent* de;
+    while ((de = readdir(dir.get()))) {
+        char* fileName = de->d_name;
+        if (fileName[0] == '.') continue;
+
+        size_t fileNameLength = strlen(fileName);
+        if (fileNameLength >= trainName.length()) {
+            if (0 == strncmp(fileName + fileNameLength - trainName.length(), trainName.c_str(),
+                             trainName.length())) {
+                return fileName;
+            }
+        }
+    }
+
+    return nullptr;
+}
+
 // Returns array of int64_t which contains timestamp in seconds, uid,
 // configID and whether the file is a local history file.
 static void parseFileName(char* name, FileName* output) {
@@ -123,20 +146,25 @@ void StorageManager::writeFile(const char* file, const void* buffer, int numByte
     close(fd);
 }
 
-bool StorageManager::writeTrainInfo(int64_t trainVersionCode, const std::string& trainName,
-                                    int32_t status, const std::vector<int64_t>& experimentIds) {
+bool StorageManager::writeTrainInfo(const InstallTrainInfo& trainInfo) {
     std::lock_guard<std::mutex> lock(sTrainInfoMutex);
 
-    deleteAllFiles(TRAIN_INFO_DIR);
+    if (trainInfo.trainName.empty()) {
+      return false;
+    }
+    deleteSuffixedFiles(TRAIN_INFO_DIR, trainInfo.trainName.c_str());
 
-    int fd = open(TRAIN_INFO_PATH, O_WRONLY | O_CREAT | O_CLOEXEC, S_IRUSR | S_IWUSR);
+    std::string fileName =
+            StringPrintf("%s/%ld_%s", TRAIN_INFO_DIR, (long) getWallClockSec(),
+                         trainInfo.trainName.c_str());
+
+    int fd = open(fileName.c_str(), O_WRONLY | O_CREAT | O_CLOEXEC, S_IRUSR | S_IWUSR);
     if (fd == -1) {
-        VLOG("Attempt to access %s but failed", TRAIN_INFO_PATH);
+        VLOG("Attempt to access %s but failed", fileName.c_str());
         return false;
     }
 
     size_t result;
-
     // Write the magic word
     result = write(fd, &TRAIN_INFO_FILE_MAGIC, sizeof(TRAIN_INFO_FILE_MAGIC));
     if (result != sizeof(TRAIN_INFO_FILE_MAGIC)) {
@@ -146,8 +174,8 @@ bool StorageManager::writeTrainInfo(int64_t trainVersionCode, const std::string&
     }
 
     // Write the train version
-    const size_t trainVersionCodeByteCount = sizeof(trainVersionCode);
-    result = write(fd, &trainVersionCode, trainVersionCodeByteCount);
+    const size_t trainVersionCodeByteCount = sizeof(trainInfo.trainVersionCode);
+    result = write(fd, &trainInfo.trainVersionCode, trainVersionCodeByteCount);
     if (result != trainVersionCodeByteCount) {
         VLOG("Failed to wrtie train version code");
         close(fd);
@@ -155,7 +183,7 @@ bool StorageManager::writeTrainInfo(int64_t trainVersionCode, const std::string&
     }
 
     // Write # of bytes in trainName to file
-    const size_t trainNameSize = trainName.size();
+    const size_t trainNameSize = trainInfo.trainName.size();
     const size_t trainNameSizeByteCount = sizeof(trainNameSize);
     result = write(fd, (uint8_t*)&trainNameSize, trainNameSizeByteCount);
     if (result != trainNameSizeByteCount) {
@@ -165,7 +193,7 @@ bool StorageManager::writeTrainInfo(int64_t trainVersionCode, const std::string&
     }
 
     // Write trainName to file
-    result = write(fd, trainName.c_str(), trainNameSize);
+    result = write(fd, trainInfo.trainName.c_str(), trainNameSize);
     if (result != trainNameSize) {
         VLOG("Failed to write train name");
         close(fd);
@@ -173,8 +201,8 @@ bool StorageManager::writeTrainInfo(int64_t trainVersionCode, const std::string&
     }
 
     // Write status to file
-    const size_t statusByteCount = sizeof(status);
-    result = write(fd, (uint8_t*)&status, statusByteCount);
+    const size_t statusByteCount = sizeof(trainInfo.status);
+    result = write(fd, (uint8_t*)&trainInfo.status, statusByteCount);
     if (result != statusByteCount) {
         VLOG("Failed to write status");
         close(fd);
@@ -182,7 +210,7 @@ bool StorageManager::writeTrainInfo(int64_t trainVersionCode, const std::string&
     }
 
     // Write experiment id count to file.
-    const size_t experimentIdsCount = experimentIds.size();
+    const size_t experimentIdsCount = trainInfo.experimentIds.size();
     const size_t experimentIdsCountByteCount = sizeof(experimentIdsCount);
     result = write(fd, (uint8_t*) &experimentIdsCount, experimentIdsCountByteCount);
     if (result != experimentIdsCountByteCount) {
@@ -193,7 +221,7 @@ bool StorageManager::writeTrainInfo(int64_t trainVersionCode, const std::string&
 
     // Write experimentIds to file
     for (size_t i = 0; i < experimentIdsCount; i++) {
-        const int64_t experimentId = experimentIds[i];
+        const int64_t experimentId = trainInfo.experimentIds[i];
         const size_t experimentIdByteCount = sizeof(experimentId);
         result = write(fd, &experimentId, experimentIdByteCount);
         if (result == experimentIdByteCount) {
@@ -205,23 +233,47 @@ bool StorageManager::writeTrainInfo(int64_t trainVersionCode, const std::string&
         }
     }
 
-    result = fchown(fd, AID_STATSD, AID_STATSD);
-    if (result) {
-        VLOG("Failed to chown train info file to statsd");
-        close(fd);
-        return false;
+    // Write bools to file
+    const size_t boolByteCount = sizeof(trainInfo.requiresStaging);
+    result = write(fd, (uint8_t*)&trainInfo.requiresStaging, boolByteCount);
+    if (result != boolByteCount) {
+      VLOG("Failed to write requires staging");
+      close(fd);
+      return false;
+    }
+
+    result = write(fd, (uint8_t*)&trainInfo.rollbackEnabled, boolByteCount);
+    if (result != boolByteCount) {
+      VLOG("Failed to write rollback enabled");
+      close(fd);
+      return false;
+    }
+
+    result = write(fd, (uint8_t*)&trainInfo.requiresLowLatencyMonitor, boolByteCount);
+    if (result != boolByteCount) {
+      VLOG("Failed to write requires log latency monitor");
+      close(fd);
+      return false;
     }
 
     close(fd);
     return true;
 }
 
-bool StorageManager::readTrainInfo(InstallTrainInfo& trainInfo) {
+bool StorageManager::readTrainInfo(const std::string& trainName, InstallTrainInfo& trainInfo) {
     std::lock_guard<std::mutex> lock(sTrainInfoMutex);
+    return readTrainInfoLocked(trainName, trainInfo);
+}
 
-    int fd = open(TRAIN_INFO_PATH, O_RDONLY | O_CLOEXEC);
+bool StorageManager::readTrainInfoLocked(const std::string& trainName, InstallTrainInfo& trainInfo) {
+    trimToFit(TRAIN_INFO_DIR, /*parseTimestampOnly=*/ true);
+    const char* fileName = findTrainInfoFileNameLocked(trainName);
+    if (fileName == nullptr) {
+        return false;
+    }
+    int fd = open(StringPrintf("%s/%s", TRAIN_INFO_DIR, fileName).c_str(), O_RDONLY | O_CLOEXEC);
     if (fd == -1) {
-        VLOG("Failed to open train-info.bin");
+        VLOG("Failed to open %s", fileName);
         return false;
     }
 
@@ -297,6 +349,29 @@ bool StorageManager::readTrainInfo(InstallTrainInfo& trainInfo) {
         trainInfo.experimentIds.push_back(experimentId);
     }
 
+    // Read bools
+    const size_t boolByteCount = sizeof(trainInfo.requiresStaging);
+    result = read(fd, &trainInfo.requiresStaging, boolByteCount);
+    if (result != boolByteCount) {
+        VLOG("Failed to read requires requires staging from train info file");
+        close(fd);
+        return false;
+    }
+
+    result = read(fd, &trainInfo.rollbackEnabled, boolByteCount);
+    if (result != boolByteCount) {
+        VLOG("Failed to read requires rollback enabled from train info file");
+        close(fd);
+        return false;
+    }
+
+    result = read(fd, &trainInfo.requiresLowLatencyMonitor, boolByteCount);
+    if (result != boolByteCount) {
+        VLOG("Failed to read requires requires low latency monitor from train info file");
+        close(fd);
+        return false;
+    }
+
     // Expect to be at EOF.
     char c;
     result = read(fd, &c, 1);
@@ -311,6 +386,32 @@ bool StorageManager::readTrainInfo(InstallTrainInfo& trainInfo) {
     return true;
 }
 
+vector<InstallTrainInfo> StorageManager::readAllTrainInfo() {
+    std::lock_guard<std::mutex> lock(sTrainInfoMutex);
+    vector<InstallTrainInfo> trainInfoList;
+    unique_ptr<DIR, decltype(&closedir)> dir(opendir(TRAIN_INFO_DIR), closedir);
+    if (dir == NULL) {
+        VLOG("Directory does not exist: %s", TRAIN_INFO_DIR);
+        return trainInfoList;
+    }
+
+    dirent* de;
+    while ((de = readdir(dir.get()))) {
+        char* name = de->d_name;
+        if (name[0] == '.') {
+            continue;
+        }
+
+        InstallTrainInfo trainInfo;
+        bool readSuccess = StorageManager::readTrainInfoLocked(name, trainInfo);
+        if (!readSuccess) {
+            continue;
+        }
+        trainInfoList.push_back(trainInfo);
+    }
+    return trainInfoList;
+}
+
 void StorageManager::deleteFile(const char* file) {
     if (remove(file) != 0) {
         VLOG("Attempt to delete %s but is not found", file);
@@ -574,7 +675,7 @@ void StorageManager::sortFiles(vector<FileInfo>* fileNames) {
     });
 }
 
-void StorageManager::trimToFit(const char* path) {
+void StorageManager::trimToFit(const char* path, bool parseTimestampOnly) {
     unique_ptr<DIR, decltype(&closedir)> dir(opendir(path), closedir);
     if (dir == NULL) {
         VLOG("Path %s does not exist", path);
@@ -589,7 +690,12 @@ void StorageManager::trimToFit(const char* path) {
         if (name[0] == '.') continue;
 
         FileName output;
-        parseFileName(name, &output);
+        if (parseTimestampOnly) {
+            output.mTimestampSec = StrToInt64(strtok(name, "_"));
+            output.mIsHistory = false;
+        } else {
+            parseFileName(name, &output);
+        }
         if (output.mTimestampSec == -1) continue;
         string file_name = output.getFullFileName(path);
 
diff --git a/cmds/statsd/src/storage/StorageManager.h b/cmds/statsd/src/storage/StorageManager.h
index 69b41c2cb974..d59046dfbb99 100644
--- a/cmds/statsd/src/storage/StorageManager.h
+++ b/cmds/statsd/src/storage/StorageManager.h
@@ -52,13 +52,22 @@ public:
     /**
      * Writes train info.
      */
-    static bool writeTrainInfo(int64_t trainVersionCode, const std::string& trainName,
-                               int32_t status, const std::vector<int64_t>& experimentIds);
+    static bool writeTrainInfo(const InstallTrainInfo& trainInfo);
 
     /**
      * Reads train info.
      */
-    static bool readTrainInfo(InstallTrainInfo& trainInfo);
+    static bool readTrainInfo(const std::string& trainName, InstallTrainInfo& trainInfo);
+
+    /**
+     * Reads train info assuming lock is obtained.
+     */
+    static bool readTrainInfoLocked(const std::string& trainName, InstallTrainInfo& trainInfo);
+
+    /**
+     * Reads all train info and returns a vector of train info.
+     */
+    static vector<InstallTrainInfo> readAllTrainInfo();
 
     /**
      * Reads the file content to the buffer.
@@ -124,7 +133,7 @@ public:
      * Trims files in the provided directory to limit the total size, number of
      * files, accumulation of outdated files.
      */
-    static void trimToFit(const char* dir);
+    static void trimToFit(const char* dir, bool parseTimestampOnly = false);
 
     /**
      * Returns true if there already exists identical configuration on device.
diff --git a/cmds/statsd/tests/storage/StorageManager_test.cpp b/cmds/statsd/tests/storage/StorageManager_test.cpp
index b91e5a0ad3a1..27a86e4230d8 100644
--- a/cmds/statsd/tests/storage/StorageManager_test.cpp
+++ b/cmds/statsd/tests/storage/StorageManager_test.cpp
@@ -40,40 +40,12 @@ TEST(StorageManagerTest, TrainInfoReadWriteTest) {
 
     bool result;
 
-    result = StorageManager::writeTrainInfo(trainInfo.trainVersionCode, trainInfo.trainName,
-                                            trainInfo.status, trainInfo.experimentIds);
+    result = StorageManager::writeTrainInfo(trainInfo);
 
     EXPECT_TRUE(result);
 
     InstallTrainInfo trainInfoResult;
-    result = StorageManager::readTrainInfo(trainInfoResult);
-    EXPECT_TRUE(result);
-
-    EXPECT_EQ(trainInfo.trainVersionCode, trainInfoResult.trainVersionCode);
-    EXPECT_EQ(trainInfo.trainName.size(), trainInfoResult.trainName.size());
-    EXPECT_EQ(trainInfo.trainName, trainInfoResult.trainName);
-    EXPECT_EQ(trainInfo.status, trainInfoResult.status);
-    EXPECT_EQ(trainInfo.experimentIds.size(), trainInfoResult.experimentIds.size());
-    EXPECT_EQ(trainInfo.experimentIds, trainInfoResult.experimentIds);
-}
-
-TEST(StorageManagerTest, TrainInfoReadWriteEmptyTrainNameTest) {
-    InstallTrainInfo trainInfo;
-    trainInfo.trainVersionCode = 12345;
-    trainInfo.trainName = "";
-    trainInfo.status = 1;
-    const char* expIds = "test_ids";
-    trainInfo.experimentIds.assign(expIds, expIds + strlen(expIds));
-
-    bool result;
-
-    result = StorageManager::writeTrainInfo(trainInfo.trainVersionCode, trainInfo.trainName,
-                                            trainInfo.status, trainInfo.experimentIds);
-
-    EXPECT_TRUE(result);
-
-    InstallTrainInfo trainInfoResult;
-    result = StorageManager::readTrainInfo(trainInfoResult);
+    result = StorageManager::readTrainInfo(trainInfo.trainName, trainInfoResult);
     EXPECT_TRUE(result);
 
     EXPECT_EQ(trainInfo.trainVersionCode, trainInfoResult.trainVersionCode);
@@ -94,13 +66,12 @@ TEST(StorageManagerTest, TrainInfoReadWriteTrainNameSizeOneTest) {
 
     bool result;
 
-    result = StorageManager::writeTrainInfo(trainInfo.trainVersionCode, trainInfo.trainName,
-                                            trainInfo.status, trainInfo.experimentIds);
+    result = StorageManager::writeTrainInfo(trainInfo);
 
     EXPECT_TRUE(result);
 
     InstallTrainInfo trainInfoResult;
-    result = StorageManager::readTrainInfo(trainInfoResult);
+    result = StorageManager::readTrainInfo(trainInfo.trainName, trainInfoResult);
     EXPECT_TRUE(result);
 
     EXPECT_EQ(trainInfo.trainVersionCode, trainInfoResult.trainVersionCode);
diff --git a/services/core/java/com/android/server/rollback/WatchdogRollbackLogger.java b/services/core/java/com/android/server/rollback/WatchdogRollbackLogger.java
index f3f14a95eac6..1be6f227b44b 100644
--- a/services/core/java/com/android/server/rollback/WatchdogRollbackLogger.java
+++ b/services/core/java/com/android/server/rollback/WatchdogRollbackLogger.java
@@ -38,9 +38,9 @@ import android.content.rollback.PackageRollbackInfo;
 import android.content.rollback.RollbackInfo;
 import android.util.ArraySet;
 import android.util.Slog;
-import android.util.StatsLog;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.util.FrameworkStatsLog;
 import com.android.server.PackageWatchdog;
 
 import java.util.ArrayList;
@@ -206,12 +206,25 @@ public final class WatchdogRollbackLogger {
                 + " rollbackReason: " + rollbackReasonToString(rollbackReason)
                 + " failedPackageName: " + failingPackageName);
         if (logPackage != null) {
-            StatsLog.logWatchdogRollbackOccurred(type, logPackage.getPackageName(),
-                    logPackage.getVersionCode(), rollbackReason, failingPackageName);
+            FrameworkStatsLog.write(
+                    FrameworkStatsLog.WATCHDOG_ROLLBACK_OCCURRED,
+                    type,
+                    logPackage.getPackageName(),
+                    logPackage.getVersionCode(),
+                    rollbackReason,
+                    failingPackageName,
+                    new byte[]{});
         } else {
             // In the case that the log package is null, still log an empty string as an
             // indication that retrieving the logging parent failed.
-            StatsLog.logWatchdogRollbackOccurred(type, "", 0, rollbackReason, failingPackageName);
+            FrameworkStatsLog.write(
+                    FrameworkStatsLog.WATCHDOG_ROLLBACK_OCCURRED,
+                    type,
+                    "",
+                    0,
+                    rollbackReason,
+                    failingPackageName,
+                    new byte[]{});
         }
     }