From b96ab38fcf11da453a63dbea8910b1ea08266c46 Mon Sep 17 00:00:00 2001
From: NagyZoltanPeter <113987313+NagyZoltanPeter@users.noreply.github.com>
Date: Thu, 21 Mar 2024 16:40:24 +0100
Subject: [PATCH] Add meterics to lightpush rate limits

---
 waku/common/tokenbucket.nim          |  9 ++++++---
 waku/common/waku_service_metrics.nim | 12 ++++++++++++
 waku/waku_lightpush/protocol.nim     | 13 ++++++++-----
 3 files changed, 26 insertions(+), 8 deletions(-)
 create mode 100644 waku/common/waku_service_metrics.nim

diff --git a/waku/common/tokenbucket.nim b/waku/common/tokenbucket.nim
index 83c602e648..0eea7fc459 100644
--- a/waku/common/tokenbucket.nim
+++ b/waku/common/tokenbucket.nim
@@ -5,6 +5,10 @@ import chronos
 ## This is an extract from chronos/ratelimit.nim due to the found bug in the original implementation.
 ## Unfortunately that bug cannot be solved without harm the original features of TokenBucket class.
 ## So, this current shortcut is used to enable move ahead with nwaku rate limiter implementation.
+
+type BucketObserver* =
+  proc(consumeResult: bool, ): Future[void] {.async, closure.}
+
 type TokenBucket* = ref object
   budget*: int
   budgetCap: int
@@ -39,14 +43,13 @@ proc tryConsume*(bucket: TokenBucket, tokens: int, now = Moment.now()): bool =
     bucket.budget -= tokens
     return true
 
-  # bucket.altUpdate(now)
   bucket.update(now)
 
   if bucket.budget >= tokens:
     bucket.budget -= tokens
-    true
+    return true
   else:
-    false
+    return false
 
 proc replenish*(bucket: TokenBucket, tokens: int, now = Moment.now()) =
   ## Add `tokens` to the budget (capped to the bucket capacity)
diff --git a/waku/common/waku_service_metrics.nim b/waku/common/waku_service_metrics.nim
new file mode 100644
index 0000000000..a6dcec8549
--- /dev/null
+++ b/waku/common/waku_service_metrics.nim
@@ -0,0 +1,12 @@
+when (NimMajor, NimMinor) < (1, 4):
+  {.push raises: [Defect].}
+else:
+  {.push raises: [].}
+
+import metrics
+
+declarePublicCounter waku_service_requests,
+  "number of non-relay service requests received", ["service"]
+declarePublicCounter waku_service_requests_rejected,
+  "number of non-relay service requests received being rejected due to limit overdue",
+  ["service"]
diff --git a/waku/waku_lightpush/protocol.nim b/waku/waku_lightpush/protocol.nim
index 700ae1e9d6..b031b1e4be 100644
--- a/waku/waku_lightpush/protocol.nim
+++ b/waku/waku_lightpush/protocol.nim
@@ -12,7 +12,8 @@ import
   ./rpc,
   ./rpc_codec,
   ./protocol_metrics,
-  ../common/ratelimit
+  ../common/ratelimit,
+  ../common/waku_service_metrics
 
 export ratelimit
 
@@ -31,6 +32,7 @@ proc handleRequest*(
   let reqDecodeRes = PushRPC.decode(buffer)
   var
     isSuccess = false
+    isRejectedDueRateLimit = false
     pushResponseInfo = ""
     requestId = ""
 
@@ -39,13 +41,14 @@ proc handleRequest*(
   elif reqDecodeRes.get().request.isNone():
     pushResponseInfo = emptyRequestBodyFailure
   elif wl.requestRateLimiter.isSome() and not wl.requestRateLimiter.get().tryConsume(1):
+    isRejectedDueRateLimit = true
     let pushRpcRequest = reqDecodeRes.get()
-    debug "lightpush request rejected due rate limit exceeded",
+    trace "lightpush request rejected due rate limit exceeded",
       peerId = $peerId, requestId = pushRpcRequest.requestId
     pushResponseInfo = TooManyRequestsMessage
-    ##TODO: add metrics of above limit requests per minute or period defined
+    waku_service_requests_rejected.inc(labelValues = ["Lightpush"])
   else:
-    ##TODO: add metrics of below limit requests per minute or period defined
+    waku_service_requests.inc(labelValues = ["Lightpush"])
 
     let pushRpcRequest = reqDecodeRes.get()
 
@@ -67,7 +70,7 @@ proc handleRequest*(
     isSuccess = handleRes.isOk()
     pushResponseInfo = (if isSuccess: "OK" else: handleRes.error)
 
-  if not isSuccess:
+  if not isSuccess and not isRejectedDueRateLimit:
     waku_lightpush_errors.inc(labelValues = [pushResponseInfo])
     error "failed to push message", error = pushResponseInfo
   let response = PushResponse(isSuccess: isSuccess, info: some(pushResponseInfo))