[7.x][ML] Recalculate memory usage before upgrading model state

include/model/CResourceMonitor.h

@@ -102,6 +102,9 @@ class MODEL_EXPORT CResourceMonitor {
     //! Recalculate the memory usage regardless of whether there is a memory limit
     void forceRefresh(CMonitoredResource& resource);
 
+    //! Recalculate the memory usage for all monitored resources
+    void forceRefreshAll();
+
     //! Set the internal memory limit, as specified in a limits config file
     void memoryLimit(std::size_t limitMBs);
 

lib/api/CAnomalyJob.cc

@@ -411,6 +411,11 @@ bool CAnomalyJob::parsePersistControlMessageArgs(const std::string& controlMessa
 
 void CAnomalyJob::processPersistControlMessage(const std::string& controlMessageArgs) {
     if (m_PersistenceManager != nullptr) {
+        // There is a subtle difference between these two cases.  When there
+        // are no control message arguments this triggers persistence of all
+        // chained processors, i.e. maybe the categorizer as well as the anomaly
+        // detector if there is one.  But when control message arguments are
+        // passed, ONLY the persistence of the anomaly detector is triggered.
         if (controlMessageArgs.empty()) {
             if (this->isPersistenceNeeded("state")) {
                 m_PersistenceManager->startPersist(core::CTimeUtils::now());
@@ -421,6 +426,10 @@ void CAnomalyJob::processPersistControlMessage(const std::string& controlMessage
             std::string snapshotDescription;
             if (parsePersistControlMessageArgs(controlMessageArgs, snapshotTimestamp,
                                                snapshotId, snapshotDescription)) {
+                // Since this is not going through the full persistence call
+                // chain, make sure model size stats are up to date before
+                // persisting
+                m_Limits.resourceMonitor().forceRefreshAll();
                 if (m_PersistenceManager->doForegroundPersist(
                         [this, &snapshotDescription, &snapshotId,
                          &snapshotTimestamp](core::CDataAdder& persister) {

lib/model/CResourceMonitor.cc

@@ -122,6 +122,14 @@ void CResourceMonitor::forceRefresh(CMonitoredResource& resource) {
     this->updateAllowAllocations();
 }
 
+void CResourceMonitor::forceRefreshAll() {
+    for (auto& resource : m_Resources) {
+        this->memUsage(resource.first);
+    }
+
+    this->updateAllowAllocations();
+}
+
 void CResourceMonitor::updateAllowAllocations() {
     std::size_t total{this->totalMemory()};
     core::CProgramCounters::counter(counter_t::E_TSADMemoryUsage) = total;

lib/model/unittest/CResourceMonitorTest.cc

@@ -264,6 +264,23 @@ BOOST_FIXTURE_TEST_CASE(testMonitor, CTestFixture) {
         BOOST_REQUIRE_EQUAL(mem, m_ReportedModelSizeStats.s_Usage);
         BOOST_REQUIRE_EQUAL(model_t::E_MemoryStatusOk, m_ReportedModelSizeStats.s_MemoryStatus);
     }
+    {
+        // As above but refreshing all resources in one call
+        CResourceMonitor mon;
+        mon.registerComponent(categorizer);
+        mon.registerComponent(detector1);
+        mon.registerComponent(detector2);
+
+        mon.memoryUsageReporter(std::bind(&CTestFixture::reportCallback, this,
+                                          std::placeholders::_1));
+        m_ReportedModelSizeStats.s_Usage = 0;
+        BOOST_REQUIRE_EQUAL(std::size_t(0), m_ReportedModelSizeStats.s_Usage);
+
+        mon.forceRefreshAll();
+        mon.sendMemoryUsageReportIfSignificantlyChanged(0);
+        BOOST_REQUIRE_EQUAL(mem, m_ReportedModelSizeStats.s_Usage);
+        BOOST_REQUIRE_EQUAL(model_t::E_MemoryStatusOk, m_ReportedModelSizeStats.s_MemoryStatus);
+    }
     {
         // Test the report callback for allocation failures
         CResourceMonitor mon;