From c5fc98a59a3a19aa696cd66abd7f841e9b51b4c8 Mon Sep 17 00:00:00 2001
From: yombo <40785594+xyombo@users.noreply.github.com>
Date: Tue, 8 Aug 2023 15:16:37 +0800
Subject: [PATCH] feat(remoting):  random load balance for getty sessions #598

---
 .../loadbalance/random_loadbalance.go         | 23 ++++++++++++++-----
 .../loadbalance/random_loadbalance_test.go    | 18 +++++++++------
 2 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/pkg/remoting/loadbalance/random_loadbalance.go b/pkg/remoting/loadbalance/random_loadbalance.go
index 9700f919a..639ebca20 100644
--- a/pkg/remoting/loadbalance/random_loadbalance.go
+++ b/pkg/remoting/loadbalance/random_loadbalance.go
@@ -18,20 +18,31 @@
 package loadbalance
 
 import (
+	"math/rand"
 	"sync"
+	"time"
 
 	getty "github.com/apache/dubbo-getty"
 )
 
 func RandomLoadBalance(sessions *sync.Map, xid string) getty.Session {
-	var session getty.Session
+	//collect sync.Map keys
+	//filted out closed session instance
+	var keys []getty.Session
 	sessions.Range(func(key, value interface{}) bool {
-		session = key.(getty.Session)
+		session := key.(getty.Session)
 		if session.IsClosed() {
-			sessions.Delete(session)
-			return true
+			sessions.Delete(key)
+		} else {
+			keys = append(keys, session)
 		}
-		return false
+		return true
 	})
-	return session
+	//keys eq 0 means there are no available session
+	if len(keys) == 0 {
+		return nil
+	}
+	//random in keys
+	randomIndex := rand.New(rand.NewSource(time.Now().UnixNano())).Intn(len(keys))
+	return keys[randomIndex]
 }
diff --git a/pkg/remoting/loadbalance/random_loadbalance_test.go b/pkg/remoting/loadbalance/random_loadbalance_test.go
index 2a3791365..5db9c8825 100644
--- a/pkg/remoting/loadbalance/random_loadbalance_test.go
+++ b/pkg/remoting/loadbalance/random_loadbalance_test.go
@@ -27,11 +27,11 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestRandomLoadBalance_Nomal(t *testing.T) {
+func TestRandomLoadBalance_Normal(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	sessions := &sync.Map{}
 
-	for i := 0; i < 3; i++ {
+	for i := 0; i < 10; i++ {
 		session := mock.NewMockTestSession(ctrl)
 		session.EXPECT().IsClosed().Return(i == 2).AnyTimes()
 		sessions.Store(session, fmt.Sprintf("session-%d", (i+1)))
@@ -44,20 +44,24 @@ func TestRandomLoadBalance_Nomal(t *testing.T) {
 }
 
 func TestRandomLoadBalance_All_Closed(t *testing.T) {
-
 	ctrl := gomock.NewController(t)
 	sessions := &sync.Map{}
-
 	//mock  closed  sessions
 	for i := 0; i < 10; i++ {
 		session := mock.NewMockTestSession(ctrl)
 		session.EXPECT().IsClosed().Return(true).AnyTimes()
 		sessions.Store(session, fmt.Sprintf("session-%d", (i+1)))
 	}
-	result := RandomLoadBalance(sessions, "some_xid")
+	if result := RandomLoadBalance(sessions, "some_xid"); result != nil {
+		t.Errorf("Expected nil, actual got %+v", result)
+	}
+}
 
-	assert.NotNil(t, result)
-	assert.True(t, result.IsClosed(), "found one un-closed session instance in ALL_CLOSED_SESSION_MAP")
+func TestRandomLoadBalance_Empty(t *testing.T) {
+	sessions := &sync.Map{}
+	if result := RandomLoadBalance(sessions, "some_xid"); result != nil {
+		t.Errorf("Expected nil, actual got %+v", result)
+	}
 }
 
 func TestRandomLoadBalance_All_Opening(t *testing.T) {