diff --git a/internals/Service/redisDB.go b/internals/Service/redisDB.go index 0a6cb47..d1e95b2 100644 --- a/internals/Service/redisDB.go +++ b/internals/Service/redisDB.go @@ -2,70 +2,68 @@ package Service import ( "context" + _ "embed" "log" - "sync" config "Load-Pulse/Config" redis "github.com/redis/go-redis/v9" ) -var client *redis.Client; +var client *redis.Client -var ctx = context.Background(); +var ctx = context.Background() -var mu sync.Mutex; +//go:embed scripts/increment.lua +var kaimScript string -func IncrementRequestCount() { - mu.Lock() - defer mu.Unlock() +// Lua script to atomically check and increment +var incrementScript = redis.NewScript(kaimScript) + +func TryIncrementRequestCount(limit int) (bool, error) { cfg := config.GetConfig() - err := client.Incr(ctx, cfg.RedisKey).Err() + result, err := incrementScript.Run(ctx, client, []string{cfg.RedisKey}, limit).Int64() if err != nil { - log.Fatal("[ERR]: Error in Incrementing Concurrent Request Count from Redis !!", err); + return false, err } + // If result is -1, limit reached. Otherwise (new count), success. + return result != -1, nil } func DecrementRequestCount() { - mu.Lock() - defer mu.Unlock() - cfg := config.GetConfig(); - err := client.Decr(ctx, cfg.RedisKey).Err(); + cfg := config.GetConfig() + err := client.Decr(ctx, cfg.RedisKey).Err() if err != nil { - log.Fatal("[ERR]: Error in Decrementing Concurrent Request Count from Redis !!", err); + log.Fatal("[ERR]: Error in Decrementing Concurrent Request Count from Redis !!", err) } } func GetRequestCount() int64 { - mu.Lock() - defer mu.Unlock() - cfg := config.GetConfig(); - currentCount, err := client.Get(ctx, cfg.RedisKey).Int64(); + cfg := config.GetConfig() + currentCount, err := client.Get(ctx, cfg.RedisKey).Int64() if err != nil { - log.Fatal("[ERR]: Error in Fetching Concurrent Requests Count from Redis !!", err); - return 0; + log.Fatal("[ERR]: Error in Fetching Concurrent Requests Count from Redis !!", err) + return 0 } - return currentCount; + return currentCount } func ResetRequestCount() { - mu.Lock() - defer mu.Unlock() - cfg := config.GetConfig(); - err := client.Set(ctx, cfg.RedisKey, 0, 0).Err(); + cfg := config.GetConfig() + err := client.Set(ctx, cfg.RedisKey, 0, 0).Err() if err != nil { - log.Fatal("[ERR]: Error in Resetting Concurrent Requests Count in Redis !!", err); + log.Fatal("[ERR]: Error in Resetting Concurrent Requests Count in Redis !!", err) } - LogServer("[LOG]: Concurrency Count Reset Done\n"); + LogServer("[LOG]: Concurrency Count Reset Done\n") } func InitRedisClient() { - cfg := config.GetConfig(); + cfg := config.GetConfig() client = redis.NewClient(&redis.Options{ Addr: cfg.RedisURL, Password: cfg.RedisPassword, DB: 0, Protocol: 2, }) - LogServer("[LOG]: Redis Client Initialized\n"); -} \ No newline at end of file + LogServer("[LOG]: Redis Client Initialized\n") +} diff --git a/internals/Service/scripts/increment.lua b/internals/Service/scripts/increment.lua new file mode 100644 index 0000000..9160496 --- /dev/null +++ b/internals/Service/scripts/increment.lua @@ -0,0 +1,8 @@ +local key = KEYS[1] +local limit = tonumber(ARGV[1]) +local current = tonumber(redis.call("GET", key) or "0") +if current < limit then + return redis.call("INCR", key) +else + return -1 +end