Golang engine.NewBatch函数代码示例

// PERSIST key
func (s *Store) Persist(db uint32, args [][]byte) (int64, error) {
	if len(args) != 1 {
		return 0, errArguments("len(args) = %d, expect = 1", len(args))
	}

	key := args[0]

	if err := s.acquire(); err != nil {
		return 0, errors.Trace(err)
	}
	defer s.release()

	o, err := s.loadStoreRow(db, key)
	if err != nil || o == nil {
		return 0, errors.Trace(err)
	}
	if o.GetExpireAt() == 0 {
		return 0, nil
	}

	fw := &Forward{DB: db, Op: "Persist", Args: args}
	bt := engine.NewBatch()
	o.SetExpireAt(0)
	bt.Set(o.MetaKey(), o.MetaValue())
	return 1, s.commit(bt, fw)
}

func (s *Store) incrInt(db uint32, key []byte, delta int64) (int64, error) {
	o, err := s.loadStringRow(db, key)
	if err != nil {
		return 0, errors.Trace(err)
	}

	bt := engine.NewBatch()
	if o != nil {
		_, err := o.LoadDataValue(s)
		if err != nil {
			return 0, errors.Trace(err)
		}
		v, err := ParseInt(o.Value)
		if err != nil {
			return 0, errors.Trace(err)
		}
		delta += v
	} else {
		o = newStringRow(db, key)
		bt.Set(o.MetaKey(), o.MetaValue())
	}

	o.Value = FormatInt(delta)
	bt.Set(o.DataKey(), o.DataValue())

	fw := &Forward{DB: db, Op: "IncrBy", Args: [][]byte{key, FormatInt(delta)}}
	return delta, s.commit(bt, fw)
}

// SPOP key
func (s *Store) SPop(db uint32, args [][]byte) ([]byte, error) {
	if len(args) != 1 {
		return nil, errArguments("len(args) = %d, expect = 1", len(args))
	}

	key := args[0]

	if err := s.acquire(); err != nil {
		return nil, errors.Trace(err)
	}
	defer s.release()

	o, err := s.loadSetRow(db, key)
	if err != nil || o == nil {
		return nil, errors.Trace(err)
	}

	members, err := o.getMembers(s, 1)
	if err != nil || len(members) == 0 {
		return nil, errors.Trace(err)
	}
	o.Member = members[0]

	bt := engine.NewBatch()
	bt.Del(o.DataKey())
	if o.Size--; o.Size > 0 {
		bt.Set(o.MetaKey(), o.MetaValue())
	} else {
		bt.Del(o.MetaKey())
	}
	fw := &Forward{DB: db, Op: "SRem", Args: [][]byte{key, members[0]}}
	return o.Member, s.commit(bt, fw)
}

// MSETNX key value [key value ...]
func (s *Store) MSetNX(db uint32, args [][]byte) (int64, error) {
	if len(args) == 0 || len(args)%2 != 0 {
		return 0, errArguments("len(args) = %d, expect != 0 && mod 2 = 0", len(args))
	}

	if err := s.acquire(); err != nil {
		return 0, errors.Trace(err)
	}
	defer s.release()

	for i := 0; i < len(args); i += 2 {
		o, err := s.loadStoreRow(db, args[i])
		if err != nil || o != nil {
			return 0, errors.Trace(err)
		}
	}

	ms := &markSet{}
	bt := engine.NewBatch()
	for i := len(args)/2 - 1; i >= 0; i-- {
		key, value := args[i*2], args[i*2+1]
		if !ms.Has(key) {
			o := newStringRow(db, key)
			o.Value = value
			bt.Set(o.DataKey(), o.DataValue())
			bt.Set(o.MetaKey(), o.MetaValue())
			ms.Set(key)
		}
	}

	fw := &Forward{DB: db, Op: "MSet", Args: args}
	return 1, s.commit(bt, fw)
}

