fix: LB handling non /invoke requests

Author: grussorusso

internal/lb/architectureAwareLb.go

@@ -24,8 +24,10 @@ type ArchitectureAwareBalancer struct {
 	armRing *HashRing
 	x86Ring *HashRing
 
-	mode      string
-	rrIndices map[string]int
+	mode        string
+	archRRIndex int
+	armRRIndex  int // for non-invoke requests using RR
+	x86RRIndex  int // for non-invoke requests using RR
 }
 
 // NewArchitectureAwareBalancer Constructor
@@ -37,9 +39,11 @@ func NewArchitectureAwareBalancer(targets []*middleware.ProxyTarget) *Architectu
 	log.Printf("Running ArchitectureAwareLB with %d replicas per node in the hash rings\n", REPLICAS)
 
 	b := &ArchitectureAwareBalancer{
-		armRing:   NewHashRing(REPLICAS),
-		x86Ring:   NewHashRing(REPLICAS),
-		rrIndices: make(map[string]int),
+		armRing:     NewHashRing(REPLICAS),
+		x86Ring:     NewHashRing(REPLICAS),
+		archRRIndex: 0,
+		armRRIndex:  0,
+		x86RRIndex:  0,
 	}
 
 	b.mode = config.GetString(config.LB_MODE, RR)
@@ -64,6 +68,29 @@ func (b *ArchitectureAwareBalancer) Next(c echo.Context) *middleware.ProxyTarget
 	b.mu.Lock()
 	defer b.mu.Unlock()
 
+	if !isInvoke(c) {
+		log.Printf("c NOT INVOKE: %s\n", c.Path())
+		// fallback to round-robin
+		var candidate *middleware.ProxyTarget
+		var arch string
+		if len(b.armRing.targetList) == 0 {
+			arch = container.X86
+		} else if len(b.x86Ring.targetList) == 0 {
+			arch = container.ARM
+		} else {
+			arch = b.selectArchitectureRR()
+		}
+
+		if arch == container.ARM {
+			b.armRRIndex = (b.armRRIndex + 1) % len(b.armRing.targetList)
+			candidate = b.armRing.targetList[b.armRRIndex]
+		} else {
+			b.x86RRIndex = (b.x86RRIndex + 1) % len(b.x86Ring.targetList)
+			candidate = b.x86Ring.targetList[b.x86RRIndex]
+		}
+		return candidate
+	}
+
 	funcName := extractFunctionName(c)        // get function's name from request's URL
 	fun, ok := function.GetFunction(funcName) // we use this to leverage cache before asking etcd
 	if !ok {
@@ -88,7 +115,7 @@ func (b *ArchitectureAwareBalancer) Next(c echo.Context) *middleware.ProxyTarget
 		bandit := mab.GlobalBanditManager.GetBandit(funcName)
 		targetArch = bandit.SelectArm(ctx)
 	} else if b.mode == RR { // RoundRobin
-		targetArch = b.selectArchitectureRR(funcName) // here the load balancer decides what architecture to use for this function
+		targetArch = b.selectArchitectureRR() // here the load balancer decides what architecture to use for this function
 	} else { // Random
 		targetArch = b.selectArchitectureRandom() // random load balancer for testing purposes
 	}
@@ -196,14 +223,14 @@ func (b *ArchitectureAwareBalancer) selectArchitecture(fun *function.Function) (
 }
 
 // selectArchitectureRR selects the architecture using a Round Robin policy.
