Adds mirroring service

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

@@ -0,0 +1,104 @@
+package mirror
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"sync"
+
+	"github.com/containous/traefik/v2/pkg/safe"
+)
+
+// Mirroring is an http.Handler that can mirror requests.
+type Mirroring struct {
+	handler        http.Handler
+	mirrorHandlers []*mirrorHandler
+	rw             http.ResponseWriter
+	routinePool    *safe.Pool
+
+	lock  sync.RWMutex
+	total uint64
+}
+
+// New returns a new instance of *Mirroring.
+func New(handler http.Handler, pool *safe.Pool) *Mirroring {
+	return &Mirroring{
+		routinePool: pool,
+		handler:     handler,
+		rw:          blackholeResponseWriter{},
+	}
+}
+
+func (m *Mirroring) inc() uint64 {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+	m.total++
+	return m.total
+}
+
+type mirrorHandler struct {
+	http.Handler
+	percent int
+
+	lock  sync.RWMutex
+	count uint64
+}
+
+func (m *Mirroring) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
+	m.handler.ServeHTTP(rw, req)
+
+	select {
+	case <-req.Context().Done():
+		// No mirroring if request has been canceled during main handler ServeHTTP
+		return
+	default:
+	}
+
+	m.routinePool.GoCtx(func(_ context.Context) {
+		total := m.inc()
+		for _, handler := range m.mirrorHandlers {
+			handler.lock.Lock()
+			if handler.count*100 < total*uint64(handler.percent) {
+				handler.count++
+				handler.lock.Unlock()
+				// When a request served by m.handler is successful, req.Context will be cancelled,
+				// which would trigger a cancellation of the ongoing mirrored requests.
+				// Therefore, we give a new, non-cancellable context  to each of the mirrored calls,
+				// so they can terminate by themselves.
+				handler.ServeHTTP(m.rw, req.WithContext(contextStopPropagation{req.Context()}))
+			} else {
+				handler.lock.Unlock()
+			}
+		}
+	})
+}
+
+// AddMirror adds an httpHandler to mirror to.
+func (m *Mirroring) AddMirror(handler http.Handler, percent int) error {
+	if percent < 0 || percent >= 100 {
+		return errors.New("percent must be between 0 and 100")
+	}
+	m.mirrorHandlers = append(m.mirrorHandlers, &mirrorHandler{Handler: handler, percent: percent})
+	return nil
+}
+
+type blackholeResponseWriter struct{}
+
+func (b blackholeResponseWriter) Header() http.Header {
+	return http.Header{}
+}
+
+func (b blackholeResponseWriter) Write(bytes []byte) (int, error) {
+	return len(bytes), nil
+}
+
+func (b blackholeResponseWriter) WriteHeader(statusCode int) {
+}
+
+type contextStopPropagation struct {
+	context.Context
+}
+
+func (c contextStopPropagation) Done() <-chan struct{} {
+	return make(chan struct{})
+}

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

@@ -0,0 +1,79 @@
+package mirror
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+
+	"github.com/containous/traefik/v2/pkg/safe"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMirroringOn100(t *testing.T) {
+	var countMirror1, countMirror2 int32
+	handler := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.WriteHeader(http.StatusOK)
+	})
+	pool := safe.NewPool(context.Background())
+	mirror := New(handler, pool)
+	err := mirror.AddMirror(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		atomic.AddInt32(&countMirror1, 1)
+	}), 10)
+	assert.NoError(t, err)
+
+	err = mirror.AddMirror(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		atomic.AddInt32(&countMirror2, 1)
+	}), 50)
+	assert.NoError(t, err)
+
+	for i := 0; i < 100; i++ {
+		mirror.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/", nil))
+	}
+
+	pool.Stop()
+
+	val1 := atomic.LoadInt32(&countMirror1)
+	val2 := atomic.LoadInt32(&countMirror2)
+	assert.Equal(t, 10, int(val1))
+	assert.Equal(t, 50, int(val2))
+}
+
+func TestMirroringOn10(t *testing.T) {
+	var countMirror1, countMirror2 int32
+	handler := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		rw.WriteHeader(http.StatusOK)
+	})
+	pool := safe.NewPool(context.Background())
+	mirror := New(handler, pool)
+	err := mirror.AddMirror(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		atomic.AddInt32(&countMirror1, 1)
+	}), 10)
+	assert.NoError(t, err)
+
+	err = mirror.AddMirror(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		atomic.AddInt32(&countMirror2, 1)
+	}), 50)
+	assert.NoError(t, err)
+
+	for i := 0; i < 10; i++ {
+		mirror.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, "/", nil))
+	}
+
+	pool.Stop()
+
+	val1 := atomic.LoadInt32(&countMirror1)
+	val2 := atomic.LoadInt32(&countMirror2)
+	assert.Equal(t, 1, int(val1))
+	assert.Equal(t, 5, int(val2))
+}
+
+func TestInvalidPercent(t *testing.T) {
+	mirror := New(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {}), safe.NewPool(context.Background()))
+	err := mirror.AddMirror(nil, -1)
+	assert.Error(t, err)
+
+	err = mirror.AddMirror(nil, 101)
+	assert.Error(t, err)
+}