mirror of
https://github.com/logos-messaging/nim-chat-sdk.git
synced 2026-01-05 23:53:09 +00:00
fix: usage of nwaku's token_bucket
This commit is contained in:
pablo
parent
f04e6ae7aa
commit
1d8ad697c2
@ -12,10 +12,11 @@ srcDir = "src"
|
||||
requires "nim >= 2.2.4",
|
||||
"chronicles",
|
||||
"chronos",
|
||||
"db_connector"
|
||||
"db_connector",
|
||||
"waku"
|
||||
|
||||
task buildSharedLib, "Build shared library for C bindings":
|
||||
exec "nim c --app:lib --out:../library/c-bindings/libchatsdk.so chat_sdk/chat_sdk.nim"
|
||||
exec "nim c --mm:refc --app:lib --out:../library/c-bindings/libchatsdk.so chat_sdk/chat_sdk.nim"
|
||||
|
||||
task buildStaticLib, "Build static library for C bindings":
|
||||
exec "nim c --app:staticLib --out:../library/c-bindings/libchatsdk.a chat_sdk/chat_sdk.nim"
|
||||
exec "nim c --mm:refc --app:staticLib --out:../library/c-bindings/libchatsdk.a chat_sdk/chat_sdk.nim"
|
||||
|
||||
@ -1,183 +1,192 @@
|
||||
import std/[times, deques]
|
||||
import std/[times, deques, options]
|
||||
# import ./token_bucket
|
||||
import waku/common/rate_limit/token_bucket
|
||||
import chronos
|
||||
|
||||
type
|
||||
Serializable* = concept x
|
||||
x.toBytes() is seq[byte]
|
||||
CapacityState = enum
|
||||
Normal = 0
|
||||
AlmostNone = 1
|
||||
None = 2
|
||||
|
||||
MessagePriority* = enum
|
||||
SendResult* = enum
|
||||
PassedToSender = 0
|
||||
Enqueued = 1
|
||||
Dropped = 2
|
||||
DroppedBatchTooLarge = 3
|
||||
|
||||
Priority* = enum
|
||||
Critical = 0
|
||||
Normal = 1
|
||||
Optional = 2
|
||||
|
||||
QueuedMessage*[T: Serializable] = ref object of RootObj
|
||||
messageId*: string
|
||||
msg*: T
|
||||
priority*: MessagePriority
|
||||
queuedAt*: float
|
||||
Serializable* =
|
||||
concept x
|
||||
x.toBytes() is seq[byte]
|
||||
|
||||
MessageSender*[T: Serializable] = proc(messageId: string, msg: T): Future[bool] {.async.}
|
||||
MessageSender*[T: Serializable] =
|
||||
proc(msgs: seq[tuple[msgId: string, msg: T]]): Future[void] {.async.}
|
||||
|
||||
RateLimitManager*[T: Serializable] = ref object
|
||||
messageCount*: int = 100 # Default to 100 messages
|
||||
epochDurationSec*: int = 600 # Default to 10 minutes
|
||||
currentCount*: int
|
||||
lastResetTime*: float
|
||||
criticalQueue*: Deque[QueuedMessage[T]]
|
||||
normalQueue*: Deque[QueuedMessage[T]]
|
||||
optionalQueue*: Deque[QueuedMessage[T]]
|
||||
sender*: MessageSender[T]
|
||||
isRunning*: bool
|
||||
sendTask*: Future[void]
|
||||
bucket: TokenBucket
|
||||
sender: MessageSender[T]
|
||||
running: bool
|
||||
queueCritical: Deque[seq[tuple[msgId: string, msg: T]]]
|
||||
queueNormal: Deque[seq[tuple[msgId: string, msg: T]]]
|
||||
sleepDuration: chronos.Duration
|
||||
|
||||
proc newRateLimitManager*[T: Serializable](messageCount: int, epochDurationSec: int, sender: MessageSender[T]): RateLimitManager[T] =
|
||||
proc newRateLimitManager*[T: Serializable](
|
||||
sender: MessageSender[T],
|
||||
capacity: int = 100,
|
||||
duration: chronos.Duration = chronos.minutes(10),
|
||||
sleepDuration: chronos.Duration = chronos.milliseconds(1000),
|
||||
): RateLimitManager[T] =
|
||||
RateLimitManager[T](
|
||||
messageCount: messageCount,
|
||||
epochDurationSec: epochDurationSec,
|
||||
currentCount: 0,
|
||||
lastResetTime: epochTime(),
|
||||
criticalQueue: initDeque[QueuedMessage[T]](),
|
||||
normalQueue: initDeque[QueuedMessage[T]](),
|
||||
optionalQueue: initDeque[QueuedMessage[T]](),
|
||||
bucket: TokenBucket.newStrict(capacity, duration),
|
||||
sender: sender,
|
||||
isRunning: false
|
||||
running: false,
|
||||
queueCritical: Deque[seq[tuple[msgId: string, msg: T]]](),
|
||||
queueNormal: Deque[seq[tuple[msgId: string, msg: T]]](),
|
||||
sleepDuration: sleepDuration,
|
||||
)
|
||||
|
||||
proc updateEpochIfNeeded[T: Serializable](manager: RateLimitManager[T]) =
|
||||
let now = epochTime()
|
||||
if now - manager.lastResetTime >= float(manager.epochDurationSec):
|
||||
manager.lastResetTime = now
|
||||
manager.currentCount = 0
|
||||
|
||||
proc getUsagePercent[T: Serializable](manager: RateLimitManager[T]): float =
|
||||
if manager.messageCount == 0:
|
||||
return 1.0
|
||||
float(manager.currentCount) / float(manager.messageCount)
|
||||
|
||||
proc queueForSend*[T: Serializable](manager: RateLimitManager[T], messageId: string, msg: T, priority: MessagePriority) =
|
||||
manager.updateEpochIfNeeded()
|
||||
|
||||
let queuedMsg = QueuedMessage[T](
|
||||
messageId: messageId,
|
||||
msg: msg,
|
||||
priority: priority,
|
||||
queuedAt: epochTime()
|
||||
)
|
||||
|
||||
let usagePercent = manager.getUsagePercent()
|
||||
|
||||
if usagePercent >= 1.0:
|
||||
# Quota exhausted - queue critical on top, queue normal, drop optional
|
||||
case priority:
|
||||
of Critical:
|
||||
manager.criticalQueue.addLast(queuedMsg)
|
||||
of Normal:
|
||||
manager.normalQueue.addLast(queuedMsg)
|
||||
of Optional:
|
||||
discard # Drop optional messages when quota exhausted
|
||||
elif usagePercent >= 0.7:
|
||||
# Low quota - send critical, queue normal, drop optional
|
||||
case priority:
|
||||
of Critical:
|
||||
manager.criticalQueue.addLast(queuedMsg)
|
||||
of Normal:
|
||||
manager.normalQueue.addLast(queuedMsg)
|
||||
of Optional:
|
||||
discard # Drop optional messages when quota low
|
||||
proc getCapacityState[T: Serializable](
|
||||
manager: RateLimitManager[T], now: Moment, count: int = 1
|
||||
): CapacityState =
|
||||
let (budget, budgetCap) = manager.bucket.getAvailableCapacity(now)
|
||||
let countAfter = budget - count
|
||||
let ratio = countAfter.float / budgetCap.float
|
||||
if ratio < 0.0:
|
||||
return CapacityState.None
|
||||
elif ratio < 0.3:
|
||||
return CapacityState.AlmostNone
|
||||
else:
|
||||
# Normal operation - queue all messages
|
||||
case priority:
|
||||
of Critical:
|
||||
manager.criticalQueue.addLast(queuedMsg)
|
||||
of Normal:
|
||||
manager.normalQueue.addLast(queuedMsg)
|
||||
of Optional:
|
||||
manager.optionalQueue.addLast(queuedMsg)
|
||||
return CapacityState.Normal
|
||||
|
||||
proc getNextMessage[T: Serializable](manager: RateLimitManager[T]): QueuedMessage[T] =
|
||||
# Priority order: Critical -> Normal -> Optional
|
||||
if manager.criticalQueue.len > 0:
|
||||
return manager.criticalQueue.popFirst()
|
||||
elif manager.normalQueue.len > 0:
|
||||
return manager.normalQueue.popFirst()
|
||||
elif manager.optionalQueue.len > 0:
|
||||
return manager.optionalQueue.popFirst()
|
||||
proc passToSender[T: Serializable](
|
||||
manager: RateLimitManager[T],
|
||||
msgs: seq[tuple[msgId: string, msg: T]],
|
||||
now: Moment,
|
||||
priority: Priority,
|
||||
): Future[SendResult] {.async.} =
|
||||
let count = msgs.len
|
||||
let capacity = manager.bucket.tryConsume(count, now)
|
||||
if not capacity:
|
||||
case priority
|
||||
of Priority.Critical:
|
||||
manager.queueCritical.addLast(msgs)
|
||||
return SendResult.Enqueued
|
||||
of Priority.Normal:
|
||||
manager.queueNormal.addLast(msgs)
|
||||
return SendResult.Enqueued
|
||||
of Priority.Optional:
|
||||
return SendResult.Dropped
|
||||
await manager.sender(msgs)
|
||||
return SendResult.PassedToSender
|
||||
|
||||
proc processCriticalQueue[T: Serializable](
|
||||
manager: RateLimitManager[T], now: Moment
|
||||
): Future[void] {.async.} =
|
||||
while manager.queueCritical.len > 0:
|
||||
let msgs = manager.queueCritical.popFirst()
|
||||
let capacityState = manager.getCapacityState(now, msgs.len)
|
||||
if capacityState == CapacityState.Normal:
|
||||
discard await manager.passToSender(msgs, now, Priority.Critical)
|
||||
elif capacityState == CapacityState.AlmostNone:
|
||||
discard await manager.passToSender(msgs, now, Priority.Critical)
|
||||
else:
|
||||
raise newException(ValueError, "No messages in queue")
|
||||
|
||||
proc hasMessages[T: Serializable](manager: RateLimitManager[T]): bool =
|
||||
manager.criticalQueue.len > 0 or manager.normalQueue.len > 0 or manager.optionalQueue.len > 0
|
||||
|
||||
proc sendLoop*[T: Serializable](manager: RateLimitManager[T]): Future[void] {.async.} =
|
||||
manager.isRunning = true
|
||||
|
||||
while manager.isRunning:
|
||||
try:
|
||||
manager.updateEpochIfNeeded()
|
||||
|
||||
while manager.hasMessages() and manager.currentCount < manager.messageCount:
|
||||
let msg = manager.getNextMessage()
|
||||
|
||||
try:
|
||||
let sent = await manager.sender(msg.messageId, msg.msg)
|
||||
if sent:
|
||||
manager.currentCount += 1
|
||||
else:
|
||||
# Re-queue failed message at beginning of appropriate queue
|
||||
case msg.priority:
|
||||
of Critical:
|
||||
manager.criticalQueue.addFirst(msg)
|
||||
of Normal:
|
||||
manager.normalQueue.addFirst(msg)
|
||||
of Optional:
|
||||
manager.optionalQueue.addFirst(msg)
|
||||
break # Stop trying to send more messages if one fails
|
||||
except:
|
||||
# Re-queue on exception
|
||||
case msg.priority:
|
||||
of Critical:
|
||||
manager.criticalQueue.addFirst(msg)
|
||||
of Normal:
|
||||
manager.normalQueue.addFirst(msg)
|
||||
of Optional:
|
||||
manager.optionalQueue.addFirst(msg)
|
||||
# add back to critical queue
|
||||
manager.queueCritical.addFirst(msgs)
|
||||
break
|
||||
|
||||
# Small delay between messages TODO not sure if needed
|
||||
await sleepAsync(chronos.milliseconds(10))
|
||||
|
||||
# Wait before next processing cycle TODO not completely sure if this is the way
|
||||
await sleepAsync(chronos.milliseconds(100))
|
||||
except CancelledError:
|
||||
proc processNormalQueue[T: Serializable](
|
||||
manager: RateLimitManager[T], now: Moment
|
||||
): Future[void] {.async.} =
|
||||
while manager.queueNormal.len > 0:
|
||||
let msgs = manager.queueNormal.popFirst()
|
||||
let capacityState = manager.getCapacityState(now, msgs.len)
|
||||
if capacityState == CapacityState.Normal:
|
||||
discard await manager.passToSender(msgs, now, Priority.Normal)
|
||||
else:
|
||||
# add back to critical queue
|
||||
manager.queueNormal.addFirst(msgs)
|
||||
break
|
||||
except:
|
||||
await sleepAsync(chronos.seconds(1)) # Wait longer on error
|
||||
|
||||
proc start*[T: Serializable](manager: RateLimitManager[T]): Future[void] {.async.} =
|
||||
if not manager.isRunning:
|
||||
manager.sendTask = manager.sendLoop()
|
||||
proc sendOrEnqueue*[T: Serializable](
|
||||
manager: RateLimitManager[T],
|
||||
msgs: seq[tuple[msgId: string, msg: T]],
|
||||
priority: Priority,
|
||||
now: Moment = Moment.now(),
|
||||
): Future[SendResult] {.async.} =
|
||||
let (_, budgetCap) = manager.bucket.getAvailableCapacity(now)
|
||||
if msgs.len.float / budgetCap.float >= 0.3:
|
||||
# drop batch if it's too large to avoid starvation
|
||||
return SendResult.DroppedBatchTooLarge
|
||||
|
||||
proc stop*[T: Serializable](manager: RateLimitManager[T]): Future[void] {.async.} =
|
||||
if manager.isRunning:
|
||||
manager.isRunning = false
|
||||
if not manager.sendTask.isNil:
|
||||
manager.sendTask.cancelSoon()
|
||||
let capacityState = manager.getCapacityState(now, msgs.len)
|
||||
case capacityState
|
||||
of CapacityState.Normal:
|
||||
return await manager.passToSender(msgs, now, priority)
|
||||
of CapacityState.AlmostNone:
|
||||
case priority
|
||||
of Priority.Critical:
|
||||
return await manager.passToSender(msgs, now, priority)
|
||||
of Priority.Normal:
|
||||
manager.queueNormal.addLast(msgs)
|
||||
return SendResult.Enqueued
|
||||
of Priority.Optional:
|
||||
return SendResult.Dropped
|
||||
of CapacityState.None:
|
||||
case priority
|
||||
of Priority.Critical:
|
||||
manager.queueCritical.addLast(msgs)
|
||||
return SendResult.Enqueued
|
||||
of Priority.Normal:
|
||||
manager.queueNormal.addLast(msgs)
|
||||
return SendResult.Enqueued
|
||||
of Priority.Optional:
|
||||
return SendResult.Dropped
|
||||
|
||||
proc getEnqueued*[T: Serializable](
|
||||
manager: RateLimitManager[T]
|
||||
): tuple[
|
||||
critical: seq[tuple[msgId: string, msg: T]], normal: seq[tuple[msgId: string, msg: T]]
|
||||
] =
|
||||
var criticalMsgs: seq[tuple[msgId: string, msg: T]]
|
||||
var normalMsgs: seq[tuple[msgId: string, msg: T]]
|
||||
|
||||
for batch in manager.queueCritical:
|
||||
criticalMsgs.add(batch)
|
||||
|
||||
for batch in manager.queueNormal:
|
||||
normalMsgs.add(batch)
|
||||
|
||||
return (criticalMsgs, normalMsgs)
|
||||
|
||||
proc start*[T: Serializable](
|
||||
manager: RateLimitManager[T],
|
||||
nowProvider: proc(): Moment {.gcsafe.} = proc(): Moment {.gcsafe.} =
|
||||
Moment.now(),
|
||||
): Future[void] {.async.} =
|
||||
manager.running = true
|
||||
while manager.running:
|
||||
try:
|
||||
await manager.sendTask
|
||||
except CancelledError:
|
||||
discard
|
||||
let now = nowProvider()
|
||||
await manager.processCriticalQueue(now)
|
||||
await manager.processNormalQueue(now)
|
||||
|
||||
proc getQueueStatus*[T: Serializable](manager: RateLimitManager[T]): tuple[critical: int, normal: int, optional: int, total: int] =
|
||||
(
|
||||
critical: manager.criticalQueue.len,
|
||||
normal: manager.normalQueue.len,
|
||||
optional: manager.optionalQueue.len,
|
||||
total: manager.criticalQueue.len + manager.normalQueue.len + manager.optionalQueue.len
|
||||
)
|
||||
# configurable sleep duration for processing queued messages
|
||||
await sleepAsync(manager.sleepDuration)
|
||||
except Exception as e:
|
||||
echo "Error in queue processing: ", e.msg
|
||||
await sleepAsync(manager.sleepDuration)
|
||||
|
||||
proc getCurrentQuota*[T: Serializable](manager: RateLimitManager[T]): tuple[total: int, used: int, remaining: int] =
|
||||
(
|
||||
total: manager.messageCount,
|
||||
used: manager.currentCount,
|
||||
remaining: max(0, manager.messageCount - manager.currentCount)
|
||||
)
|
||||
proc stop*[T: Serializable](manager: RateLimitManager[T]) =
|
||||
manager.running = false
|
||||
|
||||
func `$`*[T: Serializable](b: RateLimitManager[T]): string {.inline.} =
|
||||
if isNil(b):
|
||||
return "nil"
|
||||
return
|
||||
"RateLimitManager(running: " & $b.running & ", critical: " & $b.queueCritical.len &
|
||||
", normal: " & $b.queueNormal.len & ")"
|
||||
|
||||
@ -1,244 +1,311 @@
|
||||
{.used.}
|
||||
|
||||
import std/[sequtils, times, random], testutils/unittests
|
||||
import testutils/unittests
|
||||
import ../ratelimit/ratelimit
|
||||
import chronos
|
||||
import strutils
|
||||
|
||||
randomize()
|
||||
|
||||
# Test message types
|
||||
type
|
||||
TestMessage = object
|
||||
content: string
|
||||
id: int
|
||||
|
||||
# Implement Serializable for test types
|
||||
# Implement the Serializable concept for string
|
||||
proc toBytes*(s: string): seq[byte] =
|
||||
cast[seq[byte]](s)
|
||||
|
||||
proc toBytes*(msg: TestMessage): seq[byte] =
|
||||
result = toBytes(msg.content)
|
||||
result.add(cast[seq[byte]](@[byte(msg.id)]))
|
||||
|
||||
proc toBytes*(i: int): seq[byte] =
|
||||
cast[seq[byte]](@[byte(i)])
|
||||
|
||||
suite "Rate Limit Manager":
|
||||
suite "Queue RateLimitManager":
|
||||
setup:
|
||||
## Given
|
||||
let epochDuration = 60
|
||||
var sentMessages: seq[string]
|
||||
var sentMessages: seq[tuple[msgId: string, msg: string]]
|
||||
var senderCallCount: int = 0
|
||||
|
||||
proc testSender(messageId: string, msg: string): Future[bool] {.async.} =
|
||||
# Create a mock sender
|
||||
proc mockSender(
|
||||
msgs: seq[tuple[msgId: string, msg: string]]
|
||||
): Future[void] {.async.} =
|
||||
senderCallCount.inc()
|
||||
for msg in msgs:
|
||||
sentMessages.add(msg)
|
||||
await sleepAsync(chronos.milliseconds(10))
|
||||
return true
|
||||
|
||||
asyncTest "basic message queueing and sending":
|
||||
asyncTest "sendOrEnqueue - immediate send when capacity available":
|
||||
## Given
|
||||
let capacity = 10
|
||||
var manager = newRateLimitManager(capacity, epochDuration, testSender)
|
||||
await manager.start()
|
||||
let manager = newRateLimitManager[string](
|
||||
mockSender, capacity = 10, duration = chronos.milliseconds(100)
|
||||
)
|
||||
let testMsg = "Hello World"
|
||||
|
||||
## When
|
||||
manager.queueForSend("msg1", "Hello", Critical)
|
||||
manager.queueForSend("msg2", "World", Normal)
|
||||
|
||||
await sleepAsync(chronos.milliseconds(200))
|
||||
await manager.stop()
|
||||
let res = await manager.sendOrEnqueue(@[("msg1", testMsg)], Critical)
|
||||
|
||||
## Then
|
||||
check:
|
||||
res == PassedToSender
|
||||
senderCallCount == 1
|
||||
sentMessages.len == 1
|
||||
sentMessages[0].msgId == "msg1"
|
||||
sentMessages[0].msg == "Hello World"
|
||||
|
||||
asyncTest "sendOrEnqueue - multiple messages":
|
||||
## Given
|
||||
let manager = newRateLimitManager[string](
|
||||
mockSender, capacity = 10, duration = chronos.milliseconds(100)
|
||||
)
|
||||
|
||||
## When
|
||||
let res =
|
||||
await manager.sendOrEnqueue(@[("msg1", "First"), ("msg2", "Second")], Normal)
|
||||
|
||||
## Then
|
||||
check:
|
||||
res == PassedToSender
|
||||
senderCallCount == 1
|
||||
sentMessages.len == 2
|
||||
sentMessages[0] == "Hello"
|
||||
sentMessages[1] == "World"
|
||||
sentMessages[0].msgId == "msg1"
|
||||
sentMessages[0].msg == "First"
|
||||
sentMessages[1].msgId == "msg2"
|
||||
sentMessages[1].msg == "Second"
|
||||
|
||||
asyncTest "priority ordering - critical first, then normal, then optional":
|
||||
asyncTest "start and stop - basic functionality":
|
||||
## Given
|
||||
let capacity = 10
|
||||
var manager = newRateLimitManager(capacity, epochDuration, testSender)
|
||||
await manager.start()
|
||||
let manager = newRateLimitManager[string](
|
||||
mockSender,
|
||||
capacity = 10,
|
||||
duration = chronos.milliseconds(100),
|
||||
sleepDuration = chronos.milliseconds(50),
|
||||
)
|
||||
|
||||
## When - queue messages in mixed priority order
|
||||
manager.queueForSend("normal1", "Normal1", Normal)
|
||||
manager.queueForSend("critical1", "Critical1", Critical)
|
||||
manager.queueForSend("optional1", "Optional1", Optional)
|
||||
manager.queueForSend("critical2", "Critical2", Critical)
|
||||
manager.queueForSend("normal2", "Normal2", Normal)
|
||||
## When - start the manager
|
||||
let startFut = manager.start()
|
||||
|
||||
await sleepAsync(chronos.milliseconds(300))
|
||||
await manager.stop()
|
||||
# Give it some time to start
|
||||
await sleepAsync(chronos.milliseconds(20))
|
||||
|
||||
## Then - critical messages should be sent first
|
||||
## Then - manager should be running (check via string representation)
|
||||
check:
|
||||
sentMessages.len == 5
|
||||
sentMessages[0] == "Critical1" # First critical
|
||||
sentMessages[1] == "Critical2" # Second critical
|
||||
sentMessages[2] == "Normal1" # Then normal messages
|
||||
sentMessages[3] == "Normal2"
|
||||
sentMessages[4] == "Optional1" # Finally optional
|
||||
"running: true" in $manager
|
||||
not startFut.finished() # should still be running
|
||||
|
||||
asyncTest "rate limiting - respects message quota per epoch":
|
||||
## When - stop the manager
|
||||
manager.stop()
|
||||
|
||||
# Give it time to stop (should happen within one sleep cycle)
|
||||
await sleepAsync(chronos.milliseconds(60))
|
||||
|
||||
## Then - manager should be stopped
|
||||
check:
|
||||
"running: false" in $manager
|
||||
startFut.finished() # should have completed
|
||||
|
||||
asyncTest "start and stop - drop large batch":
|
||||
## Given
|
||||
let capacity = 3 # Only 3 messages allowed
|
||||
var manager = newRateLimitManager(capacity, epochDuration, testSender)
|
||||
await manager.start()
|
||||
let manager = newRateLimitManager[string](
|
||||
mockSender,
|
||||
capacity = 2,
|
||||
duration = chronos.milliseconds(100),
|
||||
sleepDuration = chronos.milliseconds(30),
|
||||
)
|
||||
|
||||
## When - queue more messages than the limit
|
||||
for i in 1..5:
|
||||
manager.queueForSend("msg" & $i, "Message" & $i, Normal)
|
||||
let largeBatch1 = @[("msg1", "First"), ("msg2", "Second"), ("msg3", "Third")]
|
||||
let largeBatch2 = @[("msg4", "Fourth"), ("msg5", "Fifth")]
|
||||
|
||||
await sleepAsync(chronos.milliseconds(300))
|
||||
discard await manager.sendOrEnqueue(largeBatch1, Normal)
|
||||
discard await manager.sendOrEnqueue(largeBatch2, Critical)
|
||||
|
||||
## Then
|
||||
let quota = manager.getCurrentQuota()
|
||||
let queueStatus = manager.getQueueStatus()
|
||||
asyncTest "enqueue - enqueue critical only when exceeded":
|
||||
## Given
|
||||
let manager = newRateLimitManager[string](
|
||||
mockSender,
|
||||
capacity = 10,
|
||||
duration = chronos.milliseconds(100),
|
||||
sleepDuration = chronos.milliseconds(100),
|
||||
)
|
||||
var now = Moment.now()
|
||||
let startFut = manager.start(
|
||||
nowProvider = proc(): Moment =
|
||||
now
|
||||
)
|
||||
|
||||
let r1 = await manager.sendOrEnqueue(@[("msg1", "1")], Critical, now)
|
||||
let r2 = await manager.sendOrEnqueue(@[("msg2", "2")], Critical, now)
|
||||
let r3 = await manager.sendOrEnqueue(@[("msg3", "3")], Critical, now)
|
||||
let r4 = await manager.sendOrEnqueue(@[("msg4", "4")], Critical, now)
|
||||
let r5 = await manager.sendOrEnqueue(@[("msg5", "5")], Critical, now)
|
||||
let r6 = await manager.sendOrEnqueue(@[("msg6", "6")], Critical, now)
|
||||
let r7 = await manager.sendOrEnqueue(@[("msg7", "7")], Critical, now)
|
||||
let r8 = await manager.sendOrEnqueue(@[("msg8", "8")], Critical, now)
|
||||
let r9 = await manager.sendOrEnqueue(@[("msg9", "9")], Critical, now)
|
||||
let r10 = await manager.sendOrEnqueue(@[("msg10", "10")], Critical, now)
|
||||
# will enqueue to critical queue
|
||||
let r11 = await manager.sendOrEnqueue(@[("msg11", "11")], Critical, now)
|
||||
|
||||
check:
|
||||
quota.used == 3
|
||||
quota.remaining == 0
|
||||
sentMessages.len == 3
|
||||
queueStatus.total == 2 # 2 messages should be queued
|
||||
r1 == PassedToSender
|
||||
r2 == PassedToSender
|
||||
r3 == PassedToSender
|
||||
r4 == PassedToSender
|
||||
r5 == PassedToSender
|
||||
r6 == PassedToSender
|
||||
r7 == PassedToSender
|
||||
r8 == PassedToSender
|
||||
r9 == PassedToSender
|
||||
r10 == PassedToSender
|
||||
r11 == Enqueued
|
||||
|
||||
await manager.stop()
|
||||
let (critical, normal) = manager.getEnqueued()
|
||||
check:
|
||||
critical.len == 1
|
||||
normal.len == 0
|
||||
critical[0].msgId == "msg11"
|
||||
|
||||
asyncTest "optional message dropping at high usage":
|
||||
# Give it time to stop
|
||||
manager.stop()
|
||||
await sleepAsync(chronos.milliseconds(150))
|
||||
## Then - should be stopped
|
||||
check:
|
||||
"running: false" in $manager
|
||||
startFut.finished()
|
||||
|
||||
asyncTest "enqueue - enqueue normal on 70% capacity":
|
||||
## Given
|
||||
let capacity = 10
|
||||
var manager = newRateLimitManager(capacity, epochDuration, testSender)
|
||||
await manager.start()
|
||||
let manager = newRateLimitManager[string](
|
||||
mockSender,
|
||||
capacity = 10,
|
||||
duration = chronos.milliseconds(100),
|
||||
sleepDuration = chronos.milliseconds(100),
|
||||
)
|
||||
var now = Moment.now()
|
||||
let startFut = manager.start(
|
||||
nowProvider = proc(): Moment =
|
||||
now
|
||||
)
|
||||
|
||||
# Fill to 80% capacity to trigger optional dropping
|
||||
for i in 1..8:
|
||||
manager.queueForSend("fill" & $i, "Fill" & $i, Normal)
|
||||
|
||||
await sleepAsync(chronos.milliseconds(200))
|
||||
|
||||
## When - add messages at high usage
|
||||
manager.queueForSend("critical", "Critical", Critical)
|
||||
manager.queueForSend("normal", "Normal", Normal)
|
||||
manager.queueForSend("optional", "Optional", Optional) # Should be dropped
|
||||
|
||||
await sleepAsync(chronos.milliseconds(200))
|
||||
|
||||
## Then
|
||||
let quota = manager.getCurrentQuota()
|
||||
let optionalSent = sentMessages.anyIt(it == "Optional")
|
||||
let r1 = await manager.sendOrEnqueue(@[("msg1", "1")], Normal, now)
|
||||
let r2 = await manager.sendOrEnqueue(@[("msg2", "2")], Normal, now)
|
||||
let r3 = await manager.sendOrEnqueue(@[("msg3", "3")], Normal, now)
|
||||
let r4 = await manager.sendOrEnqueue(@[("msg4", "4")], Normal, now)
|
||||
let r5 = await manager.sendOrEnqueue(@[("msg5", "5")], Normal, now)
|
||||
let r6 = await manager.sendOrEnqueue(@[("msg6", "6")], Normal, now)
|
||||
let r7 = await manager.sendOrEnqueue(@[("msg7", "7")], Normal, now)
|
||||
let r8 = await manager.sendOrEnqueue(@[("msg8", "8")], Normal, now)
|
||||
let r9 = await manager.sendOrEnqueue(@[("msg9", "9")], Normal, now)
|
||||
let r10 = await manager.sendOrEnqueue(@[("msg10", "10")], Normal, now)
|
||||
let r11 = await manager.sendOrEnqueue(@[("msg11", "11")], Critical, now)
|
||||
let r12 = await manager.sendOrEnqueue(@[("msg12", "12")], Optional, now)
|
||||
|
||||
check:
|
||||
quota.used == 10
|
||||
not optionalSent # Optional message should not be sent
|
||||
r1 == PassedToSender
|
||||
r2 == PassedToSender
|
||||
r3 == PassedToSender
|
||||
r4 == PassedToSender
|
||||
r5 == PassedToSender
|
||||
r6 == PassedToSender
|
||||
r7 == PassedToSender
|
||||
r8 == Enqueued
|
||||
r9 == Enqueued
|
||||
r10 == Enqueued
|
||||
r11 == PassedToSender
|
||||
r12 == Dropped
|
||||
|
||||
await manager.stop()
|
||||
|
||||
asyncTest "quota tracking - accurate total, used, and remaining counts":
|
||||
## Given
|
||||
let capacity = 5
|
||||
var manager = newRateLimitManager(capacity, epochDuration, testSender)
|
||||
await manager.start()
|
||||
|
||||
## When - initially no messages sent
|
||||
let initialQuota = manager.getCurrentQuota()
|
||||
|
||||
## Then - should show full quota available
|
||||
let (critical, normal) = manager.getEnqueued()
|
||||
check:
|
||||
initialQuota.total == 5
|
||||
initialQuota.used == 0
|
||||
initialQuota.remaining == 5
|
||||
critical.len == 0
|
||||
normal.len == 3
|
||||
normal[0].msgId == "msg8"
|
||||
normal[1].msgId == "msg9"
|
||||
normal[2].msgId == "msg10"
|
||||
|
||||
## When - send some messages
|
||||
manager.queueForSend("msg1", "Message1", Normal)
|
||||
manager.queueForSend("msg2", "Message2", Normal)
|
||||
|
||||
await sleepAsync(chronos.milliseconds(200))
|
||||
|
||||
## Then - quota should be updated
|
||||
let afterQuota = manager.getCurrentQuota()
|
||||
# Give it time to stop
|
||||
manager.stop()
|
||||
await sleepAsync(chronos.milliseconds(150))
|
||||
## Then - should be stopped
|
||||
check:
|
||||
afterQuota.used == 2
|
||||
afterQuota.remaining == 3
|
||||
"running: false" in $manager
|
||||
startFut.finished()
|
||||
|
||||
await manager.stop()
|
||||
|
||||
asyncTest "queue status tracking - by priority levels":
|
||||
asyncTest "enqueue - process queued messages":
|
||||
## Given
|
||||
let capacity = 2 # Small limit to force queueing
|
||||
var manager = newRateLimitManager(capacity, epochDuration, testSender)
|
||||
await manager.start()
|
||||
let manager = newRateLimitManager[string](
|
||||
mockSender,
|
||||
capacity = 10,
|
||||
duration = chronos.milliseconds(200),
|
||||
sleepDuration = chronos.milliseconds(200),
|
||||
)
|
||||
type MomentRef = ref object
|
||||
value: Moment
|
||||
|
||||
## When - fill quota and add more messages
|
||||
manager.queueForSend("msg1", "Message1", Critical)
|
||||
manager.queueForSend("msg2", "Message2", Normal)
|
||||
manager.queueForSend("msg3", "Message3", Critical) # Should be queued
|
||||
manager.queueForSend("msg4", "Message4", Normal) # Should be queued
|
||||
manager.queueForSend("msg5", "Message5", Optional) # Should be queued
|
||||
var now = Moment.now()
|
||||
var nowRef = MomentRef(value: now)
|
||||
let startFut = manager.start(
|
||||
nowProvider = proc(): Moment =
|
||||
return nowRef.value
|
||||
)
|
||||
|
||||
await sleepAsync(chronos.milliseconds(100))
|
||||
|
||||
## Then
|
||||
let queueStatus = manager.getQueueStatus()
|
||||
let r1 = await manager.sendOrEnqueue(@[("1", "val_1")], Normal, now)
|
||||
let r2 = await manager.sendOrEnqueue(@[("2", "val_2")], Normal, now)
|
||||
let r3 = await manager.sendOrEnqueue(@[("3", "val_3")], Normal, now)
|
||||
let r4 = await manager.sendOrEnqueue(@[("4", "val_4")], Normal, now)
|
||||
let r5 = await manager.sendOrEnqueue(@[("5", "val_5")], Normal, now)
|
||||
let r6 = await manager.sendOrEnqueue(@[("6", "val_6")], Normal, now)
|
||||
let r7 = await manager.sendOrEnqueue(@[("7", "val_7")], Normal, now)
|
||||
let r8 = await manager.sendOrEnqueue(@[("8", "val_8")], Normal, now)
|
||||
let r9 = await manager.sendOrEnqueue(@[("9", "val_9")], Normal, now)
|
||||
let r10 = await manager.sendOrEnqueue(@[("10", "val_10")], Normal, now)
|
||||
let r11 = await manager.sendOrEnqueue(@[("11", "val_11")], Critical, now)
|
||||
let r12 = await manager.sendOrEnqueue(@[("12", "val_12")], Optional, now)
|
||||
let r13 = await manager.sendOrEnqueue(@[("13", "val_13")], Critical, now)
|
||||
let r14 = await manager.sendOrEnqueue(@[("14", "val_14")], Critical, now)
|
||||
let r15 = await manager.sendOrEnqueue(@[("15", "val_15")], Critical, now)
|
||||
|
||||
check:
|
||||
queueStatus.total >= 2 # At least some messages queued
|
||||
queueStatus.critical >= 1 # Critical messages in queue
|
||||
queueStatus.normal >= 1 # Normal messages in queue
|
||||
r1 == PassedToSender
|
||||
r2 == PassedToSender
|
||||
r3 == PassedToSender
|
||||
r4 == PassedToSender
|
||||
r5 == PassedToSender
|
||||
r6 == PassedToSender
|
||||
r7 == PassedToSender
|
||||
r8 == Enqueued
|
||||
r9 == Enqueued
|
||||
r10 == Enqueued
|
||||
r11 == PassedToSender
|
||||
r12 == Dropped
|
||||
r13 == PassedToSender
|
||||
r14 == PassedToSender
|
||||
r15 == Enqueued
|
||||
|
||||
await manager.stop()
|
||||
|
||||
suite "Generic Type Support":
|
||||
setup:
|
||||
## Given
|
||||
let epochDuration = 60
|
||||
var sentCustomMessages: seq[TestMessage]
|
||||
|
||||
proc testCustomSender(messageId: string, msg: TestMessage): Future[bool] {.async.} =
|
||||
sentCustomMessages.add(msg)
|
||||
await sleepAsync(chronos.milliseconds(10))
|
||||
return true
|
||||
|
||||
asyncTest "custom message types - TestMessage objects":
|
||||
## Given
|
||||
let capacity = 5
|
||||
var manager = newRateLimitManager(capacity, epochDuration, testCustomSender)
|
||||
await manager.start()
|
||||
|
||||
## When
|
||||
let testMsg = TestMessage(content: "Test content", id: 42)
|
||||
manager.queueForSend("custom1", testMsg, Normal)
|
||||
|
||||
await sleepAsync(chronos.milliseconds(200))
|
||||
await manager.stop()
|
||||
|
||||
## Then
|
||||
var (critical, normal) = manager.getEnqueued()
|
||||
check:
|
||||
sentCustomMessages.len == 1
|
||||
sentCustomMessages[0].content == "Test content"
|
||||
sentCustomMessages[0].id == 42
|
||||
critical.len == 1
|
||||
normal.len == 3
|
||||
normal[0].msgId == "8"
|
||||
normal[1].msgId == "9"
|
||||
normal[2].msgId == "10"
|
||||
critical[0].msgId == "15"
|
||||
|
||||
asyncTest "integer message types":
|
||||
## Given
|
||||
let capacity = 5
|
||||
var sentInts: seq[int]
|
||||
nowRef.value = now + chronos.milliseconds(250)
|
||||
await sleepAsync(chronos.milliseconds(250))
|
||||
|
||||
proc testIntSender(messageId: string, msg: int): Future[bool] {.async.} =
|
||||
sentInts.add(msg)
|
||||
await sleepAsync(chronos.milliseconds(10))
|
||||
return true
|
||||
|
||||
var manager = newRateLimitManager(capacity, epochDuration, testIntSender)
|
||||
await manager.start()
|
||||
|
||||
## When
|
||||
manager.queueForSend("int1", 42, Critical)
|
||||
manager.queueForSend("int2", 100, Normal)
|
||||
manager.queueForSend("int3", 999, Optional)
|
||||
|
||||
await sleepAsync(chronos.milliseconds(200))
|
||||
await manager.stop()
|
||||
|
||||
## Then
|
||||
(critical, normal) = manager.getEnqueued()
|
||||
check:
|
||||
sentInts.len == 3
|
||||
sentInts[0] == 42 # Critical sent first
|
||||
sentInts[1] == 100 # Normal sent second
|
||||
sentInts[2] == 999 # Optional sent last
|
||||
critical.len == 0
|
||||
normal.len == 0
|
||||
senderCallCount == 14
|
||||
sentMessages.len == 14
|
||||
sentMessages[0].msgId == "1"
|
||||
sentMessages[1].msgId == "2"
|
||||
sentMessages[2].msgId == "3"
|
||||
sentMessages[3].msgId == "4"
|
||||
sentMessages[4].msgId == "5"
|
||||
sentMessages[5].msgId == "6"
|
||||
sentMessages[6].msgId == "7"
|
||||
sentMessages[7].msgId == "11"
|
||||
sentMessages[8].msgId == "13"
|
||||
sentMessages[9].msgId == "14"
|
||||
sentMessages[10].msgId == "15"
|
||||
sentMessages[11].msgId == "8"
|
||||
sentMessages[12].msgId == "9"
|
||||
sentMessages[13].msgId == "10"
|
||||
|
||||
# Give it time to stop
|
||||
manager.stop()
|
||||
await sleepAsync(chronos.milliseconds(250))
|
||||
## Then - should be stopped
|
||||
check:
|
||||
"running: false" in $manager
|
||||
startFut.finished()
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user