From 0a15850dc8f989706a9747cd8c79c9e7bf0d1552 Mon Sep 17 00:00:00 2001 From: davidnewhall2 Date: Thu, 28 Nov 2019 04:14:05 -0800 Subject: [PATCH] cosolidate code --- poller/influx.go | 20 +++++++++-- poller/prometheus.go | 17 +++++++++ poller/start.go | 83 +++++++++++++++----------------------------- poller/unifi.go | 41 +++++++++------------- 4 files changed, 79 insertions(+), 82 deletions(-) diff --git a/poller/influx.go b/poller/influx.go index d3173c4c..d428bdb8 100644 --- a/poller/influx.go +++ b/poller/influx.go @@ -1,13 +1,29 @@ package poller import ( + "crypto/tls" "fmt" "github.com/davidnewhall/unifi-poller/influxunifi" "github.com/davidnewhall/unifi-poller/metrics" - client "github.com/influxdata/influxdb1-client/v2" + influx "github.com/influxdata/influxdb1-client/v2" ) +// GetInfluxDB returns an InfluxDB interface. +func (u *UnifiPoller) GetInfluxDB() (err error) { + u.Influx, err = influx.NewHTTPClient(influx.HTTPConfig{ + Addr: u.Config.InfluxURL, + Username: u.Config.InfluxUser, + Password: u.Config.InfluxPass, + TLSConfig: &tls.Config{InsecureSkipVerify: u.Config.InfxBadSSL}, + }) + if err != nil { + return fmt.Errorf("influxdb: %v", err) + } + u.Logf("Logging Measurements to InfluxDB at %s as user %s", u.Config.InfluxURL, u.Config.InfluxUser) + return nil +} + // CollectAndProcess collects measurements and then reports them to InfluxDB // Can be called once or in a ticker loop. This function and all the ones below // handle their own logging. An error is returned so the calling function may @@ -31,7 +47,7 @@ func (u *UnifiPoller) ReportMetrics(metrics *metrics.Metrics) error { m := &influxunifi.Metrics{Metrics: metrics} // Make a new Influx Points Batcher. var err error - m.BatchPoints, err = client.NewBatchPoints(client.BatchPointsConfig{Database: u.Config.InfluxDB}) + m.BatchPoints, err = influx.NewBatchPoints(influx.BatchPointsConfig{Database: u.Config.InfluxDB}) if err != nil { return fmt.Errorf("influx.NewBatchPoints: %v", err) } diff --git a/poller/prometheus.go b/poller/prometheus.go index b53f7c62..030a51b1 100644 --- a/poller/prometheus.go +++ b/poller/prometheus.go @@ -2,12 +2,29 @@ package poller import ( "fmt" + "net/http" + "strings" "time" "github.com/davidnewhall/unifi-poller/metrics" "github.com/davidnewhall/unifi-poller/promunifi" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promhttp" ) +// RunPrometheus starts the web server and registers the collector. +func (u *UnifiPoller) RunPrometheus() error { + u.Logf("Exporting Measurements at https://%s/metrics for Prometheus", u.Config.HTTPListen) + http.Handle("/metrics", promhttp.Handler()) + prometheus.MustRegister(promunifi.NewUnifiCollector(promunifi.UnifiCollectorCnfg{ + Namespace: strings.Replace(u.Config.Namespace, "-", "", -1), + CollectFn: u.ExportMetrics, + LoggingFn: u.LogExportReport, + ReportErrors: true, // XXX: Does this need to be configurable? + })) + return http.ListenAndServe(u.Config.HTTPListen, nil) +} + // ExportMetrics updates the internal metrics provided via // HTTP at /metrics for prometheus collection. // This is run by Prometheus as CollectFn. diff --git a/poller/start.go b/poller/start.go index a4f1ec34..11d70228 100644 --- a/poller/start.go +++ b/poller/start.go @@ -2,20 +2,13 @@ package poller import ( - "crypto/tls" "fmt" "log" - "net/http" "os" "strings" "time" - "github.com/davidnewhall/unifi-poller/promunifi" - influx "github.com/influxdata/influxdb1-client/v2" - "github.com/prometheus/client_golang/prometheus" - "github.com/prometheus/client_golang/prometheus/promhttp" "github.com/spf13/pflag" - "golift.io/unifi" ) // New returns a new poller struct preloaded with default values. @@ -40,7 +33,7 @@ func New() *UnifiPoller { // Start begins the application from a CLI. // Parses cli flags, parses config file, parses env vars, sets up logging, then: -// - dumps a json payload OR - authenticates unifi controller and executes Run(). +// - dumps a json payload OR - executes Run(). func (u *UnifiPoller) Start() error { log.SetFlags(log.LstdFlags) u.Flag.Parse(os.Args[1:]) @@ -75,10 +68,6 @@ func (u *UnifiPoller) Start() error { } log.Printf("[INFO] UniFi Poller v%v Starting Up! PID: %d", Version, os.Getpid()) - if err := u.GetUnifi(); err != nil { - return err - } - return u.Run() } @@ -102,6 +91,9 @@ func (f *Flag) Parse(args []string) { // 2. Run the collector one time and report the metrics to influxdb. (lambda) // 3. Start a web server and wait for Prometheus to poll the application for metrics. func (u *UnifiPoller) Run() error { + if err := u.GetUnifi(); err != nil { + return err + } u.Logf("Polling UniFi Controller at %s v%s as user %s. Sites: %v", u.Config.UnifiBase, u.Unifi.ServerVersion, u.Config.UnifiUser, u.Config.Sites) @@ -110,62 +102,43 @@ func (u *UnifiPoller) Run() error { if err := u.GetInfluxDB(); err != nil { return err } - u.Logf("Logging Measurements to InfluxDB at %s as user %s", u.Config.InfluxURL, u.Config.InfluxUser) - u.Config.Mode = "influx poller" return u.PollController() case "influxlambda", "lambdainflux", "lambda_influx", "influx_lambda": if err := u.GetInfluxDB(); err != nil { return err } - u.Logf("Logging Measurements to InfluxDB at %s as user %s one time (lambda mode)", - u.Config.InfluxURL, u.Config.InfluxUser) u.LastCheck = time.Now() return u.CollectAndProcess() case "prometheus", "exporter": - u.Logf("Exporting Measurements at https://%s/metrics for Prometheus", u.Config.HTTPListen) - http.Handle("/metrics", promhttp.Handler()) - prometheus.MustRegister(promunifi.NewUnifiCollector(promunifi.UnifiCollectorCnfg{ - Namespace: strings.Replace(u.Config.Namespace, "-", "", -1), - CollectFn: u.ExportMetrics, - LoggingFn: u.LogExportReport, - ReportErrors: true, // XXX: Does this need to be configurable? - })) - return http.ListenAndServe(u.Config.HTTPListen, nil) + return u.RunPrometheus() } } -// GetInfluxDB returns an InfluxDB interface. -func (u *UnifiPoller) GetInfluxDB() (err error) { - u.Influx, err = influx.NewHTTPClient(influx.HTTPConfig{ - Addr: u.Config.InfluxURL, - Username: u.Config.InfluxUser, - Password: u.Config.InfluxPass, - TLSConfig: &tls.Config{InsecureSkipVerify: u.Config.InfxBadSSL}, - }) - if err != nil { - return fmt.Errorf("influxdb: %v", err) +// PollController runs forever, polling UniFi and pushing to InfluxDB +// This is started by Run() after everything checks out. +func (u *UnifiPoller) PollController() error { + interval := u.Config.Interval.Round(time.Second) + log.Printf("[INFO] Everything checks out! Poller started, interval: %v", interval) + ticker := time.NewTicker(interval) + defer ticker.Stop() + for u.LastCheck = range ticker.C { + var err error + if u.Config.ReAuth { + u.LogDebugf("Re-authenticating to UniFi Controller") + // Some users need to re-auth every interval because the cookie times out. + if err = u.Unifi.Login(); err != nil { + u.LogError(err, "re-authenticating") + } + } + if err == nil { + // Only run this if the authentication procedure didn't return error. + _ = u.CollectAndProcess() + } + if u.errorCount > 0 { + return fmt.Errorf("too many errors, stopping poller") + } } - return nil } - -// GetUnifi returns a UniFi controller interface. -func (u *UnifiPoller) GetUnifi() (err error) { - // Create an authenticated session to the Unifi Controller. - u.Unifi, err = unifi.NewUnifi(&unifi.Config{ - User: u.Config.UnifiUser, - Pass: u.Config.UnifiPass, - URL: u.Config.UnifiBase, - VerifySSL: u.Config.VerifySSL, - ErrorLog: u.LogErrorf, // Log all errors. - DebugLog: u.LogDebugf, // Log debug messages. - }) - if err != nil { - return fmt.Errorf("unifi controller: %v", err) - } - u.LogDebugf("Authenticated with controller successfully") - - return u.CheckSites() -} diff --git a/poller/unifi.go b/poller/unifi.go index f5bcd9ba..54135923 100644 --- a/poller/unifi.go +++ b/poller/unifi.go @@ -2,7 +2,6 @@ package poller import ( "fmt" - "log" "strings" "time" @@ -10,31 +9,23 @@ import ( "golift.io/unifi" ) -// PollController runs forever, polling UniFi -// and pushing to influx OR exporting for prometheus. -// This is started by Run() after everything checks out. -func (u *UnifiPoller) PollController() error { - interval := u.Config.Interval.Round(time.Second) - log.Printf("[INFO] Everything checks out! Poller started in %v mode, interval: %v", u.Config.Mode, interval) - ticker := time.NewTicker(interval) - for u.LastCheck = range ticker.C { - var err error - if u.Config.ReAuth { - u.LogDebugf("Re-authenticating to UniFi Controller") - // Some users need to re-auth every interval because the cookie times out. - if err = u.Unifi.Login(); err != nil { - u.LogError(err, "re-authenticating") - } - } - if err == nil { - // Only run this if the authentication procedure didn't return error. - _ = u.CollectAndProcess() - } - if u.errorCount > 0 { - return fmt.Errorf("too many errors, stopping poller") - } +// GetUnifi returns a UniFi controller interface. +func (u *UnifiPoller) GetUnifi() (err error) { + // Create an authenticated session to the Unifi Controller. + u.Unifi, err = unifi.NewUnifi(&unifi.Config{ + User: u.Config.UnifiUser, + Pass: u.Config.UnifiPass, + URL: u.Config.UnifiBase, + VerifySSL: u.Config.VerifySSL, + ErrorLog: u.LogErrorf, // Log all errors. + DebugLog: u.LogDebugf, // Log debug messages. + }) + if err != nil { + return fmt.Errorf("unifi controller: %v", err) } - return nil + u.LogDebugf("Authenticated with controller successfully") + + return u.CheckSites() } // CheckSites makes sure the list of provided sites exists on the controller.