all: support lru and internal backup

Author: changkun

.gitignore

@@ -15,4 +15,6 @@
 
 # directories (remove the comment below to include it)
 vendor/
-build/
+build/
+data/redis
+data/backup

cache.go

@@ -0,0 +1,86 @@
+package main
+
+import (
+	"container/list"
+	"sync"
+	"time"
+)
+
+type item struct {
+	k, v string
+}
+
+// lru is a naive thread-safe lru cache
+type lru struct {
+	cap   uint
+	size  uint
+	elems *list.List // of item
+
+	mu sync.RWMutex
+}
+
+func newLRU(doexpire bool) *lru {
+	l := &lru{
+		cap:   32, // could do it with memory quota
+		size:  0,
+		elems: list.New(),
+		mu:    sync.RWMutex{},
+	}
+	if doexpire {
+		go l.clear()
+	}
+	return l
+}
+
+// clear clears the lru after a while, this is just a dirty
+// solution to prevent if the database is updated but lru is
+// not synced.
+func (l *lru) clear() {
+	t := time.NewTicker(time.Minute * 5)
+	for {
+		select {
+		case <-t.C:
+			l.mu.Lock()
+			for e := l.elems.Front(); e != nil; e = e.Next() {
+				l.elems.Remove(e)
+			}
+			l.mu.Unlock()
+		}
+	}
+}
+
+func (l *lru) Get(k string) (string, bool) {
+	l.mu.RLock()
+	defer l.mu.RUnlock()
+
+	for e := l.elems.Front(); e != nil; e = e.Next() {
+		if e.Value.(*item).k == k {
+			l.elems.MoveToFront(e)
+			return e.Value.(*item).v, true
+		}
+	}
+	return "", false
+}
+
+func (l *lru) Put(k, v string) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	// found from cache
+	i := &item{k, v}
+	for e := l.elems.Front(); e != nil; e = e.Next() {
+		if e.Value.(*item).k == k {
+			i.v = l.elems.Remove(e).(*item).v
+			l.elems.PushFront(i)
+			return
+		}
+	}
+
+	// check if cache is full
+	if l.size+1 > l.cap {
+		l.elems.Remove(l.elems.Back())
+	}
+	l.elems.PushFront(i)
+	l.size++
+	return
+}

cache_test.go

@@ -0,0 +1,83 @@
+package main
+
+import (
+	"math/rand"
+	"testing"
+)
+
+func TestLRU(t *testing.T) {
+	l := newLRU(false)
+	l.cap = 2 // for testing
+
+	if _, ok := l.Get("a"); ok {
+		t.Fatalf("Get value from empty LRU")
+	}
+
+	l.Put("a", "1") // a
+	v, ok := l.Get("a")
+	if !ok { // a -> b
+		t.Fatalf("Get value from LRU found nothing")
+	}
+
+	l.Put("b", "2") // b -> a
+	v, ok = l.Get("a")
+	if !ok { // a -> b
+		t.Fatalf("Get value after Put from LRU found nothing")
+	}
+	if v != "1" {
+		t.Fatalf("Get value from LRU want 1 got %v", v)
+	}
+	l.Put("c", "3") // c -> a
+	_, ok = l.Get("b")
+	if ok {
+		t.Fatalf("Get value success meaning LRU incorrect")
+	}
+	v, ok = l.Get("c")
+	if !ok {
+		t.Fatalf("Get value fail meaning LRU incorrect")
+	}
+	if v != "3" {
+		t.Fatalf("Get value from LRU want 3 got %v", v)
+	}
+}
+
+func rands() string {
+	var alphabet = "qazwsxedcrfvtgbyhnujmikolpQAZWSXEDCRFVTGBYHNUJMIKOLP"
+	ret := make([]byte, 5)
+	for i := 0; i < 5; i++ {
+		ret[i] = alphabet[rand.Intn(len(alphabet))]
+	}
+	return BytesToString(ret)
+}
+
+func BenchmarkLRU(b *testing.B) {
+	l := newLRU(false)
+	l.Put("a", "1")
+	b.Run("Get", func(b *testing.B) {
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				l.Get("a")
+			}
+		})
+	})
+	b.Run("Put-Same", func(b *testing.B) {
+		b.RunParallel(func(pb *testing.PB) {
+			// each goroutine put its own k/v
+			k, v := rands(), rands()
+			for pb.Next() {
+				l.Put(k, v)
+			}
+		})
+	})
+
+	// This is a very naive bench test, especially it
+	// mostly measures the rands().
+	b.Run("Put-Different", func(b *testing.B) {
+		b.RunParallel(func(pb *testing.PB) {
+			for pb.Next() {
+				// each put has a different k/v
+				l.Put(rands(), rands())
+			}
+		})
+	})
+}

config.go

