From d43e7373c6da05ab5f810c7d8b444c6b34045ee7 Mon Sep 17 00:00:00 2001
From: Rory Z <16801068+Rory-Z@users.noreply.github.com>
Date: Mon, 20 Nov 2023 22:06:32 +0800
Subject: [PATCH] feat(probe): probe support keep alive

Signed-off-by: Rory Z <16801068+Rory-Z@users.noreply.github.com>
---
 config/config.go           | 32 +++++++++++++++++++++++++-------
 config/example/config.yaml |  1 +
 prober/mqtt.go             | 10 +++++++---
 3 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/config/config.go b/config/config.go
index 18f4e2b..8515749 100644
--- a/config/config.go
+++ b/config/config.go
@@ -27,13 +27,28 @@ type Metrics struct {
 }
 
 type Probe struct {
-	Target          string           `yaml:"target"`
-	Scheme          string           `yaml:"scheme,omitempty"`
-	ClientID        string           `yaml:"client_id,omitempty"`
-	Username        string           `yaml:"username,omitempty"`
-	Password        string           `yaml:"password,omitempty"`
-	Topic           string           `yaml:"topic,omitempty"`
-	QoS             byte             `yaml:"qos,omitempty"`
+	// Target is the address of the EMQX node to probe. Required.
+	Target string `yaml:"target"`
+	// Scheme is the protocol scheme of the EMQX node to probe.
+	// Enum: [mqtt | tcp | mqtts | ssl | tls | ws | wss]
+	// Default: tcp
+	Scheme string `yaml:"scheme,omitempty"`
+	// ClientID is the MQTT client ID to use when probing.
+	// Default: emqx_exporter_probe_<index>
+	ClientID string `yaml:"client_id,omitempty"`
+	// Username is the MQTT username to use when probing.
+	Username string `yaml:"username,omitempty"`
+	// Password is the MQTT password to use when probing.
+	Password string `yaml:"password,omitempty"`
+	// Topic is the MQTT topic to use when probing.
+	// Default: emqx-exporter-probe-<index>
+	Topic string `yaml:"topic,omitempty"`
+	// QoS is the MQTT QoS to use when probing.
+	// Default: 0
+	QoS byte `yaml:"qos,omitempty"`
+	// KeepAlive is the keep alive period in seconds. Defaults to 30 seconds.
+	KeepAlive int64 `yaml:"keep_alive,omitempty"`
+	// TLSClientConfig is the TLS configuration to use when probing.
 	TLSClientConfig *TLSClientConfig `yaml:"tls_config,omitempty"`
 }
 
@@ -160,6 +175,9 @@ func (sc *SafeConfig) ReloadConfig(confFile string) (err error) {
 		if probe.Topic == "" {
 			probe.Topic = "emqx-exporter-probe-" + fmt.Sprintf("%d", index)
 		}
+		if probe.KeepAlive == 0 {
+			probe.KeepAlive = 30
+		}
 		c.Probes[index] = probe
 	}
 
diff --git a/config/example/config.yaml b/config/example/config.yaml
index 5cd7790..4bbbf88 100644
--- a/config/example/config.yaml
+++ b/config/example/config.yaml
@@ -10,3 +10,4 @@ probes:
     password:
     topic:
     qos:
+    keep_alive:
diff --git a/prober/mqtt.go b/prober/mqtt.go
index ec1ea27..18e6ca5 100644
--- a/prober/mqtt.go
+++ b/prober/mqtt.go
@@ -46,14 +46,18 @@ func init() {
 			select {
 			case <-context.Background().Done():
 				return
-			case <-time.After(5 * time.Second):
+			case <-time.After(60 * time.Second):
 			}
 		}
 	}()
 }
 
 func initMQTTProbe(probe config.Probe, logger log.Logger) (*MQTTProbe, error) {
-	opt := mqtt.NewClientOptions().AddBroker(probe.Scheme + "://" + probe.Target).SetClientID(probe.ClientID).SetUsername(probe.Username).SetPassword(probe.Password)
+	opt := mqtt.NewClientOptions().AddBroker(probe.Scheme + "://" + probe.Target)
+	opt.SetClientID(probe.ClientID)
+	opt.SetUsername(probe.Username)
+	opt.SetPassword(probe.Password)
+	opt.SetKeepAlive(time.Duration(probe.KeepAlive) * time.Second)
 	if probe.TLSClientConfig != nil {
 		opt.SetTLSConfig(probe.TLSClientConfig.ToTLSConfig())
 	}
@@ -108,7 +112,7 @@ func ProbeMQTT(probe config.Probe, logger log.Logger) bool {
 		if msg == nil {
 			return false
 		}
-	case <-time.After(5 * time.Second):
+	case <-time.After(time.Duration(probe.KeepAlive) * time.Second):
 		return false
 	}