API: expose runtime representation

Co-authored-by: Julien Salleyron <julien.salleyron@gmail.com>
Co-authored-by: Jean-Baptiste Doumenjou <jb.doumenjou@gmail.com>

pkg/healthcheck/healthcheck.go

@@ -10,12 +10,18 @@ import (
 	"sync"
 	"time"
 
+	"github.com/containous/traefik/pkg/config"
 	"github.com/containous/traefik/pkg/log"
 	"github.com/containous/traefik/pkg/safe"
 	"github.com/go-kit/kit/metrics"
 	"github.com/vulcand/oxy/roundrobin"
 )
 
+const (
+	serverUp   = "UP"
+	serverDown = "DOWN"
+)
+
 var singleton *HealthCheck
 var once sync.Once
 
@@ -221,3 +227,38 @@ func checkHealth(serverURL *url.URL, backend *BackendConfig) error {
 
 	return nil
 }
+
+// NewLBStatusUpdater returns a new LbStatusUpdater
+func NewLBStatusUpdater(bh BalancerHandler, svinfo *config.ServiceInfo) *LbStatusUpdater {
+	return &LbStatusUpdater{
+		BalancerHandler: bh,
+		serviceInfo:     svinfo,
+	}
+}
+
+// LbStatusUpdater wraps a BalancerHandler and a ServiceInfo,
+// so it can keep track of the status of a server in the ServiceInfo.
+type LbStatusUpdater struct {
+	BalancerHandler
+	serviceInfo *config.ServiceInfo // can be nil
+}
+
+// RemoveServer removes the given server from the BalancerHandler,
+// and updates the status of the server to "DOWN".
+func (lb *LbStatusUpdater) RemoveServer(u *url.URL) error {
+	err := lb.BalancerHandler.RemoveServer(u)
+	if err == nil && lb.serviceInfo != nil {
+		lb.serviceInfo.UpdateStatus(u.String(), serverDown)
+	}
+	return err
+}
+
+// UpsertServer adds the given server to the BalancerHandler,
+// and updates the status of the server to "UP".
+func (lb *LbStatusUpdater) UpsertServer(u *url.URL, options ...roundrobin.ServerOption) error {
+	err := lb.BalancerHandler.UpsertServer(u, options...)
+	if err == nil && lb.serviceInfo != nil {
+		lb.serviceInfo.UpdateStatus(u.String(), serverUp)
+	}
+	return err
+}

pkg/healthcheck/healthcheck_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/containous/traefik/pkg/config"
 	"github.com/containous/traefik/pkg/testhelpers"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -367,6 +368,8 @@ type testLoadBalancer struct {
 	numRemovedServers  int
 	numUpsertedServers int
 	servers            []*url.URL
+	// options is just to make sure that LBStatusUpdater forwards options on Upsert to its BalancerHandler
+	options []roundrobin.ServerOption
 }
 
 func (lb *testLoadBalancer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
@@ -386,6 +389,7 @@ func (lb *testLoadBalancer) UpsertServer(u *url.URL, options ...roundrobin.Serve
 	defer lb.Unlock()
 	lb.numUpsertedServers++
 	lb.servers = append(lb.servers, u)
+	lb.options = append(lb.options, options...)
 	return nil
 }
 
@@ -393,14 +397,23 @@ func (lb *testLoadBalancer) Servers() []*url.URL {
 	return lb.servers
 }
 
+func (lb *testLoadBalancer) Options() []roundrobin.ServerOption {
+	return lb.options
+}
+
 func (lb *testLoadBalancer) removeServer(u *url.URL) {
 	var i int
 	var serverURL *url.URL
+	found := false
 	for i, serverURL = range lb.servers {
 		if *serverURL == *u {
+			found = true
 			break
 		}
 	}
+	if !found {
+		return
+	}
 
 	lb.servers = append(lb.servers[:i], lb.servers[i+1:]...)
 }
@@ -427,3 +440,32 @@ func (th *testHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		th.done()
 	}
 }
+
+func TestLBStatusUpdater(t *testing.T) {
+	lb := &testLoadBalancer{RWMutex: &sync.RWMutex{}}
+	svInfo := &config.ServiceInfo{}
+	lbsu := NewLBStatusUpdater(lb, svInfo)
+	newServer, err := url.Parse("http://foo.com")
+	assert.Nil(t, err)
+	err = lbsu.UpsertServer(newServer, roundrobin.Weight(1))
+	assert.Nil(t, err)
+	assert.Equal(t, len(lbsu.Servers()), 1)
+	assert.Equal(t, len(lbsu.BalancerHandler.(*testLoadBalancer).Options()), 1)
+	statuses := svInfo.GetAllStatus()
+	assert.Equal(t, len(statuses), 1)
+	for k, v := range statuses {
+		assert.Equal(t, k, newServer.String())
+		assert.Equal(t, v, serverUp)
+		break
+	}
+	err = lbsu.RemoveServer(newServer)
+	assert.Nil(t, err)
+	assert.Equal(t, len(lbsu.Servers()), 0)
+	statuses = svInfo.GetAllStatus()
+	assert.Equal(t, len(statuses), 1)
+	for k, v := range statuses {
+		assert.Equal(t, k, newServer.String())
+		assert.Equal(t, v, serverDown)
+		break
+	}
+}