-func (b *ArchitectureAwareBalancer) selectArchitectureRR(funcName string) string {
+func (b *ArchitectureAwareBalancer) selectArchitectureRR() string {
 
 	// This is just a function to use as a baseline for the LB. It should actually implement checks over the rings dimension.
 	// i.e.: it cannot select ARM/X86 "blindly", it should check if we have at least one node for that architecture.
 	archs := []string{container.ARM, container.X86}
-	index := b.rrIndices[funcName]
+	index := b.archRRIndex
 	selected := archs[index]
-	b.rrIndices[funcName] = (index + 1) % len(archs)
+	b.archRRIndex = (index + 1) % len(archs)
 	return selected
 }
 

internal/lb/architectureAwareLb_test.go

@@ -157,6 +157,7 @@ func TestGetNodeFromRing(t *testing.T) {
 	req := httptest.NewRequest(http.MethodPost, "/invoke/testGetNodeFromRingFunc", nil)
 	rec := httptest.NewRecorder()
 	c := e.NewContext(req, rec)
+	c.SetPath("/invoke/testGetNodeFromRingFunc")
 
 	// First call
 	firstTarget := b.Next(c)
@@ -207,6 +208,7 @@ func TestGetArchFallback(t *testing.T) {
 	req := httptest.NewRequest(http.MethodPost, "/invoke/testGetArchFallbackFunc", nil)
 	rec := httptest.NewRecorder()
 	c := e.NewContext(req, rec)
+	c.SetPath("/invoke/testGetArchFallbackFunc")
 
 	// First call
 	firstTarget := b.Next(c)
@@ -256,6 +258,7 @@ func TestGetArchFallbackNotPossible(t *testing.T) {
 	req := httptest.NewRequest(http.MethodPost, "/invoke/testGetArchFallbackNotPossibleFunc", nil)
 	rec := httptest.NewRecorder()
 	c := e.NewContext(req, rec)
+	c.SetPath("/invoke/testGetArchFallbackNotPossibleFunc")
 
 	// First call
 	firstTarget := b.Next(c)

internal/lb/architectureUnawareLb.go

@@ -2,6 +2,7 @@ package lb
 
 import (
 	"log"
+	"strings"
 	"sync"
 	"time"
 
@@ -18,6 +19,9 @@ type ArchitectureUnawareBalancer struct {
 
 	// instead of classic lists we will use hashRings (see hashRing.go) to implement a consistent hashing technique
 	hashRing *HashRing
+
+	// round-robin: used for non-invocation requests
+	rrIndex int
 }
 
 // NewArchitectureUnawareBalancer Constructor
@@ -53,6 +57,14 @@ func (b *ArchitectureUnawareBalancer) Next(c echo.Context) *middleware.ProxyTarg
 	b.mu.Lock()
 	defer b.mu.Unlock()
 
+	if !isInvoke(c) {
+		// fallback to round-robin
+		b.rrIndex = (b.rrIndex + 1) % len(b.hashRing.targetList)
+		candidate := b.hashRing.targetList[b.rrIndex]
+		log.Printf("Forwarding %s request to target %s\n", c.Path(), candidate.Name)
+		return candidate
+	}
+
 	funcName := extractFunctionName(c)        // get function's name from request's URL
 	fun, ok := function.GetFunction(funcName) // we use this to leverage cache before asking etcd
 	if !ok {
@@ -75,6 +87,10 @@ func (b *ArchitectureUnawareBalancer) Next(c echo.Context) *middleware.ProxyTarg
 	return candidate
 }
 
+func isInvoke(c echo.Context) bool {
+	return strings.HasPrefix(c.Path(), "/invoke/")
+}
+
 // AddTarget Echo requires this method for dynamic load-balancing. It simply inserts a new node in the respective ring.
 func (b *ArchitectureUnawareBalancer) AddTarget(t *middleware.ProxyTarget) bool {
 	b.mu.Lock()

internal/lb/lb.go

@@ -24,7 +24,7 @@ var currentTargets []*middleware.ProxyTarget
 
 func newBalancer(targets []*middleware.ProxyTarget) (middleware.ProxyBalancer, bool) {
 	// old Load Balancer: return middleware.NewRoundRobinBalancer(targets)
-	isArchAware := config.GetBool(config.Arch_AWARENESS, true)
+	isArchAware := config.GetBool(config.Arch_AWARENESS, false)
 
 	if isArchAware {
 		return NewArchitectureAwareBalancer(targets), true