// SLOTSRESTORE key ttlms value [key ttlms value ...]
func (s *Store) SlotsRestore(db uint32, args [][]byte) error {
	if len(args) == 0 || len(args)%3 != 0 {
		return errArguments("len(args) = %d, expect != 0 && mod 3 = 0", len(args))
	}

	objs := make([]*rdb.ObjEntry, len(args)/3)
	for i := 0; i < len(objs); i++ {
		key := args[i*3]
		ttlms, err := ParseInt(args[i*3+1])
		if err != nil {
			return errArguments("parse args failed - %s", err)
		}
		value := args[i*3+2]

		expireat := int64(0)
		if ttlms != 0 {
			if v, ok := TTLmsToExpireAt(ttlms); ok && v > 0 {
				expireat = v
			} else {
				return errArguments("parse args[%d] ttlms = %d", i*3+1, ttlms)
			}
		}

		obj, err := rdb.DecodeDump(value)
		if err != nil {
			return errArguments("decode args[%d] failed, %s", i*3+2, err)
		}

		objs[i] = &rdb.ObjEntry{
			DB:       db,
			Key:      key,
			ExpireAt: uint64(expireat),
			Value:    obj,
		}
	}

	if err := s.acquire(); err != nil {
		return errors.Trace(err)
	}
	defer s.release()

	ms := &markSet{}
	bt := engine.NewBatch()
	for i := len(objs) - 1; i >= 0; i-- {
		e := objs[i]
		if ms.Has(e.Key) {
			log.Debugf("[%d] restore batch, db = %d, key = %v, ignore", i, e.DB, e.Key)
			continue
		} else {
			log.Debugf("[%d] restore batch, db = %d, key = %v", i, e.DB, e.Key)
		}
		if err := s.restore(bt, e.DB, e.Key, int64(e.ExpireAt), e.Value); err != nil {
			log.Warningf("restore object failed, db = %d, key = %v, err = %s", e.DB, e.Key, err)
			return errors.Trace(err)
		}
		ms.Set(e.Key)
	}
	fw := &Forward{DB: db, Op: "SlotsRestore", Args: args}
	return s.commit(bt, fw)
}

// RPOP key
func (s *Store) RPop(db uint32, args [][]byte) ([]byte, error) {
	if len(args) != 1 {
		return nil, errArguments("len(args) = %d, expect = 1", len(args))
	}

	key := args[0]

	if err := s.acquire(); err != nil {
		return nil, errors.Trace(err)
	}
	defer s.release()

	o, err := s.loadListRow(db, key)
	if err != nil || o == nil {
		return nil, errors.Trace(err)
	}

	o.Index = o.Rindex - 1
	if _, err := o.LoadDataValue(s); err != nil {
		return nil, errors.Trace(err)
	} else {
		bt := engine.NewBatch()
		bt.Del(o.DataKey())
		if o.Rindex--; o.Lindex < o.Rindex {
			bt.Set(o.MetaKey(), o.MetaValue())
		} else {
			bt.Del(o.MetaKey())
		}
		fw := &Forward{DB: db, Op: "RPop", Args: args}
		return o.Value, s.commit(bt, fw)
	}
}

// DEL key [key ...]
func (s *Store) Del(db uint32, args [][]byte) (int64, error) {
	if len(args) == 0 {
		return 0, errArguments("len(args) = %d, expect != 0", len(args))
	}

	keys := args

	if err := s.acquire(); err != nil {
		return 0, errors.Trace(err)
	}
	defer s.release()

	ms := &markSet{}
	bt := engine.NewBatch()
	for _, key := range keys {
		if !ms.Has(key) {
			exists, err := s.deleteIfExists(bt, db, key)
			if err != nil {
				return 0, errors.Trace(err)
			}
			if exists {
				ms.Set(key)
			}
		}
	}
	fw := &Forward{DB: db, Op: "Del", Args: args}
	return ms.Len(), s.commit(bt, fw)
}

func (s *Store) incrFloat(db uint32, key []byte, delta float64) (float64, error) {
	o, err := s.loadStringRow(db, key)
	if err != nil {
		return 0, errors.Trace(err)
	}

	bt := engine.NewBatch()
	if o != nil {
		_, err := o.LoadDataValue(s)
		if err != nil {
			return 0, errors.Trace(err)
		}
		v, err := ParseFloat(o.Value)
		if err != nil {
			return 0, errors.Trace(err)
		}
		delta += v
	} else {
		o = newStringRow(db, key)
		bt.Set(o.MetaKey(), o.MetaValue())
	}

	if math.IsNaN(delta) || math.IsInf(delta, 0) {
		return 0, errors.New("increment would produce NaN or Infinity")
	}

	o.Value = FormatFloat(delta)
	bt.Set(o.DataKey(), o.DataValue())

	fw := &Forward{DB: db, Op: "IncrByFloat", Args: [][]byte{key, FormatFloat(delta)}}
	return delta, s.commit(bt, fw)
}

