Merge pull request #15 from LyricTian/develop

Token storage optimization

Author: LyricTian

.gitignore

@@ -25,3 +25,4 @@ _testmain.go
 
 # OSX
 *.DS_Store
+*.db

example/server/main.go

@@ -27,7 +27,7 @@ func init() {
 func main() {
 	manager := manage.NewDefaultManager()
 	// token store
-	manager.MapTokenStorage(store.NewMemoryTokenStore(0))
+	manager.MustTokenStorage(store.NewMemoryTokenStore())
 	// client store
 	manager.MapClientStorage(store.NewTestClientStore(&models.Client{
 		ID:     "222222",

generates/access.go

@@ -6,7 +6,7 @@ import (
 	"strconv"
 	"strings"
 
-	"github.com/LyricTian/go.uuid"
+	"github.com/satori/go.uuid"
 	"gopkg.in/oauth2.v3"
 )
 

generates/authorize.go

@@ -5,7 +5,7 @@ import (
 	"encoding/base64"
 	"strings"
 
-	"github.com/LyricTian/go.uuid"
+	"github.com/satori/go.uuid"
 	"gopkg.in/oauth2.v3"
 )
 

manage/manage_test.go

@@ -22,7 +22,7 @@ func TestManager(t *testing.T) {
 		})
 
 		Convey("Memory store test", func() {
-			manager.MapTokenStorage(store.NewMemoryTokenStore(0))
+			manager.MustTokenStorage(store.NewMemoryTokenStore())
 			testManager(manager)
 		})
 	})

manage/manager.go

@@ -1,12 +1,10 @@
 package manage
 
 import (
-	"time"
-
-	"github.com/LyricTian/inject"
-
 	"reflect"
+	"time"
 
+	"github.com/codegangsta/inject"
 	"gopkg.in/oauth2.v3"
 	"gopkg.in/oauth2.v3/errors"
 	"gopkg.in/oauth2.v3/generates"

server/server_test.go

@@ -27,7 +27,7 @@ var (
 
 func init() {
 	manager = manage.NewDefaultManager()
-	manager.MapTokenStorage(store.NewMemoryTokenStore(0))
+	manager.MustTokenStorage(store.NewMemoryTokenStore())
 }
 
 func clientStore(domain string) oauth2.ClientStore {

store/token.go

@@ -1,170 +1,126 @@
 package store
 
 import (
-	"container/list"
-	"strconv"
-	"sync"
+	"encoding/json"
 	"time"
 
+	"github.com/satori/go.uuid"
+	"github.com/tidwall/buntdb"
 	"gopkg.in/oauth2.v3"
+	"gopkg.in/oauth2.v3/models"
 )
 
 // NewMemoryTokenStore Create a token store instance based on memory
-// gcInterval Perform garbage collection intervals(The default is 30 seconds)
-func NewMemoryTokenStore(gcInterval time.Duration) oauth2.TokenStore {
-	if gcInterval == 0 {
-		gcInterval = time.Second * 30
-	}
-	store := &MemoryTokenStore{
-		gcInterval: gcInterval,
-		basicList:  list.New(),
-		data:       make(map[string]oauth2.TokenInfo),
-		access:     make(map[string]string),
-		refresh:    make(map[string]string),
+func NewMemoryTokenStore() (store oauth2.TokenStore, err error) {
+	store, err = NewFileTokenStore(":memory:")
+	return
+}
+
+// NewFileTokenStore Create a token store instance based on file
+func NewFileTokenStore(filename string) (store oauth2.TokenStore, err error) {
+	db, err := buntdb.Open(filename)
+	if err != nil {
+		return
 	}
-	go store.gc()
-	return store
+	store = &TokenStore{db: db}
+	return
 }
 
-// MemoryTokenStore Memory storage for token
-type MemoryTokenStore struct {
-	gcInterval time.Duration
-	globalID   int64
-	lock       sync.RWMutex
-	data       map[string]oauth2.TokenInfo
-	access     map[string]string
-	refresh    map[string]string
-	basicList  *list.List
-	listLock   sync.RWMutex
+// TokenStore Token storage based on buntdb(https://github.com/tidwall/buntdb)
+type TokenStore struct {
+	db *buntdb.DB
 }
 
-func (mts *MemoryTokenStore) gc() {
-	time.AfterFunc(mts.gcInterval, func() {
-		defer mts.gc()
-		rmeles := make([]*list.Element, 0, 32)
-		mts.listLock.RLock()
-		ele := mts.basicList.Front()
-		mts.listLock.RUnlock()
-		for ele != nil {
-			if rm := mts.gcElement(ele); rm {
-				rmeles = append(rmeles, ele)
+// Create Create and store the new token information
+func (ts *TokenStore) Create(info oauth2.TokenInfo) (err error) {
+	ct := time.Now()
+	jv, err := json.Marshal(info)
+	if err != nil {
+		return
+	}
+	basicID := uuid.NewV4().String()
+	aexp := info.GetAccessExpiresIn()
+	rexp := aexp
+
+	err = ts.db.Update(func(tx *buntdb.Tx) (err error) {
+		if refresh := info.GetRefresh(); refresh != "" {
+			rexp = info.GetRefreshCreateAt().Add(info.GetRefreshExpiresIn()).Sub(ct)
+			if aexp.Seconds() > rexp.Seconds() {
+				aexp = rexp
+			}
+			_, _, err = tx.Set(refresh, basicID, &buntdb.SetOptions{Expires: true, TTL: rexp})
+			if err != nil {
+				return
 			}
-			mts.listLock.RLock()
-			ele = ele.Next()
-			mts.listLock.RUnlock()
 		}
-
-		for _, e := range rmeles {
-			mts.listLock.Lock()
-			mts.basicList.Remove(e)
-			mts.listLock.Unlock()
+		_, _, err = tx.Set(basicID, string(jv), &buntdb.SetOptions{Expires: true, TTL: rexp})
+		if err != nil {
+			return
 		}
+		_, _, err = tx.Set(info.GetAccess(), basicID, &buntdb.SetOptions{Expires: true, TTL: aexp})
+		return
 	})
+	return
 }
 
-func (mts *MemoryTokenStore) gcElement(ele *list.Element) (rm bool) {
-	basicID := ele.Value.(string)
-	mts.lock.RLock()
-	ti, ok := mts.data[basicID]
-	mts.lock.RUnlock()
-	if !ok {
-		rm = true
+// remove key
+func (ts *TokenStore) remove(key string) (err error) {
+	verr := ts.db.Update(func(tx *buntdb.Tx) (err error) {
+		_, err = tx.Delete(key)
+		return
+	})
+	if verr == buntdb.ErrNotFound {
 		return
 	}
-	ct := time.Now()
-	if refresh := ti.GetRefresh(); refresh != "" &&
-		ti.GetRefreshCreateAt().Add(ti.GetRefreshExpiresIn()).Before(ct) {
-		mts.lock.RLock()
-		delete(mts.access, ti.GetAccess())
-		delete(mts.refresh, refresh)
-		delete(mts.data, basicID)
-		mts.lock.RUnlock()
-		rm = true
-	} else if ti.GetAccessCreateAt().Add(ti.GetAccessExpiresIn()).Before(ct) {
-		mts.lock.RLock()
-		delete(mts.access, ti.GetAccess())
-		if refresh := ti.GetRefresh(); refresh == "" {
-			delete(mts.data, basicID)
-			rm = true
-		}
-		mts.lock.RUnlock()
-	}
+	err = verr
 	return
 }
 
-func (mts *MemoryTokenStore) getBasicID(id int64, info oauth2.TokenInfo) string {
-	return info.GetClientID() + "_" + strconv.FormatInt(id, 10)
+// RemoveByAccess Use the access token to delete the token information
+func (ts *TokenStore) RemoveByAccess(access string) (err error) {
+	err = ts.remove(access)
+	return
 }
 
-// Create Create and store the new token information
-func (mts *MemoryTokenStore) Create(info oauth2.TokenInfo) (err error) {
-	mts.lock.Lock()
-	defer mts.lock.Unlock()
-	mts.globalID++
-	basicID := mts.getBasicID(mts.globalID, info)
-	mts.data[basicID] = info
-	mts.access[info.GetAccess()] = basicID
-	if refresh := info.GetRefresh(); refresh != "" {
-		mts.refresh[refresh] = basicID
-	}
-
-	mts.listLock.Lock()
-	mts.basicList.PushBack(basicID)
-	mts.listLock.Unlock()
+// RemoveByRefresh Use the refresh token to delete the token information
+func (ts *TokenStore) RemoveByRefresh(refresh string) (err error) {
+	err = ts.remove(refresh)
 	return
 }
 
-// RemoveByAccess Use the access token to delete the token information
-func (mts *MemoryTokenStore) RemoveByAccess(access string) (err error) {
-	mts.lock.RLock()
-	v, ok := mts.access[access]
-	if !ok {
-		mts.lock.RUnlock()
+func (ts *TokenStore) get(key string) (ti oauth2.TokenInfo, err error) {
+	verr := ts.db.View(func(tx *buntdb.Tx) (err error) {
+		basicID, err := tx.Get(key)
+		if err != nil {
+			return
+		}
+		jv, err := tx.Get(basicID)
+		if err != nil {
+			return
+		}
+		var tm models.Token
+		err = json.Unmarshal([]byte(jv), &tm)
+		if err != nil {
+			return
+		}
+		ti = &tm
+		return
+	})
+	if verr == buntdb.ErrNotFound {
 		return
 	}
-	info := mts.data[v]
-	mts.lock.RUnlock()
-
-	mts.lock.Lock()
-	defer mts.lock.Unlock()
-	delete(mts.access, access)
-	if refresh := info.GetRefresh(); refresh == "" {
-		delete(mts.data, v)
-	}
-	return
-}
-
-// RemoveByRefresh Use the refresh token to delete the token information
-func (mts *MemoryTokenStore) RemoveByRefresh(refresh string) (err error) {
-	mts.lock.Lock()
-	defer mts.lock.Unlock()
-	delete(mts.refresh, refresh)
-
+	err = verr
 	return
 }
 
 // GetByAccess Use the access token for token information data
-func (mts *MemoryTokenStore) GetByAccess(access string) (ti oauth2.TokenInfo, err error) {
-	mts.lock.RLock()
-	v, ok := mts.access[access]
-	if !ok {
-		mts.lock.RUnlock()
-		return
-	}
-	ti = mts.data[v]
-	mts.lock.RUnlock()
+func (ts *TokenStore) GetByAccess(access string) (ti oauth2.TokenInfo, err error) {
+	ti, err = ts.get(access)
 	return
 }
 
 // GetByRefresh Use the refresh token for token information data
-func (mts *MemoryTokenStore) GetByRefresh(refresh string) (ti oauth2.TokenInfo, err error) {
-	mts.lock.RLock()
-	v, ok := mts.refresh[refresh]
-	if !ok {
-		mts.lock.RUnlock()
-		return
-	}
-	ti = mts.data[v]
-	mts.lock.RUnlock()
+func (ts *TokenStore) GetByRefresh(refresh string) (ti oauth2.TokenInfo, err error) {
+	ti, err = ts.get(refresh)
 	return
 }