From 868b959c6fddc2771f02f8af2203650ed18ee4f4 Mon Sep 17 00:00:00 2001 From: zhaoyupeng Date: Tue, 15 Jul 2025 18:40:13 +0800 Subject: [PATCH] feat: add memory cache --- database/cache/cache.go | 1 + database/cache/memory.go | 155 +++++++++++++++++++++++++++++++++++++ database/cache/new.go | 42 ++++++---- database/cache/new_test.go | 29 +++++++ database/cache/option.go | 23 ++++-- go.mod | 7 +- go.sum | 11 ++- 7 files changed, 239 insertions(+), 29 deletions(-) create mode 100644 database/cache/memory.go diff --git a/database/cache/cache.go b/database/cache/cache.go index 0a224a4..4e696f7 100644 --- a/database/cache/cache.go +++ b/database/cache/cache.go @@ -62,6 +62,7 @@ var ( marshaler func(data any) ([]byte, error) = json.Marshal unmarshaler func(data []byte, model any) error = json.Unmarshal ErrorKeyNotFound = errors.New("key not found") + ErrorStoreFailed = errors.New("store failed") Default Cache ) diff --git a/database/cache/memory.go b/database/cache/memory.go new file mode 100644 index 0000000..7200e17 --- /dev/null +++ b/database/cache/memory.go @@ -0,0 +1,155 @@ +package cache + +import ( + "context" + "errors" + "time" + + "github.com/dgraph-io/ristretto/v2" +) + +var _ Cache = (*_mem)(nil) + +type _mem struct { + ctx context.Context + cache *ristretto.Cache[string, []byte] +} + +func newMemory(ctx context.Context, ins *ristretto.Cache[string, []byte]) Cache { + return &_mem{ + ctx: ctx, + cache: ins, + } +} + +func (m *_mem) Client() any { + return m.cache +} + +func (c *_mem) Close() { + c.cache.Close() +} + +func (c *_mem) Del(ctx context.Context, keys ...string) error { + for _, key := range keys { + c.cache.Del(key) + } + + return nil +} + +func (c *_mem) Get(ctx context.Context, key string) ([]byte, error) { + val, ok := c.cache.Get(key) + if !ok { + return val, ErrorKeyNotFound + } + + return val, nil +} + +func (c *_mem) GetDel(ctx context.Context, key string) ([]byte, error) { + val, err := c.Get(ctx, key) + if err != nil { + return val, err + } + + c.cache.Del(key) + + return val, err +} + +func (c *_mem) GetDelScan(ctx context.Context, key string) Scanner { + val, err := c.GetDel(ctx, key) + return newScan(val, err) +} + +func (c *_mem) GetEx(ctx context.Context, key string, duration time.Duration) ([]byte, error) { + val, err := c.Get(ctx, key) + if err != nil { + return val, err + } + + c.cache.SetWithTTL(key, val, 1, duration) + + return val, err +} + +func (m *_mem) GetExScan(ctx context.Context, key string, duration time.Duration) Scanner { + val, err := m.GetEx(ctx, key, duration) + return newScan(val, err) +} + +func (m *_mem) GetScan(ctx context.Context, key string) Scanner { + val, err := m.Get(ctx, key) + return newScan(val, err) +} + +func (m *_mem) Gets(ctx context.Context, keys ...string) ([][]byte, error) { + vals := make([][]byte, 0, len(keys)) + + for _, key := range keys { + val, err := m.Get(ctx, key) + if err != nil { + if errors.Is(err, ErrorKeyNotFound) { + continue + } + + return vals, err + } + + vals = append(vals, val) + } + + if len(vals) != len(keys) { + return vals, ErrorKeyNotFound + } + + return vals, nil +} + +func (m *_mem) Set(ctx context.Context, key string, value any) error { + val, err := handleValue(value) + if err != nil { + return err + } + + if ok := m.cache.Set(key, val, 1); !ok { + return ErrorStoreFailed + } + + m.cache.Wait() + + return nil +} + +func (m *_mem) SetEx(ctx context.Context, key string, value any, duration time.Duration) error { + val, err := handleValue(value) + if err != nil { + return err + } + + if ok := m.cache.SetWithTTL(key, val, 1, duration); !ok { + return ErrorStoreFailed + } + + m.cache.Wait() + + return nil +} + +func (m *_mem) Sets(ctx context.Context, vm map[string]any) error { + for key, value := range vm { + val, err := handleValue(value) + if err != nil { + return err + } + + if ok := m.cache.Set(key, val, 1); !ok { + return ErrorStoreFailed + } + } + + m.cache.Wait() + + return nil +} diff --git a/database/cache/new.go b/database/cache/new.go index b92b50a..221b922 100644 --- a/database/cache/new.go +++ b/database/cache/new.go @@ -4,35 +4,31 @@ import ( "context" "fmt" "gitea.loveuer.com/yizhisec/packages/tool" + "github.com/dgraph-io/ristretto/v2" "github.com/go-redis/redis/v8" _ "github.com/go-redis/redis/v8" "net/url" ) -var ( - defaultRedis = "redis://127.0.0.1:6379" -) - -func New(opts ...OptionFn) (Cache, error) { +func New(opts ...Option) (Cache, error) { var ( err error - cfg = &config{ - ctx: context.Background(), - redis: &defaultRedis, + opt = &option{ + ctx: context.Background(), } ) - for _, opt := range opts { - opt(cfg) + for _, fn := range opts { + fn(opt) } - if cfg.redis != nil { + if opt.redis != nil { var ( ins *url.URL client *redis.Client ) - if ins, err = url.Parse(*cfg.redis); err != nil { + if ins, err = url.Parse(*opt.redis); err != nil { return nil, err } @@ -45,17 +41,33 @@ func New(opts ...OptionFn) (Cache, error) { Password: password, }) - if err = client.Ping(tool.TimeoutCtx(cfg.ctx, 5)).Err(); err != nil { + if err = client.Ping(tool.TimeoutCtx(opt.ctx, 5)).Err(); err != nil { return nil, err } - return newRedis(cfg.ctx, client), nil + return newRedis(opt.ctx, client), nil + } + + if opt.memory { + var ( + ins *ristretto.Cache[string, []byte] + ) + + if ins, err = ristretto.NewCache(&ristretto.Config[string, []byte]{ + NumCounters: 1e7, + MaxCost: 1 << 30, + BufferItems: 64, + }); err != nil { + return nil, err + } + + return newMemory(opt.ctx, ins), nil } return nil, fmt.Errorf("invalid cache config") } -func Init(opts ...OptionFn) (err error) { +func Init(opts ...Option) (err error) { Default, err = New(opts...) return err } diff --git a/database/cache/new_test.go b/database/cache/new_test.go index 76cea18..9a17498 100644 --- a/database/cache/new_test.go +++ b/database/cache/new_test.go @@ -1,6 +1,7 @@ package cache import ( + "errors" "testing" ) @@ -71,3 +72,31 @@ func TestNoAuth(t *testing.T) { // t.Fatal(err) //} } + +func TestMemory(t *testing.T) { + c, err := New(WithMemory()) + if err != nil { + t.Fatal(err) + } + + bs, err := c.Get(t.Context(), "haha") + if err != nil { + if !errors.Is(err, ErrorKeyNotFound) { + t.Fatal(err) + } + + t.Logf("key not found") + } + + t.Logf("haha = %s", string(bs)) + + if err = c.Set(t.Context(), "haha", "haha"); err != nil { + t.Fatal(err) + } + + if bs, err = c.Get(t.Context(), "haha"); err != nil { + t.Fatal(err) + } + + t.Logf("haha = %s", string(bs)) +} diff --git a/database/cache/option.go b/database/cache/option.go index bc3cf96..62f882f 100644 --- a/database/cache/option.go +++ b/database/cache/option.go @@ -5,23 +5,24 @@ import ( "fmt" ) -type config struct { - ctx context.Context - redis *string +type option struct { + ctx context.Context + redis *string + memory bool } -type OptionFn func(*config) +type Option func(*option) -func WithCtx(ctx context.Context) OptionFn { - return func(c *config) { +func WithCtx(ctx context.Context) Option { + return func(c *option) { if ctx != nil { c.ctx = ctx } } } -func WithRedis(host string, port int, username, password string) OptionFn { - return func(c *config) { +func WithRedis(host string, port int, username, password string) Option { + return func(c *option) { uri := fmt.Sprintf("redis://%s:%d", host, port) if username != "" || password != "" { uri = fmt.Sprintf("redis://%s:%s@%s:%d", username, password, host, port) @@ -30,3 +31,9 @@ func WithRedis(host string, port int, username, password string) OptionFn { c.redis = &uri } } + +func WithMemory() Option { + return func(c *option) { + c.memory = true + } +} diff --git a/go.mod b/go.mod index 35929c5..f9609bf 100644 --- a/go.mod +++ b/go.mod @@ -1,10 +1,11 @@ module gitea.loveuer.com/yizhisec/packages -go 1.23 +go 1.23.0 toolchain go1.24.3 require ( + github.com/dgraph-io/ristretto/v2 v2.2.0 github.com/fatih/color v1.18.0 github.com/gin-gonic/gin v1.10.1 github.com/glebarez/sqlite v1.11.0 @@ -23,7 +24,7 @@ require ( filippo.io/edwards25519 v1.1.0 // indirect github.com/bytedance/sonic v1.11.6 // indirect github.com/bytedance/sonic/loader v0.1.1 // indirect - github.com/cespare/xxhash/v2 v2.1.1 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/cloudwego/base64x v0.1.4 // indirect github.com/cloudwego/iasm v0.2.0 // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect @@ -60,7 +61,7 @@ require ( golang.org/x/arch v0.8.0 // indirect golang.org/x/net v0.27.0 // indirect golang.org/x/sync v0.11.0 // indirect - golang.org/x/sys v0.30.0 // indirect + golang.org/x/sys v0.31.0 // indirect golang.org/x/text v0.22.0 // indirect google.golang.org/protobuf v1.34.1 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect diff --git a/go.sum b/go.sum index 57d0934..c214889 100644 --- a/go.sum +++ b/go.sum @@ -6,8 +6,9 @@ github.com/bytedance/sonic v1.11.6 h1:oUp34TzMlL+OY1OUWxHqsdkgC/Zfc85zGqw9siXjrc github.com/bytedance/sonic v1.11.6/go.mod h1:LysEHSvpvDySVdC2f87zGWf6CIKJcAvqab1ZaiQtds4= github.com/bytedance/sonic/loader v0.1.1 h1:c+e5Pt1k/cy5wMveRDyk2X4B9hF4g7an8N3zCYjJFNM= github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU= -github.com/cespare/xxhash/v2 v2.1.1 h1:6MnRN8NT7+YBpUIWxHtefFZOKTAPgGjpQSxqLNn0+qY= github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cloudwego/base64x v0.1.4 h1:jwCgWpFanWmN8xoIUHa2rtzmkd5J2plF/dnLS6Xd/0Y= github.com/cloudwego/base64x v0.1.4/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w= github.com/cloudwego/iasm v0.2.0 h1:1KNIy1I1H9hNNFEEH3DVnI4UujN+1zjpuk6gwHLTssg= @@ -19,6 +20,10 @@ github.com/creack/pty v1.1.7/go.mod h1:lj5s0c3V2DBrqTV7llrYr5NG6My20zk30Fl46Y7Do github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dgraph-io/ristretto/v2 v2.2.0 h1:bkY3XzJcXoMuELV8F+vS8kzNgicwQFAaGINAEJdWGOM= +github.com/dgraph-io/ristretto/v2 v2.2.0/go.mod h1:RZrm63UmcBAaYWC1DotLYBmTvgkrs0+XhBd7Npn7/zI= +github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da h1:aIftn67I1fkbMa512G+w+Pxci9hJPB8oMnkcP3iZF38= +github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= @@ -321,8 +326,8 @@ golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= -golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc= -golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.31.0 h1:ioabZlmFYtWhL+TRYpcnNlLwhyxaM9kWTDEmfnprqik= +golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=