feat: max requests

conf/config.yaml

@@ -2,5 +2,6 @@ app:
   port: 8194
   debug: True
   key: dify-sandbox
-max_workers: 100
+max_workers: 3
+max_requests: 100
 worker_timeout: 5

internal/controller/router.go

@@ -3,10 +3,13 @@ package controller
 import (
 	"github.com/gin-gonic/gin"
 	"github.com/langgenius/dify-sandbox/internal/middleware"
+	"github.com/langgenius/dify-sandbox/internal/static"
 )
 
 func Setup(eng *gin.Engine) {
+	eng.Use(middleware.MaxRequest(static.GetCoshubGlobalConfigurations().MaxRequests))
 	eng.Use(middleware.Auth())
+	eng.Use(middleware.MaxWoker(static.GetCoshubGlobalConfigurations().MaxWorkers))
 
 	eng.POST("/v1/sandbox/run", RunSandboxController)
 }

internal/middleware/cocrrent.go

@@ -0,0 +1,58 @@
+package middleware
+
+import (
+	"net/http"
+	"sync"
+
+	"github.com/gin-gonic/gin"
+	"github.com/langgenius/dify-sandbox/internal/types"
+)
+
+func MaxWoker(max int) gin.HandlerFunc {
+	queue := make(chan *gin.Context, max)
+
+	for i := 0; i < max; i++ {
+		go func() {
+			for {
+				select {
+				case c := <-queue:
+					c.Next()
+				}
+			}
+		}()
+	}
+
+	return func(c *gin.Context) {
+		queue <- c
+	}
+}
+
+type MaxRequestIface struct {
+	current int
+	lock    *sync.RWMutex
+}
+
+func MaxRequest(max int) gin.HandlerFunc {
+	m := &MaxRequestIface{
+		current: 0,
+		lock:    &sync.RWMutex{},
+	}
+
+	return func(c *gin.Context) {
+		m.lock.RLock()
+		if m.current >= max {
+			m.lock.RUnlock()
+			c.JSON(http.StatusServiceUnavailable, types.ErrorResponse(-503, "Too many requests"))
+			c.Abort()
+			return
+		}
+		m.lock.RUnlock()
+		m.lock.Lock()
+		m.current++
+		m.lock.Unlock()
+		c.Next()
+		m.lock.Lock()
+		m.current--
+		m.lock.Unlock()
+	}
+}

internal/static/config.go

@@ -33,6 +33,11 @@ func InitConfig(path string) error {
 		difySandboxGlobalConfigurations.MaxWorkers, _ = strconv.Atoi(max_workers)
 	}
 
+	max_requests := os.Getenv("MAX_REQUESTS")
+	if max_requests != "" {
+		difySandboxGlobalConfigurations.MaxRequests, _ = strconv.Atoi(max_requests)
+	}
+
 	port := os.Getenv("SANDBOX_PORT")
 	if port != "" {
 		difySandboxGlobalConfigurations.App.Port, _ = strconv.Atoi(port)

internal/types/config.go

@@ -7,5 +7,6 @@ type DifySandboxGlobalConfigurations struct {
 		Key   string `yaml:"key"`
 	} `yaml:"app"`
 	MaxWorkers    int `yaml:"max_workers"`
+	MaxRequests   int `yaml:"max_requests"`
 	WorkerTimeout int `yaml:"worker_timeout"`
 }