Timeouts

(cherry picked from commit 52ae87ca911fb7f5794e64ab7213c5b791ce085d)

JetHerald/Configs.cs

@@ -19,4 +19,11 @@
     {
         public string Token { get; set; }
     }
+
+    public class Timeout
+    {
+        public int DebtLimitSeconds { get; set; }
+        public int HeartbeatCost { get; set; }
+        public int ReportCost { get; set; }
+    }
 }

JetHerald/Controllers/HeartbeatController.cs

@@ -1,12 +1,10 @@
 using System;
-using System.IO;
-using System.Runtime.Serialization;
-using System.Text;
 using System.Text.Json;
 using System.Text.Json.Serialization;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Mvc;
-using Microsoft.AspNetCore.Mvc.ModelBinding;
+using Microsoft.Extensions.Options;
 
 namespace JetHerald.Controllers
 {
@@ -15,11 +13,15 @@ namespace JetHerald.Controllers
     {
         Db Db { get; }
         JetHeraldBot Herald { get; }
+        LeakyBucket Timeouts { get; }
+        Options.Timeout Config { get; }
 
-        public HeartbeatController(Db db, JetHeraldBot herald)
+        public HeartbeatController(Db db, JetHeraldBot herald, LeakyBucket timeouts, IOptions<Options.Timeout> cfgOptions)
         {
             Herald = herald;
+            Timeouts = timeouts;
             Db = db;
+            Config = cfgOptions.Value;
         }
 
         // tfw when you need to manually parse body and query
@@ -71,12 +73,16 @@ namespace JetHerald.Controllers
             else if (!t.WriteToken.Equals(args.WriteToken, StringComparison.Ordinal))
                 return StatusCode(403);
 
+            if (Timeouts.IsTimedOut(t.TopicId))
+                return StatusCode(StatusCodes.Status429TooManyRequests);
 
             var affected = await Db.ReportHeartbeat(t.TopicId, heart, args.ExpiryTimeout);
 
             if (affected == 1)
                 await Herald.HeartbeatSent(t, heart);
 
+            Timeouts.ApplyCost(t.TopicId, Config.HeartbeatCost);
+
             return new OkResult();
         }
 

JetHerald/Controllers/ReportController.cs

@@ -1,7 +1,9 @@
 using System;
 using System.Text.Json;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Mvc;
+using Microsoft.Extensions.Options;
 
 namespace JetHerald.Controllers
 {
@@ -10,11 +12,15 @@ namespace JetHerald.Controllers
     {
         Db Db { get; }
         JetHeraldBot Herald { get; }
+        LeakyBucket Timeouts { get; }
+        Options.Timeout Config { get; }
 
-        public ReportController(Db db, JetHeraldBot herald)
+        public ReportController(Db db, JetHeraldBot herald, LeakyBucket timeouts, IOptions<Options.Timeout> cfgOptions)
         {
             Herald = herald;
+            Timeouts = timeouts;
             Db = db;
+            Config = cfgOptions.Value;
         }
 
         [Route("api/report")]
@@ -55,7 +61,13 @@ namespace JetHerald.Controllers
             else if (!t.WriteToken.Equals(args.WriteToken, StringComparison.OrdinalIgnoreCase))
                 return StatusCode(403);
 
+            if (Timeouts.IsTimedOut(t.TopicId))
+                return StatusCode(StatusCodes.Status429TooManyRequests);
+
             await Herald.PublishMessage(t, args.Message);
+
+            Timeouts.ApplyCost(t.TopicId, Config.ReportCost);
+
             return new OkResult();
         }
 

JetHerald/LeakyBucket.cs

@@ -0,0 +1,39 @@
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using Microsoft.Extensions.Options;
+
+namespace JetHerald
+{
+    public class LeakyBucket
+    {
+        private readonly ConcurrentDictionary<uint, DateTime> expiryDates = new();
+        private readonly Options.Timeout config;
+
+        public LeakyBucket(IOptions<Options.Timeout> cfgOptions)
+        {
+            config = cfgOptions.Value;
+        }
+
+        public bool IsTimedOut(uint key)
+        {
+            var debtLimit = DateTime.Now.AddSeconds(config.DebtLimitSeconds);
+            var time = expiryDates.GetValueOrDefault(key, DateTime.Now);
+            Console.WriteLine(time);
+            return time > debtLimit;
+        }
+
+        public void ApplyCost(uint key, int cost)
+        {
+            expiryDates.AddOrUpdate(key,
+                key => DateTime.Now.AddSeconds(cost),
+                (key, oldDebt) =>
+                {
+                    if (oldDebt < DateTime.Now)
+                        return DateTime.Now.AddSeconds(cost);
+
+                    return oldDebt.AddSeconds(cost);
+                });
+        }
+    }
+}

JetHerald/Startup.cs

@@ -22,8 +22,10 @@ namespace JetHerald
             services.Configure<Options.ConnectionStrings>(Configuration.GetSection("ConnectionStrings"));
             services.Configure<Options.Telegram>(Configuration.GetSection("Telegram"));
             services.Configure<Options.Discord>(Configuration.GetSection("Discord"));
+            services.Configure<Options.Timeout>(Configuration.GetSection("Timeout"));
             services.AddSingleton<Db>();
             services.AddSingleton<JetHeraldBot>();
+            services.AddSingleton<LeakyBucket>();
             services.AddMvc().SetCompatibilityVersion(CompatibilityVersion.Version_3_0);
 
         }

JetHerald/appsettings.json

@@ -7,5 +7,10 @@
   "Telegram": {
     "UseProxy": "true"
   },
+  "Timeout": {
+    "DebtLimitSeconds": 600,
+    "HeartbeatCost": 60,
+    "ReportCost": 60
+  },
   "PathBase": "/"
 }