feat: Add simple healthcheck

cli/serve.go

@@ -11,9 +11,15 @@ import (
 	"gopkg.in/yaml.v2"
 )
 
+type healthCheck struct {
+	Attemps  uint
+	Endpoint string
+}
+
 type target struct {
-	Address string
-	Weight  int
+	Address     string
+	Weight      int
+	HealthCheck *healthCheck
 }
 
 type config struct {
@@ -62,7 +68,17 @@ func (a *Application) serve(c *cli.Context) error {
 		if err != nil {
 			return err
 		}
-		proxy.Add(targetURL, item.Weight)
+
+		var itemHealthCheck *rp.HealthCheck = nil
+		if item.HealthCheck != nil {
+			healthCheckEndpoint, err := url.Parse(item.HealthCheck.Endpoint)
+			if err != nil {
+				return err
+			}
+
+			itemHealthCheck = rp.NewHealthCheck(healthCheckEndpoint, item.HealthCheck.Attemps)
+		}
+		proxy.Add(targetURL, item.Weight, itemHealthCheck)
 	}
 
 	log.Printf("Starting reverse-proxy server on %s ...", cfg.ListenAddr)

reverse-proxy/health-check.go

@@ -0,0 +1,57 @@
+package rp
+
+import (
+	"net"
+	"net/http"
+	"net/url"
+	"time"
+)
+
+type HealthCheck struct {
+	endpoint  *url.URL
+	attemps   uint
+	client    *http.Client
+	sleepTime time.Duration
+}
+
+func NewHealthCheck(endpoint *url.URL, attemps uint) *HealthCheck {
+	return &HealthCheck{
+		endpoint:  endpoint,
+		attemps:   attemps,
+		sleepTime: 1 * time.Second,
+		client: &http.Client{
+			Transport: &http.Transport{
+				Proxy: http.ProxyFromEnvironment,
+				Dial: (&net.Dialer{
+					Timeout:   10 * time.Second,
+					KeepAlive: 0,
+				}).Dial,
+				TLSHandshakeTimeout:   5 * time.Second,
+				ExpectContinueTimeout: 1 * time.Second,
+				MaxIdleConnsPerHost:   1,
+				DisableCompression:    true,
+				DisableKeepAlives:     true,
+				ResponseHeaderTimeout: 5 * time.Second,
+			},
+			CheckRedirect: func(req *http.Request, via []*http.Request) error {
+				return http.ErrUseLastResponse
+			},
+			Timeout: 10 * time.Second,
+		},
+	}
+}
+
+func (h *HealthCheck) check() bool {
+	req, err := http.NewRequest("GET", h.endpoint.String(), nil)
+	if err != nil {
+		return false
+	}
+	for i := uint(0); i < h.attemps; i++ {
+		resp, err := h.client.Do(req)
+		if err == nil && resp.StatusCode >= 200 && resp.StatusCode < 300 {
+			return true
+		}
+		time.Sleep(h.sleepTime)
+	}
+	return false
+}

reverse-proxy/proxy-connection.go

@@ -0,0 +1,51 @@
+package rp
+
+import (
+	"net"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"time"
+)
+
+type proxyConnection struct {
+	reverseProxy *httputil.ReverseProxy
+	target       *url.URL
+	weight       int
+	healthCheck  *HealthCheck
+	isHealthy    bool
+}
+
+func newProxyConnection(target *url.URL, weight int, healthCheck *HealthCheck) *proxyConnection {
+	return &proxyConnection{
+		reverseProxy: httputil.NewSingleHostReverseProxy(target),
+		target:       target,
+		weight:       weight,
+		isHealthy:    true,
+		healthCheck:  healthCheck,
+	}
+}
+
+func (p *proxyConnection) String() string {
+	return p.target.String()
+}
+
+func (p *proxyConnection) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	r.URL.Host = p.target.Host
+	r.URL.Scheme = p.target.Scheme
+	r.Host = p.target.Host
+	r.Header.Set("X-Forwarded-Host", r.Header.Get("Host"))
+	p.reverseProxy.Transport = &http.Transport{
+		DialContext: (&net.Dialer{
+			Timeout:   10 * time.Second,
+			KeepAlive: 10 * time.Second,
+			DualStack: true,
+		}).DialContext,
+		MaxIdleConns:          10,
+		IdleConnTimeout:       30 * time.Second,
+		TLSHandshakeTimeout:   5 * time.Second,
+		ExpectContinueTimeout: 1 * time.Second,
+	}
+
+	p.reverseProxy.ServeHTTP(w, r)
+}

reverse-proxy/reverse-proxy.go

@@ -6,52 +6,13 @@ import (
 	"log"
 	"net"
 	"net/http"
-	"net/http/httputil"
 	"net/url"
 	"strings"
 	"time"
 
 	"github.com/bogdanovich/dns_resolver"
 )
 
-type proxyConnection struct {
-	reverseProxy *httputil.ReverseProxy
-	target       *url.URL
-	weight       int
-}
-
-func (p *proxyConnection) String() string {
-	return p.target.String()
-}
-
-func (p *proxyConnection) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	r.URL.Host = p.target.Host
-	r.URL.Scheme = p.target.Scheme
-	r.Host = p.target.Host
-	r.Header.Set("X-Forwarded-Host", r.Header.Get("Host"))
-	p.reverseProxy.Transport = &http.Transport{
-		DialContext: (&net.Dialer{
-			Timeout:   10 * time.Second,
-			KeepAlive: 10 * time.Second,
-			DualStack: true,
-		}).DialContext,
-		MaxIdleConns:          10,
-		IdleConnTimeout:       30 * time.Second,
-		TLSHandshakeTimeout:   5 * time.Second,
-		ExpectContinueTimeout: 1 * time.Second,
-	}
-
-	p.reverseProxy.ServeHTTP(w, r)
-}
-
-func newProxyConnection(target *url.URL, weight int) *proxyConnection {
-	return &proxyConnection{
-		reverseProxy: httputil.NewSingleHostReverseProxy(target),
-		target:       target,
-		weight:       weight,
-	}
-}
-
 type ReverseProxy struct {
 	rr  *roundRobin
 	log bool
@@ -95,8 +56,8 @@ func (rp *ReverseProxy) Log(mode bool) {
 	rp.log = mode
 }
 
-func (rp *ReverseProxy) Add(target *url.URL, weight int) {
-	rp.rr.Add(newProxyConnection(target, weight))
+func (rp *ReverseProxy) Add(target *url.URL, weight int, healthCheck *HealthCheck) {
+	rp.rr.Add(newProxyConnection(target, weight, healthCheck))
 }
 
 type statusWriter struct {