Adds mirroring service
This commit is contained in:
parent
fd24b1898e
commit
602a2ea541
10 changed files with 465 additions and 10 deletions
104
pkg/server/service/loadbalancer/mirror/mirror.go
Normal file
104
pkg/server/service/loadbalancer/mirror/mirror.go
Normal file
|
@ -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{})
|
||||
}
|
79
pkg/server/service/loadbalancer/mirror/mirror_test.go
Normal file
79
pkg/server/service/loadbalancer/mirror/mirror_test.go
Normal file
|
@ -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)
|
||||
}
|
Loading…
Add table
Add a link
Reference in a new issue