From 18dced44713b375742df0a547ee8197b6f173fdd Mon Sep 17 00:00:00 2001
From: Wangjunqi123 <wangjunqi@kylinos.cn>
Date: Fri, 13 Dec 2024 10:29:46 +0800
Subject: [PATCH] server: limit cluster scale

---
 cmd/server/app/agentmanager/agentmanager.go | 15 +++++++++++++++
 pkg/global/global.go                        |  3 +++
 2 files changed, 18 insertions(+)

diff --git a/cmd/server/app/agentmanager/agentmanager.go b/cmd/server/app/agentmanager/agentmanager.go
index ba8c26a8..9c332b2b 100644
--- a/cmd/server/app/agentmanager/agentmanager.go
+++ b/cmd/server/app/agentmanager/agentmanager.go
@@ -18,6 +18,7 @@ import (
 	"gitee.com/openeuler/PilotGo/pkg/global"
 	commonSDK "gitee.com/openeuler/PilotGo/sdk/common"
 	"gitee.com/openeuler/PilotGo/sdk/logger"
+	"github.com/pkg/errors"
 )
 
 // 用于管理server连接的agent
@@ -47,6 +48,15 @@ func (am *AgentManager) Stop() {
 	// stop server here
 }
 
+func (am *AgentManager) ReturnAgentNumber() int {
+	agentnum := 0
+	globalAgentManager.agentMap.Range(func(key, value any) bool {
+		agentnum++
+		return true
+	})
+	return agentnum
+}
+
 func AddAgent(a *Agent) {
 	globalAgentManager.agentMap.Store(a.UUID, a)
 }
@@ -84,6 +94,11 @@ func DeleteAgent(uuid string) {
 }
 
 func AddandRunAgent(c net.Conn) error {
+	if globalAgentManager.ReturnAgentNumber() == global.ClusterSize {
+		c.Close()
+		return errors.Errorf("cluster size is full, can't add new agent")
+	}
+
 	agent, err := NewAgent(c)
 	if err != nil {
 		logger.Warn("create agent from conn error, error:%s , remote addr is:%s",
diff --git a/pkg/global/global.go b/pkg/global/global.go
index cc80ac4e..9b44403a 100644
--- a/pkg/global/global.go
+++ b/pkg/global/global.go
@@ -13,4 +13,7 @@ const (
 	// 是否为部门根节点
 	Departroot   = 0
 	DepartUnroot = 1
+
+	// 集群规模
+	ClusterSize = 10
 )
-- 
Gitee