From 41a05db1ba5098a39423bc88dfc6791a7938fc7d Mon Sep 17 00:00:00 2001
From: zhangqianze <zhangqianze@agora.io>
Date: Fri, 13 Dec 2024 17:43:09 +0800
Subject: [PATCH] feat: add worker limit for gemini

---
 server/internal/config.go      |  2 ++
 server/internal/http_server.go | 20 ++++++++++++++++++++
 server/internal/worker.go      |  1 +
 3 files changed, 23 insertions(+)

diff --git a/server/internal/config.go b/server/internal/config.go
index 8d7bd20d..b2801e80 100644
--- a/server/internal/config.go
+++ b/server/internal/config.go
@@ -21,6 +21,8 @@ const (
 	tokenExpirationInSeconds = uint32(86400)
 
 	WORKER_TIMEOUT_INFINITY = -1
+
+	MAX_GEMINI_WORKER_COUNT = 3
 )
 
 var (
diff --git a/server/internal/http_server.go b/server/internal/http_server.go
index e1752952..a333a40a 100644
--- a/server/internal/http_server.go
+++ b/server/internal/http_server.go
@@ -212,6 +212,25 @@ func (s *HttpServer) handlerStart(c *gin.Context) {
 		return
 	}
 
+	// Check if the graphName contains "gemini"
+	if strings.Contains(req.GraphName, "gemini") {
+		// Count existing workers with the same graphName
+		graphNameCount := 0
+		for _, channelName := range workers.Keys() {
+			worker := workers.Get(channelName).(*Worker)
+			if worker.GraphName == req.GraphName {
+				graphNameCount++
+			}
+		}
+
+		// Reject if more than 3 workers are using the same graphName
+		if graphNameCount >= MAX_GEMINI_WORKER_COUNT {
+			slog.Error("handlerStart graphName workers exceed limit", "graphName", req.GraphName, "graphNameCount", graphNameCount, logTag)
+			s.output(c, codeErrWorkersLimit, http.StatusTooManyRequests)
+			return
+		}
+	}
+
 	req.WorkerHttpServerPort = getHttpServerPort()
 	propertyJsonFile, logFile, err := s.processProperty(&req)
 	if err != nil {
@@ -222,6 +241,7 @@ func (s *HttpServer) handlerStart(c *gin.Context) {
 
 	worker := newWorker(req.ChannelName, logFile, s.config.Log2Stdout, propertyJsonFile)
 	worker.HttpServerPort = req.WorkerHttpServerPort
+	worker.GraphName = req.GraphName // Save graphName in the Worker instance
 
 	if req.QuitTimeoutSeconds > 0 {
 		worker.QuitTimeoutSeconds = req.QuitTimeoutSeconds
diff --git a/server/internal/worker.go b/server/internal/worker.go
index 8e4dac9c..0222b5ee 100644
--- a/server/internal/worker.go
+++ b/server/internal/worker.go
@@ -26,6 +26,7 @@ type Worker struct {
 	LogFile            string
 	Log2Stdout         bool
 	PropertyJsonFile   string
+	GraphName          string // New field to store the graphName
 	Pid                int
 	QuitTimeoutSeconds int
 	CreateTs           int64