+4

CHANGELOG.md

··· 7 7 8 8 ## [Unreleased] 9 9 10 10 + ### Added 11 11 + 12 12 + - Added the `per_second_fn` and `burst_limit_fn` functions to dynamically set limits based on the identifier. 13 13 + 10 14 ### Changed 11 15 12 16 - Removed the need for `glimit.build()` and `glimit.try_build()` functions. Now, the `glimit.apply()` function can be used directly with the limiter configuration.

+22 -4

src/glimit.gleam

··· 55 55 /// 56 56 pub type RateLimiterBuilder(a, b, id) { 57 57 RateLimiterBuilder( 58 58 - per_second: Option(Int), 59 59 - burst_limit: Option(Int), 58 58 + per_second: Option(fn(id) -> Int), 59 59 + burst_limit: Option(fn(id) -> Int), 60 60 identifier: Option(fn(a) -> id), 61 61 on_limit_exceeded: Option(fn(a) -> b), 62 62 ) ··· 83 83 limiter: RateLimiterBuilder(a, b, id), 84 84 limit: Int, 85 85 ) -> RateLimiterBuilder(a, b, id) { 86 86 - RateLimiterBuilder(..limiter, per_second: Some(limit)) 86 86 + RateLimiterBuilder(..limiter, per_second: Some(fn(_) { limit })) 87 87 + } 88 88 + 89 89 + /// Set the rate limit per second, based on the identifier. 90 90 + /// 91 91 + pub fn per_second_fn( 92 92 + limiter: RateLimiterBuilder(a, b, id), 93 93 + limit_fn: fn(id) -> Int, 94 94 + ) -> RateLimiterBuilder(a, b, id) { 95 95 + RateLimiterBuilder(..limiter, per_second: Some(limit_fn)) 87 96 } 88 97 89 98 /// Set the maximum number of available tokens. ··· 95 104 limiter: RateLimiterBuilder(a, b, id), 96 105 burst_limit: Int, 97 106 ) -> RateLimiterBuilder(a, b, id) { 98 98 - RateLimiterBuilder(..limiter, burst_limit: Some(burst_limit)) 107 107 + RateLimiterBuilder(..limiter, burst_limit: Some(fn(_) { burst_limit })) 108 108 + } 109 109 + 110 110 + /// Set the maximum number of available tokens, based on the identifier. 111 111 + /// 112 112 + pub fn burst_limit_fn( 113 113 + limiter: RateLimiterBuilder(a, b, id), 114 114 + burst_limit_fn: fn(id) -> Int, 115 115 + ) -> RateLimiterBuilder(a, b, id) { 116 116 + RateLimiterBuilder(..limiter, burst_limit: Some(burst_limit_fn)) 99 117 } 100 118 101 119 /// Set the handler to be called when the rate limit is reached.

+6 -6

src/glimit/registry.gleam

