db_str.go

package rosedb

import (
	"bytes"
	"sync"
	"time"

	"github.com/roseduan/rosedb/index"
	"github.com/roseduan/rosedb/storage"
	"github.com/roseduan/rosedb/utils"
)

// StrIdx string index.
type StrIdx struct {
	mu      *sync.RWMutex
	idxList *index.SkipList
}

func newStrIdx() *StrIdx {
	return &StrIdx{
		idxList: index.NewSkipList(), mu: new(sync.RWMutex),
	}
}

// Set set key to hold the string value. If key already holds a value, it is overwritten.
// Any previous time to live associated with the key is discarded on successful Set operation.
func (db *RoseDB) Set(key, value interface{}) error {
	encKey, encVal, err := db.encode(key, value)
	if err != nil {
		return err
	}
	return db.setVal(encKey, encVal)
}

// SetNx is short for "Set if not exists", set key to hold string value if key does not exist.
// In that case, it is equal to Set. When key already holds a value, no operation is performed.
func (db *RoseDB) SetNx(key, value interface{}) (ok bool, err error) {
	encKey, encVal, err := db.encode(key, value)
	if err != nil {
		return false, err
	}
	if exist := db.StrExists(encKey); exist {
		return
	}

	if err = db.Set(encKey, encVal); err == nil {
		ok = true
	}
	return
}

// SetEx set key to hold the string value and set key to timeout after a given number of seconds.
func (db *RoseDB) SetEx(key, value interface{}, duration int64) (err error) {
	if duration <= 0 {
		return ErrInvalidTTL
	}

	encKey, encVal, err := db.encode(key, value)
	if err != nil {
		return err
	}

	db.strIndex.mu.Lock()
	defer db.strIndex.mu.Unlock()

	deadline := time.Now().Unix() + duration
	e := storage.NewEntryWithExpire(encKey, encVal, deadline, String, StringExpire)
	if err = db.store(e); err != nil {
		return
	}

	// set String index info, stored at skip list.
	if err = db.setIndexer(e); err != nil {
		return
	}
	// set expired info.
	db.expires[String][string(encKey)] = deadline
	return
}

// Get get the value of key. If the key does not exist an error is returned.
func (db *RoseDB) Get(key, dest interface{}) error {
	encKey, err := utils.EncodeKey(key)
	if err != nil {
		return err
	}

	if err := db.checkKeyValue(encKey, nil); err != nil {
		return err
	}

	db.strIndex.mu.RLock()
	defer db.strIndex.mu.RUnlock()

	val, err := db.getVal(encKey)
	if err != nil {
		return err
	}

	if len(val) > 0 {
		err = utils.DecodeValue(val, dest)
	}
	return err
}

// GetSet set key to value and returns the old value stored at key.
// If the key not exist, return an err.
func (db *RoseDB) GetSet(key, value, dest interface{}) (err error) {
	err = db.Get(key, dest)
	if err != nil && err != ErrKeyNotExist && err != ErrKeyExpired {
		return
	}
	return db.Set(key, value)
}

// MSet set multiple keys to multiple values
func (db *RoseDB) MSet(values ...interface{}) error {
	if len(values)%2 != 0 {
		return ErrWrongNumberOfArgs
	}

	var keys [][]byte
	var vals [][]byte

	for i := 0; i < len(values); i += 2 {
		encKey, encVal, err := db.encode(values[i], values[i+1])
		if err != nil {
			return err
		}
		if err := db.checkKeyValue(encKey, encVal); err != nil {
			return err
		}

		var skip bool
		if db.config.IdxMode == KeyValueMemMode {
			existVal, err := db.getVal(encKey)
			if err != nil && err != ErrKeyExpired && err != ErrKeyNotExist {
				return err
			}

			// if the existed value is the same as the set value, pass this key and value.
			if bytes.Compare(existVal, encVal) == 0 {
				skip = true
			}
		}

		if !skip {
			keys = append(keys, encKey)
			vals = append(vals, encVal)
		}
	}
	if len(keys) == 0 || len(vals) == 0 {
		return nil
	}

	db.strIndex.mu.Lock()
	defer db.strIndex.mu.Unlock()

	for i := 0; i < len(keys); i++ {
		e := storage.NewEntryNoExtra(keys[i], vals[i], String, StringSet)
		if err := db.store(e); err != nil {
			return err
		}

		// clear expire time.
		delete(db.expires[String], string(keys[i]))

		// set String index info, stored at skip list.
		if err := db.setIndexer(e); err != nil {
			return err
		}

		// set into cache if necessary.
		db.cache.Set(keys[i], vals[i])
	}
	return nil
}

