github.com/hdt3213/godis@v1.2.9/database/list.go (about)

     1  package database
     2  
     3  import (
     4  	"fmt"
     5  	"strconv"
     6  	"strings"
     7  
     8  	List "github.com/hdt3213/godis/datastruct/list"
     9  	"github.com/hdt3213/godis/interface/database"
    10  	"github.com/hdt3213/godis/interface/redis"
    11  	"github.com/hdt3213/godis/lib/utils"
    12  	"github.com/hdt3213/godis/redis/protocol"
    13  )
    14  
    15  func (db *DB) getAsList(key string) (List.List, protocol.ErrorReply) {
    16  	entity, ok := db.GetEntity(key)
    17  	if !ok {
    18  		return nil, nil
    19  	}
    20  	list, ok := entity.Data.(List.List)
    21  	if !ok {
    22  		return nil, &protocol.WrongTypeErrReply{}
    23  	}
    24  	return list, nil
    25  }
    26  
    27  func (db *DB) getOrInitList(key string) (list List.List, isNew bool, errReply protocol.ErrorReply) {
    28  	list, errReply = db.getAsList(key)
    29  	if errReply != nil {
    30  		return nil, false, errReply
    31  	}
    32  	isNew = false
    33  	if list == nil {
    34  		list = List.NewQuickList()
    35  		db.PutEntity(key, &database.DataEntity{
    36  			Data: list,
    37  		})
    38  		isNew = true
    39  	}
    40  	return list, isNew, nil
    41  }
    42  
    43  // execLIndex gets element of list at given list
    44  func execLIndex(db *DB, args [][]byte) redis.Reply {
    45  	// parse args
    46  	key := string(args[0])
    47  	index64, err := strconv.ParseInt(string(args[1]), 10, 64)
    48  	if err != nil {
    49  		return protocol.MakeErrReply("ERR value is not an integer or out of range")
    50  	}
    51  	index := int(index64)
    52  
    53  	// get entity
    54  	list, errReply := db.getAsList(key)
    55  	if errReply != nil {
    56  		return errReply
    57  	}
    58  	if list == nil {
    59  		return &protocol.NullBulkReply{}
    60  	}
    61  
    62  	size := list.Len() // assert: size > 0
    63  	if index < -1*size {
    64  		return &protocol.NullBulkReply{}
    65  	} else if index < 0 {
    66  		index = size + index
    67  	} else if index >= size {
    68  		return &protocol.NullBulkReply{}
    69  	}
    70  
    71  	val, _ := list.Get(index).([]byte)
    72  	return protocol.MakeBulkReply(val)
    73  }
    74  
    75  // execLLen gets length of list
    76  func execLLen(db *DB, args [][]byte) redis.Reply {
    77  	// parse args
    78  	key := string(args[0])
    79  
    80  	list, errReply := db.getAsList(key)
    81  	if errReply != nil {
    82  		return errReply
    83  	}
    84  	if list == nil {
    85  		return protocol.MakeIntReply(0)
    86  	}
    87  
    88  	size := int64(list.Len())
    89  	return protocol.MakeIntReply(size)
    90  }
    91  
    92  // execLPop removes the first element of list, and return it
    93  func execLPop(db *DB, args [][]byte) redis.Reply {
    94  	// parse args
    95  	key := string(args[0])
    96  
    97  	// get data
    98  	list, errReply := db.getAsList(key)
    99  	if errReply != nil {
   100  		return errReply
   101  	}
   102  	if list == nil {
   103  		return &protocol.NullBulkReply{}
   104  	}
   105  
   106  	val, _ := list.Remove(0).([]byte)
   107  	if list.Len() == 0 {
   108  		db.Remove(key)
   109  	}
   110  	db.addAof(utils.ToCmdLine3("lpop", args...))
   111  	return protocol.MakeBulkReply(val)
   112  }
   113  
   114  var lPushCmd = []byte("LPUSH")
   115  
   116  func undoLPop(db *DB, args [][]byte) []CmdLine {
   117  	key := string(args[0])
   118  	list, errReply := db.getAsList(key)
   119  	if errReply != nil {
   120  		return nil
   121  	}
   122  	if list == nil || list.Len() == 0 {
   123  		return nil
   124  	}
   125  	element, _ := list.Get(0).([]byte)
   126  	return []CmdLine{
   127  		{
   128  			lPushCmd,
   129  			args[0],
   130  			element,
   131  		},
   132  	}
   133  }
   134  
   135  // execLPush inserts element at head of list
   136  func execLPush(db *DB, args [][]byte) redis.Reply {
   137  	key := string(args[0])
   138  	values := args[1:]
   139  
   140  	// get or init entity
   141  	list, _, errReply := db.getOrInitList(key)
   142  	if errReply != nil {
   143  		return errReply
   144  	}
   145  
   146  	// insert
   147  	for _, value := range values {
   148  		list.Insert(0, value)
   149  	}
   150  
   151  	db.addAof(utils.ToCmdLine3("lpush", args...))
   152  	return protocol.MakeIntReply(int64(list.Len()))
   153  }
   154  
   155  func undoLPush(db *DB, args [][]byte) []CmdLine {
   156  	key := string(args[0])
   157  	count := len(args) - 1
   158  	cmdLines := make([]CmdLine, 0, count)
   159  	for i := 0; i < count; i++ {
   160  		cmdLines = append(cmdLines, utils.ToCmdLine("LPOP", key))
   161  	}
   162  	return cmdLines
   163  }
   164  
   165  // execLPushX inserts element at head of list, only if list exists
   166  func execLPushX(db *DB, args [][]byte) redis.Reply {
   167  	key := string(args[0])
   168  	values := args[1:]
   169  
   170  	// get or init entity
   171  	list, errReply := db.getAsList(key)
   172  	if errReply != nil {
   173  		return errReply
   174  	}
   175  	if list == nil {
   176  		return protocol.MakeIntReply(0)
   177  	}
   178  
   179  	// insert
   180  	for _, value := range values {
   181  		list.Insert(0, value)
   182  	}
   183  	db.addAof(utils.ToCmdLine3("lpushx", args...))
   184  	return protocol.MakeIntReply(int64(list.Len()))
   185  }
   186  
   187  // execLRange gets elements of list in given range
   188  func execLRange(db *DB, args [][]byte) redis.Reply {
   189  	// parse args
   190  	key := string(args[0])
   191  	start64, err := strconv.ParseInt(string(args[1]), 10, 64)
   192  	if err != nil {
   193  		return protocol.MakeErrReply("ERR value is not an integer or out of range")
   194  	}
   195  	start := int(start64)
   196  	stop64, err := strconv.ParseInt(string(args[2]), 10, 64)
   197  	if err != nil {
   198  		return protocol.MakeErrReply("ERR value is not an integer or out of range")
   199  	}
   200  	stop := int(stop64)
   201  
   202  	// get data
   203  	list, errReply := db.getAsList(key)
   204  	if errReply != nil {
   205  		return errReply
   206  	}
   207  	if list == nil {
   208  		return &protocol.EmptyMultiBulkReply{}
   209  	}
   210  
   211  	// compute index
   212  	size := list.Len() // assert: size > 0
   213  	if start < -1*size {
   214  		start = 0
   215  	} else if start < 0 {
   216  		start = size + start
   217  	} else if start >= size {
   218  		return &protocol.EmptyMultiBulkReply{}
   219  	}
   220  	if stop < -1*size {
   221  		stop = 0
   222  	} else if stop < 0 {
   223  		stop = size + stop + 1
   224  	} else if stop < size {
   225  		stop = stop + 1
   226  	} else {
   227  		stop = size
   228  	}
   229  	if stop < start {
   230  		stop = start
   231  	}
   232  
   233  	// assert: start in [0, size - 1], stop in [start, size]
   234  	slice := list.Range(start, stop)
   235  	result := make([][]byte, len(slice))
   236  	for i, raw := range slice {
   237  		bytes, _ := raw.([]byte)
   238  		result[i] = bytes
   239  	}
   240  	return protocol.MakeMultiBulkReply(result)
   241  }
   242  
   243  // execLRem removes element of list at specified index
   244  func execLRem(db *DB, args [][]byte) redis.Reply {
   245  	// parse args
   246  	key := string(args[0])
   247  	count64, err := strconv.ParseInt(string(args[1]), 10, 64)
   248  	if err != nil {
   249  		return protocol.MakeErrReply("ERR value is not an integer or out of range")
   250  	}
   251  	count := int(count64)
   252  	value := args[2]
   253  
   254  	// get data entity
   255  	list, errReply := db.getAsList(key)
   256  	if errReply != nil {
   257  		return errReply
   258  	}
   259  	if list == nil {
   260  		return protocol.MakeIntReply(0)
   261  	}
   262  
   263  	var removed int
   264  	if count == 0 {
   265  		removed = list.RemoveAllByVal(func(a interface{}) bool {
   266  			return utils.Equals(a, value)
   267  		})
   268  	} else if count > 0 {
   269  		removed = list.RemoveByVal(func(a interface{}) bool {
   270  			return utils.Equals(a, value)
   271  		}, count)
   272  	} else {
   273  		removed = list.ReverseRemoveByVal(func(a interface{}) bool {
   274  			return utils.Equals(a, value)
   275  		}, -count)
   276  	}
   277  
   278  	if list.Len() == 0 {
   279  		db.Remove(key)
   280  	}
   281  	if removed > 0 {
   282  		db.addAof(utils.ToCmdLine3("lrem", args...))
   283  	}
   284  
   285  	return protocol.MakeIntReply(int64(removed))
   286  }
   287  
   288  // execLSet puts element at specified index of list
   289  func execLSet(db *DB, args [][]byte) redis.Reply {
   290  	// parse args
   291  	key := string(args[0])
   292  	index64, err := strconv.ParseInt(string(args[1]), 10, 64)
   293  	if err != nil {
   294  		return protocol.MakeErrReply("ERR value is not an integer or out of range")
   295  	}
   296  	index := int(index64)
   297  	value := args[2]
   298  
   299  	// get data
   300  	list, errReply := db.getAsList(key)
   301  	if errReply != nil {
   302  		return errReply
   303  	}
   304  	if list == nil {
   305  		return protocol.MakeErrReply("ERR no such key")
   306  	}
   307  
   308  	size := list.Len() // assert: size > 0
   309  	if index < -1*size {
   310  		return protocol.MakeErrReply("ERR index out of range")
   311  	} else if index < 0 {
   312  		index = size + index
   313  	} else if index >= size {
   314  		return protocol.MakeErrReply("ERR index out of range")
   315  	}
   316  
   317  	list.Set(index, value)
   318  	db.addAof(utils.ToCmdLine3("lset", args...))
   319  	return &protocol.OkReply{}
   320  }
   321  
   322  func undoLSet(db *DB, args [][]byte) []CmdLine {
   323  	key := string(args[0])
   324  	index64, err := strconv.ParseInt(string(args[1]), 10, 64)
   325  	if err != nil {
   326  		return nil
   327  	}
   328  	index := int(index64)
   329  	list, errReply := db.getAsList(key)
   330  	if errReply != nil {
   331  		return nil
   332  	}
   333  	if list == nil {
   334  		return nil
   335  	}
   336  	size := list.Len() // assert: size > 0
   337  	if index < -1*size {
   338  		return nil
   339  	} else if index < 0 {
   340  		index = size + index
   341  	} else if index >= size {
   342  		return nil
   343  	}
   344  	value, _ := list.Get(index).([]byte)
   345  	return []CmdLine{
   346  		{
   347  			[]byte("LSET"),
   348  			args[0],
   349  			args[1],
   350  			value,
   351  		},
   352  	}
   353  }
   354  
   355  // execRPop removes last element of list then return it
   356  func execRPop(db *DB, args [][]byte) redis.Reply {
   357  	// parse args
   358  	key := string(args[0])
   359  
   360  	// get data
   361  	list, errReply := db.getAsList(key)
   362  	if errReply != nil {
   363  		return errReply
   364  	}
   365  	if list == nil {
   366  		return &protocol.NullBulkReply{}
   367  	}
   368  
   369  	val, _ := list.RemoveLast().([]byte)
   370  	if list.Len() == 0 {
   371  		db.Remove(key)
   372  	}
   373  	db.addAof(utils.ToCmdLine3("rpop", args...))
   374  	return protocol.MakeBulkReply(val)
   375  }
   376  
   377  var rPushCmd = []byte("RPUSH")
   378  
   379  func undoRPop(db *DB, args [][]byte) []CmdLine {
   380  	key := string(args[0])
   381  	list, errReply := db.getAsList(key)
   382  	if errReply != nil {
   383  		return nil
   384  	}
   385  	if list == nil || list.Len() == 0 {
   386  		return nil
   387  	}
   388  	element, _ := list.Get(list.Len() - 1).([]byte)
   389  	return []CmdLine{
   390  		{
   391  			rPushCmd,
   392  			args[0],
   393  			element,
   394  		},
   395  	}
   396  }
   397  
   398  func prepareRPopLPush(args [][]byte) ([]string, []string) {
   399  	return []string{
   400  		string(args[0]),
   401  		string(args[1]),
   402  	}, nil
   403  }
   404  
   405  // execRPopLPush pops last element of list-A then insert it to the head of list-B
   406  func execRPopLPush(db *DB, args [][]byte) redis.Reply {
   407  	sourceKey := string(args[0])
   408  	destKey := string(args[1])
   409  
   410  	// get source entity
   411  	sourceList, errReply := db.getAsList(sourceKey)
   412  	if errReply != nil {
   413  		return errReply
   414  	}
   415  	if sourceList == nil {
   416  		return &protocol.NullBulkReply{}
   417  	}
   418  
   419  	// get dest entity
   420  	destList, _, errReply := db.getOrInitList(destKey)
   421  	if errReply != nil {
   422  		return errReply
   423  	}
   424  
   425  	// pop and push
   426  	val, _ := sourceList.RemoveLast().([]byte)
   427  	destList.Insert(0, val)
   428  
   429  	if sourceList.Len() == 0 {
   430  		db.Remove(sourceKey)
   431  	}
   432  
   433  	db.addAof(utils.ToCmdLine3("rpoplpush", args...))
   434  	return protocol.MakeBulkReply(val)
   435  }
   436  
   437  func undoRPopLPush(db *DB, args [][]byte) []CmdLine {
   438  	sourceKey := string(args[0])
   439  	list, errReply := db.getAsList(sourceKey)
   440  	if errReply != nil {
   441  		return nil
   442  	}
   443  	if list == nil || list.Len() == 0 {
   444  		return nil
   445  	}
   446  	element, _ := list.Get(list.Len() - 1).([]byte)
   447  	return []CmdLine{
   448  		{
   449  			rPushCmd,
   450  			args[0],
   451  			element,
   452  		},
   453  		{
   454  			[]byte("LPOP"),
   455  			args[1],
   456  		},
   457  	}
   458  }
   459  
   460  // execRPush inserts element at last of list
   461  func execRPush(db *DB, args [][]byte) redis.Reply {
   462  	// parse args
   463  	key := string(args[0])
   464  	values := args[1:]
   465  
   466  	// get or init entity
   467  	list, _, errReply := db.getOrInitList(key)
   468  	if errReply != nil {
   469  		return errReply
   470  	}
   471  
   472  	// put list
   473  	for _, value := range values {
   474  		list.Add(value)
   475  	}
   476  	db.addAof(utils.ToCmdLine3("rpush", args...))
   477  	return protocol.MakeIntReply(int64(list.Len()))
   478  }
   479  
   480  func undoRPush(db *DB, args [][]byte) []CmdLine {
   481  	key := string(args[0])
   482  	count := len(args) - 1
   483  	cmdLines := make([]CmdLine, 0, count)
   484  	for i := 0; i < count; i++ {
   485  		cmdLines = append(cmdLines, utils.ToCmdLine("RPOP", key))
   486  	}
   487  	return cmdLines
   488  }
   489  
   490  // execRPushX inserts element at last of list only if list exists
   491  func execRPushX(db *DB, args [][]byte) redis.Reply {
   492  	if len(args) < 2 {
   493  		return protocol.MakeErrReply("ERR wrong number of arguments for 'rpush' command")
   494  	}
   495  	key := string(args[0])
   496  	values := args[1:]
   497  
   498  	// get or init entity
   499  	list, errReply := db.getAsList(key)
   500  	if errReply != nil {
   501  		return errReply
   502  	}
   503  	if list == nil {
   504  		return protocol.MakeIntReply(0)
   505  	}
   506  
   507  	// put list
   508  	for _, value := range values {
   509  		list.Add(value)
   510  	}
   511  	db.addAof(utils.ToCmdLine3("rpushx", args...))
   512  
   513  	return protocol.MakeIntReply(int64(list.Len()))
   514  }
   515  
   516  // execLTrim removes elements from both ends a list. delete the list if all elements were trimmmed.
   517  func execLTrim(db *DB, args [][]byte) redis.Reply {
   518  	n := len(args)
   519  	if n != 3 {
   520  		return protocol.MakeErrReply(fmt.Sprintf("ERR wrong number of arguments (given %d, expected 3)", n))
   521  	}
   522  	key := string(args[0])
   523  	start, err := strconv.Atoi(string(args[1]))
   524  	if err != nil {
   525  		return protocol.MakeErrReply("ERR value is not an integer or out of range")
   526  	}
   527  	end, err := strconv.Atoi(string(args[2]))
   528  	if err != nil {
   529  		return protocol.MakeErrReply("ERR value is not an integer or out of range")
   530  	}
   531  
   532  	// get or init entity
   533  	list, errReply := db.getAsList(key)
   534  	if errReply != nil {
   535  		return errReply
   536  	}
   537  	if list == nil {
   538  		return protocol.MakeOkReply()
   539  	}
   540  
   541  	length := list.Len()
   542  	if start < 0 {
   543  		start += length
   544  	}
   545  	if end < 0 {
   546  		end += length
   547  	}
   548  
   549  	leftCount := start
   550  	rightCount := length - end - 1
   551  
   552  	for i := 0; i < leftCount && list.Len() > 0; i++ {
   553  		list.Remove(0)
   554  	}
   555  	for i := 0; i < rightCount && list.Len() > 0; i++ {
   556  		list.RemoveLast()
   557  	}
   558  
   559  	db.addAof(utils.ToCmdLine3("ltrim", args...))
   560  
   561  	return protocol.MakeOkReply()
   562  }
   563  
   564  func execLInsert(db *DB, args [][]byte) redis.Reply {
   565  	n := len(args)
   566  	if n != 4 {
   567  		return protocol.MakeErrReply("ERR wrong number of arguments for 'linsert' command")
   568  	}
   569  	key := string(args[0])
   570  	list, errReply := db.getAsList(key)
   571  	if errReply != nil {
   572  		return errReply
   573  	}
   574  	if list == nil {
   575  		return protocol.MakeIntReply(0)
   576  	}
   577  
   578  	dir := strings.ToLower(string(args[1]))
   579  	if dir != "before" && dir != "after" {
   580  		return protocol.MakeErrReply("ERR syntax error")
   581  	}
   582  
   583  	pivot := string(args[2])
   584  	index := -1
   585  	list.ForEach(func(i int, v interface{}) bool {
   586  		if string(v.([]byte)) == pivot {
   587  			index = i
   588  			return false
   589  		}
   590  		return true
   591  	})
   592  	if index == -1 {
   593  		return protocol.MakeIntReply(-1)
   594  	}
   595  
   596  	val := args[3]
   597  	if dir == "before" {
   598  		list.Insert(index, val)
   599  	} else {
   600  		list.Insert(index+1, val)
   601  	}
   602  
   603  	db.addAof(utils.ToCmdLine3("linsert", args...))
   604  
   605  	return protocol.MakeIntReply(int64(list.Len()))
   606  }
   607  
   608  func init() {
   609  	registerCommand("LPush", execLPush, writeFirstKey, undoLPush, -3, flagWrite).
   610  		attachCommandExtra([]string{redisFlagWrite, redisFlagDenyOOM, redisFlagFast}, 1, 1, 1)
   611  	registerCommand("LPushX", execLPushX, writeFirstKey, undoLPush, -3, flagWrite).
   612  		attachCommandExtra([]string{redisFlagWrite, redisFlagDenyOOM, redisFlagFast}, 1, 1, 1)
   613  	registerCommand("RPush", execRPush, writeFirstKey, undoRPush, -3, flagWrite).
   614  		attachCommandExtra([]string{redisFlagWrite, redisFlagDenyOOM, redisFlagFast}, 1, 1, 1)
   615  	registerCommand("RPushX", execRPushX, writeFirstKey, undoRPush, -3, flagWrite).
   616  		attachCommandExtra([]string{redisFlagWrite, redisFlagDenyOOM, redisFlagFast}, 1, 1, 1)
   617  	registerCommand("LPop", execLPop, writeFirstKey, undoLPop, 2, flagWrite).
   618  		attachCommandExtra([]string{redisFlagWrite, redisFlagFast}, 1, 1, 1)
   619  	registerCommand("RPop", execRPop, writeFirstKey, undoRPop, 2, flagWrite).
   620  		attachCommandExtra([]string{redisFlagWrite, redisFlagFast}, 1, 1, 1)
   621  	registerCommand("RPopLPush", execRPopLPush, prepareRPopLPush, undoRPopLPush, 3, flagWrite).
   622  		attachCommandExtra([]string{redisFlagWrite, redisFlagDenyOOM}, 1, 1, 1)
   623  	registerCommand("LRem", execLRem, writeFirstKey, rollbackFirstKey, 4, flagWrite).
   624  		attachCommandExtra([]string{redisFlagWrite}, 1, 1, 1)
   625  	registerCommand("LLen", execLLen, readFirstKey, nil, 2, flagReadOnly).
   626  		attachCommandExtra([]string{redisFlagReadonly, redisFlagFast}, 1, 1, 1)
   627  	registerCommand("LIndex", execLIndex, readFirstKey, nil, 3, flagReadOnly).
   628  		attachCommandExtra([]string{redisFlagReadonly}, 1, 1, 1)
   629  	registerCommand("LSet", execLSet, writeFirstKey, undoLSet, 4, flagWrite).
   630  		attachCommandExtra([]string{redisFlagWrite, redisFlagDenyOOM}, 1, 1, 1)
   631  	registerCommand("LRange", execLRange, readFirstKey, nil, 4, flagReadOnly).
   632  		attachCommandExtra([]string{redisFlagReadonly}, 1, 1, 1)
   633  	registerCommand("LTrim", execLTrim, writeFirstKey, rollbackFirstKey, 4, flagWrite).
   634  		attachCommandExtra([]string{redisFlagWrite}, 1, 1, 1)
   635  	registerCommand("LInsert", execLInsert, writeFirstKey, rollbackFirstKey, 5, flagWrite).
   636  		attachCommandExtra([]string{redisFlagWrite, redisFlagDenyOOM}, 1, 1, 1)
   637  }