From ca4ca3e4b1aef6e622d52f29ce320c4f35f15bbf Mon Sep 17 00:00:00 2001
From: Chien-Fu Chen <chienfuchen32@gmail.com>
Date: Thu, 6 Dec 2018 16:13:20 +0800
Subject: [PATCH] add global mongo keep-connection session option

---
 collector/mongodb_collector.go | 34 +++++++++++++++++++++++++++++++---
 mongodb_exporter.go            |  2 ++
 shared/connection.go           |  5 +++++
 3 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/collector/mongodb_collector.go b/collector/mongodb_collector.go
index 97da59cc..6a037557 100644
--- a/collector/mongodb_collector.go
+++ b/collector/mongodb_collector.go
@@ -1,6 +1,7 @@
 package collector
 
 import (
+	"errors"
 	"time"
 
 	"github.com/dcu/mongodb_exporter/shared"
@@ -40,6 +41,7 @@ type MongodbCollectorOpts struct {
 	UserName                 string
 	AuthMechanism            string
 	SocketTimeout            time.Duration
+	KeepConnection           bool
 }
 
 func (in MongodbCollectorOpts) toSessionOps() shared.MongoSessionOpts {
@@ -52,6 +54,7 @@ func (in MongodbCollectorOpts) toSessionOps() shared.MongoSessionOpts {
 		UserName:              in.UserName,
 		AuthMechanism:         in.AuthMechanism,
 		SocketTimeout:         in.SocketTimeout,
+		KeepConnection:        in.KeepConnection,
 	}
 }
 
@@ -81,14 +84,39 @@ func (exporter *MongodbCollector) Describe(ch chan<- *prometheus.Desc) {
 	}
 }
 
+var mongoGlobalSession *mgo.Session
+
+func checkGlobalSession(exporter *MongodbCollector) error {
+	if mongoGlobalSession == nil {
+		mongoGlobalSession = shared.MongoSession(exporter.Opts.toSessionOps())
+		if mongoGlobalSession == nil {
+			return errors.New("init mongo global session failed")
+		}
+	}
+	err := mongoGlobalSession.Ping()
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 // Collect collects all mongodb's metrics.
 func (exporter *MongodbCollector) Collect(ch chan<- prometheus.Metric) {
-	mongoSess := shared.MongoSession(exporter.Opts.toSessionOps())
-	if mongoSess != nil {
+	var mongoSess *mgo.Session
+	var errCheckSession error
+	if exporter.Opts.KeepConnection {
+		mongoSess = mongoGlobalSession
+		errCheckSession = checkGlobalSession(exporter)
+	} else {
+		mongoSess = shared.MongoSession(exporter.Opts.toSessionOps())
+	}
+	if mongoSess != nil && errCheckSession == nil {
 		upGauge.WithLabelValues().Set(float64(1))
 		upGauge.Collect(ch)
 		upGauge.Reset()
-		defer mongoSess.Close()
+		if !exporter.Opts.KeepConnection {
+			defer mongoSess.Close()
+		}
 		glog.Info("Collecting Server Status")
 		exporter.collectServerStatus(mongoSess, ch)
 		if exporter.Opts.CollectReplSet {
diff --git a/mongodb_exporter.go b/mongodb_exporter.go
index 9d0607fa..bd6998a0 100644
--- a/mongodb_exporter.go
+++ b/mongodb_exporter.go
@@ -58,6 +58,7 @@ var (
 	mongodbCollectProfileMetrics        = flag.Bool("mongodb.collect.profile", false, "Collect MongoDB profile metrics")
 	mongodbCollectConnPoolStats         = flag.Bool("mongodb.collect.connpoolstats", false, "Collect MongoDB connpoolstats")
 	mongodbSocketTimeout                = flag.Duration("mongodb.socket-timeout", 0, "timeout for socket operations to mongodb")
+	mongodbKeepConnection               = flag.Bool("mongodb.keep-connection", false, "keep one client connection to mongodb")
 	version                             = flag.Bool("version", false, "Print mongodb_exporter version")
 )
 
@@ -164,6 +165,7 @@ func registerCollector() {
 		UserName:                 *mongodbUserName,
 		AuthMechanism:            *mongodbAuthMechanism,
 		SocketTimeout:            *mongodbSocketTimeout,
+		KeepConnection:           *mongodbKeepConnection,
 	})
 	prometheus.MustRegister(mongodbCollector)
 }
diff --git a/shared/connection.go b/shared/connection.go
index c51fbd0c..eaea8503 100644
--- a/shared/connection.go
+++ b/shared/connection.go
@@ -27,6 +27,7 @@ type MongoSessionOpts struct {
 	UserName              string
 	AuthMechanism         string
 	SocketTimeout         time.Duration
+	KeepConnection        bool
 }
 
 // MongoSession creates a Mongo session
@@ -57,6 +58,10 @@ func MongoSession(opts MongoSessionOpts) *mgo.Session {
 	session.SetMode(mgo.Eventual, true)
 	session.SetSyncTimeout(syncMongodbTimeout)
 	session.SetSocketTimeout(opts.SocketTimeout)
+	if opts.KeepConnection {
+		session.SetMode(mgo.Eventual, false) // refresh doesn't work well while reconnect
+		session.SetPoolLimit(1)              // keep 1 global Connection
+	}
 	return session
 }