nim-chat-sdk/ratelimit/store/memory.nim

import std/[times, options, deques, tables]
import ./store
import chronos

# Memory Implementation
type MemoryRateLimitStore*[T] = ref object
  bucketState: Option[BucketState]
  criticalQueue: Deque[seq[tuple[msgId: string, msg: T]]]
  normalQueue: Deque[seq[tuple[msgId: string, msg: T]]]
  criticalLength: int
  normalLength: int

proc new*[T](M: type[MemoryRateLimitStore[T]]): M =
  return M(
    bucketState: none(BucketState),
    criticalQueue: initDeque[seq[tuple[msgId: string, msg: T]]](),
    normalQueue: initDeque[seq[tuple[msgId: string, msg: T]]](),
    criticalLength: 0,
    normalLength: 0
  )

proc saveBucketState*[T](
    store: MemoryRateLimitStore[T], bucketState: BucketState
): Future[bool] {.async.} =
  store.bucketState = some(bucketState)
  return true

proc loadBucketState*[T](
    store: MemoryRateLimitStore[T]
): Future[Option[BucketState]] {.async.} =
  return store.bucketState

proc addToQueue*[T](
    store: MemoryRateLimitStore[T], 
    queueType: QueueType, 
    msgs: seq[tuple[msgId: string, msg: T]]
): Future[bool] {.async.} =
  case queueType
  of QueueType.Critical:
    store.criticalQueue.addLast(msgs)
    inc store.criticalLength
  of QueueType.Normal:
    store.normalQueue.addLast(msgs)
    inc store.normalLength
  return true

proc popFromQueue*[T](
    store: MemoryRateLimitStore[T], 
    queueType: QueueType
): Future[Option[seq[tuple[msgId: string, msg: T]]]] {.async.} =
  case queueType
  of QueueType.Critical:
    if store.criticalQueue.len > 0:
      dec store.criticalLength
      return some(store.criticalQueue.popFirst())
  of QueueType.Normal:
    if store.normalQueue.len > 0:
      dec store.normalLength
      return some(store.normalQueue.popFirst())
  return none(seq[tuple[msgId: string, msg: T]])

proc getQueueLength*[T](
    store: MemoryRateLimitStore[T], 
    queueType: QueueType
): int =
  case queueType
  of QueueType.Critical:
    return store.criticalLength
  of QueueType.Normal:
    return store.normalLength