1 files changed, 77 insertions, 4 deletions

diff --git a/ml/ml.cc b/ml/ml.cc
index cfda3d28f..b56401b07 100644
--- a/ml/ml.cc
+++ b/ml/ml.cc
@@ -4,8 +4,24 @@
 #include "Dimension.h"
 #include "Host.h"
 
+#include <random>
+
 using namespace ml;
 
+bool ml_capable() {
+    return true;
+}
+
+bool ml_enabled(RRDHOST *RH) {
+    if (!Cfg.EnableAnomalyDetection)
+        return false;
+
+    if (simple_pattern_matches(Cfg.SP_HostsToSkip, RH->hostname))
+        return false;
+
+    return true;
+}
+
 /*
  * Assumptions:
  *  1) hosts outlive their sets, and sets outlive their dimensions,
@@ -13,14 +29,24 @@ using namespace ml;
  */
 
 void ml_init(void) {
+    // Read config values
     Cfg.readMLConfig();
-}
 
-void ml_new_host(RRDHOST *RH) {
     if (!Cfg.EnableAnomalyDetection)
         return;
 
-    if (simple_pattern_matches(Cfg.SP_HostsToSkip, RH->hostname))
+    // Generate random numbers to efficiently sample the features we need
+    // for KMeans clustering.
+    std::random_device RD;
+    std::mt19937 Gen(RD());
+
+    Cfg.RandomNums.reserve(Cfg.MaxTrainSamples);
+    for (size_t Idx = 0; Idx != Cfg.MaxTrainSamples; Idx++)
+        Cfg.RandomNums.push_back(Gen());
+}
+
+void ml_new_host(RRDHOST *RH) {
+    if (!ml_enabled(RH))
         return;
 
     Host *H = new Host(RH);
@@ -64,7 +90,10 @@ void ml_delete_dimension(RRDDIM *RD) {
         return;
 
     Host *H = static_cast<Host *>(RD->rrdset->rrdhost->ml_host);
-    H->removeDimension(D);
+    if (!H)
+        delete D;
+    else
+        H->removeDimension(D);
 
     RD->state->ml_dimension = nullptr;
 }
@@ -150,6 +179,48 @@ char *ml_get_anomaly_event_info(RRDHOST *RH, const char *AnomalyDetectorName,
     return strdup(Json.dump(4, '\t').c_str());
 }
 
+void ml_process_rrdr(RRDR *R, int MaxAnomalyRates) {
+    if (R->rows != 1)
+        return;
+
+    if (MaxAnomalyRates < 1 || MaxAnomalyRates >= R->d)
+        return;
+
+    calculated_number *CNs = R->v;
+    RRDR_DIMENSION_FLAGS *DimFlags = R->od;
+
+    std::vector<std::pair<calculated_number, int>> V;
+
+    V.reserve(R->d);
+    for (int Idx = 0; Idx != R->d; Idx++)
+        V.emplace_back(CNs[Idx], Idx);
+
+    std::sort(V.rbegin(), V.rend());
+
+    for (int Idx = MaxAnomalyRates; Idx != R->d; Idx++) {
+        int UnsortedIdx = V[Idx].second;
+
+        int OldFlags = static_cast<int>(DimFlags[UnsortedIdx]);
+        int NewFlags = OldFlags | RRDR_DIMENSION_HIDDEN;
+
+        DimFlags[UnsortedIdx] = static_cast<rrdr_dimension_flag>(NewFlags);
+    }
+}
+
+void ml_dimension_update_name(RRDSET *RS, RRDDIM *RD, const char *Name) {
+    (void) RS;
+
+    Dimension *D = static_cast<Dimension *>(RD->state->ml_dimension);
+    if (!D)
+        return;
+
+    D->setAnomalyRateRDName(Name);
+}
+
+bool ml_streaming_enabled() {
+    return Cfg.StreamADCharts;
+}
+
 #if defined(ENABLE_ML_TESTS)
 
 #include "gtest/gtest.h"
@@ -163,3 +234,5 @@ int test_ml(int argc, char *argv[]) {
 }
 
 #endif // ENABLE_ML_TESTS
+
+#include "ml-private.h"