No More Redis... simplify!

This commit is contained in:
Atridad Lahiji 2024-01-03 14:32:37 -07:00
parent 8655019a4f
commit 06074fb84a
No known key found for this signature in database
5 changed files with 52 additions and 64 deletions

View file

@ -2,7 +2,5 @@
DISCORD_TOKEN=""
OPENAI_API_KEY=""
REPLICATE_API_TOKEN=""
REDIS_HOST=""
REDIS_PASSWORD=""
# Comma separated
COOLDOWN_ALLOW_LIST=""

41
lib/cooldowns.go Normal file
View file

@ -0,0 +1,41 @@
package lib
import (
"sync"
"time"
)
type CooldownManager struct {
cooldowns map[string]time.Time
mu sync.Mutex
}
func NewCooldownManager() *CooldownManager {
return &CooldownManager{
cooldowns: make(map[string]time.Time),
}
}
func (m *CooldownManager) StartCooldown(key string, duration time.Duration) {
m.mu.Lock()
defer m.mu.Unlock()
m.cooldowns[key] = time.Now().Add(duration)
}
func (m *CooldownManager) IsOnCooldown(key string) bool {
m.mu.Lock()
defer m.mu.Unlock()
cooldownEnd, exists := m.cooldowns[key]
if !exists {
return false
}
if time.Now().After(cooldownEnd) {
delete(m.cooldowns, key)
return false
}
return true
}

View file

@ -3,10 +3,13 @@ package lib
import (
"os"
"strings"
"time"
"github.com/diamondburned/arikawa/v3/discord"
)
var manager = NewCooldownManager()
// Userish is an interface that captures the common methods you may want to call
// on either a discord.Member or discord.User, including a display name.
type Userish interface {
@ -64,7 +67,7 @@ func GetUserObject(event discord.InteractionEvent) Userish {
}
}
func CooldownHandler(event discord.InteractionEvent, command string, cooldownMinutes int) bool {
func CooldownHandler(event discord.InteractionEvent, key string, duration time.Duration) bool {
user := GetUserObject(event)
allowList := strings.Split(os.Getenv("COOLDOWN_ALLOW_LIST"), ",")
@ -75,10 +78,10 @@ func CooldownHandler(event discord.InteractionEvent, command string, cooldownMin
}
}
cachedVal := GetCache(user.ID().String() + ":" + command)
if cachedVal != "nil" {
if manager.IsOnCooldown(key) {
return false
}
SetCache(user.ID().String()+":"+"hdpic", user.ID().String()+":"+command, cooldownMinutes)
manager.StartCooldown(key, duration)
return true
}

View file

@ -1,55 +0,0 @@
package lib
import (
"context"
"os"
"time"
"github.com/redis/go-redis/v9"
)
var redis_host = os.Getenv("REDIS_HOST")
var redis_password = os.Getenv("REDIS_PASSWORD")
var rdb = redis.NewClient(&redis.Options{
Addr: redis_host,
Password: redis_password,
DB: 0,
})
func SetCache(key string, value string, ttlMinutes int) bool {
println("Setting the Cache")
rdb := redis.NewClient(&redis.Options{
Addr: redis_host,
Password: redis_password,
DB: 0,
})
if rdb == nil {
panic("Failed to create Redis client")
}
err := rdb.Set(context.Background(), key, value, time.Minute*time.Duration(ttlMinutes)).Err()
if err != nil {
panic(err)
}
return err != nil
}
func GetCache(key string) string {
println("Fetching From Cache")
if rdb == nil {
panic("Failed to create Redis client")
}
val, err := rdb.Get(context.Background(), key).Result()
if err != nil {
println("Cache Miss")
return "nil"
}
println("Cache Hit")
return val
}

View file

@ -11,6 +11,7 @@ import (
"net/http"
"os"
"os/signal"
"time"
"github.com/diamondburned/arikawa/v3/api"
"github.com/diamondburned/arikawa/v3/api/cmdroute"
@ -132,7 +133,7 @@ func (h *handler) cmdPing(ctx context.Context, data cmdroute.CommandData) *api.I
func (h *handler) cmdAsk(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData {
// Cooldown Logic
allowed := lib.CooldownHandler(*data.Event, "ask", 1)
allowed := lib.CooldownHandler(*data.Event, "ask", time.Minute)
if !allowed {
return errorResponse(errors.New("please wait for the cooldown"))
@ -194,7 +195,7 @@ func (h *handler) cmdAsk(ctx context.Context, data cmdroute.CommandData) *api.In
func (h *handler) cmdPic(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData {
// Cooldown Logic
allowed := lib.CooldownHandler(*data.Event, "pic", 1)
allowed := lib.CooldownHandler(*data.Event, "pic", time.Minute)
if !allowed {
return errorResponse(errors.New("please wait for the cooldown"))
@ -270,7 +271,7 @@ func (h *handler) cmdPic(ctx context.Context, data cmdroute.CommandData) *api.In
func (h *handler) cmdHDPic(ctx context.Context, data cmdroute.CommandData) *api.InteractionResponseData {
// Cooldown Logic
allowed := lib.CooldownHandler(*data.Event, "hdPic", 1)
allowed := lib.CooldownHandler(*data.Event, "hdPic", time.Minute*10)
if !allowed {
return errorResponse(errors.New("please wait for the cooldown"))