179 lines
7.7 KiB
Go
179 lines
7.7 KiB
Go
// Package redis supplies atomic shared middleware state in one Redis hash slot.
|
|
// Use a dedicated namespace and a noeviction Redis deployment for correctness.
|
|
package redis
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"git.yuntuo.hk/k2safe/go-omni/middleware"
|
|
goredis "github.com/redis/go-redis/v9"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
type IdempotencyConfig struct {
|
|
MaxEntries int `json:"maxEntries"`
|
|
MaxRecordBytes int `json:"maxRecordBytes"`
|
|
}
|
|
type IdempotencyStore struct {
|
|
client goredis.Scripter
|
|
keys []string
|
|
cfg IdempotencyConfig
|
|
signature string
|
|
}
|
|
|
|
func namespace(prefix, kind string) (string, error) {
|
|
if prefix == "" || len(prefix) > 256 || strings.ContainsAny(prefix, "{}\r\n") {
|
|
return "", errors.New("middleware/redis: nonempty bounded namespace without braces required")
|
|
}
|
|
sum := sha256.Sum256([]byte(prefix))
|
|
return "goOmni:{" + hex.EncodeToString(sum[:]) + "}:" + kind, nil
|
|
}
|
|
func NewIdempotencyStore(client goredis.Scripter, prefix string, cfg IdempotencyConfig) (*IdempotencyStore, error) {
|
|
if client == nil || cfg.MaxEntries < 0 || cfg.MaxRecordBytes < 0 {
|
|
return nil, errors.New("middleware/redis: invalid configuration")
|
|
}
|
|
if cfg.MaxEntries == 0 {
|
|
cfg.MaxEntries = 128
|
|
}
|
|
if cfg.MaxRecordBytes == 0 {
|
|
cfg.MaxRecordBytes = 128 << 10
|
|
}
|
|
if cfg.MaxEntries > 100000 || cfg.MaxRecordBytes < 512 || cfg.MaxRecordBytes > 1<<20 {
|
|
return nil, errors.New("middleware/redis: capacity outside supported bounds")
|
|
}
|
|
base, err := namespace(prefix, "idempotency")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &IdempotencyStore{client: client, keys: []string{base + ":records", base + ":expiry", base + ":config"}, cfg: cfg, signature: fmt.Sprintf("%d:%d", cfg.MaxEntries, cfg.MaxRecordBytes)}, nil
|
|
}
|
|
|
|
var claim = goredis.NewScript(`
|
|
local config=redis.call('GET',KEYS[3])
|
|
if config and config~=ARGV[5] then return redis.error_reply('configuration mismatch') end
|
|
if not config then redis.call('SET',KEYS[3],ARGV[5]) end
|
|
local clock=redis.call('TIME');local now=clock[1]*1000+math.floor(clock[2]/1000)
|
|
local expired=redis.call('ZRANGEBYSCORE',KEYS[2],'-inf',now,'LIMIT',0,128)
|
|
for _,key in ipairs(expired) do redis.call('HDEL',KEYS[1],key);redis.call('ZREM',KEYS[2],key) end
|
|
local expiry=redis.call('ZSCORE',KEYS[2],ARGV[1]);if expiry and tonumber(expiry)<=now then redis.call('HDEL',KEYS[1],ARGV[1]);redis.call('ZREM',KEYS[2],ARGV[1]) end
|
|
local value=redis.call('HGET',KEYS[1],ARGV[1]);if value then return {0,value} end
|
|
if redis.call('HLEN',KEYS[1])>=tonumber(ARGV[4]) then return redis.error_reply('idempotency capacity reached') end
|
|
value=cjson.encode({fingerprint=ARGV[2],owner=ARGV[3],state='running'})
|
|
redis.call('HSET',KEYS[1],ARGV[1],value);return {1,value}
|
|
`)
|
|
|
|
func (s *IdempotencyStore) Claim(ctx context.Context, key, fp, owner string) (middleware.IdempotencyRecord, bool, error) {
|
|
if len(key) != 64 || len(fp) != 64 || len(owner) != 32 {
|
|
return middleware.IdempotencyRecord{}, false, errors.New("middleware/redis: invalid claim")
|
|
}
|
|
result, err := claim.Run(ctx, s.client, s.keys, key, fp, owner, s.cfg.MaxEntries, s.signature).Slice()
|
|
if err != nil {
|
|
return middleware.IdempotencyRecord{}, false, err
|
|
}
|
|
if len(result) != 2 {
|
|
return middleware.IdempotencyRecord{}, false, errors.New("middleware/redis: invalid claim result")
|
|
}
|
|
value, ok := result[1].(string)
|
|
if !ok || len(value) > s.cfg.MaxRecordBytes {
|
|
return middleware.IdempotencyRecord{}, false, errors.New("middleware/redis: invalid record")
|
|
}
|
|
var record middleware.IdempotencyRecord
|
|
err = json.Unmarshal([]byte(value), &record)
|
|
return record, result[0] == int64(1), err
|
|
}
|
|
|
|
var complete = goredis.NewScript(`
|
|
local current=redis.call('HGET',KEYS[1],ARGV[1]);if not current then return redis.error_reply('claim missing') end
|
|
local value=cjson.decode(current)
|
|
if value.owner~=ARGV[2] or value.state~='running' then return redis.error_reply('claim owner/state changed') end
|
|
local response=ARGV[3]
|
|
if response~='' then value.state='completed';value.response=cjson.decode(response) else value.state='failed' end
|
|
local encoded=cjson.encode(value)
|
|
if string.len(encoded)>tonumber(ARGV[5]) then value.state='failed';value.response=nil;encoded=cjson.encode(value) end
|
|
local clock=redis.call('TIME');local now=clock[1]*1000+math.floor(clock[2]/1000)
|
|
redis.call('HSET',KEYS[1],ARGV[1],encoded);redis.call('ZADD',KEYS[2],now+tonumber(ARGV[4]),ARGV[1]);return 1
|
|
`)
|
|
|
|
func (s *IdempotencyStore) Complete(ctx context.Context, key, owner string, response *middleware.StoredResponse, ttl time.Duration) error {
|
|
if ttl <= 0 || ttl > 30*24*time.Hour {
|
|
return errors.New("middleware/redis: TTL must be positive and <=30d")
|
|
}
|
|
var encoded []byte
|
|
var err error
|
|
if response != nil {
|
|
encoded, err = json.Marshal(response)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if len(encoded) > s.cfg.MaxRecordBytes {
|
|
encoded = nil
|
|
}
|
|
}
|
|
return complete.Run(ctx, s.client, s.keys, key, owner, string(encoded), max(1, ttl.Milliseconds()), s.cfg.MaxRecordBytes).Err()
|
|
}
|
|
|
|
// Resolve removes only a matching non-replayable claim after the operator has
|
|
// reconciled the business outcome. Never invoke automatically on a timeout.
|
|
func (s *IdempotencyStore) Resolve(ctx context.Context, key, owner string) error {
|
|
return goredis.NewScript(`local v=redis.call('HGET',KEYS[1],ARGV[1]);if not v then return 0 end;local r=cjson.decode(v);if r.owner~=ARGV[2] or r.state=='completed' then return redis.error_reply('claim cannot be resolved') end;redis.call('HDEL',KEYS[1],ARGV[1]);redis.call('ZREM',KEYS[2],ARGV[1]);return 1`).Run(ctx, s.client, s.keys, key, owner).Err()
|
|
}
|
|
|
|
type limiter struct {
|
|
client goredis.Scripter
|
|
keys []string
|
|
cfg middleware.RateLimitConfig
|
|
signature string
|
|
}
|
|
|
|
func NewRateLimiter(client goredis.Scripter, prefix string, cfg middleware.RateLimitConfig) (*middleware.RateLimiter, error) {
|
|
if client == nil {
|
|
return nil, errors.New("middleware/redis: client required")
|
|
}
|
|
if cfg.MaxKeys == 0 {
|
|
cfg.MaxKeys = 10000
|
|
}
|
|
if cfg.IdleTTL == 0 {
|
|
cfg.IdleTTL = 24 * time.Hour
|
|
}
|
|
if _, err := middleware.NewRateLimiter(cfg); err != nil {
|
|
return nil, err
|
|
}
|
|
base, err := namespace(prefix, "rate")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
backend := &limiter{client: client, keys: []string{base + ":buckets", base + ":seen", base + ":config"}, cfg: cfg, signature: fmt.Sprintf("%g:%d:%d:%d", cfg.Rate, cfg.Burst, cfg.MaxKeys, cfg.IdleTTL.Milliseconds())}
|
|
cfg.Backend = backend
|
|
return middleware.NewRateLimiter(cfg)
|
|
}
|
|
|
|
var take = goredis.NewScript(`
|
|
local config=redis.call('GET',KEYS[3]);if config and config~=ARGV[6] then return redis.error_reply('configuration mismatch') end
|
|
if not config then redis.call('SET',KEYS[3],ARGV[6]) end
|
|
local clock=redis.call('TIME');local now=clock[1]*1000+math.floor(clock[2]/1000)
|
|
local expired=redis.call('ZRANGEBYSCORE',KEYS[2],'-inf',now-tonumber(ARGV[5]),'LIMIT',0,128)
|
|
for _,key in ipairs(expired) do redis.call('HDEL',KEYS[1],key);redis.call('ZREM',KEYS[2],key) end
|
|
local value=redis.call('HGET',KEYS[1],ARGV[1]);local tokens=tonumber(ARGV[3])
|
|
if value then local b=cjson.decode(value);tokens=math.min(tokens,b.tokens+math.max(0,now-b.time)*tonumber(ARGV[2])/1000)
|
|
elseif redis.call('HLEN',KEYS[1])>=tonumber(ARGV[4]) then return -1 end
|
|
local delay=0;if tokens>=1 then tokens=tokens-1 else delay=math.ceil((1-tokens)*1000/tonumber(ARGV[2])) end
|
|
redis.call('HSET',KEYS[1],ARGV[1],cjson.encode({tokens=tokens,time=now}));redis.call('ZADD',KEYS[2],now,ARGV[1]);return delay
|
|
`)
|
|
|
|
func (l *limiter) Take(ctx context.Context, key string) (time.Duration, error) {
|
|
sum := sha256.Sum256([]byte(key))
|
|
delay, err := take.Run(ctx, l.client, l.keys, hex.EncodeToString(sum[:]), l.cfg.Rate, l.cfg.Burst, l.cfg.MaxKeys, l.cfg.IdleTTL.Milliseconds(), l.signature).Int64()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if delay < 0 {
|
|
return 0, middleware.ErrLimitCapacity
|
|
}
|
|
return time.Duration(delay) * time.Millisecond, nil
|
|
}
|