From 4ddece299009bf57563dbc2e19f6eafa66465ffc Mon Sep 17 00:00:00 2001
From: "sweep-ai[bot]" <128439645+sweep-ai[bot]@users.noreply.github.com>
Date: Fri, 22 Mar 2024 15:46:32 +0000
Subject: [PATCH] feat: Updated gateway/internal/router/priority.go

---
 gateway/internal/router/priority.go | 25 ++++++++++++++++++-------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/gateway/internal/router/priority.go b/gateway/internal/router/priority.go
index 7b0f80e..71a5136 100644
--- a/gateway/internal/router/priority.go
+++ b/gateway/internal/router/priority.go
@@ -2,6 +2,8 @@ package router
 
 import (
 	"sync/atomic"
+	"log"
+	"gateway/internal/router" // Importing to use HealthChecker
 )
 
 const (
@@ -21,11 +23,20 @@ func NewPriorityRouter(providers []RouterConfig) *PriorityRouter {
 }
 
 func (r *PriorityRouter) Next() (*RouterConfig, error) {
-	idx := int(r.idx.Load())
-
-	// Todo: make a check for healthy provider
-	model := &r.providers[idx]
-	r.idx.Add(1)
-
-	return model, nil
+	providerLen := len(r.providers)
+	originalIdx := r.idx.Load()
+	var healthyProvider *RouterConfig
+	for i := 0; i < providerLen; i++ {
+		idx := (originalIdx + uint64(i)) % uint64(providerLen)
+		if router.DefaultHealthChecker{}.IsHealthy(r.providers[idx].Name) {
+			healthyProvider = &r.providers[idx]
+			r.idx.Store(idx + 1)
+			break
+		}
+	}
+	if healthyProvider == nil {
+		log.Println("Error: No healthy providers available.")
+		return nil, fmt.Errorf("no healthy providers available")
+	}
+	return healthyProvider, nil
 }