From b8d0567889c35788f19eb143a8c25de57664232f Mon Sep 17 00:00:00 2001
From: Rafael Vasquez <raf.vasquez@ibm.com>
Date: Wed, 8 Feb 2023 17:01:28 -0500
Subject: [PATCH] feat: Allow custom prometheus info metric

Signed-off-by: Rafael Vasquez <raf.vasquez@ibm.com>
---
 .../com/ibm/watson/modelmesh/Metrics.java     | 19 ++++++++++++++++
 .../com/ibm/watson/modelmesh/ModelMesh.java   | 22 ++++++++++++++++++-
 .../watson/modelmesh/ModelMeshEnvVars.java    |  1 +
 3 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/ibm/watson/modelmesh/Metrics.java b/src/main/java/com/ibm/watson/modelmesh/Metrics.java
index 6db8430b6..a9895a958 100644
--- a/src/main/java/com/ibm/watson/modelmesh/Metrics.java
+++ b/src/main/java/com/ibm/watson/modelmesh/Metrics.java
@@ -159,6 +159,7 @@ public PrometheusMetrics(Map<String, String> params) throws Exception {
             int port = 2112;
             boolean shortNames = true;
             boolean https = true;
+            boolean customMetrics = false;
             String memMetrics = "all"; // default to all
             for (Entry<String, String> ent : params.entrySet()) {
                 switch (ent.getKey()) {
@@ -183,6 +184,9 @@ public PrometheusMetrics(Map<String, String> params) throws Exception {
                 case "mem_detail":
                     memMetrics = ent.getValue();
                     break;
+                case "metric_name":
+                    customMetrics = true;
+                    break;
                 default:
                     throw new Exception("Unrecognized metrics config parameter: " + ent.getKey());
                 }
@@ -230,6 +234,21 @@ public PrometheusMetrics(Map<String, String> params) throws Exception {
                 }
             }
 
+            // Custom info metrics
+            if (customMetrics){
+                @SuppressWarnings("rawtypes")
+                SimpleCollector.Builder builder;
+                builder = Gauge.build();
+                Collector collector = (
+                    builder
+                    .name(params.get("metric_name"))
+                    .help("Custom Info Metrics")
+                    .labelNames(params.get("deployment"), params.get("slot"), params.get("component"), params.get("group"))
+                    .create()
+                );
+                registry.register(collector);
+            }
+
             this.metricServer = new NettyServer(registry, port, https);
             this.shortNames = shortNames;
 
diff --git a/src/main/java/com/ibm/watson/modelmesh/ModelMesh.java b/src/main/java/com/ibm/watson/modelmesh/ModelMesh.java
index f4f964ec5..aa720709c 100644
--- a/src/main/java/com/ibm/watson/modelmesh/ModelMesh.java
+++ b/src/main/java/com/ibm/watson/modelmesh/ModelMesh.java
@@ -883,7 +883,9 @@ protected final TProcessor initialize() throws Exception {
 //    }
 
     // "type" or "type:p1=v1;p2=v2;...;pn=vn"
-    private static final Pattern METRICS_CONFIG_PATT = Pattern.compile("([a-z]+)(:\\w+=[^;]+(?:;\\w+=[^;]+)*)?");
+    private static final Pattern METRICS_CONFIG_PATT = Pattern.compile("([a-z;]+)(:\\w+=[^;]+(?:;\\w+=[^;]+)*)?");
+    // "metric_name" or "metric:name;l1=v1,l2=v2,...,ln=vn,"
+    private static final Pattern CUSTOM_METRIC_CONFIG_PATT = Pattern.compile("([a-z_:]+);(\\w+=[^;]+(?:;\\w+=[^,]+)*)?");
 
     private static Metrics setUpMetrics() throws Exception {
         if (System.getenv("MM_METRICS_STATSD_PORT") != null || System.getenv("MM_METRICS_PROMETHEUS_PORT") != null) {
@@ -916,6 +918,24 @@ private static Metrics setUpMetrics() throws Exception {
                 params.put(kv[0], kv[1]);
             }
         }
+        String customMetricConfig = getStringParameter(MMESH_CUSTOM_ENV_VAR, null);
+        if (customMetricConfig == null) {
+            logger.info("{} returned null", MMESH_CUSTOM_ENV_VAR);
+        } else {
+            logger.info("{} set to \"{}\"", MMESH_CUSTOM_ENV_VAR, customMetricConfig);
+            Matcher customMetricMatcher = CUSTOM_METRIC_CONFIG_PATT.matcher(customMetricConfig);
+            if (!customMetricMatcher.matches()) {
+                throw new Exception("Invalid metrics configuration provided in env var " + MMESH_CUSTOM_ENV_VAR + ": \""
+                                    + customMetricConfig + "\"");
+            }
+            String name = customMetricMatcher.group(1);
+            String customParamString = customMetricMatcher.group(2);
+            params.put("metric_name", name);
+            for (String customParm : customParamString.substring(1).split(",")) {
+                String[] kv = customParm.split("=");
+                params.put(kv[0], kv[1]);
+            }
+        }
         try {
             switch (type.toLowerCase()) {
             case "statsd":
diff --git a/src/main/java/com/ibm/watson/modelmesh/ModelMeshEnvVars.java b/src/main/java/com/ibm/watson/modelmesh/ModelMeshEnvVars.java
index 6351f5f64..b295dcf2b 100644
--- a/src/main/java/com/ibm/watson/modelmesh/ModelMeshEnvVars.java
+++ b/src/main/java/com/ibm/watson/modelmesh/ModelMeshEnvVars.java
@@ -46,6 +46,7 @@ private ModelMeshEnvVars() {}
     public static final String LOAD_FAILURE_EXPIRY_ENV_VAR = "MM_LOAD_FAILURE_EXPIRY_TIME_MS";
 
     public static final String MMESH_METRICS_ENV_VAR = "MM_METRICS";
+    public static final String MMESH_CUSTOM_ENV_VAR = "MM_INFO_METRICS";
 
     public static final String LOG_EACH_INVOKE_ENV_VAR = "MM_LOG_EACH_INVOKE";
     public static final String SEND_DEST_ID_ENV_VAR = "MM_SEND_DEST_ID";