From 91df8f413e1a65752eb84f5df793943def370508 Mon Sep 17 00:00:00 2001
From: Caleb Lloyd <2414837+caleblloyd@users.noreply.github.com>
Date: Wed, 12 Jun 2024 17:45:30 -0400
Subject: [PATCH] Allow passing natsOpts from service observations and js
 advisories (#203)

* allow passing natsOpts from service observations and js advisories

Signed-off-by: Caleb Lloyd <caleb@synadia.com>

* Lint fix

---------

Signed-off-by: Caleb Lloyd <caleb@synadia.com>
Co-authored-by: Samuel Attwood <sam@synadia.com>
---
 go.mod                           |  2 +-
 surveyor/conn_pool.go            |  5 ++++-
 surveyor/jetstream_advisories.go | 13 +++++++++----
 surveyor/observation.go          | 13 +++++++++----
 surveyor/surveyor.go             |  3 +--
 5 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/go.mod b/go.mod
index e3859d0..5f66db0 100644
--- a/go.mod
+++ b/go.mod
@@ -18,7 +18,6 @@ require (
 	github.com/spf13/viper v1.18.2
 	github.com/stretchr/testify v1.8.4
 	golang.org/x/crypto v0.19.0
-	golang.org/x/exp v0.0.0-20240213143201-ec583247a57a
 	golang.org/x/sync v0.6.0
 )
 
@@ -46,6 +45,7 @@ require (
 	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/subosito/gotenv v1.6.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
+	golang.org/x/exp v0.0.0-20240213143201-ec583247a57a // indirect
 	golang.org/x/sys v0.17.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
 	golang.org/x/time v0.5.0 // indirect
diff --git a/surveyor/conn_pool.go b/surveyor/conn_pool.go
index c3b47b2..a9579f1 100644
--- a/surveyor/conn_pool.go
+++ b/surveyor/conn_pool.go
@@ -26,6 +26,9 @@ type natsContext struct {
 	TLSCA       string `json:"tls_ca"`
 	TLSCert     string `json:"tls_cert"`
 	TLSKey      string `json:"tls_key"`
+
+	// only passed programmatically
+	NatsOpts []nats.Option `json:"-"`
 }
 
 func (c *natsContext) copy() *natsContext {
@@ -199,7 +202,7 @@ func (cp *natsConnPool) getPooledConn(key string, cfg *natsContext) (*pooledNats
 		}
 		cp.Unlock()
 
-		opts := cp.natsOpts
+		opts := append(cp.natsOpts, cfg.NatsOpts...)
 		opts = append(opts, func(options *nats.Options) error {
 			if cfg.Name != "" {
 				options.Name = cfg.Name
diff --git a/surveyor/jetstream_advisories.go b/surveyor/jetstream_advisories.go
index 8c15f2d..7512122 100644
--- a/surveyor/jetstream_advisories.go
+++ b/surveyor/jetstream_advisories.go
@@ -235,10 +235,6 @@ type JSAdvisoryConfig struct {
 	// account name
 	AccountName string `json:"name"`
 
-	// optional configuration for importing JS metrics and advisories from other accounts
-	// it can only be set via JSAdvisoryConfig directly (not from config file)
-	ExternalAccountConfig *JSAdvisoriesExternalAccountConfig `json:"-"`
-
 	// connection options
 	JWT         string `json:"jwt"`
 	Seed        string `json:"seed"`
@@ -250,6 +246,14 @@ type JSAdvisoryConfig struct {
 	TLSCA       string `json:"tls_ca"`
 	TLSCert     string `json:"tls_cert"`
 	TLSKey      string `json:"tls_key"`
+
+	// additional opts available via JSAdvisoryConfig directly (not from config file)
+
+	// optional configuration for importing JS metrics and advisories from other accounts
+	ExternalAccountConfig *JSAdvisoriesExternalAccountConfig `json:"-"`
+
+	// nats options appended to base surveyor options
+	NatsOpts []nats.Option `json:"-"`
 }
 
 // JSAdvisoriesExternalAccountConfig is used to configure external accounts from which
@@ -387,6 +391,7 @@ func (o *jsAdvisoryListener) natsContext() *natsContext {
 		TLSCA:       o.config.TLSCA,
 		TLSCert:     o.config.TLSCert,
 		TLSKey:      o.config.TLSKey,
+		NatsOpts:    o.config.NatsOpts,
 	}
 
 	return natsCtx
diff --git a/surveyor/observation.go b/surveyor/observation.go
index 408a10c..6c405ed 100644
--- a/surveyor/observation.go
+++ b/surveyor/observation.go
@@ -125,10 +125,6 @@ type ServiceObsConfig struct {
 	ServiceName string `json:"name"`
 	Topic       string `json:"topic"`
 
-	// optional configuration for importing observations from other accounts
-	// it can only be set via ServiceObsConfig directly (not from config file)
-	ExternalAccountConfig *ServiceObservationExternalAccountConfig `json:"-"`
-
 	// connection options
 	JWT         string `json:"jwt"`
 	Seed        string `json:"seed"`
@@ -140,6 +136,14 @@ type ServiceObsConfig struct {
 	TLSCA       string `json:"tls_ca"`
 	TLSCert     string `json:"tls_cert"`
 	TLSKey      string `json:"tls_key"`
+
+	// additional opts available via ServiceObsConfig directly (not from config file)
+
+	// optional configuration for importing observations from other accounts
+	ExternalAccountConfig *ServiceObservationExternalAccountConfig `json:"-"`
+
+	// nats options appended to base surveyor options
+	NatsOpts []nats.Option `json:"-"`
 }
 
 type ServiceObservationExternalAccountConfig struct {
@@ -262,6 +266,7 @@ func (o *serviceObsListener) natsContext() *natsContext {
 		TLSCA:       o.config.TLSCA,
 		TLSCert:     o.config.TLSCert,
 		TLSKey:      o.config.TLSKey,
+		NatsOpts:    o.config.NatsOpts,
 	}
 
 	// legacy Credentials field
diff --git a/surveyor/surveyor.go b/surveyor/surveyor.go
index b536c27..b1d5c7f 100644
--- a/surveyor/surveyor.go
+++ b/surveyor/surveyor.go
@@ -197,8 +197,7 @@ func (s *Surveyor) createStatszCollector() error {
 	}
 
 	s.statzC = NewStatzCollector(s.sysAcctPC.nc, s.logger, s.opts.ExpectedServers, s.opts.ServerResponseWait, s.opts.PollTimeout, s.opts.Accounts, s.opts.ConstLabels)
-	s.promRegistry.MustRegister(s.statzC)
-	return nil
+	return s.promRegistry.Register(s.statzC)
 }
 
 // generates the TLS config for https