extproc: fixes race condition with RNG (#417)

**Commit Message**

Previously, router instance shared the rng among multiple goroutines
which is not a thread safe. This fixes it by having a goroutine local
RNG.


**Related Issues/PRs (if applicable)**

Contributes to #53

---------

Signed-off-by: Takeshi Yoneda <[email protected]>

Author: mathetake

internal/extproc/router/router.go

@@ -17,12 +17,11 @@ import (
 // router implements [x.Router].
 type router struct {
 	rules []filterapi.RouteRule
-	rng   *rand.Rand
 }
 
 // New creates a new [x.Router] implementation for the given config.
 func New(config *filterapi.Config, newCustomFn x.NewCustomRouterFn) (x.Router, error) {
-	r := &router{rules: config.Rules, rng: rand.New(rand.NewSource(uint64(time.Now().UnixNano())))} //nolint:gosec
+	r := &router{rules: config.Rules}
 	if newCustomFn != nil {
 		customRouter := newCustomFn(r, config)
 		return customRouter, nil
@@ -44,26 +43,32 @@ func (r *router) Calculate(headers map[string]string) (backend *filterapi.Backen
 			}
 		}
 	}
-	if rule == nil {
+	if rule == nil || len(rule.Backends) == 0 {
 		return nil, x.ErrNoMatchingRule
 	}
 	return r.selectBackendFromRule(rule), nil
 }
 
+// selectBackendFromRule selects a backend from the given rule. Precondition: len(rule.Backends) > 0.
 func (r *router) selectBackendFromRule(rule *filterapi.RouteRule) (backend *filterapi.Backend) {
+	if len(rule.Backends) == 1 {
+		return &rule.Backends[0]
+	}
+
 	// Each backend has a weight, so we randomly select depending on the weight.
 	// This is a pretty naive implementation and can be buggy, so fix it later.
 	totalWeight := 0
 	for _, b := range rule.Backends {
 		totalWeight += b.Weight
 	}
 
+	rng := rand.New(rand.NewSource(uint64(time.Now().UnixNano()))) // nolint:gosec
 	// Pick a random backend if none of them have a weight.
 	if totalWeight == 0 {
-		return &rule.Backends[r.rng.Intn(len(rule.Backends))]
+		return &rule.Backends[rng.Intn(len(rule.Backends))]
 	}
 
-	selected := r.rng.Intn(totalWeight)
+	selected := rng.Intn(totalWeight)
 	for i := range rule.Backends {
 		b := &rule.Backends[i]
 		if selected < b.Weight {

internal/extproc/router/router_test.go

@@ -6,6 +6,8 @@
 package router
 
 import (
+	"sync"
+	"sync/atomic"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -110,6 +112,33 @@ func TestRouter_Calculate(t *testing.T) {
 		require.Greater(t, chosenNames["bar"], 700)
 		require.Greater(t, chosenNames["foo"], 200)
 	})
+
+	t.Run("concurrent access", func(t *testing.T) {
+		var wg sync.WaitGroup
+		wg.Add(1000)
+
+		var foo atomic.Int32
+		var bar atomic.Int32
+		for range 1000 {
+			go func() {
+				defer wg.Done()
+				b, err := r.Calculate(map[string]string{"x-model-name": "llama3.3333"})
+				require.NoError(t, err)
+				require.NotNil(t, b)
+
+				if b.Name == "foo" {
+					foo.Add(1)
+				} else {
+					bar.Add(1)
+				}
+			}()
+		}
+		wg.Wait()
+		require.Equal(t, int32(1000), bar.Load()+foo.Load())
+		require.Greater(t, bar.Load(), foo.Load())
+		require.Greater(t, bar.Load(), int32(700))
+		require.Greater(t, foo.Load(), int32(200))
+	})
 }
 
 func TestRouter_selectBackendFromRule(t *testing.T) {