[ML] Recalculate memory usage before upgrading model state (#1585)

When a persist control message with arguments is received
by the anomaly detector it doesn't go through the standard
chain of persistence calls, as it unconditionally rewrites
the state (even if no data has been seen) and includes only
the anomaly detector state rather than the categorizer state
too.  Because of this the memory usage was not being
recalculated prior to persisting the state as would normally
happen.  This PR rectifies that omission.

Fixes one of the problems detailed in
elastic/elasticsearch#64665 (review)

Author: droberts195

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;