package ledis import ( "encoding/binary" "errors" "fmt" "github.com/siddontang/go/num" "github.com/siddontang/ledisdb/store" "strings" "time" ) type KVPair struct { Key []byte Value []byte } var errKVKey = errors.New("invalid encode kv key") func checkKeySize(key []byte) error { if len(key) > MaxKeySize || len(key) == 0 { return errKeySize } return nil } func checkValueSize(value []byte) error { if len(value) > MaxValueSize { return errValueSize } return nil } func (db *DB) encodeKVKey(key []byte) []byte { ek := make([]byte, len(key)+1+len(db.indexVarBuf)) pos := copy(ek, db.indexVarBuf) ek[pos] = KVType pos++ copy(ek[pos:], key) return ek } func (db *DB) decodeKVKey(ek []byte) ([]byte, error) { pos, err := db.checkKeyIndex(ek) if err != nil { return nil, err } if pos+1 > len(ek) || ek[pos] != KVType { return nil, errKVKey } pos++ return ek[pos:], nil } func (db *DB) encodeKVMinKey() []byte { ek := db.encodeKVKey(nil) return ek } func (db *DB) encodeKVMaxKey() []byte { ek := db.encodeKVKey(nil) ek[len(ek)-1] = KVType + 1 return ek } func (db *DB) incr(key []byte, delta int64) (int64, error) { if err := checkKeySize(key); err != nil { return 0, err } var err error key = db.encodeKVKey(key) t := db.kvBatch t.Lock() defer t.Unlock() var n int64 n, err = StrInt64(db.bucket.Get(key)) if err != nil { return 0, err } n += delta t.Put(key, num.FormatInt64ToSlice(n)) err = t.Commit() return n, err } // ps : here just focus on deleting the key-value data, // any other likes expire is ignore. func (db *DB) delete(t *batch, key []byte) int64 { key = db.encodeKVKey(key) t.Delete(key) return 1 } func (db *DB) setExpireAt(key []byte, when int64) (int64, error) { t := db.kvBatch t.Lock() defer t.Unlock() if exist, err := db.Exists(key); err != nil || exist == 0 { return 0, err } else { db.expireAt(t, KVType, key, when) if err := t.Commit(); err != nil { return 0, err } } return 1, nil } func (db *DB) Decr(key []byte) (int64, error) { return db.incr(key, -1) } func (db *DB) DecrBy(key []byte, decrement int64) (int64, error) { return db.incr(key, -decrement) } func (db *DB) Del(keys ...[]byte) (int64, error) { if len(keys) == 0 { return 0, nil } codedKeys := make([][]byte, len(keys)) for i, k := range keys { codedKeys[i] = db.encodeKVKey(k) } t := db.kvBatch t.Lock() defer t.Unlock() for i, k := range keys { t.Delete(codedKeys[i]) db.rmExpire(t, KVType, k) } err := t.Commit() return int64(len(keys)), err } func (db *DB) Exists(key []byte) (int64, error) { if err := checkKeySize(key); err != nil { return 0, err } var err error key = db.encodeKVKey(key) var v []byte v, err = db.bucket.Get(key) if v != nil && err == nil { return 1, nil } return 0, err } func (db *DB) Get(key []byte) ([]byte, error) { if err := checkKeySize(key); err != nil { return nil, err } key = db.encodeKVKey(key) return db.bucket.Get(key) } func (db *DB) GetSlice(key []byte) (store.Slice, error) { if err := checkKeySize(key); err != nil { return nil, err } key = db.encodeKVKey(key) return db.bucket.GetSlice(key) } func (db *DB) GetSet(key []byte, value []byte) ([]byte, error) { if err := checkKeySize(key); err != nil { return nil, err } else if err := checkValueSize(value); err != nil { return nil, err } key = db.encodeKVKey(key) t := db.kvBatch t.Lock() defer t.Unlock() oldValue, err := db.bucket.Get(key) if err != nil { return nil, err } t.Put(key, value) err = t.Commit() return oldValue, err } func (db *DB) Incr(key []byte) (int64, error) { return db.incr(key, 1) } func (db *DB) IncrBy(key []byte, increment int64) (int64, error) { return db.incr(key, increment) } func (db *DB) MGet(keys ...[]byte) ([][]byte, error) { values := make([][]byte, len(keys)) it := db.bucket.NewIterator() defer it.Close() for i := range keys { if err := checkKeySize(keys[i]); err != nil { return nil, err } values[i] = it.Find(db.encodeKVKey(keys[i])) } return values, nil } func (db *DB) MSet(args ...KVPair) error { if len(args) == 0 { return nil } t := db.kvBatch var err error var key []byte var value []byte t.Lock() defer t.Unlock() for i := 0; i < len(args); i++ { if err := checkKeySize(args[i].Key); err != nil { return err } else if err := checkValueSize(args[i].Value); err != nil { return err } key = db.encodeKVKey(args[i].Key) value = args[i].Value t.Put(key, value) } err = t.Commit() return err } func (db *DB) Set(key []byte, value []byte) error { if err := checkKeySize(key); err != nil { return err } else if err := checkValueSize(value); err != nil { return err } var err error key = db.encodeKVKey(key) t := db.kvBatch t.Lock() defer t.Unlock() t.Put(key, value) err = t.Commit() return err } func (db *DB) SetNX(key []byte, value []byte) (int64, error) { if err := checkKeySize(key); err != nil { return 0, err } else if err := checkValueSize(value); err != nil { return 0, err } var err error key = db.encodeKVKey(key) var n int64 = 1 t := db.kvBatch t.Lock() defer t.Unlock() if v, err := db.bucket.Get(key); err != nil { return 0, err } else if v != nil { n = 0 } else { t.Put(key, value) err = t.Commit() } return n, err } func (db *DB) SetEX(key []byte, duration int64, value []byte) error { if err := checkKeySize(key); err != nil { return err } else if err := checkValueSize(value); err != nil { return err } else if duration <= 0 { return errExpireValue } ek := db.encodeKVKey(key) t := db.kvBatch t.Lock() defer t.Unlock() t.Put(ek, value) db.expireAt(t, KVType, key, time.Now().Unix()+duration) if err := t.Commit(); err != nil { return err } return nil } func (db *DB) flush() (drop int64, err error) { t := db.kvBatch t.Lock() defer t.Unlock() return db.flushType(t, KVType) } func (db *DB) Expire(key []byte, duration int64) (int64, error) { if duration <= 0 { return 0, errExpireValue } return db.setExpireAt(key, time.Now().Unix()+duration) } func (db *DB) ExpireAt(key []byte, when int64) (int64, error) { if when <= time.Now().Unix() { return 0, errExpireValue } return db.setExpireAt(key, when) } func (db *DB) TTL(key []byte) (int64, error) { if err := checkKeySize(key); err != nil { return -1, err } return db.ttl(KVType, key) } func (db *DB) Persist(key []byte) (int64, error) { if err := checkKeySize(key); err != nil { return 0, err } t := db.kvBatch t.Lock() defer t.Unlock() n, err := db.rmExpire(t, KVType, key) if err != nil { return 0, err } err = t.Commit() return n, err } func (db *DB) SetRange(key []byte, offset int, value []byte) (int64, error) { if len(value) == 0 { return 0, nil } if err := checkKeySize(key); err != nil { return 0, err } else if len(value)+offset > MaxValueSize { return 0, errValueSize } key = db.encodeKVKey(key) t := db.kvBatch t.Lock() defer t.Unlock() oldValue, err := db.bucket.Get(key) if err != nil { return 0, err } extra := offset + len(value) - len(oldValue) if extra > 0 { oldValue = append(oldValue, make([]byte, extra)...) } copy(oldValue[offset:], value) t.Put(key, oldValue) if err := t.Commit(); err != nil { return 0, err } return int64(len(oldValue)), nil } func getRange(start int, end int, valLen int) (int, int) { if start < 0 { start = valLen + start } if end < 0 { end = valLen + end } if start < 0 { start = 0 } if end < 0 { end = 0 } if end >= valLen { end = valLen - 1 } return start, end } func (db *DB) GetRange(key []byte, start int, end int) ([]byte, error) { if err := checkKeySize(key); err != nil { return nil, err } key = db.encodeKVKey(key) value, err := db.bucket.Get(key) if err != nil { return nil, err } valLen := len(value) start, end = getRange(start, end, valLen) if start > end { return nil, nil } return value[start : end+1], nil } func (db *DB) StrLen(key []byte) (int64, error) { s, err := db.GetSlice(key) if err != nil { return 0, err } n := s.Size() s.Free() return int64(n), nil } func (db *DB) Append(key []byte, value []byte) (int64, error) { if len(value) == 0 { return 0, nil } if err := checkKeySize(key); err != nil { return 0, err } key = db.encodeKVKey(key) t := db.kvBatch t.Lock() defer t.Unlock() oldValue, err := db.bucket.Get(key) if err != nil { return 0, err } if len(oldValue)+len(value) > MaxValueSize { return 0, errValueSize } oldValue = append(oldValue, value...) t.Put(key, oldValue) if err := t.Commit(); err != nil { return 0, nil } return int64(len(oldValue)), nil } func (db *DB) BitOP(op string, destKey []byte, srcKeys ...[]byte) (int64, error) { if err := checkKeySize(destKey); err != nil { return 0, err } op = strings.ToLower(op) if len(srcKeys) == 0 { return 0, nil } else if op == BitNot && len(srcKeys) > 1 { return 0, fmt.Errorf("BITOP NOT has only one srckey") } else if len(srcKeys) < 2 { return 0, nil } key := db.encodeKVKey(srcKeys[0]) value, err := db.bucket.Get(key) if err != nil { return 0, err } if op == BitNot { for i := 0; i < len(value); i++ { value[i] ^= value[i] } } else { for j := 1; j < len(srcKeys); j++ { if err := checkKeySize(srcKeys[j]); err != nil { return 0, err } key = db.encodeKVKey(srcKeys[j]) ovalue, err := db.bucket.Get(key) if err != nil { return 0, err } if len(value) < len(ovalue) { value, ovalue = ovalue, value } for i := 0; i < len(ovalue); i++ { switch op { case BitAND: value[i] &= ovalue[i] case BitOR: value[i] |= ovalue[i] case BitXOR: value[i] ^= ovalue[i] default: return 0, fmt.Errorf("invalid op type: %s", op) } } for i := len(ovalue); i < len(value); i++ { switch op { case BitAND: value[i] &= 0 case BitOR: value[i] |= 0 case BitXOR: value[i] ^= 0 } } } } key = db.encodeKVKey(destKey) t := db.kvBatch t.Lock() defer t.Unlock() t.Put(key, value) if err := t.Commit(); err != nil { return 0, err } return int64(len(value)), nil } var bitsInByte = [256]int32{0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7, 4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8} func numberBitCount(i uint32) uint32 { i = i - ((i >> 1) & 0x55555555) i = (i & 0x33333333) + ((i >> 2) & 0x33333333) return (((i + (i >> 4)) & 0x0F0F0F0F) * 0x01010101) >> 24 } func (db *DB) BitCount(key []byte, start int, end int) (int64, error) { if err := checkKeySize(key); err != nil { return 0, err } key = db.encodeKVKey(key) value, err := db.bucket.Get(key) if err != nil { return 0, err } start, end = getRange(start, end, len(value)) value = value[start : end+1] var n int64 = 0 pos := 0 for ; pos+4 <= len(value); pos = pos + 4 { n += int64(numberBitCount(binary.BigEndian.Uint32(value[pos : pos+4]))) } for ; pos < len(value); pos++ { n += int64(bitsInByte[value[pos]]) } return n, nil } func (db *DB) BitPos(key []byte, on int, start int, end int) (int64, error) { if err := checkKeySize(key); err != nil { return 0, err } if (on & ^1) != 0 { return 0, fmt.Errorf("bit must be 0 or 1, not %d", on) } var skipValue uint8 = 0 if on == 0 { skipValue = 0xFF } key = db.encodeKVKey(key) value, err := db.bucket.Get(key) if err != nil { return 0, err } start, end = getRange(start, end, len(value)) value = value[start : end+1] for i, v := range value { if uint8(v) != skipValue { for j := 0; j < 8; j++ { isNull := uint8(v)&(1<> 3) extra := byteOffset + 1 - len(value) if extra > 0 { value = append(value, make([]byte, extra)...) } byteVal := value[byteOffset] bit := 7 - uint8(uint32(offset)&0x7) bitVal := byteVal & (1 << bit) byteVal &= ^(1 << bit) byteVal |= (uint8(on&0x1) << bit) value[byteOffset] = byteVal t.Put(key, value) if err := t.Commit(); err != nil { return 0, err } if bitVal > 0 { return 1, nil } else { return 0, nil } } func (db *DB) GetBit(key []byte, offset int) (int64, error) { if err := checkKeySize(key); err != nil { return 0, err } key = db.encodeKVKey(key) value, err := db.bucket.Get(key) if err != nil { return 0, err } byteOffset := uint32(offset) >> 3 bit := 7 - uint8(uint32(offset)&0x7) if byteOffset >= uint32(len(value)) { return 0, nil } bitVal := value[byteOffset] & (1 << bit) if bitVal > 0 { return 1, nil } else { return 0, nil } }