@@ -21,11 +21,13 @@ var (
 )
 
 type config struct {
-	Host  string `yaml:"host"`
-	Addr  string `yaml:"addr"`
-	Store string `yaml:"store"`
-	Log   string `yaml:"log"`
-	S     struct {
+	Host      string `yaml:"host"`
+	Addr      string `yaml:"addr"`
+	Store     string `yaml:"store"`
+	BackupMin int    `yaml:"backup_min"`
+	BackupDir string `yaml:"backup_dir"`
+	Log       string `yaml:"log"`
+	S         struct {
 		Prefix string `yaml:"prefix"`
 	} `yaml:"s"`
 	X struct {

config.yml

@@ -6,6 +6,8 @@
 host: https://golang.design
 addr: :8080
 store: redis://localhost:6379/9
+backup_min: 10
+backup_dir: ./data/backup
 log: "golang.design/redir: "
 s:
   prefix: /s/

data/container.yml

@@ -6,6 +6,8 @@
 host: https://golang.design/
 addr: :8080
 store: redis://redis:6379/9
+backup_min: 1440 # 24h
+backup_dir: ./data/backup
 log: "golang.design/redir: "
 s:
   prefix: /s/

db_test.go

@@ -47,7 +47,7 @@ func check(ctx context.Context, t *testing.T, s *store, key string, rr interface
 	if err != nil {
 		t.Fatalf("Fetch failure, err: %v\n", err)
 	}
-	err = json.Unmarshal([]byte(r), rr)
+	err = json.Unmarshal(StringToBytes(r), rr)
 	if err != nil {
 		t.Fatalf("Unmarshal failure, err: %v\n", err)
 	}

docker/docker-compose.yml

@@ -24,7 +24,7 @@ services:
     container_name: redis
     restart: always
     volumes:
-      - ../data:/data
+      - ../data/redis:/data
     image: redis:3.2
     ports:
       - "6379:6379"

handler.go

@@ -7,9 +7,15 @@ package main
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"html/template"
+	"io/ioutil"
 	"log"
 	"net/http"
+	"os"
+	"time"
+
+	"gopkg.in/yaml.v3"
 )
 
 type visit struct {
@@ -18,7 +24,8 @@ type visit struct {
 }
 
 type server struct {
-	db      *store // *store
+	db      *store
+	cache   *lru
 	visitCh chan visit
 }
 
@@ -37,6 +44,7 @@ func newServer(ctx context.Context) *server {
 	}
 	s := &server{
 		db:      db,
+		cache:   newLRU(true),
 		visitCh: make(chan visit, 100),
 	}
 	go s.counting(ctx)
@@ -69,8 +77,59 @@ func (s *server) registerHandler() {
 	http.Handle(conf.X.Prefix, s.xHandler(conf.X.VCS, conf.X.ImportPath, conf.X.RepoPath))
 }
 
-// backup tries to backup the data store to local files every week.
-// it will keeps the latest 10 backups of the data read from data store.
+// backup tries to backup the data store to local files.
 func (s *server) backup(ctx context.Context) {
-	// TODO: do self-backups
+	if _, err := os.Stat(conf.BackupDir); os.IsNotExist(err) {
+		err := os.Mkdir(conf.BackupDir, os.ModePerm)
+		if err != nil {
+			log.Fatalf("cannot create backup directory, err: %v\n", err)
+		}
+	}
+
+	t := time.NewTicker(time.Minute * time.Duration(conf.BackupMin))
+	log.Printf("internal backup is running...")
+	for {
+		select {
+		case <-t.C:
+			r, err := s.db.Keys(ctx, "*")
+			if err != nil {
+				log.Printf("backup failure, err: %v\n", err)
+				continue
+			}
+			if len(r) == 0 { // no keys for backup
+				continue
+			}
+
+			d := make(map[string]interface{}, len(r))
+			for _, k := range r {
+				v, err := s.db.Fetch(ctx, k)
+				if err != nil {
+					log.Printf("backup failed because of key %v, err: %v\n", k, err)
+					continue
+				}
+				var vv interface{}
+				err = json.Unmarshal(StringToBytes(v), &vv)
+				if err != nil {
+					log.Printf("backup failed because unmarshal of key %v, err: %v\n", k, err)
+					continue
+				}
+				d[k] = vv
+			}
+
+			b, err := yaml.Marshal(d)
+			if err != nil {
+				log.Printf("backup failed when converting to yaml, err: %v\n", err)
+				continue
+			}
+
+			name := fmt.Sprintf("/backup-%s.yml", time.Now().Format(time.RFC3339))
+			err = ioutil.WriteFile(conf.BackupDir+name, b, os.ModePerm)
+			if err != nil {
+				log.Printf("backup failed when saving the file, err: %v\n", err)
+				continue
+			}
+		case <-ctx.Done():
+			return
+		}
+	}
 }

short.go

@@ -105,22 +105,36 @@ func (s *server) sHandler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// TODO: use LRU to optimize fetch speed in the future.
-	raw, err := s.db.FetchAlias(ctx, alias)
-	if err != nil {
-		return
+	checkdb := func(url string) (string, error) {
+		raw, err := s.db.FetchAlias(ctx, alias)
+		if err != nil {
+			return "", err
+		}
+		c := arecord{}
+		err = json.Unmarshal(StringToBytes(raw), &c)
+		if err != nil {
+			return "", err
+		}
+		if url != c.URL {
+			s.cache.Put(alias, c.URL)
+			url = c.URL
+		}
+		return url, nil
 	}
-	c := arecord{}
-	err = json.Unmarshal([]byte(raw), &c)
-	if err != nil {
-		return
+
+	url, ok := s.cache.Get(alias)
+	if !ok {
+		url, err = checkdb(url)
+		if err != nil {
+			return
+		}
 	}
 
 	// redirect the user immediate, but run pv/uv count in background
-	http.Redirect(w, r, c.URL, http.StatusTemporaryRedirect)
+	http.Redirect(w, r, url, http.StatusTemporaryRedirect)
 
-	// count after the forwarding
-	s.visitCh <- visit{s.readIP(r), alias}
+	// count visit in another goroutine so it won't block the redirect.
+	go func() { s.visitCh <- visit{s.readIP(r), alias} }()
 }
 
 // readIP implements a best effort approach to return the real client IP,
@@ -165,7 +179,7 @@ func (s *server) stats(ctx context.Context, w http.ResponseWriter) (retErr error
 			retErr = err
 			return
 		}
-		err = json.Unmarshal([]byte(raw), &ars.Records[i])
+		err = json.Unmarshal(StringToBytes(raw), &ars.Records[i])
 		if err != nil {
 			retErr = err
 			return