From 461e25ab958038f2b656cf7fc7b3c23eb60e930f Mon Sep 17 00:00:00 2001
From: Michael Martin <michael.martin@konghq.com>
Date: Sun, 1 Dec 2024 19:40:53 -0800
Subject: [PATCH] fix(prometheus): add a toggle switch for wasm metrics export

---
 kong/clustering/compat/removed_fields.lua     |  7 +++++
 kong/plugins/prometheus/exporter.lua          | 14 +++++++++-
 kong/plugins/prometheus/schema.lua            |  1 +
 kong/plugins/prometheus/wasmx.lua             |  7 ++++-
 .../09-hybrid_mode/09-config-compat_spec.lua  | 28 +++++++++++++++++--
 .../26-prometheus/09-wasmx_spec.lua           |  5 +++-
 6 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/kong/clustering/compat/removed_fields.lua b/kong/clustering/compat/removed_fields.lua
index 00bafc650f8..9127a4d3dea 100644
--- a/kong/clustering/compat/removed_fields.lua
+++ b/kong/clustering/compat/removed_fields.lua
@@ -228,4 +228,11 @@ return {
       "queue.concurrency_limit",
     },
   },
+
+  -- Any dataplane older than 3.9.0
+  [3009000000] = {
+    prometheus = {
+      "wasm_metrics",
+    },
+  },
 }
diff --git a/kong/plugins/prometheus/exporter.lua b/kong/plugins/prometheus/exporter.lua
index 83e3d05a917..d629b4a20e3 100644
--- a/kong/plugins/prometheus/exporter.lua
+++ b/kong/plugins/prometheus/exporter.lua
@@ -206,7 +206,19 @@ end
 
 
 local function configure(configs)
-  IS_PROMETHEUS_ENABLED = configs ~= nil
+  IS_PROMETHEUS_ENABLED = false
+  wasm.set_enabled(false)
+
+  if configs ~= nil then
+    IS_PROMETHEUS_ENABLED = true
+
+    for i = 1, #configs do
+      if configs[i].wasm_metrics then
+        wasm.set_enabled(true)
+        break
+      end
+    end
+  end
 end
 
 
diff --git a/kong/plugins/prometheus/schema.lua b/kong/plugins/prometheus/schema.lua
index a23e3b3fc5e..df9d3dc8cbb 100644
--- a/kong/plugins/prometheus/schema.lua
+++ b/kong/plugins/prometheus/schema.lua
@@ -22,6 +22,7 @@ return {
           { latency_metrics = { description = "A boolean value that determines if latency metrics should be collected. If enabled, `kong_latency_ms`, `upstream_latency_ms` and `request_latency_ms` metrics will be exported.", type = "boolean", default = false }, },
           { bandwidth_metrics = { description = "A boolean value that determines if bandwidth metrics should be collected. If enabled, `bandwidth_bytes` and `stream_sessions_total` metrics will be exported.", type = "boolean", default = false }, },
           { upstream_health_metrics = { description = "A boolean value that determines if upstream metrics should be collected. If enabled, `upstream_target_health` metric will be exported.", type = "boolean", default = false }, },
+          { wasm_metrics = { description = "A boolean value that determines if Wasm metrics should be collected.", type = "boolean", default = false }, },
         },
         custom_validator = validate_shared_dict,
     }, },
diff --git a/kong/plugins/prometheus/wasmx.lua b/kong/plugins/prometheus/wasmx.lua
index 3f3e36546fd..c053819ca0d 100644
--- a/kong/plugins/prometheus/wasmx.lua
+++ b/kong/plugins/prometheus/wasmx.lua
@@ -18,6 +18,7 @@ local _M = {}
 local FLUSH_EVERY = 100
 local GET_METRIC_OPTS = { prefix = false }
 
+local export_enabled = false
 
 local metrics_data_buf = buf_new()
 local labels_serialization_buf = buf_new()
@@ -181,7 +182,7 @@ end
 
 
 _M.metrics_data = function()
-  if not wasm.enabled() then
+  if not export_enabled or not wasm.enabled() then
     return
   end
 
@@ -231,4 +232,8 @@ _M.metrics_data = function()
 end
 
 
+function _M.set_enabled(enabled)
+  export_enabled = enabled
+end
+
 return _M
diff --git a/spec/02-integration/09-hybrid_mode/09-config-compat_spec.lua b/spec/02-integration/09-hybrid_mode/09-config-compat_spec.lua
index 8c2b26fba41..56d81f22e82 100644
--- a/spec/02-integration/09-hybrid_mode/09-config-compat_spec.lua
+++ b/spec/02-integration/09-hybrid_mode/09-config-compat_spec.lua
@@ -1013,13 +1013,37 @@ describe("CP/DP config compat transformations #" .. strategy, function()
         }
         -- ]]
 
+        finally(function()
+          admin.plugins:remove({ id = prometheus.id })
+        end)
+
         local expected_prometheus_prior_38 = cycle_aware_deep_copy(prometheus)
         expected_prometheus_prior_38.config.ai_metrics = nil
+        expected_prometheus_prior_38.config.wasm_metrics = nil
 
         do_assert(uuid(), "3.7.0", expected_prometheus_prior_38)
+      end)
 
-        -- cleanup
-        admin.plugins:remove({ id = prometheus.id })
+      it("[prometheus] remove wasm_metrics property for versions below 3.9", function()
+        -- [[ 3.9.x ]] --
+        local prometheus = admin.plugins:insert {
+          name = "prometheus",
+          enabled = true,
+          config = {
+            wasm_metrics = true, -- becomes nil
+          },
+        }
+        -- ]]
+
+        finally(function()
+          admin.plugins:remove({ id = prometheus.id })
+        end)
+
+
+        local expected_prometheus_prior_39 = cycle_aware_deep_copy(prometheus)
+        expected_prometheus_prior_39.config.wasm_metrics = nil
+
+        do_assert(uuid(), "3.8.0", expected_prometheus_prior_39)
       end)
     end)
 
diff --git a/spec/03-plugins/26-prometheus/09-wasmx_spec.lua b/spec/03-plugins/26-prometheus/09-wasmx_spec.lua
index d52d27bb5bd..4684636bcee 100644
--- a/spec/03-plugins/26-prometheus/09-wasmx_spec.lua
+++ b/spec/03-plugins/26-prometheus/09-wasmx_spec.lua
@@ -90,7 +90,10 @@ for _, strategy in helpers.each_strategy() do
       add_filter_to_service(bp, "tests", service)
       add_filter_to_service(bp, "tests", service2)
 
-      bp.plugins:insert({ name = "prometheus" })
+      assert(bp.plugins:insert({
+        name = "prometheus",
+        config = { wasm_metrics = true },
+      }))
 
       assert(helpers.start_kong({
         nginx_conf = "spec/fixtures/custom_nginx.template",