From 34f7cad9277a3a274a062c6c6e641c918af3e513 Mon Sep 17 00:00:00 2001
From: Jianbo Sun <jianbosun@microsoft.com>
Date: Wed, 18 Dec 2024 10:52:06 +0800
Subject: [PATCH] support expose metrics for plugins

---
 api/pkg/filtermanager/config.go       |  6 ++++++
 api/pkg/plugins/plugins.go            | 12 ++++++++++++
 api/pkg/plugins/plugins_test.go       |  6 ++++++
 api/tests/integration/test_plugins.go | 16 +++++++++++++++-
 4 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/api/pkg/filtermanager/config.go b/api/pkg/filtermanager/config.go
index 6cd411fe2..44ba24c90 100644
--- a/api/pkg/filtermanager/config.go
+++ b/api/pkg/filtermanager/config.go
@@ -216,6 +216,12 @@ func (p *FilterManagerConfigParser) Parse(any *anypb.Any, callbacks capi.ConfigC
 
 	for _, proto := range plugins {
 		name := proto.Name
+
+		if registerMetrics := pkgPlugins.LoadMetricsCallback(name); registerMetrics != nil {
+			api.LogInfof("register metrics for plugin %s", name)
+			registerMetrics(callbacks)
+		}
+
 		if plugin := pkgPlugins.LoadHTTPFilterFactoryAndParser(name); plugin != nil {
 			config, err := plugin.ConfigParser.Parse(proto.Config)
 			if err != nil {
diff --git a/api/pkg/plugins/plugins.go b/api/pkg/plugins/plugins.go
index 5ab2742ab..6bcbbda29 100644
--- a/api/pkg/plugins/plugins.go
+++ b/api/pkg/plugins/plugins.go
@@ -20,6 +20,8 @@ import (
 	"errors"
 	"runtime/debug"
 
+	capi "github.com/envoyproxy/envoy/contrib/golang/common/go/api"
+
 	"mosn.io/htnn/api/internal/proto"
 	"mosn.io/htnn/api/pkg/filtermanager/api"
 	"mosn.io/htnn/api/pkg/log"
@@ -31,6 +33,7 @@ var (
 	pluginTypes                = map[string]Plugin{}
 	plugins                    = map[string]Plugin{}
 	httpFilterFactoryAndParser = map[string]*FilterFactoryAndParser{}
+	metricsRegister            = map[string]func(capi.ConfigCallbacks){}
 )
 
 // Here we introduce extra struct to avoid cyclic import between pkg/filtermanager and pkg/plugins
@@ -188,6 +191,15 @@ func (cp *PluginConfigParser) Parse(any interface{}) (res interface{}, err error
 	return conf, nil
 }
 
+func RegisterMetricsCallback(pluginName string, registerMetricFunc func(capi.ConfigCallbacks)) {
+	logger.Info("register metrics for plugin", "name", pluginName)
+	metricsRegister[pluginName] = registerMetricFunc
+}
+
+func LoadMetricsCallback(pluginName string) func(capi.ConfigCallbacks) {
+	return metricsRegister[pluginName]
+}
+
 // PluginMethodDefaultImpl provides reasonable implementation for optional methods
 type PluginMethodDefaultImpl struct{}
 
diff --git a/api/pkg/plugins/plugins_test.go b/api/pkg/plugins/plugins_test.go
index 103e6572f..235d68687 100644
--- a/api/pkg/plugins/plugins_test.go
+++ b/api/pkg/plugins/plugins_test.go
@@ -20,6 +20,7 @@ import (
 	"testing"
 
 	"github.com/agiledragon/gomonkey/v2"
+	capi "github.com/envoyproxy/envoy/contrib/golang/common/go/api"
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/protobuf/types/known/anypb"
 
@@ -285,3 +286,8 @@ func TestRegisterPluginWithType(t *testing.T) {
 	assert.NotNil(t, LoadPlugin("mock"))
 	assert.NotNil(t, LoadPluginType("mock"))
 }
+
+func TestRegisterPluginMetrics(t *testing.T) {
+	RegisterMetricsCallback("mock", func(cc capi.ConfigCallbacks) {})
+	assert.NotNil(t, LoadMetricsCallback("mock"))
+}
diff --git a/api/tests/integration/test_plugins.go b/api/tests/integration/test_plugins.go
index dac15c4cb..78b842ff3 100644
--- a/api/tests/integration/test_plugins.go
+++ b/api/tests/integration/test_plugins.go
@@ -21,6 +21,8 @@ import (
 	"strconv"
 	"strings"
 
+	capi "github.com/envoyproxy/envoy/contrib/golang/common/go/api"
+
 	"mosn.io/htnn/api/pkg/filtermanager/api"
 	"mosn.io/htnn/api/pkg/plugins"
 )
@@ -200,6 +202,7 @@ func (p *bufferPlugin) Factory() api.FilterFactory {
 type localReplyPlugin struct {
 	plugins.PluginMethodDefaultImpl
 	basePlugin
+	usageCounter capi.CounterMetric
 }
 
 func localReplyFactory(c interface{}, callbacks api.FilterCallbackHandler) api.Filter {
@@ -242,6 +245,7 @@ func (f *localReplyFilter) DecodeRequest(headers api.RequestHeaderMap, buf api.B
 	f.reqHdr = headers
 	f.runFilters = headers.Values("run")
 	if f.config.Decode {
+		lrp.usageCounter.Increment(1)
 		return f.NewLocalResponse("reply", true)
 	}
 	return api.Continue
@@ -309,6 +313,11 @@ func (p *localReplyPlugin) Factory() api.FilterFactory {
 	return localReplyFactory
 }
 
+func (p *localReplyPlugin) MetricsDefinition(c capi.ConfigCallbacks) {
+	p.usageCounter = c.DefineCounterMetric("localreply.usage.counter")
+	// Define more metrics here
+}
+
 type badPlugin struct {
 	plugins.PluginMethodDefaultImpl
 }
@@ -619,10 +628,15 @@ func (f *onLogFilter) OnLog(reqHeaders api.RequestHeaderMap, reqTrailers api.Req
 	api.LogWarnf("receive request trailers: %+v", trailers)
 }
 
+var lrp = &localReplyPlugin{}
+
 func init() {
 	plugins.RegisterPlugin("stream", &streamPlugin{})
 	plugins.RegisterPlugin("buffer", &bufferPlugin{})
-	plugins.RegisterPlugin("localReply", &localReplyPlugin{})
+
+	plugins.RegisterPlugin("localReply", lrp)
+	plugins.RegisterMetricsCallback("localReply", lrp.MetricsDefinition)
+
 	plugins.RegisterPlugin("bad", &badPlugin{})
 	plugins.RegisterPlugin("consumer", &consumerPlugin{})
 	plugins.RegisterPlugin("init", &initPlugin{})