From 0ab0bdf818c43e5fab2e11e4d90d730a6afcc042 Mon Sep 17 00:00:00 2001
From: Julien Salleyron <juliens@php.net>
Date: Sat, 26 Nov 2016 19:48:49 +0100
Subject: [PATCH] feat: HealthCheck

---
 healthcheck/healthcheck.go | 94 ++++++++++++++++++++++++++++++++++++++
 server.go                  | 11 +++++
 types/types.go             |  6 +++
 3 files changed, 111 insertions(+)
 create mode 100644 healthcheck/healthcheck.go

diff --git a/healthcheck/healthcheck.go b/healthcheck/healthcheck.go
new file mode 100644
index 000000000..66c31c9ce
--- /dev/null
+++ b/healthcheck/healthcheck.go
@@ -0,0 +1,94 @@
+package healthcheck
+
+import (
+	"github.com/containous/traefik/log"
+	"github.com/vulcand/oxy/roundrobin"
+	"net/http"
+	"net/url"
+	"sync"
+	"time"
+)
+
+var singleton *HealthCheck
+var once sync.Once
+
+// GetHealthCheck Get HealtchCheck Singleton
+func GetHealthCheck() *HealthCheck {
+	once.Do(func() {
+		singleton = newHealthCheck()
+		singleton.execute()
+	})
+	return singleton
+}
+
+// BackendHealthCheck HealthCheck configuration for a backend
+type BackendHealthCheck struct {
+	URL          string
+	DisabledURLs []*url.URL
+	lb           loadBalancer
+}
+
+var launch = false
+
+//HealthCheck struct
+type HealthCheck struct {
+	Backends map[string]*BackendHealthCheck
+}
+
+type loadBalancer interface {
+	RemoveServer(u *url.URL) error
+	UpsertServer(u *url.URL, options ...roundrobin.ServerOption) error
+	Servers() []*url.URL
+}
+
+func newHealthCheck() *HealthCheck {
+	return &HealthCheck{make(map[string]*BackendHealthCheck)}
+}
+
+// NewBackendHealthCheck Instantiate a new BackendHealthCheck
+func NewBackendHealthCheck(URL string, lb loadBalancer) *BackendHealthCheck {
+	return &BackendHealthCheck{URL, nil, lb}
+}
+
+//SetBackends set backends configuration
+func (hc *HealthCheck) SetBackendsConfiguration(backends map[string]*BackendHealthCheck) {
+	hc.Backends = backends
+}
+
+func (hc *HealthCheck) execute() {
+	ticker := time.NewTicker(time.Second * 30)
+	go func() {
+		for t := range ticker.C {
+			for backendID, backend := range hc.Backends {
+				log.Debugf("Refreshing Healthcheck %s for backend %s ", t.String(), backendID)
+				enabledURLs := backend.lb.Servers()
+				var newDisabledURLs []*url.URL
+				for _, url := range backend.DisabledURLs {
+					if testHealth(url, backend.URL) {
+						log.Debugf("Upsert %s", url.String())
+						backend.lb.UpsertServer(url, roundrobin.Weight(1))
+					} else {
+						newDisabledURLs = append(newDisabledURLs, url)
+					}
+				}
+				backend.DisabledURLs = newDisabledURLs
+
+				for _, url := range enabledURLs {
+					if !testHealth(url, backend.URL) {
+						log.Debugf("Remove %s", url.String())
+						backend.lb.RemoveServer(url)
+						backend.DisabledURLs = append(backend.DisabledURLs, url)
+					}
+				}
+			}
+		}
+	}()
+}
+
+func testHealth(serverURL *url.URL, checkURL string) bool {
+	resp, err := http.Get(serverURL.String() + checkURL)
+	if err != nil || resp.StatusCode != 200 {
+		return false
+	}
+	return true
+}
diff --git a/server.go b/server.go
index 3694a65b0..b968a30bf 100644
--- a/server.go
+++ b/server.go
@@ -23,6 +23,7 @@ import (
 	"github.com/codegangsta/negroni"
 	"github.com/containous/mux"
 	"github.com/containous/traefik/cluster"
+	"github.com/containous/traefik/healthcheck"
 	"github.com/containous/traefik/log"
 	"github.com/containous/traefik/middlewares"
 	"github.com/containous/traefik/provider"
@@ -551,6 +552,9 @@ func (server *Server) loadConfig(configurations configs, globalConfiguration Glo
 	redirectHandlers := make(map[string]http.Handler)
 
 	backends := map[string]http.Handler{}
+
+	backendsHealcheck := map[string]*healthcheck.BackendHealthCheck{}
+
 	backend2FrontendMap := map[string]string{}
 	for _, configuration := range configurations {
 		frontendNames := sortedFrontendNamesForConfig(configuration)
@@ -650,6 +654,9 @@ func (server *Server) loadConfig(configurations configs, globalConfiguration Glo
 									log.Errorf("Skipping frontend %s...", frontendName)
 									continue frontend
 								}
+								if configuration.Backends[frontend.Backend].HealthCheck != nil {
+									backendsHealcheck[frontend.Backend] = healthcheck.NewBackendHealthCheck(configuration.Backends[frontend.Backend].HealthCheck.URL, rebalancer)
+								}
 							}
 						case types.Wrr:
 							log.Debugf("Creating load-balancer wrr")
@@ -673,6 +680,9 @@ func (server *Server) loadConfig(configurations configs, globalConfiguration Glo
 									continue frontend
 								}
 							}
+							if configuration.Backends[frontend.Backend].HealthCheck != nil {
+								backendsHealcheck[frontend.Backend] = healthcheck.NewBackendHealthCheck(configuration.Backends[frontend.Backend].HealthCheck.URL, rr)
+							}
 						}
 						maxConns := configuration.Backends[frontend.Backend].MaxConn
 						if maxConns != nil && maxConns.Amount != 0 {
@@ -735,6 +745,7 @@ func (server *Server) loadConfig(configurations configs, globalConfiguration Glo
 			}
 		}
 	}
+	healthcheck.GetHealthCheck().SetBackendsConfiguration(backendsHealcheck)
 	middlewares.SetBackend2FrontendMap(&backend2FrontendMap)
 	//sort routes
 	for _, serverEntryPoint := range serverEntryPoints {
diff --git a/types/types.go b/types/types.go
index b2d764fd5..d3937e156 100644
--- a/types/types.go
+++ b/types/types.go
@@ -16,6 +16,7 @@ type Backend struct {
 	CircuitBreaker *CircuitBreaker   `json:"circuitBreaker,omitempty"`
 	LoadBalancer   *LoadBalancer     `json:"loadBalancer,omitempty"`
 	MaxConn        *MaxConn          `json:"maxConn,omitempty"`
+	HealthCheck    *HealthCheck      `json:"healthCheck,omitempty"`
 }
 
 // MaxConn holds maximum connection configuration
@@ -35,6 +36,11 @@ type CircuitBreaker struct {
 	Expression string `json:"expression,omitempty"`
 }
 
+// HealthCheck holds healthchk configuration
+type HealthCheck struct {
+	URL string `json:"url,omitempty"`
+}
+
 // Server holds server configuration.
 type Server struct {
 	URL    string `json:"url,omitempty"`