// ZINCRBY key delta member
func (s *Store) ZIncrBy(db uint32, args [][]byte) (float64, error) {
	if len(args) != 3 {
		return 0, errArguments("len(args) = %d, expect = 3", len(args))
	}

	key := args[0]
	delta, err := ParseFloat(args[1])
	if err != nil {
		return 0, errArguments("parse args failed - %s", err)
	}
	member := args[2]

	if err := s.acquire(); err != nil {
		return 0, errors.Trace(err)
	}
	defer s.release()

	o, err := s.loadZSetRow(db, key)
	if err != nil {
		return 0, errors.Trace(err)
	}

	bt := engine.NewBatch()

	var exists bool = false
	if o != nil {
		o.Member = member
		exists, err = o.LoadDataValue(s)
		if err != nil {
			return 0, errors.Trace(err)
		} else if exists {
			bt.Del(o.IndexKey())
		}
	} else {
		o = newZSetRow(db, key)
		o.Member = member
	}

	if exists {
		delta += o.Score
	} else {
		o.Size++
		bt.Set(o.MetaKey(), o.MetaValue())
	}
	o.Score = delta
	if math.IsNaN(delta) {
		return 0, errors.New("invalid nan score")
	}

	bt.Set(o.DataKey(), o.DataValue())
	bt.Set(o.IndexKey(), o.IndexValue())

	fw := &Forward{DB: db, Op: "ZIncrBy", Args: args}
	return delta, s.commit(bt, fw)
}

// HINCRBYFLOAT key field delta
func (s *Store) HIncrByFloat(db uint32, args [][]byte) (float64, error) {
	if len(args) != 3 {
		return 0, errArguments("len(args) = %d, expect = 2", len(args))
	}

	key := args[0]
	field := args[1]
	delta, err := ParseFloat(args[2])
	if err != nil {
		return 0, errArguments("parse args failed - %s", err)
	}

	if err := s.acquire(); err != nil {
		return 0, errors.Trace(err)
	}
	defer s.release()

	o, err := s.loadHashRow(db, key)
	if err != nil {
		return 0, errors.Trace(err)
	}

	var exists bool = false
	if o != nil {
		o.Field = field
		exists, err = o.LoadDataValue(s)
		if err != nil {
			return 0, errors.Trace(err)
		}
	} else {
		o = newHashRow(db, key)
		o.Field = field
	}

	bt := engine.NewBatch()
	if exists {
		v, err := ParseFloat(o.Value)
		if err != nil {
			return 0, errors.Trace(err)
		}
		delta += v
	} else {
		o.Size++
		bt.Set(o.MetaKey(), o.MetaValue())
	}

	if math.IsNaN(delta) || math.IsInf(delta, 0) {
		return 0, errors.New("increment would produce NaN or Infinity")
	}

	o.Value = FormatFloat(delta)
	bt.Set(o.DataKey(), o.DataValue())
	fw := &Forward{DB: db, Op: "HIncrByFloat", Args: args}
	return delta, s.commit(bt, fw)
}

// ZREMRANGEBYSCORE key min max
func (s *Store) ZRemRangeByScore(db uint32, args [][]byte) (int64, error) {
	if len(args) != 3 {
		return 0, errArguments("len(args) = %d, expect 3", len(args))
	}

	key := args[0]
	min := args[1]
	max := args[2]

	r, err := parseRangeSpec(min, max)
	if err != nil {
		return 0, errors.Trace(err)
	}

	if err := s.acquire(); err != nil {
		return 0, errors.Trace(err)
	}
	defer s.release()

	o, err := s.loadZSetRow(db, key)
	if err != nil {
		return 0, errors.Trace(err)
	} else if o == nil {
		return 0, nil
	}

	bt := engine.NewBatch()
	n := int64(0)

	f := func(o *zsetRow) error {
		bt.Del(o.DataKey())
		bt.Del(o.IndexKey())
		n++
		return nil
	}

	if err := o.travelInRange(s, r, f); err != nil {
		return 0, errors.Trace(err)
	}

	if n > 0 {
		if o.Size -= n; o.Size > 0 {
			bt.Set(o.MetaKey(), o.MetaValue())
		} else {
			bt.Del(o.MetaKey())
		}
	}

	fw := &Forward{DB: db, Op: "ZRemRangeByScore", Args: args}
	return n, s.commit(bt, fw)
}

