Merge pull request #10

FIX: reset and refill tokens in distributed mode

Author: mehrantsi

src/native/nats_storage.hpp

@@ -108,7 +108,7 @@ class NatsStorage : public DistributedStorage {
 
         kvConf.bucket = bucket.c_str();
         kvConf.history = 1;
-        kvConf.ttl = 3600000000000;  // 1 hour TTL in nanoseconds (3600 * 1e9)
+        kvConf.ttl = 86400000000000;  // 24 hours TTL in nanoseconds (86400 * 1e9)
 
         s = g_natsLoader.js_CreateKeyValue(&kv, js, &kvConf);
         if (s != NATS_OK) {
@@ -133,7 +133,7 @@ class NatsStorage : public DistributedStorage {
         if (nc) g_natsLoader.natsConnection_Destroy(nc);
     }
 
-    bool tryAcquire(const std::string& key, int64_t maxTokens) override {
+    bool tryAcquire(const std::string& key, int64_t tokens) override {
         if (!kv) return false; // Safety check
 
         // NATS JetStream KV Store does not allow colons in key names
@@ -148,7 +148,7 @@ class NatsStorage : public DistributedStorage {
 
         if (s == NATS_NOT_FOUND) {
             // Key doesn't exist, initialize it with maxTokens, then decrement by 1
-            std::string value = std::to_string(maxTokens);
+            std::string value = std::to_string(tokens);
             uint64_t rev;
             s = g_natsLoader.kvStore_CreateString(&rev, kv, fullKey.c_str(), value.c_str());
 
@@ -157,8 +157,8 @@ class NatsStorage : public DistributedStorage {
             }
 
             // Now decrement by 1 (acquiring 1 token)
-            if (maxTokens > 0) {
-                std::string newValue = std::to_string(maxTokens - 1);
+            if (tokens > 0) {
+                std::string newValue = std::to_string(tokens - 1);
                 uint64_t newRev;
                 s = g_natsLoader.kvStore_UpdateString(&newRev, kv, fullKey.c_str(), newValue.c_str(), rev);
                 return s == NATS_OK;
@@ -245,4 +245,18 @@ class NatsStorage : public DistributedStorage {
         uint64_t newRev;
         g_natsLoader.kvStore_UpdateString(&newRev, kv, fullKey.c_str(), newValue.c_str(), revision);
     }
+    
+    void reset(const std::string& key, int64_t maxTokens) override {
+        if (!kv) return; // Safety check
+        
+        // NATS JetStream KV Store does not allow colons in key names
+        std::string sanitizedKey = prefix + key;
+        std::replace(sanitizedKey.begin(), sanitizedKey.end(), ':', '_');
+        const std::string fullKey = sanitizedKey;
+        
+        // Set the value to maxTokens regardless of current value
+        std::string value = std::to_string(maxTokens);
+        uint64_t rev;
+        g_natsLoader.kvStore_Put(&rev, kv, fullKey.c_str(), value.c_str(), value.length());
+    }
 };

src/native/ratelimiter.hpp

@@ -34,6 +34,7 @@ class DistributedStorage {
     virtual ~DistributedStorage() = default;
     virtual bool tryAcquire(const std::string& key, int64_t tokens) = 0;
     virtual void release(const std::string& key, int64_t tokens) = 0;
+    virtual void reset(const std::string& key, int64_t maxTokens) = 0;
 };
 
 class RateLimiter {
@@ -257,6 +258,17 @@ class RateLimiter {
                     std::memory_order_acq_rel, std::memory_order_acquire)) {
                     entry.dynamicMaxTokens.store(dynamicLimit, std::memory_order_release);
                     entry.tokens.store(newTokens, std::memory_order_release);
+                    
+                    // Sync sliding window refill with distributed storage
+                    if (distributedStorage && !entry.distributedKey.empty() && tokensToAdd > 0) {
+                        try {
+                            // Release tokens back to distributed storage (effectively adding them)
+                            distributedStorage->release(entry.distributedKey, tokensToAdd);
+                        } catch (...) {
+                            // Ignore errors - distributed storage might be temporarily unavailable
+                        }
+                    }
+                    
                     return;
                 }
             } else {
@@ -265,6 +277,16 @@ class RateLimiter {
                     std::memory_order_acq_rel, std::memory_order_acquire)) {
                     entry.dynamicMaxTokens.store(dynamicLimit, std::memory_order_release);
                     entry.tokens.store(dynamicLimit, std::memory_order_release);
+                    
+                    // Reset distributed storage for fixed window
+                    if (distributedStorage && !entry.distributedKey.empty()) {
+                        try {
+                            distributedStorage->reset(entry.distributedKey, dynamicLimit);
+                        } catch (...) {
+                            // Ignore errors - distributed storage might be temporarily unavailable
+                        }
+                    }
+                    
                     return;
                 }
             }
@@ -512,7 +534,7 @@ class RateLimiter {
         // If we have distributed storage and a distributed key is set, check it first
         if (distributedStorage && !entry->distributedKey.empty()) {
             try {
-                if (!distributedStorage->tryAcquire(entry->distributedKey, entry->baseMaxTokens)) {
+                if (!distributedStorage->tryAcquire(entry->distributedKey, entry->dynamicMaxTokens.load(std::memory_order_acquire))) {
                     metrics.blockedRequests.fetch_add(1, std::memory_order_relaxed);
                     return false;
                 }

src/native/redis_storage.hpp

@@ -80,4 +80,19 @@ class RedisStorage : public DistributedStorage {
 
         g_redisLoader.freeReplyObject(reply);
     }
+    
+    void reset(const std::string& key, int64_t maxTokens) override {
+        std::string fullKey = prefix + key;
+        
+        // SET key maxTokens
+        redisReply* reply = (redisReply*)g_redisLoader.redisCommand(redis,
+            "SET %s %lld",
+            fullKey.c_str(), maxTokens);
+
+        if (!reply) {
+            throw std::runtime_error("Redis command failed");
+        }
+
+        g_redisLoader.freeReplyObject(reply);
+    }
 };