make shorturl persist

src/routes/shorturl.go

@@ -1,19 +1,22 @@
 package routes
 
 import (
+	"bytes"
+	"encoding/json"
 	"net/http"
 	"net/url"
 	"regexp"
-	"time"
 
+	"github.com/coreos/bbolt"
 	"github.com/gorilla/mux"
-	"github.com/patrickmn/go-cache"
 	"sour.is/x/toolbox/httpsrv"
+	"sour.is/x/toolbox/log"
 	"sour.is/x/toolbox/uuid"
 )
 
 func init() {
-	s := NewShortManager(365 * 24 * time.Hour)
+	s := &shortDB{}
+	httpsrv.RegisterModule("short", s.config)
 
 	httpsrv.HttpRegister("short", httpsrv.HttpRoutes{
 		{Name: "getShort", Method: "GET", Pattern: "/s/{id}", HandlerFunc: s.getShort},
@@ -21,38 +24,40 @@ func init() {
 	})
 }
 
-type shortManager struct {
-	defaultExpire time.Duration
-	db            *cache.Cache
+type shortDB struct {
+	path   string
+	bucket string
 }
 
-type shortURL struct {
-	ID     string
-	URL    string
-	Secret string
-}
-
-func NewShortManager(defaultExpire time.Duration) *shortManager {
-	return &shortManager{
-		defaultExpire: defaultExpire,
-		db:            cache.New(defaultExpire, defaultExpire/10),
+func (s *shortDB) config(config map[string]string) {
+	s.bucket = "shortURL"
+	if config["bucket"] != "" {
+		s.bucket = config["bucket"]
 	}
-}
 
-func (s *shortManager) GetURL(id string) *shortURL {
-	if u, ok := s.db.Get(id); ok {
-		if url, ok := u.(*shortURL); ok {
-			return url
+	s.path = "data/meta.db"
+	if config["store"] != "" {
+		s.path = config["store"]
+	}
+
+	db, err := bbolt.Open(s.path, 0666, nil)
+	if err != nil {
+		log.Fatalf("ShortURL: failed to open db at [%s]", s.path)
+	}
+	defer db.Close()
+
+	db.Update(func(tx *bbolt.Tx) error {
+		_, err := tx.CreateBucketIfNotExists([]byte(s.bucket))
+		if err != nil {
+			log.Fatalf("ShortURL: create bucket: %s", err)
 		}
-	}
+		return nil
+	})
 
-	return nil
-}
-func (s *shortManager) PutURL(id string, url *shortURL) {
-	s.db.SetDefault(id, url)
+	log.Noticef("ShortURL: opened db at [%s] bucket [%s]", s.path, s.bucket)
 }
 
-func (s *shortManager) getShort(w http.ResponseWriter, r *http.Request) {
+func (s *shortDB) getShort(w http.ResponseWriter, r *http.Request) {
 	vars := mux.Vars(r)
 
 	id := vars["id"]
@@ -67,7 +72,7 @@ func (s *shortManager) getShort(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(http.StatusFound)
 }
 
-func (s *shortManager) putShort(w http.ResponseWriter, r *http.Request) {
+func (s *shortDB) putShort(w http.ResponseWriter, r *http.Request) {
 	defer r.Body.Close()
 
 	vars := mux.Vars(r)
@@ -107,6 +112,12 @@ func (s *shortManager) putShort(w http.ResponseWriter, r *http.Request) {
 	httpsrv.WriteObject(w, 200, short)
 }
 
+type shortURL struct {
+	ID     string
+	URL    string
+	Secret string
+}
+
 func newshort(id, secret, u string) *shortURL {
 	m, err := regexp.MatchString("[a-z-]{1,64}", id)
 	if id == "" || !m || err != nil {
@@ -118,3 +129,64 @@ func newshort(id, secret, u string) *shortURL {
 	}
 	return &shortURL{ID: id, Secret: secret, URL: u}
 }
+
+func (s *shortURL) Bytes() []byte {
+	if s == nil {
+		return nil
+	}
+
+	var w bytes.Buffer
+	json.NewEncoder(&w).Encode(*s)
+	return w.Bytes()
+}
+
+func URLFromBytes(b []byte) *shortURL {
+	if len(b) == 0 {
+		return nil
+	}
+	var s shortURL
+	json.Unmarshal(b, &s)
+
+	log.Debug(s)
+	return &s
+}
+
+func (s *shortDB) GetURL(id string) *shortURL {
+	db, err := bbolt.Open(s.path, 0666, nil)
+	if err != nil {
+		log.Errorf("ShortURL: failed to open db at [%s]", s.path)
+		return nil
+	}
+	defer db.Close()
+
+	var url *shortURL
+
+	err = db.View(func(tx *bbolt.Tx) error {
+		b := tx.Bucket([]byte(s.bucket))
+		v := b.Get([]byte(id))
+
+		url = URLFromBytes(v)
+
+		return nil
+	})
+	if err != nil {
+		log.Errorf("ShortURL: failed to open db at [%s]", s.path)
+	}
+	return url
+}
+func (s *shortDB) PutURL(id string, url *shortURL) {
+	db, err := bbolt.Open(s.path, 0666, nil)
+	if err != nil {
+		log.Errorf("ShortURL: failed to open db at [%s]", s.path)
+		return
+	}
+	defer db.Close()
+
+	err = db.Update(func(tx *bbolt.Tx) error {
+		b := tx.Bucket([]byte(s.bucket))
+		return b.Put([]byte(id), url.Bytes())
+	})
+	if err != nil {
+		log.Errorf("ShortURL: failed to write db at [%s]", s.path)
+	}
+}