Cache and Optimization update

go.mod

@@ -5,11 +5,11 @@ go 1.22.0
 toolchain go1.22.2
 
 require (
+	github.com/fluffy-melli/korcen-go v1.1.0
 	github.com/gin-gonic/gin v1.10.0
 	github.com/swaggo/files v1.0.1
 	github.com/swaggo/gin-swagger v1.6.0
 	github.com/swaggo/swag v1.16.4
-	github.com/fluffy-melli/korcen-go v1.1.0
 )
 
 require (
@@ -22,7 +22,6 @@ require (
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/cloudwego/base64x v0.1.4 // indirect
 	github.com/emirpasic/gods v1.18.1 // indirect
-	github.com/fluffy-melli/korcen-go v1.1.0 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.8 // indirect
 	github.com/gin-contrib/sse v1.0.0 // indirect
 	github.com/go-logr/logr v1.3.0 // indirect

internal/middleware/middleware.go

@@ -15,30 +15,37 @@ type TokenBucket struct {
 	tokens     float64
 	refillRate float64
 	lastRefill time.Time
+	lastSeen   time.Time
 	mutex      sync.Mutex
 }
 
+type MiddlewareConfig struct {
+	Capacity   int
+	RefillRate float64
+}
+
 func NewTokenBucket(capacity int, refillRate float64) *TokenBucket {
+	now := time.Now()
 	return &TokenBucket{
 		capacity:   capacity,
 		tokens:     float64(capacity),
 		refillRate: refillRate,
-		lastRefill: time.Now(),
+		lastRefill: now,
+		lastSeen:   now,
 	}
 }
 
-func (tb *TokenBucket) AllowRequest() bool {
+func (tb *TokenBucket) AllowRequest(currentTime time.Time) bool {
 	tb.mutex.Lock()
 	defer tb.mutex.Unlock()
 
-	now := time.Now()
-	elapsed := now.Sub(tb.lastRefill).Seconds()
-
+	elapsed := currentTime.Sub(tb.lastRefill).Seconds()
 	tb.tokens += elapsed * tb.refillRate
 	if tb.tokens > float64(tb.capacity) {
 		tb.tokens = float64(tb.capacity)
 	}
-	tb.lastRefill = now
+	tb.lastRefill = currentTime
+	tb.lastSeen = currentTime
 
 	if tb.tokens >= 1 {
 		tb.tokens -= 1
@@ -47,24 +54,81 @@ func (tb *TokenBucket) AllowRequest() bool {
 	return false
 }
 
+func (tb *TokenBucket) Reset(capacity int, refillRate float64) {
+	tb.mutex.Lock()
+	defer tb.mutex.Unlock()
+
+	tb.capacity = capacity
+	tb.refillRate = refillRate
+	tb.tokens = float64(capacity)
+	now := time.Now()
+	tb.lastRefill = now
+	tb.lastSeen = now
+}
+
+var TokenBucketPool = sync.Pool{
+	New: func() interface{} {
+		return NewTokenBucket(100, 100.0/60.0)
+	},
+}
+
 var tokenBuckets sync.Map
 
+const CleanupInterval = 5 * time.Minute
+
+const InactiveDuration = 10 * time.Minute
+
+func init() {
+	go cleanupTokenBuckets()
+}
+
+func cleanupTokenBuckets() {
+	ticker := time.NewTicker(CleanupInterval)
+	defer ticker.Stop()
+
+	for {
+		<-ticker.C
+		now := time.Now()
+		tokenBuckets.Range(func(key, value interface{}) bool {
+			tb := value.(*TokenBucket)
+
+			tb.mutex.Lock()
+			lastSeen := tb.lastSeen
+			tb.mutex.Unlock()
+
+			if now.Sub(lastSeen) > InactiveDuration {
+				tokenBuckets.Delete(key)
+				TokenBucketPool.Put(tb)
+			}
+			return true
+		})
+	}
+}
+
 func getOrCreateTokenBucket(ip string, capacity int, refillRate float64) *TokenBucket {
 	if v, ok := tokenBuckets.Load(ip); ok {
 		return v.(*TokenBucket)
 	}
-	newBucket := NewTokenBucket(capacity, refillRate)
-	actual, _ := tokenBuckets.LoadOrStore(ip, newBucket)
+
+	newBucket := TokenBucketPool.Get().(*TokenBucket)
+	newBucket.Reset(capacity, refillRate)
+
+	actual, loaded := tokenBuckets.LoadOrStore(ip, newBucket)
+	if loaded {
+		TokenBucketPool.Put(newBucket)
+		return actual.(*TokenBucket)
+	}
 	return actual.(*TokenBucket)
 }
 
-func TokenBucketMiddleware() gin.HandlerFunc {
+func TokenBucketMiddleware(config MiddlewareConfig) gin.HandlerFunc {
 	return func(c *gin.Context) {
 		ip := c.ClientIP()
+		currentTime := time.Now()
 
-		tb := getOrCreateTokenBucket(ip, 100, 100.0/60.0) // 1분당 100개 (초당 약 1.67개 리필)
+		tb := getOrCreateTokenBucket(ip, config.Capacity, config.RefillRate)
 
-		if !tb.AllowRequest() {
+		if !tb.AllowRequest(currentTime) {
 			c.JSON(http.StatusTooManyRequests, gin.H{"error": "Too many requests"})
 			c.Abort()
 			return

internal/router/routes.go

@@ -1,24 +1,34 @@
 // internal/router/routes.go
+
 package router
 
 import (
 	"net/http"
 	"strings"
 	"time"
 
-	_ "github.com/fluffy-melli/korcen-api/docs"
+	"github.com/asynkron/protoactor-go/actor"
 	"github.com/fluffy-melli/korcen-api/internal/middleware"
 	"github.com/fluffy-melli/korcen-api/pkg/check"
 	"github.com/gin-gonic/gin"
-
-	"github.com/asynkron/protoactor-go/actor"
 )
 
-func SetupRouter(system *actor.ActorSystem, korcenPID *actor.PID) *gin.Engine {
+func respond(c *gin.Context, status int, isXML bool, payload interface{}) {
+	if isXML {
+		c.XML(status, payload)
+	} else {
+		c.JSON(status, payload)
+	}
+}
+
+func SetupRouter(system *actor.ActorSystem, korcenPID *actor.PID, config middleware.MiddlewareConfig) *gin.Engine {
 	gin.SetMode(gin.ReleaseMode)
-	r := gin.Default()
+	r := gin.New()
+
+	r.Use(gin.Logger())
+	r.Use(gin.Recovery())
 
-	r.Use(middleware.TokenBucketMiddleware())
+	r.Use(middleware.TokenBucketMiddleware(config))
 
 	APIGroup := r.Group("/api/v1")
 	{
@@ -29,13 +39,13 @@ func SetupRouter(system *actor.ActorSystem, korcenPID *actor.PID) *gin.Engine {
 			switch c.ContentType() {
 			case "text/xml", "application/xml":
 				if err := c.ShouldBindXML(&header); err != nil {
-					c.XML(http.StatusBadRequest, gin.H{"error": "Invalid XML request"})
+					respond(c, http.StatusBadRequest, true, gin.H{"error": "Invalid XML request"})
 					return
 				}
 				isXML = true
 			default:
 				if err := c.ShouldBindJSON(&header); err != nil {
-					c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid JSON request"})
+					respond(c, http.StatusBadRequest, false, gin.H{"error": "Invalid JSON request"})
 					return
 				}
 			}
@@ -49,9 +59,7 @@ func SetupRouter(system *actor.ActorSystem, korcenPID *actor.PID) *gin.Engine {
 				return
 			}
 
-			msg := &check.KorcenRequest{Header: &header}
-			future := system.Root.RequestFuture(korcenPID, msg, 5*time.Second)
-
+			future := system.Root.RequestFuture(korcenPID, &check.KorcenRequest{Header: &header}, 5*time.Second)
 			result, err := future.Result()
 			if err != nil {
 				if isXML {
@@ -82,12 +90,7 @@ func SetupRouter(system *actor.ActorSystem, korcenPID *actor.PID) *gin.Engine {
 			}
 
 			response := korcenResp.Respond
-
-			if isXML {
-				c.XML(http.StatusOK, response)
-			} else {
-				c.JSON(http.StatusOK, response)
-			}
+			respond(c, http.StatusOK, isXML, response)
 		})
 	}
 

main.go

@@ -7,10 +7,10 @@ import (
 	"log"
 	"os"
 
-	// Proto.Actor
 	"github.com/asynkron/protoactor-go/actor"
 
 	_ "github.com/fluffy-melli/korcen-api/docs"
+	"github.com/fluffy-melli/korcen-api/internal/middleware"
 	"github.com/fluffy-melli/korcen-api/internal/router"
 	"github.com/fluffy-melli/korcen-api/pkg/check"
 	swaggerFiles "github.com/swaggo/files"
@@ -25,9 +25,15 @@ func main() {
 	props := actor.PropsFromProducer(func() actor.Actor {
 		return &check.KorcenActor{}
 	})
+
 	korcenPID := system.Root.Spawn(props)
 
-	r := router.SetupRouter(system, korcenPID)
+	config := middleware.MiddlewareConfig{
+		Capacity:   100,          // 최대 토큰 수
+		RefillRate: 100.0 / 60.0, // 초당 리필 속도 (1분당 100개)
+	}
+
+	r := router.SetupRouter(system, korcenPID, config)
 
 	r.GET("/swagger/*any", ginSwagger.WrapHandler(swaggerFiles.Handler))