// MGet get the values of all the given keys
func (db *RoseDB) MGet(keys ...interface{}) ([]interface{}, error) {
	var encKeys [][]byte
	for _, key := range keys {
		encKey, err := utils.EncodeKey(key)
		if err != nil {
			return nil, err
		}
		if err := db.checkKeyValue(encKey, nil); err != nil {
			return nil, err
		}

		encKeys = append(encKeys, encKey)
	}

	db.strIndex.mu.RLock()
	defer db.strIndex.mu.RUnlock()

	var vals []interface{}
	for _, encKey := range encKeys {
		val, err := db.getVal(encKey)
		if err != nil && err != ErrKeyNotExist && err != ErrKeyExpired {
			return nil, err
		}
		vals = append(vals, val)
	}
	return vals, nil
}

// Append if key already exists and is a string, this command appends the value at the end of the string.
// If key does not exist it is created and set as an empty string, so Append will be similar to Set in this special case.
// The original type of the value must be string, otherwise you will get a wrong value after appending.
func (db *RoseDB) Append(key interface{}, value string) (err error) {
	encKey, encVal, err := db.encode(key, value)
	if err != nil {
		return err
	}
	if err := db.checkKeyValue(encKey, encVal); err != nil {
		return err
	}

	var existVal []byte
	err = db.Get(key, &existVal)
	if err != nil && err != ErrKeyNotExist && err != ErrKeyExpired {
		return err
	}

	existVal = append(existVal, []byte(value)...)
	return db.Set(encKey, existVal)
}

// StrExists check whether the key exists.
func (db *RoseDB) StrExists(key interface{}) bool {
	encKey, err := utils.EncodeKey(key)
	if err != nil {
		return false
	}
	if err := db.checkKeyValue(encKey, nil); err != nil {
		return false
	}

	db.strIndex.mu.RLock()
	defer db.strIndex.mu.RUnlock()

	exist := db.strIndex.idxList.Exist(encKey)
	if exist && !db.checkExpired(encKey, String) {
		return true
	}
	return false
}

// Remove remove the value stored at key.
func (db *RoseDB) Remove(key interface{}) error {
	encKey, err := utils.EncodeKey(key)
	if err != nil {
		return err
	}
	if err := db.checkKeyValue(encKey, nil); err != nil {
		return err
	}

	db.strIndex.mu.Lock()
	defer db.strIndex.mu.Unlock()

	e := storage.NewEntryNoExtra(encKey, nil, String, StringRem)
	if err := db.store(e); err != nil {
		return err
	}

	db.strIndex.idxList.Remove(encKey)
	delete(db.expires[String], string(encKey))
	db.cache.Remove(encKey)
	return nil
}

// Expire set the expiration time of the key.
func (db *RoseDB) Expire(key interface{}, duration int64) (err error) {
	encKey, err := utils.EncodeKey(key)
	if err != nil {
		return err
	}
	if duration <= 0 {
		return ErrInvalidTTL
	}

	db.strIndex.mu.Lock()
	defer db.strIndex.mu.Unlock()

	var value []byte
	if value, err = db.getVal(encKey); err != nil {
		return
	}

	deadline := time.Now().Unix() + duration
	e := storage.NewEntryWithExpire(encKey, value, deadline, String, StringExpire)
	if err = db.store(e); err != nil {
		return err
	}

	db.expires[String][string(encKey)] = deadline
	return
}

