180 lines
4.2 KiB
Go
180 lines
4.2 KiB
Go
|
package storage
|
||
|
|
||
|
import (
|
||
|
"context"
|
||
|
"fmt"
|
||
|
"github.com/spf13/viper"
|
||
|
"gopkg.in/square/go-jose.v2/json"
|
||
|
"log"
|
||
|
"net/http"
|
||
|
"reflect"
|
||
|
"time"
|
||
|
|
||
|
"github.com/mailgun/groupcache/v2"
|
||
|
)
|
||
|
|
||
|
type GroupCacheHandler struct {
|
||
|
KVHandler KVHandler
|
||
|
Group *groupcache.Group
|
||
|
cacheTTL int64
|
||
|
}
|
||
|
|
||
|
func NewGroupCacheHandler(cfg *viper.Viper) (*GroupCacheHandler, error) {
|
||
|
var (
|
||
|
endpoint = cfg.GetString("storage.kv.groupcache.endpoint")
|
||
|
ports = cfg.GetStringSlice("storage.kv.groupcache.ports")
|
||
|
cacheTTL = cfg.GetInt64("storage.kv.groupcache.cacheTTL")
|
||
|
cacheServers []*http.Server
|
||
|
)
|
||
|
|
||
|
var kvHandler KVHandler
|
||
|
var getterFunc groupcache.GetterFunc
|
||
|
kvType := cfg.GetString("storage.kv.groupcache.db")
|
||
|
getterFunc = func(ctx context.Context, key string, dest groupcache.Sink) error {
|
||
|
resp, err := kvHandler.Get(key)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
if resp == nil {
|
||
|
return fmt.Errorf("key not found in etcd: %s", key)
|
||
|
}
|
||
|
jsonResp, err := json.Marshal(resp.(map[string]interface{}))
|
||
|
if err != nil {
|
||
|
// Handle the error
|
||
|
}
|
||
|
dest.SetBytes(jsonResp, time.Now().Add(time.Duration(cacheTTL)*time.Minute))
|
||
|
return nil
|
||
|
|
||
|
}
|
||
|
switch kvType {
|
||
|
case "etcd":
|
||
|
etcdHandler, err := NewEtcdHandler(cfg)
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
kvHandler = etcdHandler
|
||
|
case "badger":
|
||
|
badgerHandler, err := NewBadgerHandler(cfg)
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
kvHandler = badgerHandler
|
||
|
default:
|
||
|
return nil, fmt.Errorf("unsupported kvType: %s", kvType)
|
||
|
}
|
||
|
|
||
|
pool := groupcache.NewHTTPPoolOpts("", &groupcache.HTTPPoolOptions{})
|
||
|
for _, port := range ports {
|
||
|
pool.Set("http://" + endpoint + ":" + port)
|
||
|
server := &http.Server{
|
||
|
Addr: endpoint + ":" + port,
|
||
|
Handler: pool,
|
||
|
}
|
||
|
cacheServers = append(cacheServers, server)
|
||
|
go func(srv *http.Server) {
|
||
|
log.Printf("Serving cache server %s \n", srv.Addr)
|
||
|
if err := srv.ListenAndServe(); err != nil {
|
||
|
log.Fatal(err)
|
||
|
}
|
||
|
}(server)
|
||
|
}
|
||
|
|
||
|
group := groupcache.NewGroup("data", 3000000, getterFunc)
|
||
|
handler := &GroupCacheHandler{
|
||
|
KVHandler: kvHandler,
|
||
|
Group: group,
|
||
|
cacheTTL: cacheTTL,
|
||
|
}
|
||
|
|
||
|
return handler, nil
|
||
|
}
|
||
|
|
||
|
func (h *GroupCacheHandler) Put(k string, v any) error {
|
||
|
// Update the value in the underlying key-value store
|
||
|
err := h.KVHandler.Put(k, v)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
|
||
|
// Check if the key exists in the cache
|
||
|
var value []byte
|
||
|
err = h.Group.Get(nil, k, groupcache.AllocatingByteSliceSink(&value))
|
||
|
if err == nil {
|
||
|
// Key exists in the cache, compare with the new value
|
||
|
if !reflect.DeepEqual(value, v) {
|
||
|
// Values are different, update the cache
|
||
|
jsonValue, err := json.Marshal(v)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
expireTime := time.Now().Add(time.Duration(h.cacheTTL) * time.Minute)
|
||
|
err = h.Group.Set(context.Background(), k, jsonValue, expireTime, false)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
return nil
|
||
|
}
|
||
|
|
||
|
func (h *GroupCacheHandler) PutWithTTL(k string, v any, duration time.Duration) error {
|
||
|
|
||
|
err := h.KVHandler.PutWithTTL(k, v, duration)
|
||
|
// Check if the key exists in the cache
|
||
|
var value []byte
|
||
|
err = h.Group.Get(nil, k, groupcache.AllocatingByteSliceSink(&value))
|
||
|
if err == nil {
|
||
|
// Key exists in the cache, compare with the new value
|
||
|
if !reflect.DeepEqual(value, v) {
|
||
|
// Values are different, update the cache
|
||
|
jsonValue, err := json.Marshal(v)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
expireTime := time.Now().Add(time.Duration(h.cacheTTL) * time.Minute)
|
||
|
err = h.Group.Set(context.Background(), k, jsonValue, expireTime, false)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
return nil
|
||
|
|
||
|
}
|
||
|
|
||
|
func (h *GroupCacheHandler) Get(k string) (any, error) {
|
||
|
var value []byte
|
||
|
err := h.Group.Get(nil, k, groupcache.AllocatingByteSliceSink(&value))
|
||
|
if err != nil {
|
||
|
fmt.Println("error", err)
|
||
|
return nil, err
|
||
|
}
|
||
|
|
||
|
// If the value is empty, return an empty map
|
||
|
if len(value) == 0 {
|
||
|
return make(map[string]interface{}), nil
|
||
|
}
|
||
|
|
||
|
var result map[string]interface{}
|
||
|
err = json.Unmarshal(value, &result)
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
return result, nil
|
||
|
|
||
|
}
|
||
|
|
||
|
func (h *GroupCacheHandler) Delete(k string) error {
|
||
|
err := h.KVHandler.Delete(k)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
err = h.Group.Remove(context.Background(), k)
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
return nil
|
||
|
}
|