healthcheck: add support at the load-balancers of services level

Co-authored-by: Dmitry Sharshakov <d3dx12.xx@gmail.com>
Co-authored-by: Julien Salleyron <julien.salleyron@gmail.com>
Co-authored-by: Jean-Baptiste Doumenjou <925513+jbdoumenjou@users.noreply.github.com>
Co-authored-by: Romain <rtribotte@users.noreply.github.com>
Co-authored-by: Tom Moulard <tom.moulard@traefik.io>

pkg/server/service/loadbalancer/mirror/mirror.go

@@ -11,6 +11,8 @@ import (
 	"net/http"
 	"sync"
 
+	"github.com/traefik/traefik/v2/pkg/config/dynamic"
+	"github.com/traefik/traefik/v2/pkg/healthcheck"
 	"github.com/traefik/traefik/v2/pkg/log"
 	"github.com/traefik/traefik/v2/pkg/middlewares/accesslog"
 	"github.com/traefik/traefik/v2/pkg/safe"
@@ -23,19 +25,21 @@ type Mirroring struct {
 	rw             http.ResponseWriter
 	routinePool    *safe.Pool
 
-	maxBodySize int64
+	maxBodySize      int64
+	wantsHealthCheck bool
 
 	lock  sync.RWMutex
 	total uint64
 }
 
 // New returns a new instance of *Mirroring.
-func New(handler http.Handler, pool *safe.Pool, maxBodySize int64) *Mirroring {
+func New(handler http.Handler, pool *safe.Pool, maxBodySize int64, hc *dynamic.HealthCheck) *Mirroring {
 	return &Mirroring{
-		routinePool: pool,
-		handler:     handler,
-		rw:          blackHoleResponseWriter{},
-		maxBodySize: maxBodySize,
+		routinePool:      pool,
+		handler:          handler,
+		rw:               blackHoleResponseWriter{},
+		maxBodySize:      maxBodySize,
+		wantsHealthCheck: hc != nil,
 	}
 }
 
@@ -134,6 +138,31 @@ func (m *Mirroring) AddMirror(handler http.Handler, percent int) error {
 	return nil
 }
 
+// RegisterStatusUpdater adds fn to the list of hooks that are run when the
+// status of handler of the Mirroring changes.
+// Not thread safe.
+func (m *Mirroring) RegisterStatusUpdater(fn func(up bool)) error {
+	// Since the status propagation is completely transparent through the
+	// mirroring (because of the recursion on the underlying service), we could maybe
+	// skip that below, and even not add HealthCheck as a field of
+	// dynamic.Mirroring. But I think it's easier to understand for the user
+	// if the HealthCheck is required absolutely everywhere in the config.
+	if !m.wantsHealthCheck {
+		return errors.New("healthCheck not enabled in config for this mirroring service")
+	}
+
+	updater, ok := m.handler.(healthcheck.StatusUpdater)
+	if !ok {
+		return fmt.Errorf("service of mirroring %T not a healthcheck.StatusUpdater", m.handler)
+	}
+
+	if err := updater.RegisterStatusUpdater(fn); err != nil {
+		return fmt.Errorf("cannot register service of mirroring as updater: %w", err)
+	}
+
+	return nil
+}
+
 type blackHoleResponseWriter struct{}
 
 func (b blackHoleResponseWriter) Flush() {}

pkg/server/service/loadbalancer/mirror/mirror_test.go

@@ -21,7 +21,7 @@ func TestMirroringOn100(t *testing.T) {
 		rw.WriteHeader(http.StatusOK)
 	})
 	pool := safe.NewPool(context.Background())
-	mirror := New(handler, pool, defaultMaxBodySize)
+	mirror := New(handler, pool, defaultMaxBodySize, nil)
 	err := mirror.AddMirror(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
 		atomic.AddInt32(&countMirror1, 1)
 	}), 10)
@@ -50,7 +50,7 @@ func TestMirroringOn10(t *testing.T) {
 		rw.WriteHeader(http.StatusOK)
 	})
 	pool := safe.NewPool(context.Background())
-	mirror := New(handler, pool, defaultMaxBodySize)
+	mirror := New(handler, pool, defaultMaxBodySize, nil)
 	err := mirror.AddMirror(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
 		atomic.AddInt32(&countMirror1, 1)
 	}), 10)
@@ -74,7 +74,7 @@ func TestMirroringOn10(t *testing.T) {
 }
 
 func TestInvalidPercent(t *testing.T) {
-	mirror := New(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), safe.NewPool(context.Background()), defaultMaxBodySize)
+	mirror := New(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), safe.NewPool(context.Background()), defaultMaxBodySize, nil)
 	err := mirror.AddMirror(nil, -1)
 	assert.Error(t, err)
 
@@ -93,7 +93,7 @@ func TestHijack(t *testing.T) {
 		rw.WriteHeader(http.StatusOK)
 	})
 	pool := safe.NewPool(context.Background())
-	mirror := New(handler, pool, defaultMaxBodySize)
+	mirror := New(handler, pool, defaultMaxBodySize, nil)
 
 	var mirrorRequest bool
 	err := mirror.AddMirror(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
@@ -117,7 +117,7 @@ func TestFlush(t *testing.T) {
 		rw.WriteHeader(http.StatusOK)
 	})
 	pool := safe.NewPool(context.Background())
-	mirror := New(handler, pool, defaultMaxBodySize)
+	mirror := New(handler, pool, defaultMaxBodySize, nil)
 
 	var mirrorRequest bool
 	err := mirror.AddMirror(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
@@ -154,7 +154,7 @@ func TestMirroringWithBody(t *testing.T) {
 		rw.WriteHeader(http.StatusOK)
 	})
 
-	mirror := New(handler, pool, defaultMaxBodySize)
+	mirror := New(handler, pool, defaultMaxBodySize, nil)
 
 	for i := 0; i < numMirrors; i++ {
 		err := mirror.AddMirror(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {