Add async metrics fetch for multiple piholes

parent 8d5586558c author Galorhallen <andrea.ponte1987@gmail.com> 1640558190 +0100 committer Galorhallen <andrea.ponte1987@gmail.com> 1640821760 +0100 Add test for multiple pihole Add async mode for multiple piholes Fixed GitHub Actions go versions Add test for multiple pihole Cleanup
2024-11-23 11:26:30 +01:00 · 2021-12-26 23:36:30 +01:00 · 2021-12-26 23:36:30 +01:00 · 7009c705bb
commit 7009c705bb
parent 8d5586558c
7 changed files with 283 additions and 59 deletions
--- a/config/configuration.go
+++ b/config/configuration.go
@ -2,6 +2,7 @@ package config

 import (
 	"context"
+	"errors"
 	"fmt"
 	"log"
 	"reflect"
@ -46,7 +47,7 @@ func getDefaultEnvConfig() *EnvConfig {
 }

 // Load method loads the configuration by using both flag or environment variables.
-func Load() (*EnvConfig, []Config) {
+func Load() (*EnvConfig, []Config, error) {
 	loaders := []backend.Backend{
 		env.NewBackend(),
 		flags.NewBackend(),
@ -62,7 +63,11 @@ func Load() (*EnvConfig, []Config) {

 	cfg.show()

-	return cfg, cfg.Split()
+	if clientsConfig, err := cfg.Split(); err != nil {
+		return cfg, nil, err
+	} else {
+		return cfg, clientsConfig, nil
+	}
 }

 func (c *Config) String() string {
@ -87,50 +92,64 @@ func (c Config) Validate() error {
 	return nil
 }

-func (c EnvConfig) Split() []Config {
-	result := make([]Config, 0, len(c.PIHoleHostname))
+func (c EnvConfig) Split() ([]Config, error) {
+	hostsCount := len(c.PIHoleHostname)
+	result := make([]Config, 0, hostsCount)

 	for i, hostname := range c.PIHoleHostname {
 		config := Config{
 			PIHoleHostname: strings.TrimSpace(hostname),
-			PIHolePort:     c.PIHolePort[i],
 		}

-		if hasData, data := extractConfig(c.PIHoleProtocol, i); hasData {
+		if len(c.PIHolePort) == 1 {
+			config.PIHolePort = c.PIHolePort[0]
+		} else if len(c.PIHolePort) == hostsCount {
+			config.PIHolePort = c.PIHolePort[i]
+		} else if len(c.PIHolePort) != 0 {
+			return nil, errors.New("Wrong number of ports. Port can be empty to use default, one value to use for all hosts, or match the number of hosts")
+		}
+
+		if hasData, data, isValid := extractStringConfig(c.PIHoleProtocol, i, hostsCount); hasData {
 			config.PIHoleProtocol = data
+		} else if !isValid {
+			return nil, errors.New("Wrong number of PIHoleProtocol. PIHoleProtocol can be empty to use default, one value to use for all hosts, or match the number of hosts")
 		}

-		if hasData, data := extractConfig(c.PIHoleApiToken, i); hasData {
+		if hasData, data, isValid := extractStringConfig(c.PIHoleApiToken, i, hostsCount); hasData {
 			config.PIHoleApiToken = data
+		} else if !isValid {
+			return nil, errors.New(fmt.Sprintf("Wrong number of PIHoleApiToken %d (Hosts: %d). PIHoleApiToken can be empty to use default, one value to use for all hosts, or match the number of hosts", len(c.PIHoleApiToken), hostsCount))
 		}

-		if hasData, data := extractConfig(c.PIHoleApiToken, i); hasData {
-			config.PIHoleApiToken = data
-		}
-
-		if hasData, data := extractConfig(c.PIHolePassword, i); hasData {
+		if hasData, data, isValid := extractStringConfig(c.PIHolePassword, i, hostsCount); hasData {
 			config.PIHolePassword = data
+		} else if !isValid {
+			return nil, errors.New("Wrong number of PIHolePassword. PIHolePassword can be empty to use default, one value to use for all hosts, or match the number of hosts")
 		}

 		result = append(result, config)
 	}
-	return result
+
+	return result, nil
 }

-func extractConfig(data []string, idx int) (bool, string) {
+func extractStringConfig(data []string, idx int, hostsCount int) (bool, string, bool) {
 	if len(data) == 1 {
 		v := strings.TrimSpace(data[0])
 		if v != "" {
-			return true, v
+			return true, v, true
 		}
-	} else if len(data) > 1 {
+	} else if len(data) == hostsCount {
 		v := strings.TrimSpace(data[idx])
 		if v != "" {
-			return true, v
+			return true, v, true
 		}
+	} else if len(data) != 0 { //Host count missmatch
+		return false, "", false
 	}

-	return false, ""
+	// Empty
+	return false, "", true
 }

 func (c Config) hostnameURL() string {
--- a/config/configuration_test.go
+++ b/config/configuration_test.go
@ -0,0 +1,154 @@
+package config
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+const (
+	PIHOLE_HOSTNAME  = "PIHOLE_HOSTNAME"
+	PIHOLE_PORT      = "PIHOLE_PORT"
+	PIHOLE_API_TOKEN = "PIHOLE_API_TOKEN"
+	PIHOLE_PROTOCOL  = "PIHOLE_PROTOCOL"
+)
+
+type EnvInitiazlier func(*testing.T)
+
+type TestCase struct {
+	Name        string
+	Initializer EnvInitiazlier
+}
+
+func TestSplitDefault(t *testing.T) {
+	assert := assert.New(t)
+
+	env := getDefaultEnvConfig()
+
+	clientConfigs, err := env.Split()
+	assert.NoError(err)
+
+	clientConfig := clientConfigs[0]
+	assert.Equal("127.0.0.1", clientConfig.PIHoleHostname)
+	assert.Equal("http", clientConfig.PIHoleProtocol)
+	assert.Equal(uint16(80), clientConfig.PIHolePort)
+	assert.Empty(clientConfig.PIHoleApiToken)
+	assert.Empty(clientConfig.PIHolePassword)
+}
+
+func TestSplitMultipleHostWithSameConfig(t *testing.T) {
+	assert := assert.New(t)
+
+	env := getDefaultEnvConfig()
+	env.PIHoleHostname = []string{"127.0.0.1", "127.0.0.2", "127.0.0.3"}
+	env.PIHoleApiToken = []string{"api-token"}
+	env.PIHolePort = []uint16{8080}
+
+	clientConfigs, err := env.Split()
+	assert.NoError(err)
+	assert.Len(clientConfigs, 3)
+
+	testCases := []struct {
+		Host     string
+		Port     uint16
+		Protocol string
+	}{
+		{
+			Host:     "127.0.0.1",
+			Port:     8080,
+			Protocol: "http",
+		},
+		{
+			Host:     "127.0.0.2",
+			Port:     8080,
+			Protocol: "http",
+		},
+		{
+			Host:     "127.0.0.3",
+			Port:     8080,
+			Protocol: "http",
+		},
+	}
+
+	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("Test %s", tc.Host), func(t *testing.T) {
+			clientConfig := clientConfigs[i]
+
+			assert.Equal(tc.Host, clientConfig.PIHoleHostname)
+			assert.Equal(tc.Protocol, clientConfig.PIHoleProtocol)
+			assert.Equal(tc.Port, clientConfig.PIHolePort)
+			assert.Equal("api-token", clientConfig.PIHoleApiToken)
+			assert.Empty(clientConfig.PIHolePassword)
+		})
+	}
+}
+
+func TestSplitMultipleHostWithMultipleConfigs(t *testing.T) {
+	assert := assert.New(t)
+
+	env := getDefaultEnvConfig()
+	env.PIHoleHostname = []string{"127.0.0.1", "127.0.0.2", "127.0.0.3"}
+	env.PIHoleApiToken = []string{"api-token1", "", "api-token3"}
+	env.PIHolePassword = []string{"", "password2", ""}
+	env.PIHolePort = []uint16{8081, 8082, 8083}
+
+	clientConfigs, err := env.Split()
+	assert.NoError(err)
+	assert.Len(clientConfigs, 3)
+
+	testCases := []struct {
+		Host     string
+		Port     uint16
+		Protocol string
+		ApiToken string
+		Password string
+	}{
+		{
+			Host:     "127.0.0.1",
+			Port:     8081,
+			Protocol: "http",
+			ApiToken: "api-token1",
+			Password: "",
+		},
+		{
+			Host:     "127.0.0.2",
+			Port:     8082,
+			Protocol: "http",
+			ApiToken: "",
+			Password: "password2",
+		},
+		{
+			Host:     "127.0.0.3",
+			Port:     8083,
+			Protocol: "http",
+			ApiToken: "api-token3",
+			Password: "",
+		},
+	}
+
+	for i, tc := range testCases {
+		t.Run(fmt.Sprintf("Test %s", tc.Host), func(t *testing.T) {
+			clientConfig := clientConfigs[i]
+
+			assert.Equal(tc.Host, clientConfig.PIHoleHostname)
+			assert.Equal(tc.Protocol, clientConfig.PIHoleProtocol)
+			assert.Equal(tc.Port, clientConfig.PIHolePort)
+			assert.Equal(tc.ApiToken, clientConfig.PIHoleApiToken)
+			assert.Equal(tc.Password, clientConfig.PIHolePassword)
+		})
+	}
+}
+
+func TestWrongParams(t *testing.T) {
+	assert := assert.New(t)
+
+	env := getDefaultEnvConfig()
+	env.PIHoleHostname = []string{"127.0.0.1", "127.0.0.2", "127.0.0.3"}
+	env.PIHoleApiToken = []string{"api-token1", "api-token2"}
+	env.PIHolePort = []uint16{808}
+
+	clientConfigs, err := env.Split()
+	assert.Errorf(err, "Wrong number of PIHoleApiToken. PIHoleApiToken can be empty to use default, one value to use for all hosts, or match the number of hosts")
+	assert.Nil(clientConfigs)
+}
--- a/go.mod
+++ b/go.mod
@ -5,6 +5,8 @@ go 1.15
 require (
 	github.com/heetch/confita v0.10.0
 	github.com/prometheus/client_golang v1.11.0
+	github.com/stretchr/testify v1.7.0
 	github.com/xonvanetta/shutdown v0.0.3
 	golang.org/x/net v0.0.0-20200625001655-4c5254603344
+	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
 )
--- a/go.sum
+++ b/go.sum
@ -123,6 +123,8 @@ github.com/konsorten/go-windows-terminal-sequences v1.0.3/go.mod h1:T0+1ngSBFLxv
 github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc=
 github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
+github.com/kr/pretty v0.2.1 h1:Fmg33tUaq4/8ym9TJN1x7sLJnHVwhP33CNkpYV/7rwI=
+github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
@ -195,12 +197,15 @@ github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6Mwd
 github.com/sirupsen/logrus v1.6.0/go.mod h1:7uNnSEd1DgxDLC74fIahvMZmmYsHGZGEOFrfsX/uA88=
 github.com/soheilhy/cmux v0.1.4/go.mod h1:IM3LyeVVIOuxMH7sFAkER9+bJ4dT7Ms6E4xg4kGIyLM=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.1.1 h1:2vfRuCMp5sSVIDSqO8oNnWJq7mPa6KVP3iPIwFBuy8A=
 github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/stretchr/testify v1.6.0 h1:jlIyCplCJFULU/01vCkhKuTyc3OorI3bJFuw6obfgho=
 github.com/stretchr/testify v1.6.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/tmc/grpc-websocket-proxy v0.0.0-20190109142713-0ad062ec5ee5/go.mod h1:ncp9v5uamzpCO7NfCPTXjqaC+bZgJeR0sMTm6dMHP7U=
 github.com/tv42/httpunix v0.0.0-20150427012821-b75d8614f926/go.mod h1:9ESjWnEqriFuLhtthL60Sar/7RFoluCcXsuvEwTV5KM=
 github.com/ugorji/go v1.1.4/go.mod h1:uQMGLiO92mf5W77hV/PUCpI3pbzQx3CRekS0kk+RGrc=
@ -241,6 +246,8 @@ golang.org/x/sync v0.0.0-20190227155943-e225da77a7e6/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201207232520-09787c993a3a/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ=
+golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180823144017-11551d06cbcc/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
@ -292,6 +299,8 @@ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
 gopkg.in/resty.v1 v1.12.0/go.mod h1:mDo4pnntr5jdWRML875a/NmxYqAlA73dVijT2AXvQQo=
 gopkg.in/square/go-jose.v2 v2.3.1/go.mod h1:M9dMgbHiYLoDGQrXy7OpJDJWiKiU//h+vD76mk0e1AI=
 gopkg.in/yaml.v2 v2.0.0-20170812160011-eb3733d160e7/go.mod h1:JAlM8MvJe8wmxCU4Bli9HhUf9+ttbYbLASfIpnQbh74=
--- a/internal/pihole/client.go
+++ b/internal/pihole/client.go
@ -16,12 +16,38 @@ import (
 	"github.com/eko/pihole-exporter/internal/metrics"
 )

+type ClientStatus byte
+
+const (
+	MetricsCollectionInProgress ClientStatus = iota
+	MetricsCollectionSuccess
+	MetricsCollectionError
+	MetricsCollectionTimeout
+)
+
+func (status ClientStatus) String() string {
+	return []string{"MetricsCollectionInProgress", "MetricsCollectionSuccess", "MetricsCollectionError", "MetricsCollectionTimeout"}[status]
+}
+
+type ClientChannel struct {
+	Status ClientStatus
+	Err    error
+}
+
+func (c *ClientChannel) String() string {
+	if c.Err != nil {
+		return fmt.Sprintf("ClientChannel<Status: %s, Err: '%s'>", c.Status, c.Err.Error())
+	} else {
+		return fmt.Sprintf("ClientChannel<Status: %s, Err: <nil>>", c.Status)
+	}
+}
+
 // Client struct is a PI-Hole client to request an instance of a PI-Hole ad blocker.
 type Client struct {
 	httpClient http.Client
 	interval   time.Duration
 	config     *config.Config
-	MetricRetrieved chan bool
+	Status     chan *ClientChannel
 }

 // NewClient method initializes a new PI-Hole client.
@ -41,6 +67,7 @@ func NewClient(config *config.Config) *Client {
 				return http.ErrUseLastResponse
 			},
 		},
+		Status: make(chan *ClientChannel, 1),
 	}
 }

@ -48,32 +75,24 @@ func (c *Client) String() string {
 	return c.config.PIHoleHostname
 }

-/*
-// Metrics scrapes pihole and sets them
-func (c *Client) Metrics() http.HandlerFunc {
-	return func(writer http.ResponseWriter, request *http.Request) {
-		stats, err := c.getStatistics()
-		if err != nil {
-			writer.WriteHeader(http.StatusBadRequest)
-			_, _ = writer.Write([]byte(err.Error()))
-			return
-		}
+func (c *Client) CollectMetricsAsync(writer http.ResponseWriter, request *http.Request) {
+	log.Printf("Collecting from %s", c.config.PIHoleHostname)
+	if stats, err := c.getStatistics(); err == nil {
 		c.setMetrics(stats)
-
-		log.Printf("New tick of statistics: %s", stats.ToString())
-		promhttp.Handler().ServeHTTP(writer, request)
+		c.Status <- &ClientChannel{Status: MetricsCollectionSuccess, Err: nil}
+		log.Printf("New tick of statistics from %s: %s", c.config.PIHoleHostname, stats)
+	} else {
+		c.Status <- &ClientChannel{Status: MetricsCollectionError, Err: err}
 	}
-}*/
+}

 func (c *Client) CollectMetrics(writer http.ResponseWriter, request *http.Request) error {
-
 	stats, err := c.getStatistics()
 	if err != nil {
 		return err
 	}
 	c.setMetrics(stats)
-
-	log.Printf("New tick of statistics from %s: %s", c, stats)
+	log.Printf("New tick of statistics from %s: %s", c.config.PIHoleHostname, stats)
 	return nil
 }

--- a/internal/server/server.go
+++ b/internal/server/server.go
@ -22,32 +22,28 @@ type Server struct {
 // the different routes that will be used by Prometheus (metrics) or for monitoring (readiness, liveness).
 func NewServer(port uint16, clients []*pihole.Client) *Server {
 	mux := http.NewServeMux()
-	httpServer := &http.Server{Addr: ":" + strconv.Itoa(int(port)), Handler: mux}
+	httpServer := &http.Server{
+		Addr:    ":" + strconv.Itoa(int(port)),
+		Handler: mux,
+	}

 	s := &Server{
 		httpServer: httpServer,
 	}

-	mux.HandleFunc("/metrics",
-		func(writer http.ResponseWriter, request *http.Request) {
-			errors := make([]string, 0)
+	mux.HandleFunc("/metrics", func(writer http.ResponseWriter, request *http.Request) {
+		log.Printf("request.Header: %v\n", request.Header)
 		
 		for _, client := range clients {
-				if err := client.CollectMetrics(writer, request); err != nil {
-					errors = append(errors, err.Error())
-					fmt.Printf("Error %s\n", err)
-				}
+			go client.CollectMetricsAsync(writer, request)
 		}

-			if len(errors) == len(clients) {
-				writer.WriteHeader(http.StatusBadRequest)
-				body := strings.Join(errors, "\n")
-				_, _ = writer.Write([]byte(body))
+		for _, client := range clients {
+			log.Printf("Received %s from %s\n", <-client.Status, client.GetHostname())
 		}

 		promhttp.Handler().ServeHTTP(writer, request)
-		},
-	)
+	})

 	mux.Handle("/readiness", s.readinessHandler())
 	mux.Handle("/liveness", s.livenessHandler())
@ -73,6 +69,27 @@ func (s *Server) Stop() {
 	s.httpServer.Shutdown(ctx)
 }

+func (s *Server) handleMetrics(clients []*pihole.Client) http.HandlerFunc {
+	return func(writer http.ResponseWriter, request *http.Request) {
+		errors := make([]string, 0)
+
+		for _, client := range clients {
+			if err := client.CollectMetrics(writer, request); err != nil {
+				errors = append(errors, err.Error())
+				fmt.Printf("Error %s\n", err)
+			}
+		}
+
+		if len(errors) == len(clients) {
+			writer.WriteHeader(http.StatusBadRequest)
+			body := strings.Join(errors, "\n")
+			_, _ = writer.Write([]byte(body))
+		}
+
+		promhttp.Handler().ServeHTTP(writer, request)
+	}
+}
+
 func (s *Server) readinessHandler() http.HandlerFunc {
 	return func(w http.ResponseWriter, req *http.Request) {
 		if s.isReady() {
--- a/main.go
+++ b/main.go
@ -2,6 +2,7 @@ package main

 import (
 	"fmt"
+	"log"

 	"github.com/eko/pihole-exporter/config"
 	"github.com/eko/pihole-exporter/internal/metrics"
@ -11,7 +12,10 @@ import (
 )

 func main() {
-	envConf, clientConfigs := config.Load()
+	envConf, clientConfigs, err := config.Load()
+	if err != nil {
+		log.Fatal(err.Error())
+	}

 	metrics.Init()