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 }