// Persist clear expiration time.
func (db *RoseDB) Persist(key interface{}) (err error) {
	encKey, err := utils.EncodeKey(key)
	if err != nil {
		return err
	}

	db.strIndex.mu.Lock()
	defer db.strIndex.mu.Unlock()

	val, err := db.getVal(encKey)
	if err != nil {
		return err
	}
	e := storage.NewEntryNoExtra(encKey, val, String, StringPersist)
	if err = db.store(e); err != nil {
		return
	}

	if err = db.setIndexer(e); err != nil {
		return
	}
	delete(db.expires[String], string(encKey))
	db.cache.Set(encKey, val)
	return
}

// TTL Time to live.
func (db *RoseDB) TTL(key interface{}) (ttl int64) {
	encKey, err := utils.EncodeKey(key)
	if err != nil {
		return
	}

	db.strIndex.mu.Lock()
	defer db.strIndex.mu.Unlock()

	deadline, exist := db.expires[String][string(encKey)]
	if !exist {
		return
	}
	if expired := db.checkExpired(encKey, String); expired {
		return
	}

	return deadline - time.Now().Unix()
}

func (db *RoseDB) setVal(key, value []byte) (err error) {
	if err = db.checkKeyValue(key, value); err != nil {
		return err
	}

	// If the existed value is the same as the set value, nothing will be done.
	if db.config.IdxMode == KeyValueMemMode {
		var existVal []byte
		existVal, err = db.getVal(key)
		if err != nil && err != ErrKeyExpired && err != ErrKeyNotExist {
			return
		}

		if bytes.Compare(existVal, value) == 0 {
			// clear expire time.
			if _, ok := db.expires[String][string(key)]; ok {
				delete(db.expires[String], string(key))
			}
			return
		}
	}

	db.strIndex.mu.Lock()
	defer db.strIndex.mu.Unlock()

	e := storage.NewEntryNoExtra(key, value, String, StringSet)
	if err := db.store(e); err != nil {
		return err
	}

	// clear expire time.
	delete(db.expires[String], string(key))
	// set String index info, stored at skip list.
	if err = db.setIndexer(e); err != nil {
		return
	}

	// set into cache if necessary.
	db.cache.Set(key, value)
	return
}

func (db *RoseDB) setIndexer(e *storage.Entry) error {
	activeFile, err := db.getActiveFile(String)
	if err != nil {
		return err
	}
	// string indexes, stored in skiplist.
	idx := &index.Indexer{
		Meta: &storage.Meta{
			Key: e.Meta.Key,
		},
		FileId: activeFile.Id,
		Offset: activeFile.Offset - int64(e.Size()),
	}

	// in KeyValueMemMode, both key and value will store in memory.
	if db.config.IdxMode == KeyValueMemMode {
		idx.Meta.Value = e.Meta.Value
	}

	db.strIndex.idxList.Put(idx.Meta.Key, idx)
	return nil
}

func (db *RoseDB) getVal(key []byte) ([]byte, error) {
	// Get index info from a skip list in memory.
	node := db.strIndex.idxList.Get(key)
	if node == nil {
		return nil, ErrKeyNotExist
	}

	idx, _ := node.Value().(*index.Indexer)
	if idx == nil {
		return nil, ErrNilIndexer
	}

	// Check if the key is expired.
	if db.checkExpired(key, String) {
		return nil, ErrKeyExpired
	}

	// In KeyValueMemMode, the value will be stored in memory.
	// So get the value from the index info.
	if db.config.IdxMode == KeyValueMemMode {
		return idx.Meta.Value, nil
	}

	// In KeyOnlyMemMode, the value not in memory.
	// So get the value from cache if exists in lru cache.
	// Otherwise, get the value from the db file at the offset.
	if db.config.IdxMode == KeyOnlyMemMode {
		// get from cache if necessary.
		if value, ok := db.cache.Get(key); ok {
			return value, nil
		}

		df, err := db.getActiveFile(String)
		if err != nil {
			return nil, err
		}
		if idx.FileId != df.Id {
			df = db.archFiles[String][idx.FileId]
		}

		e, err := df.Read(idx.Offset)
		if err != nil {
			return nil, err
		}
		value := e.Meta.Value

		// set value into cache.
		db.cache.Set(key, value)
		return value, nil
	}
	return nil, ErrKeyNotExist
}