diff --git a/README.md b/README.md index 060a70d0..1c2fc472 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,34 @@ if you'd like to see some endpoint or endpoint ideas implemented. This system is completely extensible, so I can introduce them without breaking backwards compatibility. +## Cache + +Cache is a provider for cache that provides a contract to work with external storage. +By default, is in memory, can be implemented redis, mecmached, etc. + +```go +// Cache is a provider for cache +// +// All cache providers must implement methods, which work with storage +// to store a cache of user state, inline results +// for a Telegram bot. +type Cache interface { + Get(kind CacheKind, key string) (interface{}, error) + Set(kind CacheKind, key string, value interface{}) error + Clear(kind CacheKind, key string) error + + Keys(kind CacheKind) []string +} +``` + +To store context to cache, could be used cache context middleware. + +```go +cache := tele.NewInMemoryCache() + +b.Use(middleware.CacheContext(cache)) +``` + ## Context Context is a special type that wraps a huge update structure and represents the context of the current event. It provides several helpers, which allow diff --git a/cache.go b/cache.go new file mode 100644 index 00000000..e93ff3d2 --- /dev/null +++ b/cache.go @@ -0,0 +1,104 @@ +package telebot + +import ( + "errors" + "sync" +) + +type CacheKind string + +const ( + CacheUserContext CacheKind = "user_context" +) + +// Cache is a provider for cache +// +// All cache providers must implement methods, which work with storage +// to store a cache of user state, inline results +// for a Telegram bot. +type Cache interface { + Get(kind CacheKind, key string) (interface{}, error) + Put(kind CacheKind, key string, value interface{}) error + Clear(kind CacheKind, key string) error + + Keys(kind CacheKind) []string +} + +// InMemoryCache is a provider of in memory cache. +// +// Would be enabled by default +type inMemoryCache struct { + lock sync.RWMutex + data map[CacheKind]map[string]interface{} + + keys map[CacheKind][]string // To maintain order of keys +} + +func NewInMemoryCache() Cache { + return &inMemoryCache{} +} + +func (m *inMemoryCache) Get(kind CacheKind, key string) (interface{}, error) { + m.lock.RLock() + defer m.lock.RUnlock() + + if _, ok := m.data[kind]; !ok { + return nil, errors.New("telebot: cache kind is not found") + } + + if _, ok := m.data[kind][key]; !ok { + return nil, errors.New("telebot: cache key is not found") + } + + return m.data[kind][key], nil +} + +func (m *inMemoryCache) Put(kind CacheKind, key string, value interface{}) error { + m.lock.Lock() + defer m.lock.Unlock() + + if _, ok := m.data[kind]; !ok { + m.data[kind] = make(map[string]interface{}) + m.keys[kind] = make([]string, 1) + } + + m.data[kind][key] = value + m.keys[kind] = append(m.keys[kind], key) // Update the keys slice + + return nil +} + +func (m *inMemoryCache) Clear(kind CacheKind, key string) error { + m.lock.RLock() + defer m.lock.RUnlock() + + if _, ok := m.data[kind]; !ok { + return errors.New("telebot: cache kind is not found") + } + + if _, ok := m.data[kind][key]; !ok { + return errors.New("telebot: cache key not found") + } + + delete(m.data[kind], key) // delete value + keysRemove(m.keys[kind], key) // delete key + + return nil +} + +func (m *inMemoryCache) Keys(kind CacheKind) []string { + m.lock.RLock() + defer m.lock.RUnlock() + + return m.keys[kind] +} + +func keysRemove(keys []string, remove string) []string { + for i, k := range keys { + if k == remove { + return append(keys[:i], keys[i+1:]...) + } + } + + return keys +} diff --git a/cache_test.go b/cache_test.go new file mode 100644 index 00000000..ccacb736 --- /dev/null +++ b/cache_test.go @@ -0,0 +1,31 @@ +package telebot + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestInMemoryCache_Get(t *testing.T) { + mc := NewInMemoryCache() + + test, _ := mc.Get(CacheUserContext, "test") + assert.Equal(t, nil, test) +} + +func TestInMemoryCache_Set(t *testing.T) { + mc := NewInMemoryCache() + + _ = mc.Put(CacheUserContext, "test", "test") + test, _ := mc.Get(CacheUserContext, "test") + assert.Equal(t, "test", test) +} + +func TestInMemoryCache_Keys(t *testing.T) { + mc := NewInMemoryCache() + + _ = mc.Put(CacheUserContext, "test", "test") + + keys := mc.Keys(CacheUserContext) + assert.Equal(t, []string{"test"}, keys) +} diff --git a/context.go b/context.go index 98ede2ee..d28b52f1 100644 --- a/context.go +++ b/context.go @@ -157,6 +157,9 @@ type Context interface { // Set saves data in the context. Set(key string, val interface{}) + + // Keys retrieves store keys + Keys() []string } // nativeContext is a native implementation of the Context interface. @@ -503,3 +506,15 @@ func (c *nativeContext) Get(key string) interface{} { defer c.lock.RUnlock() return c.store[key] } + +func (c *nativeContext) Keys() []string { + c.lock.RLock() + defer c.lock.RUnlock() + + keys := make([]string, 0, len(c.store)) + for k := range c.store { + keys = append(keys, k) + } + + return keys +} diff --git a/middleware/cache.go b/middleware/cache.go new file mode 100644 index 00000000..3f438358 --- /dev/null +++ b/middleware/cache.go @@ -0,0 +1,38 @@ +package middleware + +import ( + "log" + + tele "gopkg.in/telebot.v3" +) + +// CacheContext returns a middleware that store context to cache and retreive data from cache to context +// If no custom cache provided, context store will be rested each iteration. +func CacheContext(cache tele.Cache) tele.MiddlewareFunc { + return func(next tele.HandlerFunc) tele.HandlerFunc { + return func(ctx tele.Context) error { + for _, key := range cache.Keys(tele.CacheUserContext) { + value, err := cache.Get(tele.CacheUserContext, key) + if err != nil { + log.Printf("err: %s was happened, %s -> %s was not got from cache", err, value, key) + + return err + } + + ctx.Set(key, value) + } + + defer func() { + for _, key := range ctx.Keys() { + value := ctx.Get(key) + err := cache.Put(tele.CacheUserContext, key, value) + if err != nil { + log.Printf("err: %s was happened, %s -> %s was not stored from context", err, value, key) + } + } + }() + + return next(ctx) + } + } +}