Fixes the concurrency issue of calling the Next() proxy target on RRB

middleware/proxy.go

@@ -12,7 +12,6 @@ import (
 	"regexp"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/labstack/echo/v4"
@@ -79,19 +78,20 @@ type (
 
 	commonBalancer struct {
 		targets []*ProxyTarget
-		mutex   sync.RWMutex
+		mutex   sync.Mutex
 	}
 
 	// RandomBalancer implements a random load balancing technique.
 	randomBalancer struct {
-		*commonBalancer
+		commonBalancer
 		random *rand.Rand
 	}
 
 	// RoundRobinBalancer implements a round-robin load balancing technique.
 	roundRobinBalancer struct {
-		*commonBalancer
-		i uint32
+		commonBalancer
+		// tracking the index on `targets` slice for the next `*ProxyTarget` to be used
+		i int
 	}
 )
 
@@ -143,32 +143,37 @@ func proxyRaw(t *ProxyTarget, c echo.Context) http.Handler {
 
 // NewRandomBalancer returns a random proxy balancer.
 func NewRandomBalancer(targets []*ProxyTarget) ProxyBalancer {
-	b := &randomBalancer{commonBalancer: new(commonBalancer)}
+	b := randomBalancer{}
 	b.targets = targets
-	return b
+	b.random = rand.New(rand.NewSource(int64(time.Now().Nanosecond())))
+	return &b
 }
 
 // NewRoundRobinBalancer returns a round-robin proxy balancer.
 func NewRoundRobinBalancer(targets []*ProxyTarget) ProxyBalancer {
-	b := &roundRobinBalancer{commonBalancer: new(commonBalancer)}
+	b := roundRobinBalancer{}
 	b.targets = targets
-	return b
+	return &b
 }
 
-// AddTarget adds an upstream target to the list.
+// AddTarget adds an upstream target to the list and returns `true`.
+//
+// However, if a target with the same name already exists then the operation is aborted returning `false`.
 func (b *commonBalancer) AddTarget(target *ProxyTarget) bool {
+	b.mutex.Lock()
+	defer b.mutex.Unlock()
 	for _, t := range b.targets {
 		if t.Name == target.Name {
 			return false
 		}
 	}
-	b.mutex.Lock()
-	defer b.mutex.Unlock()
 	b.targets = append(b.targets, target)
 	return true
 }
 
-// RemoveTarget removes an upstream target from the list.
+// RemoveTarget removes an upstream target from the list by name.
+//
+// Returns `true` on success, `false` if no target with the name is found.
 func (b *commonBalancer) RemoveTarget(name string) bool {
 	b.mutex.Lock()
 	defer b.mutex.Unlock()
@@ -182,20 +187,36 @@ func (b *commonBalancer) RemoveTarget(name string) bool {
 }
 
 // Next randomly returns an upstream target.
+//
+// Note: `nil` is returned in case upstream target list is empty.
 func (b *randomBalancer) Next(c echo.Context) *ProxyTarget {
-	if b.random == nil {
-		b.random = rand.New(rand.NewSource(int64(time.Now().Nanosecond())))
+	b.mutex.Lock()
+	defer b.mutex.Unlock()
+	if len(b.targets) == 0 {
+		return nil
+	} else if len(b.targets) == 1 {
+		return b.targets[0]
 	}
-	b.mutex.RLock()
-	defer b.mutex.RUnlock()
 	return b.targets[b.random.Intn(len(b.targets))]
 }
 
 // Next returns an upstream target using round-robin technique.
+//
+// Note: `nil` is returned in case upstream target list is empty.
 func (b *roundRobinBalancer) Next(c echo.Context) *ProxyTarget {
-	b.i = b.i % uint32(len(b.targets))
+	b.mutex.Lock()
+	defer b.mutex.Unlock()
+	if len(b.targets) == 0 {
+		return nil
+	} else if len(b.targets) == 1 {
+		return b.targets[0]
+	}
+	// reset the index if out of bounds
+	if b.i >= len(b.targets) {
+		b.i = 0
+	}
 	t := b.targets[b.i]
-	atomic.AddUint32(&b.i, 1)
+	b.i++
 	return t
 }
 

middleware/proxy_test.go

@@ -122,7 +122,7 @@ func TestProxy(t *testing.T) {
 }
 
 type testProvider struct {
-	*commonBalancer
+	commonBalancer
 	target *ProxyTarget
 	err    error
 }
@@ -143,7 +143,7 @@ func TestTargetProvider(t *testing.T) {
 	url1, _ := url.Parse(t1.URL)
 
 	e := echo.New()
-	tp := &testProvider{commonBalancer: new(commonBalancer)}
+	tp := &testProvider{}
 	tp.target = &ProxyTarget{Name: "target 1", URL: url1}
 	e.Use(Proxy(tp))
 	rec := httptest.NewRecorder()
@@ -158,7 +158,7 @@ func TestFailNextTarget(t *testing.T) {
 	assert.Nil(t, err)
 
 	e := echo.New()
-	tp := &testProvider{commonBalancer: new(commonBalancer)}
+	tp := &testProvider{}
 	tp.target = &ProxyTarget{Name: "target 1", URL: url1}
 	tp.err = echo.NewHTTPError(http.StatusInternalServerError, "method could not select target")
 
@@ -422,3 +422,12 @@ func TestClientCancelConnectionResultsHTTPCode499(t *testing.T) {
 	timeoutStop.Done()
 	assert.Equal(t, 499, rec.Code)
 }
+
+// Assert balancer with empty targets does return `nil` on `Next()`
+func TestProxyBalancerWithNoTargets(t *testing.T) {
+	rb := NewRandomBalancer(nil)
+	assert.Nil(t, rb.Next(nil))
+
+	rrb := NewRoundRobinBalancer([]*ProxyTarget{})
+	assert.Nil(t, rrb.Next(nil))
+}