discord.w.ratelimiter source code

1 module discord.w.ratelimiter;
2 
3 import core.time;
4 
5 mixin template RateLimitObject(Duration repeat)
6 {
7 	import std.datetime.systime;
8 	import core.sync.mutex;
9 	import vibe.core.core;
10 	import core.time;
11 
12 	__gshared Mutex lock;
13 	__gshared long access;
14 
15 	void waitFor() @trusted
16 	{
17 		if (!lock)
18 			lock = new Mutex();
19 		while (true)
20 		{
21 			long t;
22 			synchronized (lock)
23 			{
24 				auto now = Clock.currStdTime;
25 				t = now - access;
26 				if (t > repeat.total!"hnsecs")
27 				{
28 					access = now;
29 					return;
30 				}
31 			}
32 			auto wait = repeat - (t + 1000).hnsecs;
33 			if (wait.total!"hnsecs" > 0)
34 				sleep(wait);
35 		}
36 	}
37 }
38 
39 mixin template DynamicRateLimitObject(size_t limit, Duration range, Duration repeat)
40 {
41 	import std.datetime.systime;
42 	import core.sync.mutex;
43 	import vibe.core.core;
44 	import core.time;
45 
46 	__gshared Mutex lock;
47 
48 	static assert(limit > 0, "Requires a limit of at least 1");
49 
50 	__gshared size_t index;
51 	__gshared long[limit] accesses;
52 
53 	void waitFor() @trusted
54 	{
55 		if (!lock)
56 			lock = new Mutex();
57 		while (true)
58 		{
59 			Duration wait;
60 			synchronized (lock)
61 			{
62 				auto now = Clock.currStdTime;
63 				auto prevAccess = accesses[(index + limit - 1) % $];
64 				auto access = accesses[index];
65 				bool hasLeft = now - access >= range.total!"hnsecs";
66 				if (hasLeft)
67 				{
68 					auto t = now - prevAccess;
69 					if (t > repeat.total!"hnsecs")
70 					{
71 						accesses[index] = now;
72 						index = (index + 1) % accesses.length;
73 						return;
74 					}
75 					wait = repeat - (t + 1000).hnsecs;
76 				}
77 				else
78 					wait = range - (now - access + 1000).hnsecs;
79 			}
80 			if (wait.total!"hnsecs" > 0)
81 				sleep(wait);
82 		}
83 	}
84 }
85 
86 unittest
87 {
88 	import std.datetime.stopwatch;
89 	import vibe.core.core;
90 
91 	auto a = runTask({
92 		mixin DynamicRateLimitObject!(4, 60.msecs, 10.msecs) testLimit;
93 		StopWatch sw;
94 		sw.start();
95 		testLimit.waitFor();
96 		assert(sw.peek >= 0.msecs && sw.peek < 10.msecs);
97 		testLimit.waitFor();
98 		assert(sw.peek >= 10.msecs && sw.peek < 20.msecs);
99 		testLimit.waitFor();
100 		assert(sw.peek >= 20.msecs && sw.peek < 30.msecs);
101 		testLimit.waitFor();
102 		assert(sw.peek >= 30.msecs && sw.peek < 40.msecs);
103 		testLimit.waitFor();
104 		assert(sw.peek >= 60.msecs && sw.peek < 70.msecs);
105 		sw.stop();
106 	});
107 
108 	auto b = runTask({
109 		mixin RateLimitObject!(10.msecs) testLimit;
110 		StopWatch sw;
111 		sw.start();
112 		testLimit.waitFor();
113 		assert(sw.peek >= 0.msecs && sw.peek < 10.msecs);
114 		testLimit.waitFor();
115 		assert(sw.peek >= 10.msecs && sw.peek < 20.msecs);
116 		testLimit.waitFor();
117 		assert(sw.peek >= 20.msecs && sw.peek < 30.msecs);
118 		testLimit.waitFor();
119 		assert(sw.peek >= 30.msecs && sw.peek < 40.msecs);
120 		sw.stop();
121 	});
122 
123 	a.join();
124 	b.join();
125 }