fix: concurrency issues

distantmagic · Jul 17, 2024 · 6874008 · 6874008
1 parent e95f486
commit 6874008
Show file tree

Hide file tree

Showing 6 changed files with 57 additions and 5 deletions.
diff --git a/loadbalancer/LlamaCppTarget.go b/loadbalancer/LlamaCppTarget.go
@@ -34,6 +34,13 @@ func (self *LlamaCppTarget) DecrementRemainingTicks() {
 	self.RemainingTicksUntilRemoved -= 1
 }
 
+func (self *LlamaCppTarget) GetSlotsStatus() (int, int) {
+	mutexToken := self.RBMutex.RLock()
+	defer self.RBMutex.RUnlock(mutexToken)
+
+	return self.LlamaCppHealthStatus.SlotsIdle, self.LlamaCppHealthStatus.SlotsProcessing
+}
+
 func (self *LlamaCppTarget) HasLessSlotsThan(other *LlamaCppTarget) bool {
 	mutexToken := self.RBMutex.RLock()
 	defer self.RBMutex.RUnlock(mutexToken)
@@ -44,6 +51,13 @@ func (self *LlamaCppTarget) HasLessSlotsThan(other *LlamaCppTarget) bool {
 	return self.LlamaCppHealthStatus.SlotsIdle < other.LlamaCppHealthStatus.SlotsIdle
 }
 
+func (self *LlamaCppTarget) HasRemainingTicks() bool {
+	mutexToken := self.RBMutex.RLock()
+	defer self.RBMutex.RUnlock(mutexToken)
+
+	return self.RemainingTicksUntilRemoved > 0
+}
+
 func (self *LlamaCppTarget) SetTickStatus(
 	lastUpdate time.Time,
 	llamaCppHealthStatus *llamacpp.LlamaCppHealthStatus,
@@ -52,8 +66,8 @@ func (self *LlamaCppTarget) SetTickStatus(
 	self.RBMutex.Lock()
 	defer self.RBMutex.Unlock()
 
-	slotsIdleDiff := self.LlamaCppHealthStatus.SlotsIdle-llamaCppHealthStatus.SlotsIdle
-	slotsProcessingDiff := self.LlamaCppHealthStatus.SlotsProcessing-llamaCppHealthStatus.SlotsProcessing
+	slotsIdleDiff := self.LlamaCppHealthStatus.SlotsIdle - llamaCppHealthStatus.SlotsIdle
+	slotsProcessingDiff := self.LlamaCppHealthStatus.SlotsProcessing - llamaCppHealthStatus.SlotsProcessing
 
 	self.LastUpdate = lastUpdate
 	self.LlamaCppHealthStatus = llamaCppHealthStatus

diff --git a/loadbalancer/LoadBalancerTargetCollection.go b/loadbalancer/LoadBalancerTargetCollection.go
@@ -24,6 +24,9 @@ func (self *LoadBalancerTargetCollection) FixTargetOrder(target *LlamaCppTarget)
 		return
 	}
 
+	self.TargetsRBMutex.Lock()
+	defer self.TargetsRBMutex.Unlock()
+
 	nextElement := element.Next()
 
 	for nextElement != nil {
@@ -62,6 +65,9 @@ func (self *LoadBalancerTargetCollection) GetTargetByConfiguration(
 }
 
 func (self *LoadBalancerTargetCollection) GetHeadTarget() *LlamaCppTarget {
+	mutexToken := self.TargetsRBMutex.RLock()
+	defer self.TargetsRBMutex.RUnlock(mutexToken)
+
 	headElement := self.Targets.Front()
 
 	if headElement == nil {
@@ -76,13 +82,19 @@ func (self *LoadBalancerTargetCollection) GetHeadTarget() *LlamaCppTarget {
 }
 
 func (self *LoadBalancerTargetCollection) Len() int {
+	mutexToken := self.TargetsRBMutex.RLock()
+	defer self.TargetsRBMutex.RUnlock(mutexToken)
+
 	return self.Targets.Len()
 }
 
 func (self *LoadBalancerTargetCollection) RegisterTarget(llamaCppTarget *LlamaCppTarget) {
 	self.setTargetByConfiguration(llamaCppTarget)
 	self.LlamaCppHealthStatusAggregate.AddSlotsFrom(llamaCppTarget)
 
+	self.TargetsRBMutex.Lock()
+	defer self.TargetsRBMutex.Unlock()
+
 	if self.Targets.Len() < 1 {
 		self.elementByTarget.Store(llamaCppTarget, self.Targets.PushFront(llamaCppTarget))
 
@@ -101,6 +113,9 @@ func (self *LoadBalancerTargetCollection) RegisterTarget(llamaCppTarget *LlamaCp
 }
 
 func (self *LoadBalancerTargetCollection) RemoveTarget(llamaCppTarget *LlamaCppTarget) {
+	self.TargetsRBMutex.Lock()
+	defer self.TargetsRBMutex.Unlock()
+
 	self.LlamaCppHealthStatusAggregate.RemoveSlotsFrom(llamaCppTarget)
 	element := self.getElementByTarget(llamaCppTarget)
 
@@ -124,6 +139,14 @@ func (self *LoadBalancerTargetCollection) UpdateTargetWithLlamaCppHealthStatus(
 
 func (self *LoadBalancerTargetCollection) UseSlot(llamaCppTarget *LlamaCppTarget) {
 	targetElement := self.getElementByTarget(llamaCppTarget)
+
+	if targetElement == nil {
+		return
+	}
+
+	self.TargetsRBMutex.Lock()
+	defer self.TargetsRBMutex.Unlock()
+
 	nextTarget := targetElement.Next()
 
 	llamaCppTarget.DecrementIdleSlots()

diff --git a/loadbalancer/LoadBalancerTemporalManager.go b/loadbalancer/LoadBalancerTemporalManager.go
@@ -23,6 +23,8 @@ func (self *LoadBalancerTemporalManager) ReduceTargetCollectionRemainingTicks()
 	var aggregatedSlotsIdle int
 	var aggregatedSlotsProcessing int
 
+	targetsMutexToken := self.LoadBalancerTargetCollection.TargetsRBMutex.RLock()
+
 	for element := self.LoadBalancerTargetCollection.Targets.Front(); element != nil; element = element.Next() {
 		if element.Value == nil {
 			continue
@@ -31,14 +33,18 @@ func (self *LoadBalancerTemporalManager) ReduceTargetCollectionRemainingTicks()
 		target := element.Value.(*LlamaCppTarget)
 		target.DecrementRemainingTicks()
 
-		if target.RemainingTicksUntilRemoved < 1 {
+		if !target.HasRemainingTicks() {
 			defer self.LoadBalancerTargetCollection.RemoveTarget(target)
 		}
 
-		aggregatedSlotsIdle += target.LlamaCppHealthStatus.SlotsIdle
-		aggregatedSlotsProcessing += target.LlamaCppHealthStatus.SlotsProcessing
+		slotsIdle, slotsProcessing := target.GetSlotsStatus()
+
+		aggregatedSlotsIdle += slotsIdle
+		aggregatedSlotsProcessing += slotsProcessing
 	}
 
+	self.LoadBalancerTargetCollection.TargetsRBMutex.RUnlock(targetsMutexToken)
+
 	self.LlamaCppHealthStatusAggregate.SetTo(
 		aggregatedSlotsIdle,
 		aggregatedSlotsProcessing,

diff --git a/loadbalancer/RespondToAggregatedHealth.go b/loadbalancer/RespondToAggregatedHealth.go
@@ -13,6 +13,9 @@ type RespondToAggregatedHealth struct {
 }
 
 func (self *RespondToAggregatedHealth) ServeHTTP(response http.ResponseWriter, request *http.Request) {
+	mutexToken := self.LlamaCppHealthStatusAggregate.RBMutex.RLock()
+	defer self.LlamaCppHealthStatusAggregate.RBMutex.RUnlock(mutexToken)
+
 	jsonLoadBalancerStatus, err := json.Marshal(self.LlamaCppHealthStatusAggregate.AggregatedHealthStatus)
 
 	if err != nil {

diff --git a/loadbalancer/StatsdReporter.go b/loadbalancer/StatsdReporter.go
@@ -12,6 +12,9 @@ func (self *StatsdReporter) ReportAggregatedHealthStatus(
 	bufferedRequestsStats *BufferedRequestsStats,
 	llamaCppHealthStatusAggregate *LlamaCppHealthStatusAggregate,
 ) error {
+	mutexToken := llamaCppHealthStatusAggregate.RBMutex.RLock()
+	defer llamaCppHealthStatusAggregate.RBMutex.RUnlock(mutexToken)
+
 	self.StatsdClient.Gauge("requests_buffered", int64(bufferedRequestsStats.RequestsBuffered))
 	self.StatsdClient.Gauge("slots_idle", int64(llamaCppHealthStatusAggregate.AggregatedHealthStatus.SlotsIdle))
 	self.StatsdClient.Gauge("slots_processing", int64(llamaCppHealthStatusAggregate.AggregatedHealthStatus.SlotsProcessing))

diff --git a/management/RespondToDashboard.go b/management/RespondToDashboard.go
@@ -30,6 +30,9 @@ func iterList[T any](loadBalancerCollection *list.List) <-chan T {
 }
 
 func (self *RespondToDashboard) ServeHTTP(response http.ResponseWriter, request *http.Request) {
+	mutexToken := self.LoadBalancer.LoadBalancerTargetCollection.TargetsRBMutex.RLock()
+	defer self.LoadBalancer.LoadBalancerTargetCollection.TargetsRBMutex.RUnlock(mutexToken)
+
 	response.Header().Set("Content-Type", "text/html")
 	response.WriteHeader(http.StatusOK)