From 1cfa55817ad5cc188ec0c87efd51d4c151c87c4d Mon Sep 17 00:00:00 2001
From: wangchuxiao <wangchuxiao97@outlook.com>
Date: Thu, 15 Sep 2022 17:50:45 +0800
Subject: [PATCH] prometheus

---
 internal/msg_gateway/gate/relay_rpc_server.go |  1 +
 internal/msg_gateway/gate/ws_server.go        |  3 +++
 pkg/common/prometheus/gather.go               | 11 +++++++++++
 pkg/common/prometheus/prometheus.go           | 16 ++++++++++++++++
 4 files changed, 31 insertions(+)

diff --git a/internal/msg_gateway/gate/relay_rpc_server.go b/internal/msg_gateway/gate/relay_rpc_server.go
index 830a6cb80..f5573eefa 100644
--- a/internal/msg_gateway/gate/relay_rpc_server.go
+++ b/internal/msg_gateway/gate/relay_rpc_server.go
@@ -39,6 +39,7 @@ func initPrometheus() {
 	promePkg.NewGetNewestSeqTotalCounter()
 	promePkg.NewPullMsgBySeqListTotalCounter()
 	promePkg.NewMsgOnlinePushSuccessCounter()
+	promePkg.NewOnlineUserCounter()
 	//promePkg.NewSingleChatMsgRecvSuccessCounter()
 	//promePkg.NewGroupChatMsgRecvSuccessCounter()
 	//promePkg.NewWorkSuperGroupChatMsgRecvSuccessCounter()
diff --git a/internal/msg_gateway/gate/ws_server.go b/internal/msg_gateway/gate/ws_server.go
index 9b61318b3..aac99a768 100644
--- a/internal/msg_gateway/gate/ws_server.go
+++ b/internal/msg_gateway/gate/ws_server.go
@@ -5,6 +5,7 @@ import (
 	"Open_IM/pkg/common/constant"
 	"Open_IM/pkg/common/db"
 	"Open_IM/pkg/common/log"
+	promePkg "Open_IM/pkg/common/prometheus"
 	"Open_IM/pkg/common/token_verify"
 	"Open_IM/pkg/grpc-etcdv3/getcdv3"
 	pbRelay "Open_IM/pkg/proto/relay"
@@ -313,6 +314,7 @@ func (ws *WServer) addUserConn(uid string, platformID int, conn *UserConn, token
 	for _, v := range ws.wsUserToConn {
 		count = count + len(v)
 	}
+	promePkg.PromeGaugeInc(promePkg.OnlineUserGauge)
 	log.Debug(operationID, "WS Add operation", "", "wsUser added", ws.wsUserToConn, "connection_uid", uid, "connection_platform", constant.PlatformIDToName(platformID), "online_user_num", len(ws.wsUserToConn), "online_conn_num", count)
 }
 
@@ -352,6 +354,7 @@ func (ws *WServer) delUserConn(conn *UserConn) {
 	if callbackResp.ErrCode != 0 {
 		log.NewError(operationID, utils.GetSelfFuncName(), "callbackUserOffline failed", callbackResp)
 	}
+	promePkg.PromeGaugeDec(promePkg.OnlineUserGauge)
 }
 
 func (ws *WServer) getUserConn(uid string, platform int) *UserConn {
diff --git a/pkg/common/prometheus/gather.go b/pkg/common/prometheus/gather.go
index 1d6b71ebb..3dd7d05a2 100644
--- a/pkg/common/prometheus/gather.go
+++ b/pkg/common/prometheus/gather.go
@@ -34,6 +34,7 @@ var (
 	SingleChatMsgRecvSuccessCounter         prometheus.Counter
 	GroupChatMsgRecvSuccessCounter          prometheus.Counter
 	WorkSuperGroupChatMsgRecvSuccessCounter prometheus.Counter
+	OnlineUserGauge                         prometheus.Gauge
 
 	//msg-msg
 	SingleChatMsgProcessSuccessCounter         prometheus.Counter
@@ -326,6 +327,16 @@ func NewWorkSuperGroupChatMsgRecvSuccessCounter() {
 	})
 }
 
+func NewOnlineUserGauges() {
+	if OnlineUserGauge != nil {
+		return
+	}
+	OnlineUserGauge = promauto.NewGauge(prometheus.GaugeOpts{
+		Name: "online_user_num",
+		Help: "The number of online user num",
+	})
+}
+
 func NewSingleChatMsgProcessSuccessCounter() {
 	if SingleChatMsgProcessSuccessCounter != nil {
 		return
diff --git a/pkg/common/prometheus/prometheus.go b/pkg/common/prometheus/prometheus.go
index e8bc87b47..7c497a78c 100644
--- a/pkg/common/prometheus/prometheus.go
+++ b/pkg/common/prometheus/prometheus.go
@@ -64,3 +64,19 @@ func PromeAdd(counter prometheus.Counter, add int) {
 		}
 	}
 }
+
+func PromeGaugeInc(gauges prometheus.Gauge) {
+	if config.Config.Prometheus.Enable {
+		if gauges != nil {
+			gauges.Inc()
+		}
+	}
+}
+
+func PromeGaugeDec(gauges prometheus.Gauge) {
+	if config.Config.Prometheus.Enable {
+		if gauges != nil {
+			gauges.Dec()
+		}
+	}
+}