升级支持一级缓存

cmd/service/service.go

@@ -52,6 +52,7 @@ var (
 	rateBucketNum = fs.Int("rate-bucket", 10, "rate bucket num")
 	maxLength     = fs.Int("max-length", -1, "code length")
 	alphabet      = fs.String("alphabet", "", "alphabet length")
+	cacheCap      = fs.Int("cache-cap", 2048, "cacheCap length")
 	err           error
 )
 
@@ -73,6 +74,7 @@ func Run() {
 	alphabet = envString("ALPHABET", alphabet)
 	rateBucketNum = envInt("RATE_BUCKET", rateBucketNum)
 	maxLength = envInt("MAX_LENGTH", maxLength)
+	cacheCap = envInt("CACHE_CAP", cacheCap)
 
 	logger = logging.SetLogging(logger, logPath, logLevel)
 
@@ -145,8 +147,7 @@ func initGrpcHandler(endpoints endpoint.Endpoints, g *group.Group) {
 
 func getServiceMiddleware(logger log.Logger) (mw []service.Middleware) {
 	mw = []service.Middleware{}
-	mw = addDefaultServiceMiddleware(logger, mw)
-
+	mw = addDefaultServiceMiddleware(logger, mw, *cacheCap)
 	return
 }
 func getEndpointMiddleware(logger log.Logger) (mw map[string][]kitendpoint.Middleware) {

cmd/service/service_gen.go

@@ -55,8 +55,8 @@ func defaultGrpcOptions(logger log.Logger) map[string][]kitgrpc.ServerOption {
 
 }
 
-func addDefaultServiceMiddleware(logger log.Logger, mw []service.Middleware) []service.Middleware {
-	mw = append(mw, service.LoggingMiddleware(logger))
+func addDefaultServiceMiddleware(logger log.Logger, mw []service.Middleware, cacheCap int) []service.Middleware {
+	mw = append(mw, service.NewLogging(logger, "traceId"), service.NewCache(logger, "traceId", cacheCap))
 	return mw
 }
 

pkg/service/cache.go

@@ -0,0 +1,124 @@
+/**
+ * @Time : 2021/11/15 5:23 PM
+ * @Author : solacowa@gmail.com
+ * @File : cache
+ * @Software: GoLand
+ */
+
+package service
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"github.com/go-kit/kit/log"
+	"github.com/go-kit/kit/log/level"
+)
+
+type lruCache struct {
+	m          map[string]*LinkNode // 指向哈希表的指针
+	cap        int                  // 长度
+	head, tail *LinkNode            // 两个哨兵
+}
+
+type LinkNode struct {
+	key       string
+	val       []byte
+	pre, next *LinkNode
+}
+
+func NewLruCache(capacity int) *lruCache {
+	head := &LinkNode{"", nil, nil, nil}
+	tail := &LinkNode{"", nil, nil, nil}
+	head.next = tail
+	tail.pre = head
+	return &lruCache{make(map[string]*LinkNode), capacity, head, tail}
+}
+
+func (s *lruCache) Get(ctx context.Context, key string) (res []byte, err error) {
+	if node, exist := s.m[key]; exist {
+		s.moveToHead(node)
+		return node.val, nil
+	}
+	return nil, errors.New("key notfound")
+}
+
+func (s *lruCache) moveToHead(node *LinkNode) {
+	head := s.head
+	// 从当前位置删除节点
+	node.pre.next = node.next
+	node.next.pre = node.pre
+	// 将节点插入头部
+	node.next = head.next
+	head.next.pre = node
+	node.pre = head
+	head.next = node
+}
+
+func (s *lruCache) Put(ctx context.Context, key string, value []byte) {
+	head := s.head
+	tail := s.tail
+	if node, exist := s.m[key]; exist { // 如果已经存在对应的节点，要将对应的元素交换到头部；并更新值（值可能变了）
+		node.val = value
+		s.moveToHead(node)
+	} else { // 节点不存在
+		// 创建节点
+		node := &LinkNode{key, value, nil, nil}
+		// 判断缓存容量是否已经满了
+		if len(s.m) == s.cap {
+			// 删除最后的元素
+			delete(s.m, tail.pre.key)
+			tail.pre.pre.next = tail
+			tail.pre = tail.pre.pre
+		}
+		// 将节点放到头部
+		node.next = head.next
+		node.pre = head
+		head.next.pre = node
+		head.next = node
+		// 参入哈希表中
+		s.m[key] = node
+	}
+}
+
+type cache struct {
+	lruCache *lruCache
+	traceId  string
+	next     Service
+	pkgName  string
+	logger   log.Logger
+}
+
+func (s *cache) Get(ctx context.Context, code string) (redirect *Redirect, err error) {
+	get, err := s.lruCache.Get(ctx, code)
+	if err == nil {
+		err = json.Unmarshal(get, &redirect)
+		if err == nil {
+			return redirect, nil
+		}
+		_ = level.Error(s.logger).Log("json", "Unmarshal", "err", err.Error())
+	}
+	defer func() {
+		b, _ := json.Marshal(redirect)
+		s.lruCache.Put(ctx, code, b)
+	}()
+	_ = level.Warn(s.logger).Log("lruCache", "Get", "err", err.Error())
+	return s.next.Get(ctx, code)
+}
+
+func (s *cache) Post(ctx context.Context, domain string) (redirect *Redirect, err error) {
+	return s.next.Post(ctx, domain)
+}
+
+func NewCache(logger log.Logger, traceId string, cacheCap int) Middleware {
+	logger = log.With(logger, "service", "cache")
+	return func(next Service) Service {
+		return &cache{
+			logger:   logger,
+			next:     next,
+			traceId:  traceId,
+			pkgName:  "service",
+			lruCache: NewLruCache(cacheCap),
+		}
+	}
+}

pkg/service/logging.go

@@ -0,0 +1,50 @@
+/**
+ * @Time : 2021/11/15 4:55 PM
+ * @Author : solacowa@gmail.com
+ * @File : logging
+ * @Software: GoLand
+ */
+
+package service
+
+import (
+	"context"
+	"github.com/go-kit/kit/log"
+	"time"
+)
+
+type logging struct {
+	traceId string
+	logger  log.Logger
+	next    Service
+}
+
+func (s *logging) Get(ctx context.Context, code string) (redirect *Redirect, err error) {
+	defer func(begin time.Time) {
+		_ = s.logger.Log(
+			s.traceId, ctx.Value(s.traceId),
+			"method", "Get", "code", code,
+			"took", time.Since(begin),
+			"err", err,
+		)
+	}(time.Now())
+	return s.next.Get(ctx, code)
+}
+
+func (s *logging) Post(ctx context.Context, domain string) (redirect *Redirect, err error) {
+	defer func(begin time.Time) {
+		_ = s.logger.Log(
+			s.traceId, ctx.Value(s.traceId),
+			"method", "Post", "domain", domain,
+			"took", time.Since(begin),
+			"err", err,
+		)
+	}(time.Now())
+	return s.next.Post(ctx, domain)
+}
+
+func NewLogging(logger log.Logger, traceId string) Middleware {
+	return func(next Service) Service {
+		return &logging{traceId, logger, next}
+	}
+}

pkg/service/middleware.go

@@ -6,35 +6,3 @@
  */
 
 package service
-
-import (
-	"context"
-	"github.com/go-kit/kit/log"
-)
-
-type Middleware func(svc Service) Service
-
-type loggingMiddleware struct {
-	logger log.Logger
-	next   Service
-}
-
-func LoggingMiddleware(logger log.Logger) Middleware {
-	return func(next Service) Service {
-		return &loggingMiddleware{logger, next}
-	}
-}
-
-func (l loggingMiddleware) Get(ctx context.Context, code string) (rs *Redirect, err error) {
-	defer func() {
-		_ = l.logger.Log("method", "Get", "s", code, "err", err)
-	}()
-	return l.next.Get(ctx, code)
-}
-
-func (l loggingMiddleware) Post(ctx context.Context, uri string) (redirect *Redirect, err error) {
-	defer func() {
-		_ = l.logger.Log("method", "Post", "uri", uri, "err", err)
-	}()
-	return l.next.Post(ctx, uri)
-}

pkg/service/service.go

@@ -21,6 +21,8 @@ var (
 	ErrRedirectInvalid  = errors.New("Redirect Invalid")
 )
 
+type Middleware func(svc Service) Service
+
 type Service interface {
 	Get(ctx context.Context, code string) (redirect *Redirect, err error)
 	Post(ctx context.Context, domain string) (redirect *Redirect, err error)