Skip to content

Commit 26eb71a

Browse files
committed
rewrite to be store-agnostic, fix a couple of bugs, docs, etc...
1 parent e888cb8 commit 26eb71a

File tree

12 files changed

+368
-376
lines changed

12 files changed

+368
-376
lines changed

README.md

Lines changed: 38 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
1-
# Golang Sliding Window Counters Rate Limiter
1+
# Swirl - Sliding Window Increment Rate Limit
22

3-
> _(I tried to come up with a nicer name...)_
3+
> Sliding Window ~~Counters~~ Increment\* rate limit implementation for Go
4+
5+
_(\*the name ["swc"](https://swc.rs/) is already taken and who doesn't love a good backronym?)_
46

57
This is a simple rate limiter built based on [this blog post](https://www.figma.com/blog/an-alternative-approach-to-rate-limiting) from Figma's engineering team.
68

@@ -11,33 +13,49 @@ See the post for information about the requirements and design of the actual alg
1113
The rate limiter satisfies this interface:
1214

1315
```go
14-
type Limiter interface {
15-
Increment(context.Context, string, int) error
16-
}
16+
Increment(ctx context.Context, key string, incr int) (*Status, bool, error)
1717
```
1818

19-
The implementation is backed by Redis and uses the go-redis library.
19+
- Status includes information you'd want to set in [`RateLimit` headers](https://datatracker.ietf.org/doc/draft-ietf-httpapi-ratelimit-headers/).
20+
- Bool is whether the limit was exceeded or not, true means reject the request.
21+
- Errors occur for cache issues, such as Redis connectivity or malformed data.
22+
23+
The implementation is store agnostic, however due to the way it works, Redis is the recommended approach due to the usage of [hash sets](https://redis.io/docs/latest/develop/data-types/hashes/).
24+
25+
The `incr` argument allows you to assign different weights to the action being rate limited. For example, a simple request may use a value of 1 and an expensive request may use a value of 10.
2026

2127
```go
22-
client := goredis.NewClient(&goredis.Options{
23-
Addr: "localhost:6379",
24-
})
28+
status, exceeded, err := m.rl.Increment(ctx, key, cost)
29+
if err != nil {
30+
http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
31+
return
32+
}
2533

26-
ratelimiter := redis.New(client, 10, time.Minute, time.Hour)
34+
limit := status.Limit
35+
remaining := status.Remaining
36+
resetTime := status.Reset.UTC().Format(time.RFC1123)
2737

28-
ratelimiter.Increment(ctx, "user_id", 1)
38+
w.Header().Set(RateLimitLimit, strconv.FormatUint(uint64(limit), 10))
39+
w.Header().Set(RateLimitRemaining, strconv.FormatUint(uint64(remaining), 10))
40+
w.Header().Set(RateLimitReset, resetTime)
41+
42+
if exceeded {
43+
// you shall not pass.
44+
w.Header().Set(RetryAfter, resetTime)
45+
http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
46+
return
47+
}
48+
49+
// all good my g come thru
50+
next.ServeHTTP(w, r)
2951
```
3052

31-
## Middleware
53+
Note that according to the IETF spec, the `X-RateLimit-*` headers are not standardised, but commonly used. See the spec for advisories on `RateLimit-Policy` etc.
3254

33-
There's a http middleware too, for convenience. Inspired by Seth Vargo's [rate limit](https://github.com/sethvargo/go-limiter) library:
55+
## `memory`
3456

35-
```go
36-
ratelimiter := redis.New(client, 10, time.Minute, time.Hour)
37-
mw := ratelimit.Middleware(ratelimiter, ratelimit.IPKeyFunc, 1)
38-
// use mw in your favourite HTTP library
39-
```
57+
This is a very basic in-memory cache that mirrors the tiny subset of Redis-based hash set APIs necessary to use the rate limiter in pure Go. You can probably use this in a very basic single-server application but it's not covered by tests and has not been extensively used in production so... beware. Treat it as a testing mock.
4058

41-
The `Middleware` function has a `weight` parameter, which allows you to give a higher increment weight to certain routes. So for example, your base rate limit can be 1000 requests per hour and each request has a weight of 1, but a particularly computationally intensive endpoint may want to have a weight of 10, so each request increments the internal counter by 10 instead of 1.
59+
## HTTP middleware
4260

43-
If you use multiple middleware instances, make sure you don't add one to the global mux/router otherwise your requests will be triggering two rate limit calculations (and thus, Redis network I/O operations). This could be solved in future by a centralised middleware controller that provides a tree of limiters with order of precedence rules etc.
61+
This package, unlike most rate limit packages, purposely does not include HTTP middleware, you probably want to write your own with your own logging, response logic, etc. anyway. It's super simple and the code above gets you most of the way already.

docker-compose.yml

Lines changed: 0 additions & 12 deletions
This file was deleted.

go.mod

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,14 @@
1-
module github.com/Southclaws/sliding-window-counters
1+
module github.com/Southclaws/swirl
22

3-
go 1.16
3+
go 1.22.1
44

55
require (
6-
github.com/go-redis/redis/v8 v8.11.4 // indirect
7-
github.com/stretchr/testify v1.7.0 // indirect
6+
github.com/puzpuzpuz/xsync/v3 v3.4.0
7+
github.com/stretchr/testify v1.7.0
8+
)
9+
10+
require (
11+
github.com/davecgh/go-spew v1.1.1 // indirect
12+
github.com/pmezard/go-difflib v1.0.0 // indirect
13+
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c // indirect
814
)

go.sum

Lines changed: 3 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -1,94 +1,14 @@
1-
github.com/cespare/xxhash/v2 v2.1.2 h1:YRXhKfTDauu4ajMg1TPgFO5jnlC2HCbmLXMcTG5cbYE=
2-
github.com/cespare/xxhash/v2 v2.1.2/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
31
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
42
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
53
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
6-
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
7-
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
8-
github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
9-
github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ=
10-
github.com/go-redis/redis/v8 v8.11.4 h1:kHoYkfZP6+pe04aFTnhDH6GDROa5yJdHJVNxV3F46Tg=
11-
github.com/go-redis/redis/v8 v8.11.4/go.mod h1:2Z2wHZXdQpCDXEGzqMockDpNyYvi2l4Pxt6RJr792+w=
12-
github.com/go-task/slim-sprig v0.0.0-20210107165309-348f09dbbbc0/go.mod h1:fyg7847qk6SyHyPtNmDHnmrv/HOrqktSC+C9fM+CJOE=
13-
github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
14-
github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8=
15-
github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA=
16-
github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs=
17-
github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w=
18-
github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0=
19-
github.com/golang/protobuf v1.4.2/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI=
20-
github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
21-
github.com/golang/protobuf v1.5.2/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
22-
github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
23-
github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
24-
github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
25-
github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
26-
github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
27-
github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
28-
github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
29-
github.com/nxadm/tail v1.4.8/go.mod h1:+ncqLTQzXmGhMZNUePPaPqPvBxHAIsmXswZKocGu+AU=
30-
github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
31-
github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk=
32-
github.com/onsi/ginkgo v1.16.4/go.mod h1:dX+/inL/fNMqNlz0e9LfyB9TswhZpCVdJM/Z6Vvnwo0=
33-
github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY=
34-
github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo=
35-
github.com/onsi/gomega v1.16.0/go.mod h1:HnhC7FXeEQY45zxNK3PPoIUhzk/80Xly9PcubAlGdZY=
364
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
375
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
6+
github.com/puzpuzpuz/xsync/v3 v3.4.0 h1:DuVBAdXuGFHv8adVXjWWZ63pJq+NRXOWVXlKDBZ+mJ4=
7+
github.com/puzpuzpuz/xsync/v3 v3.4.0/go.mod h1:VjzYrABPabuM4KyBh1Ftq6u8nhwY5tBPKP9jpmh0nnA=
388
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
39-
github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
409
github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
4110
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
42-
github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
43-
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
44-
golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
45-
golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
46-
golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
47-
golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
48-
golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
49-
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
50-
golang.org/x/net v0.0.0-20200520004742-59133d7f0dd7/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A=
51-
golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
52-
golang.org/x/net v0.0.0-20210428140749-89ef3d95e781/go.mod h1:OJAsFXCWl8Ukc7SiCT/9KSuxbyM7479/AVlXFRxuMCk=
53-
golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
54-
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
55-
golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
56-
golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
57-
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
58-
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
59-
golang.org/x/sys v0.0.0-20190904154756-749cb33beabd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
60-
golang.org/x/sys v0.0.0-20191005200804-aed5e4c7ecf9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
61-
golang.org/x/sys v0.0.0-20191120155948-bd437916bb0e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
62-
golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
63-
golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
64-
golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
65-
golang.org/x/sys v0.0.0-20210112080510-489259a85091/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
66-
golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
67-
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
68-
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
69-
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
70-
golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
71-
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
72-
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
73-
golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
74-
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
75-
golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
76-
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
77-
golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
78-
google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8=
79-
google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0=
80-
google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM=
81-
google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE=
82-
google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo=
83-
google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU=
84-
google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
85-
google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
11+
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
8612
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
87-
gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys=
88-
gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw=
89-
gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
90-
gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
91-
gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
92-
gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
9313
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
9414
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

memory/store.go

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
package memory
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"strconv"
7+
"time"
8+
9+
"github.com/puzpuzpuz/xsync/v3"
10+
)
11+
12+
var errNotFound = fmt.Errorf("not found")
13+
14+
type Entry struct {
15+
Value any
16+
Expiry *time.Time
17+
}
18+
19+
// Cache implements a very rudimentary in-memory cache that supports HSET style
20+
// operations similar to Redis. Purely for testing the rate limiter.
21+
type Cache struct {
22+
local *xsync.MapOf[string, Entry]
23+
}
24+
25+
func New() *Cache {
26+
return &Cache{
27+
local: xsync.NewMapOf[string, Entry](),
28+
}
29+
}
30+
31+
func (c *Cache) Get(ctx context.Context, key string) (string, error) {
32+
v, found := c.local.Load(key)
33+
if !found {
34+
return "", errNotFound
35+
}
36+
37+
if v.Expiry.Before(time.Now()) {
38+
c.local.Delete(key)
39+
return "", errNotFound
40+
}
41+
42+
return v.Value.(string), nil
43+
}
44+
45+
func (c *Cache) Set(ctx context.Context, key string, value string) error {
46+
c.local.Store(key, Entry{
47+
Value: value,
48+
})
49+
return nil
50+
}
51+
52+
func (c *Cache) Delete(ctx context.Context, key string) error {
53+
c.local.Delete(key)
54+
return nil
55+
}
56+
57+
func (c *Cache) HIncrBy(ctx context.Context, key string, field string, incr int64) (int, error) {
58+
ac, _ := c.local.Compute(key, func(old Entry, found bool) (Entry, bool) {
59+
if found {
60+
hash := old.Value.(map[string]string)
61+
if curr, ok := hash[field]; ok {
62+
i, err := strconv.Atoi(curr)
63+
if err != nil {
64+
return old, false
65+
}
66+
67+
i += int(incr)
68+
hash[field] = strconv.Itoa(i)
69+
old.Value = hash
70+
return old, false
71+
} else {
72+
hash[field] = strconv.Itoa(int(incr))
73+
old.Value = hash
74+
return old, false
75+
}
76+
77+
} else {
78+
hash := map[string]string{
79+
field: strconv.Itoa(int(incr)),
80+
}
81+
old.Value = hash
82+
return old, false
83+
}
84+
})
85+
86+
hash := ac.Value.(map[string]string)
87+
i, err := strconv.Atoi(hash[field])
88+
if err != nil {
89+
return 0, fmt.Errorf("failed to convert hash field to integer")
90+
}
91+
92+
return i, nil
93+
}
94+
95+
func (c *Cache) HGetAll(ctx context.Context, key string) (map[string]string, error) {
96+
v, ok := c.local.Load(key)
97+
if !ok {
98+
return map[string]string{}, nil
99+
}
100+
101+
return v.Value.(map[string]string), nil
102+
}
103+
104+
func (c *Cache) HDel(ctx context.Context, key string, field string) error {
105+
_, ok := c.local.Compute(key, func(old Entry, found bool) (Entry, bool) {
106+
if found {
107+
hash := old.Value.(map[string]string)
108+
delete(hash, field)
109+
old.Value = hash
110+
return old, false
111+
}
112+
return old, false
113+
})
114+
if !ok {
115+
return fmt.Errorf("failed to delete hash field")
116+
}
117+
118+
return nil
119+
}
120+
121+
func (c *Cache) Expire(ctx context.Context, key string, expiration time.Duration) error {
122+
c.local.Compute(key, func(old Entry, found bool) (Entry, bool) {
123+
if found {
124+
if old.Expiry != nil && old.Expiry.Before(time.Now()) {
125+
return old, true
126+
}
127+
128+
expiry := time.Now().Add(expiration)
129+
old.Expiry = &expiry
130+
return old, false
131+
}
132+
return old, false
133+
})
134+
135+
return nil
136+
}

0 commit comments

Comments
 (0)