jybase/go-xweb/httpsession/memorystore.go

package httpsession

import (
	"sync"
	"time"
)

var _ Store = NewMemoryStore(30)

type sessionNode struct {
	lock   sync.RWMutex
	kvs    map[string]interface{}
	last   time.Time
	maxAge time.Duration
}

func (node *sessionNode) Get(key string) interface{} {
	node.lock.RLock()
	v := node.kvs[key]
	node.lock.RUnlock()
	node.lock.Lock()
	node.last = time.Now()
	node.lock.Unlock()
	return v
}

func (node *sessionNode) GetMultiple() map[string]interface{} {
	m := make(map[string]interface{})
	node.lock.RLock()
	for k, v := range node.kvs {
		m[k] = v
	}
	node.lock.RUnlock()
	node.lock.Lock()
	node.last = time.Now()
	node.lock.Unlock()
	return m
}

func (node *sessionNode) Set(key string, v interface{}) {
	node.lock.Lock()
	node.kvs[key] = v
	node.last = time.Now()
	node.lock.Unlock()
}

func (node *sessionNode) SetMultiple(m map[string]interface{}) error {
	node.lock.Lock()
	for k, v := range m {
		node.kvs[k] = v
	}
	node.last = time.Now()
	node.lock.Unlock()
	return nil
}

func (node *sessionNode) Del(keys ...string) {
	node.lock.Lock()
	for _, v := range keys {
		delete(node.kvs, v)
	}
	node.last = time.Now()
	node.lock.Unlock()
}

//根据自定义字段，更新
func (node *sessionNode) UpdateByCustomField(findkey string, findvalue interface{}, setkey string, setvalue interface{}) bool {
	node.lock.Lock()
	if v2, ok := node.kvs[findkey]; ok && v2 == findvalue {
		//存在
		if setkey != "" {
			flag := setkey[:1]
			switch flag {
			case "+":
				nkey := setkey[1:]
				node.kvs[nkey] = node.kvs[nkey].(int) + setvalue.(int)
			case "-":
				nkey := setkey[1:]
				node.kvs[nkey] = node.kvs[nkey].(int) - setvalue.(int)
			default:
				node.kvs[setkey] = setvalue
			}
			node.last = time.Now()
		} else {
			mapVal := setvalue.(*map[string]interface{})
			for k, v := range *mapVal {
				node.kvs[k] = v
			}
			node.last = time.Now()
		}
	}
	node.lock.Unlock()
	return true
}

type MemoryStore struct {
	lock       sync.RWMutex
	nodes      map[Id]*sessionNode
	GcInterval time.Duration
	maxAge     time.Duration
}

func NewMemoryStore(maxAge time.Duration) *MemoryStore {
	return &MemoryStore{nodes: make(map[Id]*sessionNode),
		maxAge: maxAge, GcInterval: 10 * time.Second}
}

func (store *MemoryStore) SetMaxAge(maxAge time.Duration) {
	store.lock.Lock()
	store.maxAge = maxAge
	store.lock.Unlock()
}

func (store *MemoryStore) Get(id Id, key string) interface{} {
	store.lock.RLock()
	node, ok := store.nodes[id]
	store.lock.RUnlock()
	if !ok {
		return nil
	}

	if store.maxAge > 0 && time.Now().Sub(node.last) > node.maxAge {
		// lazy DELETE expire
		store.lock.Lock()
		delete(store.nodes, id)
		store.lock.Unlock()
		return nil
	}

	return node.Get(key)
}

func (store *MemoryStore) GetMultiple(id Id) map[string]interface{} {
	store.lock.RLock()
	node, ok := store.nodes[id]
	store.lock.RUnlock()
	if !ok {
		return make(map[string]interface{})
	}

	if store.maxAge > 0 && time.Now().Sub(node.last) > node.maxAge {
		// lazy DELETE expire
		store.lock.Lock()
		delete(store.nodes, id)
		store.lock.Unlock()
		return make(map[string]interface{})
	}

	return node.GetMultiple()
}

func (store *MemoryStore) Set(id Id, key string, value interface{}) {
	store.lock.RLock()
	node, ok := store.nodes[id]
	store.lock.RUnlock()
	if !ok {
		store.lock.Lock()
		node = store.newNode()
		node.kvs[key] = value
		store.nodes[id] = node
		store.lock.Unlock()
	}
	node.Set(key, value)
}

func (store *MemoryStore) SetMultiple(id Id, m map[string]interface{}) error {
	store.lock.RLock()
	node, ok := store.nodes[id]
	store.lock.RUnlock()
	if !ok {
		store.lock.Lock()
		node = store.newNode()
		for k, v := range m {
			node.kvs[k] = v
		}
		store.nodes[id] = node
		store.lock.Unlock()
	}
	return node.SetMultiple(m)
}

func (store *MemoryStore) newNode() *sessionNode {
	return &sessionNode{
		kvs:    make(map[string]interface{}),
		last:   time.Now(),
		maxAge: store.maxAge,
	}
}

func (store *MemoryStore) Add(id Id) {
	node := store.newNode()
	store.lock.Lock()
	store.nodes[id] = node
	store.lock.Unlock()
}

func (store *MemoryStore) Del(id Id, keys ...string) bool {
	store.lock.RLock()
	node, ok := store.nodes[id]
	store.lock.RUnlock()
	if ok {
		node.Del(keys...)
	}
	return true
}

func (store *MemoryStore) Exist(id Id) bool {
	store.lock.RLock()
	defer store.lock.RUnlock()
	_, ok := store.nodes[id]
	return ok
}

func (store *MemoryStore) Clear(id Id) bool {
	store.lock.Lock()
	defer store.lock.Unlock()
	delete(store.nodes, id)
	return true
}

func (store *MemoryStore) Run() error {
	time.AfterFunc(store.GcInterval, func() {
		store.Run()
		store.GC()
	})
	return nil
}

//根据指定字段更新Session
func (store *MemoryStore) UpdateByCustomField(findkey string, findvalue interface{}, setkey string, setvalue interface{}) bool {
	store.lock.RLock()
	for _, v := range store.nodes {
		v.UpdateByCustomField(findkey, findvalue, setkey, setvalue)
	}
	store.lock.RUnlock()
	return true
}

//
func (store *MemoryStore) GC() {
	store.lock.Lock()
	defer store.lock.Unlock()
	if store.maxAge == 0 {
		return
	}
	var i, j int
	for k, v := range store.nodes {
		if j > 20 || i > 5 {
			break
		}
		if time.Now().Sub(v.last) > v.maxAge {
			delete(store.nodes, k)
			i = i + 1
		}
		j = j + 1
	}

}