100 lines
2.2 KiB
Go
100 lines
2.2 KiB
Go
|
package storage
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/base64"
|
||
|
|
"encoding/json"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"log"
|
||
|
|
"net/http"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/mailgun/groupcache/v2"
|
||
|
|
"github.com/spf13/viper"
|
||
|
|
)
|
||
|
|
|
||
|
|
type EtcdGroupCacheHandler struct {
|
||
|
|
EtcdHandler *EtcdHandler
|
||
|
|
Group *groupcache.Group
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewEtcdGroupCacheHandler(cfg *viper.Viper) (*EtcdGroupCacheHandler, error) {
|
||
|
|
var endpoint = cfg.GetString("storage.kv.groupcache.endpoint")
|
||
|
|
etcdHandler, err := NewEtcdHandler(cfg)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
pool := groupcache.NewHTTPPoolOpts("http://"+endpoint, &groupcache.HTTPPoolOptions{})
|
||
|
|
server := http.Server{
|
||
|
|
Addr: endpoint,
|
||
|
|
Handler: pool,
|
||
|
|
}
|
||
|
|
go func() {
|
||
|
|
log.Printf("Serving....\n")
|
||
|
|
if err := server.ListenAndServe(); err != nil {
|
||
|
|
log.Fatal(err)
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
group := groupcache.NewGroup("data", 3000000, groupcache.GetterFunc(
|
||
|
|
func(ctx context.Context, key string, dest groupcache.Sink) error {
|
||
|
|
|
||
|
|
resp, err := etcdHandler.Get(key)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if resp == nil {
|
||
|
|
return errors.New("key not found in ETCD")
|
||
|
|
}
|
||
|
|
decoded, err := base64.StdEncoding.DecodeString(resp.(string))
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("failed to decode value for key %s: %v", key, err)
|
||
|
|
}
|
||
|
|
dest.SetString(string(decoded), time.Now().Add(time.Minute*5))
|
||
|
|
return nil
|
||
|
|
},
|
||
|
|
))
|
||
|
|
|
||
|
|
return &EtcdGroupCacheHandler{
|
||
|
|
EtcdHandler: etcdHandler,
|
||
|
|
Group: group,
|
||
|
|
}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (h *EtcdGroupCacheHandler) Put(k string, v interface{}) error {
|
||
|
|
value, err := json.Marshal(v)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
return h.EtcdHandler.Put(k, value)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (h *EtcdGroupCacheHandler) PutWithTTL(k string, v interface{}, duration time.Duration) error {
|
||
|
|
value, err := json.Marshal(v)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
return h.EtcdHandler.PutWithTTL(k, value, duration)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (h *EtcdGroupCacheHandler) Get(k string) (interface{}, error) {
|
||
|
|
var value []byte
|
||
|
|
err := h.Group.Get(nil, k, groupcache.AllocatingByteSliceSink(&value))
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return value, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (h *EtcdGroupCacheHandler) Delete(k string) error {
|
||
|
|
err := h.EtcdHandler.Delete(k)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
err = h.Group.Remove(context.Background(), k)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|