From 6399cc427de5a11d8660e985586b2d34cc1daf18 Mon Sep 17 00:00:00 2001
From: postmannen <postmannen@gmail.com>
Date: Fri, 17 Dec 2021 17:28:57 +0100
Subject: [PATCH] Added connection timeouts

---
 configuration_flags.go | 10 ++++++++++
 server.go              |  7 ++++++-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/configuration_flags.go b/configuration_flags.go
index 2444ad9..cc18d2c 100644
--- a/configuration_flags.go
+++ b/configuration_flags.go
@@ -31,6 +31,8 @@ type Configuration struct {
 	NodeName string
 	// the address of the message broker
 	BrokerAddress string
+	// NatsConnOptTimeout the timeout for trying the connect to nats broker
+	NatsConnOptTimeout int
 	// nats connect retry
 	NatsConnectRetryInterval int
 	// NatsReconnectJitter in milliseconds
@@ -110,6 +112,7 @@ type ConfigurationFromFile struct {
 	DatabaseFolder            *string
 	NodeName                  *string
 	BrokerAddress             *string
+	NatsConnOptTimeout        *int
 	NatsConnectRetryInterval  *int
 	NatsReconnectJitter       *int
 	NatsReconnectJitterTLS    *int
@@ -157,6 +160,7 @@ func newConfigurationDefaults() Configuration {
 		HTTPListener:             "",
 		DatabaseFolder:           "./var/lib",
 		BrokerAddress:            "127.0.0.1:4222",
+		NatsConnOptTimeout:       20,
 		NatsConnectRetryInterval: 10,
 		NatsReconnectJitter:      100,
 		NatsReconnectJitterTLS:   1,
@@ -233,6 +237,11 @@ func checkConfigValues(cf ConfigurationFromFile) Configuration {
 	} else {
 		conf.BrokerAddress = *cf.BrokerAddress
 	}
+	if cf.NatsConnOptTimeout == nil {
+		conf.NatsConnOptTimeout = cd.NatsConnOptTimeout
+	} else {
+		conf.NatsConnOptTimeout = *cf.NatsConnOptTimeout
+	}
 	if cf.NatsConnectRetryInterval == nil {
 		conf.NatsConnectRetryInterval = cd.NatsConnectRetryInterval
 	} else {
@@ -422,6 +431,7 @@ func (c *Configuration) CheckFlags() error {
 	flag.StringVar(&c.DatabaseFolder, "databaseFolder", fc.DatabaseFolder, "folder who contains the database file. Defaults to ./var/lib/. If other folder is used this flag must be specified at startup.")
 	flag.StringVar(&c.NodeName, "nodeName", fc.NodeName, "some unique string to identify this Edge unit")
 	flag.StringVar(&c.BrokerAddress, "brokerAddress", fc.BrokerAddress, "the address of the message broker")
+	flag.IntVar(&c.NatsConnOptTimeout, "natsConnOptTimeout", fc.NatsConnOptTimeout, "default nats client conn timeout")
 	flag.IntVar(&c.NatsConnectRetryInterval, "natsConnectRetryInterval", fc.NatsConnectRetryInterval, "default nats retry connect interval in seconds.")
 	flag.IntVar(&c.NatsReconnectJitter, "natsReconnectJitter", fc.NatsReconnectJitter, "default nats ReconnectJitter interval in milliseconds.")
 	flag.IntVar(&c.NatsReconnectJitterTLS, "natsReconnectJitterTLS", fc.NatsReconnectJitterTLS, "default nats ReconnectJitterTLS interval in seconds.")
diff --git a/server.go b/server.go
index c9f738f..c9df359 100644
--- a/server.go
+++ b/server.go
@@ -91,7 +91,12 @@ func NewServer(c *Configuration, version string) (*server, error) {
 	for {
 		var err error
 		// Setting MaxReconnects to -1 which equals unlimited.
-		conn, err = nats.Connect(c.BrokerAddress, opt, nats.MaxReconnects(-1), nats.ReconnectJitter(time.Duration(c.NatsReconnectJitter)*time.Millisecond, time.Duration(c.NatsReconnectJitterTLS)*time.Second))
+		conn, err = nats.Connect(c.BrokerAddress,
+			opt,
+			nats.MaxReconnects(-1),
+			nats.ReconnectJitter(time.Duration(c.NatsReconnectJitter)*time.Millisecond, time.Duration(c.NatsReconnectJitterTLS)*time.Second),
+			nats.Timeout(time.Second*time.Duration(c.NatsConnOptTimeout)),
+		)
 		// If no servers where available, we loop and retry until succesful.
 		if err != nil {
 			log.Printf("error: could not connect, waiting %v seconds, and retrying: %v\n", c.NatsConnectRetryInterval, err)