··· 19 19 State( 20 20 /// The maximum number of tokens. 21 21 /// 22 22 - max_token_count: Int, 22 22 + max_token_count: fn(id) -> Int, 23 23 /// The rate of token generation per second. 24 24 /// 25 25 - token_rate: Int, 25 25 + token_rate: fn(id) -> Int, 26 26 /// The registry of rate limiters. 27 27 /// 28 28 registry: Dict(id, Subject(rate_limiter.Message)), ··· 54 54 } 55 55 Error(_) -> { 56 56 use rate_limiter <- result.try(rate_limiter.new( 57 57 - state.max_token_count, 58 58 - state.token_rate, 57 57 + state.max_token_count(identifier), 58 58 + state.token_rate(identifier), 59 59 )) 60 60 Ok(rate_limiter) 61 61 } ··· 103 103 /// Create a new rate limiter registry. 104 104 /// 105 105 pub fn new( 106 106 - per_second: Int, 107 107 - burst_limit: Int, 106 106 + per_second: fn(id) -> Int, 107 107 + burst_limit: fn(id) -> Int, 108 108 ) -> Result(RateLimiterRegistryActor(id), Nil) { 109 109 let state = 110 110 State(

+5 -5

test/glimit_registry_test.gleam

··· 4 4 import glimit/registry 5 5 6 6 pub fn same_id_same_actor_test() { 7 7 - let assert Ok(registry) = registry.new(2, 2) 7 7 + let assert Ok(registry) = registry.new(fn(_) { 2 }, fn(_) { 2 }) 8 8 let assert Ok(rate_limiter) = registry |> registry.get_or_create("🚀") 9 9 let assert Ok(same_rate_limiter) = registry |> registry.get_or_create("🚀") 10 10 ··· 13 13 } 14 14 15 15 pub fn other_id_other_actor_test() { 16 16 - let assert Ok(registry) = registry.new(2, 2) 16 16 + let assert Ok(registry) = registry.new(fn(_) { 2 }, fn(_) { 2 }) 17 17 let assert Ok(rate_limiter) = registry |> registry.get_or_create("🚀") 18 18 let assert Ok(same_rate_limiter) = registry |> registry.get_or_create("💫") 19 19 ··· 22 22 } 23 23 24 24 pub fn sweep_full_bucket_test() { 25 25 - let assert Ok(registry) = registry.new(2, 2) 25 25 + let assert Ok(registry) = registry.new(fn(_) { 2 }, fn(_) { 2 }) 26 26 let assert Ok(rate_limiter) = registry |> registry.get_or_create("🚀") 27 27 28 28 registry |> registry.sweep(None) ··· 34 34 } 35 35 36 36 pub fn sweep_not_full_bucket_test() { 37 37 - let assert Ok(registry) = registry.new(2, 2) 37 37 + let assert Ok(registry) = registry.new(fn(_) { 2 }, fn(_) { 2 }) 38 38 let assert Ok(rate_limiter) = registry |> registry.get_or_create("🚀") 39 39 40 40 let _ = rate_limiter |> rate_limiter.hit ··· 47 47 } 48 48 49 49 pub fn sweep_after_long_time_test() { 50 50 - let assert Ok(registry) = registry.new(2, 2) 50 50 + let assert Ok(registry) = registry.new(fn(_) { 2 }, fn(_) { 2 }) 51 51 let assert Ok(rate_limiter) = registry |> registry.get_or_create("🚀") 52 52 53 53 rate_limiter |> rate_limiter.set_now(0)

+186

test/glimit_test.gleam

··· 93 93 func(Nil) |> should.equal("Stop!") 94 94 func(Nil) |> should.equal("Stop!") 95 95 } 96 96 + 97 97 + pub fn dynamic_per_second_test() { 98 98 + let limiter = 99 99 + glimit.new() 100 100 + |> glimit.per_second_fn(fn(id) { 101 101 + case id { 102 102 + "id" -> 2 103 103 + _ -> 1 104 104 + } 105 105 + }) 106 106 + |> glimit.identifier(fn(x) { x }) 107 107 + |> glimit.on_limit_exceeded(fn(_) { "Stop!" }) 108 108 + 109 109 + let func = 110 110 + fn(_) { "OK" } 111 111 + |> glimit.apply(limiter) 112 112 + 113 113 + func("id") |> should.equal("OK") 114 114 + func("id") |> should.equal("OK") 115 115 + func("id") |> should.equal("Stop!") 116 116 + func("id") |> should.equal("Stop!") 117 117 + 118 118 + func("other") |> should.equal("OK") 119 119 + func("other") |> should.equal("Stop!") 120 120 + func("other") |> should.equal("Stop!") 121 121 + } 122 122 + 123 123 + pub fn dynamic_per_second_static_burst_limit_test() { 124 124 + let assert Ok(limiter) = 125 125 + glimit.new() 126 126 + |> glimit.per_second_fn(fn(id) { 127 127 + case id { 128 128 + "id" -> 2 129 129 + _ -> 1 130 130 + } 131 131 + }) 132 132 + |> glimit.burst_limit(3) 133 133 + |> glimit.identifier(fn(x) { x }) 134 134 + |> glimit.on_limit_exceeded(fn(_) { "Stop!" }) 135 135 + |> glimit.build 136 136 + 137 137 + let func = 138 138 + fn(_) { "OK" } 139 139 + |> glimit.apply_built(limiter) 140 140 + 141 141 + let assert Ok(rate_limiter) = 142 142 + limiter.rate_limiter_registry 143 143 + |> registry.get_or_create("id") 144 144 + 145 145 + rate_limiter |> rate_limiter.set_now(0) 146 146 + func("id") |> should.equal("OK") 147 147 + func("id") |> should.equal("OK") 148 148 + func("id") |> should.equal("OK") 149 149 + func("id") |> should.equal("Stop!") 150 150 + func("id") |> should.equal("Stop!") 151 151 + 152 152 + rate_limiter |> rate_limiter.set_now(1) 153 153 + func("id") |> should.equal("OK") 154 154 + func("id") |> should.equal("OK") 155 155 + func("id") |> should.equal("Stop!") 156 156 + func("id") |> should.equal("Stop!") 157 157 + 158 158 + let assert Ok(rate_limiter) = 159 159 + limiter.rate_limiter_registry 160 160 + |> registry.get_or_create("other") 161 161 + 162 162 + rate_limiter |> rate_limiter.set_now(0) 163 163 + func("other") |> should.equal("OK") 164 164 + func("other") |> should.equal("OK") 165 165 + func("other") |> should.equal("OK") 166 166 + func("other") |> should.equal("Stop!") 167 167 + func("other") |> should.equal("Stop!") 168 168 + 169 169 + rate_limiter |> rate_limiter.set_now(1) 170 170 + func("other") |> should.equal("OK") 171 171 + func("other") |> should.equal("Stop!") 172 172 + func("other") |> should.equal("Stop!") 173 173 + } 174 174 + 175 175 + pub fn static_per_second_dynamic_burst_limit_test() { 176 176 + let assert Ok(limiter) = 177 177 + glimit.new() 178 178 + |> glimit.per_second(1) 179 179 + |> glimit.burst_limit_fn(fn(id) { 180 180 + case id { 181 181 + "id" -> 3 182 182 + _ -> 2 183 183 + } 184 184 + }) 185 185 + |> glimit.identifier(fn(x) { x }) 186 186 + |> glimit.on_limit_exceeded(fn(_) { "Stop!" }) 187 187 + |> glimit.build 188 188 + 189 189 + let func = 190 190 + fn(_) { "OK" } 191 191 + |> glimit.apply_built(limiter) 192 192 + 193 193 + let assert Ok(rate_limiter) = 194 194 + limiter.rate_limiter_registry 195 195 + |> registry.get_or_create("id") 196 196 + 197 197 + rate_limiter |> rate_limiter.set_now(0) 198 198 + func("id") |> should.equal("OK") 199 199 + func("id") |> should.equal("OK") 200 200 + func("id") |> should.equal("OK") 201 201 + func("id") |> should.equal("Stop!") 202 202 + func("id") |> should.equal("Stop!") 203 203 + 204 204 + rate_limiter |> rate_limiter.set_now(1) 205 205 + func("id") |> should.equal("OK") 206 206 + func("id") |> should.equal("Stop!") 207 207 + func("id") |> should.equal("Stop!") 208 208 + 209 209 + let assert Ok(rate_limiter) = 210 210 + limiter.rate_limiter_registry 211 211 + |> registry.get_or_create("other") 212 212 + 213 213 + rate_limiter |> rate_limiter.set_now(0) 214 214 + func("other") |> should.equal("OK") 215 215 + func("other") |> should.equal("OK") 216 216 + func("other") |> should.equal("Stop!") 217 217 + func("other") |> should.equal("Stop!") 218 218 + 219 219 + rate_limiter |> rate_limiter.set_now(1) 220 220 + func("other") |> should.equal("OK") 221 221 + func("other") |> should.equal("Stop!") 222 222 + func("other") |> should.equal("Stop!") 223 223 + } 224 224 + 225 225 + pub fn dynamic_per_second_dynamic_burst_limit_test() { 226 226 + let assert Ok(limiter) = 227 227 + glimit.new() 228 228 + |> glimit.per_second_fn(fn(id) { 229 229 + case id { 230 230 + "id" -> 2 231 231 + _ -> 1 232 232 + } 233 233 + }) 234 234 + |> glimit.burst_limit_fn(fn(id) { 235 235 + case id { 236 236 + "id" -> 4 237 237 + _ -> 3 238 238 + } 239 239 + }) 240 240 + |> glimit.identifier(fn(x) { x }) 241 241 + |> glimit.on_limit_exceeded(fn(_) { "Stop!" }) 242 242 + |> glimit.build 243 243 + 244 244 + let func = 245 245 + fn(_) { "OK" } 246 246 + |> glimit.apply_built(limiter) 247 247 + 248 248 + let assert Ok(rate_limiter) = 249 249 + limiter.rate_limiter_registry 250 250 + |> registry.get_or_create("id") 251 251 + 252 252 + rate_limiter |> rate_limiter.set_now(0) 253 253 + func("id") |> should.equal("OK") 254 254 + func("id") |> should.equal("OK") 255 255 + func("id") |> should.equal("OK") 256 256 + func("id") |> should.equal("OK") 257 257 + func("id") |> should.equal("Stop!") 258 258 + func("id") |> should.equal("Stop!") 259 259 + 260 260 + rate_limiter |> rate_limiter.set_now(1) 261 261 + func("id") |> should.equal("OK") 262 262 + func("id") |> should.equal("OK") 263 263 + func("id") |> should.equal("Stop!") 264 264 + func("id") |> should.equal("Stop!") 265 265 + 266 266 + let assert Ok(rate_limiter) = 267 267 + limiter.rate_limiter_registry 268 268 + |> registry.get_or_create("other") 269 269 + 270 270 + rate_limiter |> rate_limiter.set_now(0) 271 271 + func("other") |> should.equal("OK") 272 272 + func("other") |> should.equal("OK") 273 273 + func("other") |> should.equal("OK") 274 274 + func("other") |> should.equal("Stop!") 275 275 + func("other") |> should.equal("Stop!") 276 276 + 277 277 + rate_limiter |> rate_limiter.set_now(1) 278 278 + func("other") |> should.equal("OK") 279 279 + func("other") |> should.equal("Stop!") 280 280 + func("other") |> should.equal("Stop!") 281 281 + }