// HMSET key field value [field value ...]
func (s *Store) HMSet(db uint32, args [][]byte) error {
	if len(args) == 1 || len(args)%2 != 1 {
		return errArguments("len(args) = %d, expect != 1 && mod 2 = 1", len(args))
	}

	key := args[0]

	var eles = make([]*rdb.HashElement, len(args)/2)
	for i := 0; i < len(eles); i++ {
		e := &rdb.HashElement{}
		e.Field = args[i*2+1]
		e.Value = args[i*2+2]
		eles[i] = e
	}

	if err := s.acquire(); err != nil {
		return errors.Trace(err)
	}
	defer s.release()

	o, err := s.loadHashRow(db, key)
	if err != nil {
		return errors.Trace(err)
	}

	if o == nil {
		o = newHashRow(db, key)
	}

	ms := &markSet{}
	bt := engine.NewBatch()
	for _, e := range eles {
		o.Field, o.Value = e.Field, e.Value
		exists, err := o.TestDataValue(s)
		if err != nil {
			return errors.Trace(err)
		}
		if !exists {
			ms.Set(o.Field)
		}
		bt.Set(o.DataKey(), o.DataValue())
	}

	n := ms.Len()
	if n != 0 {
		o.Size += n
		bt.Set(o.MetaKey(), o.MetaValue())
	}
	fw := &Forward{DB: db, Op: "HMSet", Args: args}
	return s.commit(bt, fw)
}

// LTRIM key beg end
func (s *Store) LTrim(db uint32, args [][]byte) error {
	if len(args) != 3 {
		return errArguments("len(args) = %d, expect = 2", len(args))
	}

	key := args[0]
	beg, err := ParseInt(args[1])
	if err != nil {
		return errArguments("parse args failed - %s", err)
	}
	end, err := ParseInt(args[2])
	if err != nil {
		return errArguments("parse args failed - %s", err)
	}

	if err := s.acquire(); err != nil {
		return errors.Trace(err)
	}
	defer s.release()

	o, err := s.loadListRow(db, key)
	if err != nil || o == nil {
		return errors.Trace(err)
	}

	beg = maxIntValue(adjustIndex(beg, o.Lindex, o.Rindex), o.Lindex)
	end = minIntValue(adjustIndex(end, o.Lindex, o.Rindex), o.Rindex-1)
	if beg == o.Lindex && end == o.Rindex-1 {
		return nil
	}

	bt := engine.NewBatch()
	if beg <= end {
		for o.Index = o.Lindex; o.Index < beg; o.Index++ {
			bt.Del(o.DataKey())
		}
		for o.Index = o.Rindex - 1; o.Index > end; o.Index-- {
			bt.Del(o.DataKey())
		}
		o.Lindex, o.Rindex = beg, end+1
		bt.Set(o.MetaKey(), o.MetaValue())
	} else {
		for o.Index = o.Lindex; o.Index < o.Rindex; o.Index++ {
			bt.Del(o.DataKey())
		}
		bt.Del(o.MetaKey())
	}
	fw := &Forward{DB: db, Op: "LTrim", Args: args}
	return s.commit(bt, fw)
}

func (s *Store) loadStoreRow(db uint32, key []byte) (storeRow, error) {
	o, err := loadStoreRow(s, db, key)
	if err != nil || o == nil {
		return nil, errors.Trace(err)
	}

	if s.needDeleteIfExpired() && o.IsExpired() {
		bt := engine.NewBatch()
		if err := o.deleteObject(s, bt); err != nil {
			return nil, errors.Trace(err)
		}
		fw := &Forward{DB: db, Op: "Del", Args: [][]byte{key}}
		return nil, s.commit(bt, fw)
	}
	return o, nil
}

// SETRANGE key offset value
func (s *Store) SetRange(db uint32, args [][]byte) (int64, error) {
	if len(args) != 3 {
		return 0, errArguments("len(args) = %d, expect = 3", len(args))
	}

	key := args[0]
	offset, err := ParseUint(args[1])
	if err != nil {
		return 0, errArguments("parse args failed - %s", err)
	}
	value := args[2]

	if offset > maxVarbytesLen {
		return 0, errArguments("offset = %d", offset)
	}

	if err := s.acquire(); err != nil {
		return 0, errors.Trace(err)
	}
	defer s.release()

	o, err := s.loadStringRow(db, key)
	if err != nil {
		return 0, errors.Trace(err)
	}

	bt := engine.NewBatch()
	if o != nil {
		_, err := o.LoadDataValue(s)
		if err != nil {
			return 0, errors.Trace(err)
		}
	} else {
		o = newStringRow(db, key)
		bt.Set(o.MetaKey(), o.MetaValue())
	}

	if n := int(offset) + len(value); n > len(o.Value) {
		o.Value = append(o.Value, make([]byte, n-len(o.Value))...)
	}

	copy(o.Value[offset:], value)
	bt.Set(o.DataKey(), o.DataValue())

	fw := &Forward{DB: db, Op: "SetRange", Args: args}
	return int64(len(o.Value)), s.commit(bt, fw)
}