Redis 内存数据库

package dict
 
type Consumer func(key string, val interface{}) bool // function type for iterating over key-value pairs
 
type Dict interface {
	Get(key string) (val interface{}, exists bool)        // get value by key, return the value and a boolean indicating if the key exists
	Len() int                                             // get the number of key-value pairs
	Put(key string, val interface{}) (result int)         // put key-value pair, if exists, modify the value, return 0, if doesn't exist, add it, return 1
	PutIfAbsent(key string, val interface{}) (result int) // put key-value pair if absent, return 0, if exists, return 1
	PutIfExists(key string, val interface{}) (result int) // put key-value pair if exists, return 0, if absent, return 1
	Remove(key string) (result int)                       // remove key-value pair, return the count of pairs
	ForEach(consumer Consumer)                            // iterate over all key-value pairs
	Keys() []string                                       // get all keys
	RandomKeys(n int) []string                            // get n random keys
	RandomDistinctKeys(n int) []string                    // get n distinct random keys
	Clear()                                               // clear all key-value pairs
}

func printKeyValue(key string, val interface{}) {
    fmt.Printf("key: %s, value: %v\n", key, val)
}
 
// 使用 ForEach 方法遍历所有的键值对
dict.ForEach(printKeyValue)

// MakeSyncDict creates a new SyncDict instance
func MakeSyncDict() *SyncDict {
	return &SyncDict{}
}

// Get returns the value associated with the given key and a boolean indicating if the key exists
func (dict *SyncDict) Get(key string) (val interface{}, exists bool) {
	// Get the value by key
	if value, ok := dict.m.Load(key); ok {
		return value, true
	}
	return nil, false
}

// Len returns the number of key-value pairs in the dictionary
func (dict *SyncDict) Len() int {
	count := 0
	// Iterate over all key-value pairs and count them
	dict.m.Range(func(key, value interface{}) bool {
		count++
		return true
	})
	return count
}

// Put adds a key-value pair to the dictionary, if the key already exists, return 1 else 0
func (dict *SyncDict) Put(key string, val interface{}) (result int) {
	_, exists := dict.m.Load(key)
	// Store the key-value pair
	dict.m.Store(key, val)
	// Return the count of pairs
	if exists {
		return 0
	}
	return 1
}

// PutIfAbsent adds a key-value pair to the dictionary if the key does not exist, return 1 if it exists, else 0
func (dict *SyncDict) PutIfAbsent(key string, val interface{}) (result int) {
	_, exists := dict.m.Load(key)
	if exists {
		return 0
	}
	// Store the key-value pair
	dict.m.Store(key, val)
	return 1
}
 
// PutIfExists adds a key-value pair to the dictionary if the key exists, return 1 if it does not exist, else 0
func (dict *SyncDict) PutIfExists(key string, val interface{}) (result int) {
	_, exists := dict.m.Load(key)
	if !exists {
		return 0
	}
	// Store the key-value pair
	dict.m.Store(key, val)
	return 1
}

// Remove removes a key-value pair from the dictionary, return the count of pairs were removed
func (dict *SyncDict) Remove(key string) (result int) {
	_, exists := dict.m.Load(key)
	if !exists {
		return 0
	}
	// Delete the key-value pair
	dict.m.Delete(key)
	return 1
}

// ForEach iterates over all key-value pairs in the dictionary and applies the consumer function to each pair
func (dict *SyncDict) ForEach(consumer Consumer) {
	// Iterate over all key-value pairs and apply the consumer function
	dict.m.Range(func(key, value interface{}) bool {
		consumer(key.(string), value)
		// Always return true to continue iteration
		return true
	})
}

// Keys returns a slice of all keys in the dictionary
func (dict *SyncDict) Keys() []string {
	keys := make([]string, dict.Len())
	// Iterate over all key-value pairs and collect the keys
	dict.m.Range(func(key, value interface{}) bool {
		keys = append(keys, key.(string))
		return true
	})
	return keys
}

// RandomKeys returns a slice of n random keys from the dictionary
// Due to m.Range doesn't guarantee the order of iteration, we can use this feature to get random keys
// Note: This method may not be truly random, but it will give different keys each time
// Duplicate keys may be returned
func (dict *SyncDict) RandomKeys(n int) []string {
	keys := make([]string, dict.Len())
	for i := 0; i < n; i++ {
		// Randomly select a key from the dictionary
		dict.m.Range(func(key, value interface{}) bool {
			keys = append(keys, key.(string))
			return false
		})
	}
	return keys
}

// RandomDistinctKeys returns a slice of n distinct random keys from the dictionary
func (dict *SyncDict) RandomDistinctKeys(n int) []string {
	result := make([]string, dict.Len())
 
	i := 0
 
	// Iterate over all key-value pairs and collect the keys
	dict.m.Range(func(key, value interface{}) bool {
		result[i] = key.(string)
		i++
		return i != n
	})
	return result
}

// Clear clears all key-value pairs in the dictionary
func (dict *SyncDict) Clear() {
	*dict = *MakeSyncDict()
}

type DB struct {
	index int
	data  dict.Dict
}
 
// MakeDB creates a new DB instance
func MakeDB() *DB {
	return &DB{
		index: 0,
		data:  dict.MakeSyncDict(),
	}
}

// ExecFunc is a function type that takes a DB instance and a slice of byte slices as arguments and returns a resp.Reply
// All redis commands like PING, SET, GET, etc. are implemented as functions of this type
type ExecFunc func(db *DB, args [][]byte) resp.Reply

// CmdLine is a type alias for a slice of byte slices
// It is used to represent the command line arguments passed to the ExecFunc
type CmdLine = [][]byte

package database
 
import "strings"
 
// cmdTable is a map that associates command names (as strings) with their corresponding command structures
var cmdTable = make(map[string]*command)
 
type command struct {
	exec  ExecFunc // function to execute the command
	arity int      // number of arguments required for the command
}
 
// RegisterCommand registers a command with the command table
func RegisterCommand(name string, exec ExecFunc, arity int) {
	name = strings.ToLower(name)
	cmdTable[name] = &command{
		exec:  exec,
		arity: arity,
	}
}

// Exec executes a command on the DB instance
// It takes a connection and a command line as arguments
// It returns a resp.Reply which is the response to the command
func (db *DB) Exec(c resp.Connection, cmdLine CmdLine) resp.Reply {
	// The first element of cmdLine is the command name, like "PING", "SET", etc.
	// Convert it to lowercase to ensure case-insensitivity
	cmdName := strings.ToLower(string(cmdLine[0]))
	// Get the command from the command table using the command name
	// If the command is not found, return an error reply
	cmd, ok := cmdTable[cmdName]
	if !ok {
		return reply.MakeStandardErrorReply("ERR unknown command '" + cmdName + "'")
	}
	// Validate the number of arguments passed to the command
	if !ValidateArity(cmd.arity, cmdLine) {
		return reply.MakeArgNumErrReply(cmdName)
	}
	// Execute the command and return the response
	return cmd.exec(db, cmdLine[1:])
}
 
// ValidateArity checks if the number of arguments passed to a command is valid
func ValidateArity(arity int, args [][]byte) bool {
	// Check if the number of arguments is less than the required arity
	if arity >= 0 {
		return len(args) == arity
	} else {
		// If the arity is negative, it means the command takes a variable number of arguments
		// Check if the number of arguments is within the valid range
		return len(args) >= -arity
	}
}

package database
 
import (
	"redigo/interface/resp"
	"redigo/resp/reply"
)
 
func Ping(db *DB, args [][]byte) resp.Reply {
	return reply.MakePongReply()
}

// Register the PING command to the command table
func init() {
	// Register the PING command with the command table
	RegisterCommand("ping", Ping, 1)
}

// GetEntity returns DataEntity bind to the given key
func (db *DB) GetEntity(key string) (*database.DataEntity, bool) {
	raw, ok := db.data.Get(key)
	if !ok {
		return nil, false
	}
	entity, _ := raw.(*database.DataEntity)
	return entity, true
}
 
// PutEntity stores the given DataEntity in the database
func (db *DB) PutEntity(key string, entity *database.DataEntity) int {
	return db.data.Put(key, entity)
}
 
// PutIfExists edit the given DataEntity in the database
func (db *DB) PutIfExists(key string, entity *database.DataEntity) int {
	return db.data.PutIfExists(key, entity)
}
 
// PutIfAbsent stores the given DataEntity in the database if it doesn't already exist
func (db *DB) PutIfAbsent(key string, entity *database.DataEntity) int {
	return db.data.PutIfAbsent(key, entity)
}
 
// Remove deletes the DataEntity associated with the given key from the database
func (db *DB) Remove(key string) int {
	return db.data.Remove(key)
}
 
// Removes deletes the DataEntity associated with the given keys from the database
func (db *DB) Removes(keys ...string) int {
	deleted := 0
	for _, key := range keys {
		_, ok := db.data.Get(key)
		if ok {
			db.data.Remove(key)
			deleted++
		}
	}
	return deleted
}
 
// Flush clears the database by removing all DataEntity objects
func (db *DB) Flush() {
	db.data.Clear()
}

// Clear clears all key-value pairs in the dictionary
func (dict *SyncDict) Clear() {
	*dict = *MakeSyncDict()
}

type Dict interface {
	// ... other methods
	Clear()                                               // clear all key-value pairs
}

SET key1 value1
SET key2 value2
SET key3 value3

// Handle the DEL command
// It deletes the specified keys from the database
func execDel(db *DB, args [][]byte) resp.Reply {
	keys := make([]string, len(args))
	for i, arg := range args {
		keys[i] = string(arg)
	}
	deleted := db.Removes(keys...)
	return reply.MakeIntReply(int64(deleted))
}

func init() {
	RegisterCommand("DEL", execDel, -2)
}

func execExists(db *DB, args [][]byte) resp.Reply {
	result := int64(0)
	for _, arg := range args {
		key := string(arg)
		if _, ok := db.GetEntity(key); ok {
			result++
		}
	}
	return reply.MakeIntReply(result)
}

func init() {
	// ...
	RegisterCommand("EXISTS", execExists, -2)
}

func execFlushDB(db *DB, args [][]byte) resp.Reply {
	db.Flush()
	return reply.MakeOKReply()
}
 
func init() {
	// ...
	RegisterCommand("FLUSHDB", execFlushDB, -1)
}

// Handle the TYPE command
// It returns the type of the specified key
func execType(db *DB, args [][]byte) resp.Reply {
	key := string(args[0])
	if entity, ok := db.GetEntity(key); ok {
		switch entity.Data.(type) {
		// If the entity is []byte, return the type as "string"
		case []byte:
			return reply.MakeBulkReply([]byte("string"))
		}
		// TODO: Add more types as needed
	} else {
		return reply.MakeStatusReply("none")
	}
	return reply.MakeUnknownReply()
}
 
func init() {
	// ...
	RegisterCommand("TYPE", execType, 2)
}

RENAME key1 key2

// Handle the RENAME command.
// It renames a key in the database.
// RENAME key newkey
func execRename(db *DB, args [][]byte) resp.Reply {
	src := string(args[0])
	dst := string(args[1])
	entity, ok := db.GetEntity(src)
	if !ok {
		return reply.MakeStandardErrorReply("ERR no such key")
	}
	db.PutEntity(dst, entity)
	db.Remove(src)
	return reply.MakeOKReply()
}
func init() {
	// ...
	RegisterCommand("RENAME", execRename, 3)
}

RENAMENX key1 key2

// Handle the RENAMENX command.
// It renames a key in the database only if the new key does not exist.
// RENAME key newkey
func execRenameNX(db *DB, args [][]byte) resp.Reply {
	src := string(args[0])
	dst := string(args[1])
	entity, ok := db.GetEntity(src)
	if !ok {
		return reply.MakeStandardErrorReply("ERR no such key")
	}
	if _, ok := db.GetEntity(dst); ok {
		return reply.MakeIntReply(0)
	}
	db.PutEntity(dst, entity)
	db.Remove(src)
	return reply.MakeIntReply(1)
}
 
func init() {
	// ...
	RegisterCommand("RENAMENX", execRenameNX, 3)
}

KEYS *

package wildcard
 
const (
	normal     = iota
	all        // *
	any        // ?
	setSymbol  // []
	rangSymbol // [a-b]
	negSymbol  // [^a]
)
 
type item struct {
	character byte
	set       map[byte]bool
	typeCode  int
}
 
func (i *item) contains(c byte) bool {
	if i.typeCode == setSymbol {
		_, ok := i.set[c]
		return ok
	} else if i.typeCode == rangSymbol {
		if _, ok := i.set[c]; ok {
			return true
		}
		var (
			min uint8 = 255
			max uint8 = 0
		)
		for k := range i.set {
			if min > k {
				min = k
			}
			if max < k {
				max = k
			}
		}
		return c >= min && c <= max
	} else {
		_, ok := i.set[c]
		return !ok
	}
}
 
// Pattern represents a wildcard pattern
type Pattern struct {
	items []*item
}
 
// CompilePattern convert wildcard string to Pattern
func CompilePattern(src string) *Pattern {
	items := make([]*item, 0)
	escape := false
	inSet := false
	var set map[byte]bool
	for _, v := range src {
		c := byte(v)
		if escape {
			items = append(items, &item{typeCode: normal, character: c})
			escape = false
		} else if c == '*' {
			items = append(items, &item{typeCode: all})
		} else if c == '?' {
			items = append(items, &item{typeCode: any})
		} else if c == '\\' {
			escape = true
		} else if c == '[' {
			if !inSet {
				inSet = true
				set = make(map[byte]bool)
			} else {
				set[c] = true
			}
		} else if c == ']' {
			if inSet {
				inSet = false
				typeCode := setSymbol
				if _, ok := set['-']; ok {
					typeCode = rangSymbol
					delete(set, '-')
				}
				if _, ok := set['^']; ok {
					typeCode = negSymbol
					delete(set, '^')
				}
				items = append(items, &item{typeCode: typeCode, set: set})
			} else {
				items = append(items, &item{typeCode: normal, character: c})
			}
		} else {
			if inSet {
				set[c] = true
			} else {
				items = append(items, &item{typeCode: normal, character: c})
			}
		}
	}
	return &Pattern{
		items: items,
	}
}
 
// IsMatch returns whether the given string matches pattern
func (p *Pattern) IsMatch(s string) bool {
	if len(p.items) == 0 {
		return len(s) == 0
	}
	m := len(s)
	n := len(p.items)
	table := make([][]bool, m+1)
	for i := 0; i < m+1; i++ {
		table[i] = make([]bool, n+1)
	}
	table[0][0] = true
	for j := 1; j < n+1; j++ {
		table[0][j] = table[0][j-1] && p.items[j-1].typeCode == all
	}
	for i := 1; i < m+1; i++ {
		for j := 1; j < n+1; j++ {
			if p.items[j-1].typeCode == all {
				table[i][j] = table[i-1][j] || table[i][j-1]
			} else {
				table[i][j] = table[i-1][j-1] &&
					(p.items[j-1].typeCode == any ||
						(p.items[j-1].typeCode == normal && uint8(s[i-1]) == p.items[j-1].character) ||
						(p.items[j-1].typeCode >= setSymbol && p.items[j-1].contains(s[i-1])))
			}
		}
	}
	return table[m][n]
}

// Handle the KEYS command.
// It returns all keys in the database that match the specified pattern.
func execKeys(db *DB, args [][]byte) resp.Reply {
	pattern := wildcard.CompilePattern(string(args[0]))
	result := make([][]byte, 0) // Store all matching keys
	db.data.ForEach(func(key string, val interface{}) bool {
		if pattern.IsMatch(key) {
			result = append(result, []byte(key))
		}
		return true
	})
	return reply.MakeMultiBulkReply(result)
}
 
func init() {
	// ...
	RegisterCommand("KEYS", execKeys, 2)
}

type Consumer func(key string, val interface{}) bool // function type for iterating over key-value pairs

// execGet retrieves the value associated with the specified key from the database.
func execGet(db *DB, args [][]byte) resp.Reply {
	key := string(args[0])
	if entity, ok := db.GetEntity(key); ok {
		// TODO: If we have multiple types, we need to check the conversion if it's not []byte
		return reply.MakeBulkReply(entity.Data.([]byte))
	}
	return reply.MakeNullBulkReply()
}
 
func init() {
	// ...
	RegisterCommand("GET", execGet, 2)
}

// execSet stores the specified key-value pair in the database.
func execSet(db *DB, args [][]byte) resp.Reply {
	key := string(args[0])
	value := args[1]
	entity := &database.DataEntity{
		Data: value,
	}
	db.PutEntity(key, entity)
	return reply.MakeOKReply()
}
 
func init() {
	// ...
	RegisterCommand("SET", execSet, 3)
}

// execSetNX stores the specified key-value pair in the database only if the key does not already exist.
// If the key already exists, it does not modify the value and returns 0.
// If the key does not exist, it sets the value and returns 1.
// SETNX key value
func execSetNX(db *DB, args [][]byte) resp.Reply {
	key := string(args[0])
	value := args[1]
	entity := &database.DataEntity{
		Data: value,
	}
	result := db.PutIfAbsent(key, entity)
	return reply.MakeIntReply(int64(result))
}
 
func init() {
	// ...
	RegisterCommand("SETNX", execSetNX, 3)
}

// execGetSet stores the specified key-value pair in the database and returns the old value associated with the key.
func execGetSet(db *DB, args [][]byte) resp.Reply {
	key := string(args[0])
	value := args[1]
 
	entity, ok := db.GetEntity(key)
	db.PutEntity(key, &database.DataEntity{
		Data: value,
	})
	if !ok {
		return reply.MakeNullBulkReply()
	}
	return reply.MakeBulkReply(entity.Data.([]byte))
}
 
func init() {
	// ...
	RegisterCommand("GETSET", execGetSet, 3)
}

STRLEN key

// execStrLen retrieves the length of the value associated with the specified key.
func execStrLen(db *DB, args [][]byte) resp.Reply {
	key := string(args[0])
	entity, ok := db.GetEntity(key)
	if !ok {
		return reply.MakeNullBulkReply()
	}
	return reply.MakeIntReply(int64(len(entity.Data.([]byte))))
}
 
func init() {
	// ...
	RegisterCommand("STRLEN", execStrLen, 2)
}

func main() {
	// ...
	err := tcp.ListenAndServeWithSignal(
		&tcp.Config{
			Address: fmt.Sprintf("%s:%d",
				config.Properties.Bind,
				config.Properties.Port),
		},
		handler.MakeHandler())	
	if err != nil {
		logger.Error(err)
	}
}

// MakeHandler creates a RespHandler instance
func MakeHandler() *RespHandler {
	var db databaseface.Database
	db = database.NewEchoDatabase()
	return &RespHandler{
		db: db,
	}
}

package database
 
type Database struct {
	dbSet []*DB
}

type Database interface {
	Exec(client resp.Connection, args [][]byte) resp.Reply
	AfterClientClose(c resp.Connection)
	Close()
}

databases 16

// NewDatabase creates a new Database instance
func NewDatabase() *Database {
	database := &Database{}
	if config.Properties.Databases == 0 {
		config.Properties.Databases = 16
	}
	database.dbSet = make([]*DB, config.Properties.Databases)
	for i := range database.dbSet {
		db := MakeDB()
		db.index = i
		database.dbSet[i] = db
	}
	return database
}

type Connection interface {
	Write([]byte) error // Write data to the connection
	GetDBIndex() int    // Get database index
	SelectDB(int)       // Select database
}

// SelectDB selects a database
func (c *Connection) SelectDB(dbNum int) {
	c.selectedDB = dbNum
}

// execSelect sets the current database for the client connection.
// select x
func execSelect(c resp.Connection, database *Database, args [][]byte) resp.Reply {
	dbIndex, err := strconv.Atoi(string(args[0]))
	if err != nil {
		return reply.MakeStandardErrorReply("ERR invalid DB index")
	}
	if dbIndex < 0 || dbIndex >= len(database.dbSet) {
		return reply.MakeStandardErrorReply("ERR DB index out of range")
	}
	c.SelectDB(dbIndex)
	return reply.MakeIntReply(int64(dbIndex))
}

// Exec executes a command on the database
func (d *Database) Exec(client resp.Connection, args [][]byte) resp.Reply {
	defer func() {
		if err := recover(); err != nil {
			logger.Error("Database Exec panic:" + err.(error).Error())
		}
	}()
	cmdName := strings.ToLower(string(args[0]))
	if cmdName == "select" {
		if len(args) != 2 {
			return reply.MakeArgNumErrReply("select")
		}
		return execSelect(client, d, args[1:])
	}
	// Get the current database index from the client connection
	db := d.dbSet[client.GetDBIndex()]
	return db.Exec(client, args)
}

func MakeHandler() *RespHandler {
	db := database.NewDatabase()
	return &RespHandler{
		db: db,
	}
}

go run main.go

printf "*1\r\n\$4\r\nPING\r\n" | nc localhost 6379

+PONG

printf "*3\r\n\$3\r\nSET\r\n\$3\r\nkey\r\n\$5\r\nvalue\r\n" | nc localhost 6379

+OK

printf "*2\r\n\$3\r\nGET\r\n\$3\r\nkey\r\n" | nc localhost 6379

$5
value

printf "*2\r\n\$4\r\nKEYS\r\n\$1\r\n*\r\n" | nc localhost 6379

*1
$3
key

git add .
 
# 示例：可以自己写一下
git commit -m "feat: implement in-memory database
 
- Create Dict interface and SyncDict implementation
- Add database structure with multi-DB support
- Implement command pattern for Redis commands
- Add basic Redis commands (PING, SET, GET, DEL, etc.)
- Implement pattern matching for KEYS command
- Support string operations like GETSET and STRLEN"
 
git push origin main