From 2fd65a469c25edf2531f58e239c6e92a636862f6 Mon Sep 17 00:00:00 2001
From: Ming Deng <mingflycash@gmail.com>
Date: Fri, 7 Aug 2020 14:14:07 +0000
Subject: [PATCH] Support prometheus

---
 pkg/orm/filter/prometheus/filter.go      | 84 ++++++++++++++++++++++++
 pkg/orm/filter/prometheus/filter_test.go | 60 +++++++++++++++++
 2 files changed, 144 insertions(+)
 create mode 100644 pkg/orm/filter/prometheus/filter.go
 create mode 100644 pkg/orm/filter/prometheus/filter_test.go

diff --git a/pkg/orm/filter/prometheus/filter.go b/pkg/orm/filter/prometheus/filter.go
new file mode 100644
index 00000000..9f177deb
--- /dev/null
+++ b/pkg/orm/filter/prometheus/filter.go
@@ -0,0 +1,84 @@
+// Copyright 2020 beego 
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package prometheus
+
+import (
+	"context"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+
+	beego "github.com/astaxie/beego/pkg"
+	"github.com/astaxie/beego/pkg/orm"
+)
+
+// FilterChainBuilder is an extension point,
+// when we want to support some configuration,
+// please use this structure
+type FilterChainBuilder struct {
+	summaryVec prometheus.ObserverVec
+}
+
+func NewFilterChainBuilder() *FilterChainBuilder {
+	summaryVec := prometheus.NewSummaryVec(prometheus.SummaryOpts{
+		Name:      "beego",
+		Subsystem: "orm_operation",
+		ConstLabels: map[string]string{
+			"server":  beego.BConfig.ServerName,
+			"env":     beego.BConfig.RunMode,
+			"appname": beego.BConfig.AppName,
+		},
+		Help: "The statics info for orm operation",
+	}, []string{"method", "name", "duration", "insideTx", "txName"})
+
+	prometheus.MustRegister(summaryVec)
+	return &FilterChainBuilder{
+		summaryVec: summaryVec,
+	}
+}
+
+func (builder *FilterChainBuilder) FilterChain(next orm.Filter) orm.Filter {
+	return func(ctx context.Context, inv *orm.Invocation) {
+		startTime := time.Now()
+		next(ctx, inv)
+		endTime := time.Now()
+		dur := (endTime.Sub(startTime)) / time.Millisecond
+
+		// if the TPS is too large, here may be some problem
+		// thinking about using goroutine pool
+		go builder.report(ctx, inv, dur)
+	}
+}
+
+func (builder *FilterChainBuilder) report(ctx context.Context, inv *orm.Invocation, dur time.Duration) {
+	// start a transaction, we don't record it
+	if strings.HasPrefix(inv.Method, "Begin") {
+		return
+	}
+	if inv.Method == "Commit" || inv.Method == "Rollback" {
+		builder.reportTxn(ctx, inv)
+		return
+	}
+	builder.summaryVec.WithLabelValues(inv.Method, inv.GetTableName(), strconv.Itoa(int(dur)),
+		strconv.FormatBool(inv.InsideTx), inv.TxName)
+}
+
+func (builder *FilterChainBuilder) reportTxn(ctx context.Context, inv *orm.Invocation) {
+	dur := time.Now().Sub(inv.TxStartTime) / time.Millisecond
+	builder.summaryVec.WithLabelValues(inv.Method, inv.TxName, strconv.Itoa(int(dur)),
+		strconv.FormatBool(inv.InsideTx), inv.TxName)
+}
diff --git a/pkg/orm/filter/prometheus/filter_test.go b/pkg/orm/filter/prometheus/filter_test.go
new file mode 100644
index 00000000..a71e8f50
--- /dev/null
+++ b/pkg/orm/filter/prometheus/filter_test.go
@@ -0,0 +1,60 @@
+// Copyright 2020 beego 
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package prometheus
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/astaxie/beego/pkg/orm"
+)
+
+func TestFilterChainBuilder_FilterChain(t *testing.T) {
+	builder := NewFilterChainBuilder()
+	assert.NotNil(t, builder.summaryVec)
+
+	filter := builder.FilterChain(func(ctx context.Context, inv *orm.Invocation) {
+		inv.Method = "coming"
+	})
+	assert.NotNil(t, filter)
+
+	inv := &orm.Invocation{}
+	filter(context.Background(), inv)
+	assert.Equal(t, "coming", inv.Method)
+
+	inv = &orm.Invocation{
+		Method:      "Hello",
+		TxStartTime: time.Now(),
+	}
+	builder.reportTxn(context.Background(), inv)
+
+	inv = &orm.Invocation{
+		Method: "Begin",
+	}
+
+	ctx := context.Background()
+	// it will be ignored
+	builder.report(ctx, inv, time.Second)
+
+	inv.Method = "Commit"
+	builder.report(ctx, inv, time.Second)
+
+	inv.Method = "Update"
+	builder.report(ctx, inv, time.Second)
+
+}