Files
redis-go/database/server.go
2025-05-05 18:26:10 +08:00

449 lines
13 KiB
Go

package database
import (
"fmt"
"os"
"runtime/debug"
"strconv"
"strings"
"sync/atomic"
"time"
"github.com/hdt3213/godis/aof"
"github.com/hdt3213/godis/config"
"github.com/hdt3213/godis/interface/database"
"github.com/hdt3213/godis/interface/redis"
"github.com/hdt3213/godis/lib/logger"
"github.com/hdt3213/godis/lib/utils"
"github.com/hdt3213/godis/pubsub"
"github.com/hdt3213/godis/redis/protocol"
)
var godisVersion = "1.2.8" // do not modify
// Server is a redis-server with full capabilities including multiple database, rdb loader, replication
type Server struct {
dbSet []*atomic.Value // *DB
// handle publish/subscribe
hub *pubsub.Hub
// handle aof persistence
persister *aof.Persister
// for replication
role int32
slaveStatus *slaveStatus
masterStatus *masterStatus
// hooks
insertCallback database.KeyEventCallback
deleteCallback database.KeyEventCallback
}
func fileExists(filename string) bool {
info, err := os.Stat(filename)
return err == nil && !info.IsDir()
}
// NewStandaloneServer creates a standalone redis server, with multi database and all other funtions
func NewStandaloneServer() *Server {
server := &Server{}
if config.Properties.Databases == 0 {
config.Properties.Databases = 16
}
// creat tmp dir
err := os.MkdirAll(config.GetTmpDir(), os.ModePerm)
if err != nil {
panic(fmt.Errorf("create tmp dir failed: %v", err))
}
// make db set
server.dbSet = make([]*atomic.Value, config.Properties.Databases)
for i := range server.dbSet {
singleDB := makeDB()
singleDB.index = i
holder := &atomic.Value{}
holder.Store(singleDB)
server.dbSet[i] = holder
}
server.hub = pubsub.MakeHub()
// record aof
validAof := false
if config.Properties.AppendOnly {
validAof = fileExists(config.Properties.AppendFilename)
aofHandler, err := NewPersister(server,
config.Properties.AppendFilename, true, config.Properties.AppendFsync)
if err != nil {
panic(err)
}
server.bindPersister(aofHandler)
}
if config.Properties.RDBFilename != "" && !validAof {
// load rdb
err := server.loadRdbFile()
if err != nil {
logger.Error(err)
}
}
server.slaveStatus = initReplSlaveStatus()
server.initMasterStatus()
server.startReplCron()
server.role = masterRole // The initialization process does not require atomicity
return server
}
// Exec executes command
// parameter `cmdLine` contains command and its arguments, for example: "set key value"
func (server *Server) Exec(c redis.Connection, cmdLine [][]byte) (result redis.Reply) {
defer func() {
if err := recover(); err != nil {
logger.Warn(fmt.Sprintf("error occurs: %v\n%s", err, string(debug.Stack())))
result = &protocol.UnknownErrReply{}
}
}()
cmdName := strings.ToLower(string(cmdLine[0]))
// ping
if cmdName == "ping" {
return Ping(c, cmdLine[1:])
}
// authenticate
if cmdName == "auth" {
return Auth(c, cmdLine[1:])
}
if !isAuthenticated(c) {
return protocol.MakeErrReply("NOAUTH Authentication required")
}
// info
if cmdName == "info" {
return Info(server, cmdLine[1:])
}
if cmdName == "dbsize" {
return DbSize(c, server)
}
if cmdName == "slaveof" {
if c != nil && c.InMultiState() {
return protocol.MakeErrReply("cannot use slave of database within multi")
}
if len(cmdLine) != 3 {
return protocol.MakeArgNumErrReply("SLAVEOF")
}
return server.execSlaveOf(c, cmdLine[1:])
} else if cmdName == "command" {
return execCommand(cmdLine[1:])
}
// read only slave
role := atomic.LoadInt32(&server.role)
if role == slaveRole && !c.IsMaster() {
// only allow read only command, forbid all special commands except `auth` and `slaveof`
if !isReadOnlyCommand(cmdName) {
return protocol.MakeErrReply("READONLY You can't write against a read only slave.")
}
}
// special commands which cannot execute within transaction
if cmdName == "subscribe" {
if len(cmdLine) < 2 {
return protocol.MakeArgNumErrReply("subscribe")
}
return pubsub.Subscribe(server.hub, c, cmdLine[1:])
} else if cmdName == "publish" {
return pubsub.Publish(server.hub, cmdLine[1:])
} else if cmdName == "unsubscribe" {
return pubsub.UnSubscribe(server.hub, c, cmdLine[1:])
} else if cmdName == "bgrewriteaof" {
if !config.Properties.AppendOnly {
return protocol.MakeErrReply("AppendOnly is false, you can't rewrite aof file")
}
// aof.go imports router.go, router.go cannot import BGRewriteAOF from aof.go
return BGRewriteAOF(server, cmdLine[1:])
} else if cmdName == "rewriteaof" {
if !config.Properties.AppendOnly {
return protocol.MakeErrReply("AppendOnly is false, you can't rewrite aof file")
}
return RewriteAOF(server, cmdLine[1:])
} else if cmdName == "flushall" {
return server.flushAll()
} else if cmdName == "flushdb" {
if !validateArity(1, cmdLine) {
return protocol.MakeArgNumErrReply(cmdName)
}
if c.InMultiState() {
return protocol.MakeErrReply("ERR command 'FlushDB' cannot be used in MULTI")
}
return server.execFlushDB(c.GetDBIndex())
} else if cmdName == "save" {
return SaveRDB(server, cmdLine[1:])
} else if cmdName == "bgsave" {
return BGSaveRDB(server, cmdLine[1:])
} else if cmdName == "select" {
if c != nil && c.InMultiState() {
return protocol.MakeErrReply("cannot select database within multi")
}
if len(cmdLine) != 2 {
return protocol.MakeArgNumErrReply("select")
}
return execSelect(c, server, cmdLine[1:])
} else if cmdName == "copy" {
if len(cmdLine) < 3 {
return protocol.MakeArgNumErrReply("copy")
}
return execCopy(server, c, cmdLine[1:])
} else if cmdName == "replconf" {
return server.execReplConf(c, cmdLine[1:])
} else if cmdName == "psync" {
return server.execPSync(c, cmdLine[1:])
}
// todo: support multi database transaction
// normal commands
dbIndex := c.GetDBIndex()
selectedDB, errReply := server.selectDB(dbIndex)
if errReply != nil {
return errReply
}
return selectedDB.Exec(c, cmdLine)
}
// AfterClientClose does some clean after client close connection
func (server *Server) AfterClientClose(c redis.Connection) {
pubsub.UnsubscribeAll(server.hub, c)
}
// Close graceful shutdown database
func (server *Server) Close() {
// stop slaveStatus first
server.slaveStatus.close()
if server.persister != nil {
server.persister.Close()
}
server.stopMaster()
}
func execSelect(c redis.Connection, mdb *Server, args [][]byte) redis.Reply {
dbIndex, err := strconv.Atoi(string(args[0]))
if err != nil {
return protocol.MakeErrReply("ERR invalid DB index")
}
if dbIndex >= len(mdb.dbSet) || dbIndex < 0 {
return protocol.MakeErrReply("ERR DB index is out of range")
}
c.SelectDB(dbIndex)
return protocol.MakeOkReply()
}
func (server *Server) execFlushDB(dbIndex int) redis.Reply {
if server.persister != nil {
server.persister.SaveCmdLine(dbIndex, utils.ToCmdLine("FlushDB"))
}
return server.flushDB(dbIndex)
}
// flushDB flushes the selected database
func (server *Server) flushDB(dbIndex int) redis.Reply {
if dbIndex >= len(server.dbSet) || dbIndex < 0 {
return protocol.MakeErrReply("ERR DB index is out of range")
}
newDB := makeDB()
server.loadDB(dbIndex, newDB)
return &protocol.OkReply{}
}
func (server *Server) loadDB(dbIndex int, newDB *DB) redis.Reply {
if dbIndex >= len(server.dbSet) || dbIndex < 0 {
return protocol.MakeErrReply("ERR DB index is out of range")
}
oldDB := server.mustSelectDB(dbIndex)
newDB.index = dbIndex
newDB.addAof = oldDB.addAof // inherit oldDB
server.dbSet[dbIndex].Store(newDB)
return &protocol.OkReply{}
}
// flushAll flushes all databases.
func (server *Server) flushAll() redis.Reply {
for i := range server.dbSet {
server.flushDB(i)
}
if server.persister != nil {
server.persister.SaveCmdLine(0, utils.ToCmdLine("FlushAll"))
}
return &protocol.OkReply{}
}
// selectDB returns the database with the given index, or an error if the index is out of range.
func (server *Server) selectDB(dbIndex int) (*DB, *protocol.StandardErrReply) {
if dbIndex >= len(server.dbSet) || dbIndex < 0 {
return nil, protocol.MakeErrReply("ERR DB index is out of range")
}
return server.dbSet[dbIndex].Load().(*DB), nil
}
// mustSelectDB is like selectDB, but panics if an error occurs.
func (server *Server) mustSelectDB(dbIndex int) *DB {
selectedDB, err := server.selectDB(dbIndex)
if err != nil {
panic(err)
}
return selectedDB
}
// ForEach traverses all the keys in the given database
func (server *Server) ForEach(dbIndex int, cb func(key string, data *database.DataEntity, expiration *time.Time) bool) {
server.mustSelectDB(dbIndex).ForEach(cb)
}
// GetEntity returns the data entity to the given key
func (server *Server) GetEntity(dbIndex int, key string) (*database.DataEntity, bool) {
return server.mustSelectDB(dbIndex).GetEntity(key)
}
func (server *Server) GetExpiration(dbIndex int, key string) *time.Time {
raw, ok := server.mustSelectDB(dbIndex).ttlMap.Get(key)
if !ok {
return nil
}
expireTime, _ := raw.(time.Time)
return &expireTime
}
// ExecMulti executes multi commands transaction Atomically and Isolated
func (server *Server) ExecMulti(conn redis.Connection, watching map[string]uint32, cmdLines []CmdLine) redis.Reply {
selectedDB, errReply := server.selectDB(conn.GetDBIndex())
if errReply != nil {
return errReply
}
return selectedDB.ExecMulti(conn, watching, cmdLines)
}
// RWLocks lock keys for writing and reading
func (server *Server) RWLocks(dbIndex int, writeKeys []string, readKeys []string) {
server.mustSelectDB(dbIndex).RWLocks(writeKeys, readKeys)
}
// RWUnLocks unlock keys for writing and reading
func (server *Server) RWUnLocks(dbIndex int, writeKeys []string, readKeys []string) {
server.mustSelectDB(dbIndex).RWUnLocks(writeKeys, readKeys)
}
// GetUndoLogs return rollback commands
func (server *Server) GetUndoLogs(dbIndex int, cmdLine [][]byte) []CmdLine {
return server.mustSelectDB(dbIndex).GetUndoLogs(cmdLine)
}
// ExecWithLock executes normal commands, invoker should provide locks
func (server *Server) ExecWithLock(conn redis.Connection, cmdLine [][]byte) redis.Reply {
db, errReply := server.selectDB(conn.GetDBIndex())
if errReply != nil {
return errReply
}
return db.execWithLock(cmdLine)
}
// BGRewriteAOF asynchronously rewrites Append-Only-File
func BGRewriteAOF(db *Server, args [][]byte) redis.Reply {
go db.persister.Rewrite()
return protocol.MakeStatusReply("Background append only file rewriting started")
}
// RewriteAOF start Append-Only-File rewriting and blocked until it finished
func RewriteAOF(db *Server, args [][]byte) redis.Reply {
err := db.persister.Rewrite()
if err != nil {
return protocol.MakeErrReply(err.Error())
}
return protocol.MakeOkReply()
}
// SaveRDB start RDB writing and blocked until it finished
func SaveRDB(db *Server, args [][]byte) redis.Reply {
if db.persister == nil {
return protocol.MakeErrReply("please enable aof before using save")
}
rdbFilename := config.Properties.RDBFilename
if rdbFilename == "" {
rdbFilename = "dump.rdb"
}
err := db.persister.GenerateRDB(rdbFilename)
if err != nil {
return protocol.MakeErrReply(err.Error())
}
return protocol.MakeOkReply()
}
// BGSaveRDB asynchronously save RDB
func BGSaveRDB(db *Server, args [][]byte) redis.Reply {
if db.persister == nil {
return protocol.MakeErrReply("please enable aof before using save")
}
go func() {
defer func() {
if err := recover(); err != nil {
logger.Error(err)
}
}()
rdbFilename := config.Properties.RDBFilename
if rdbFilename == "" {
rdbFilename = "dump.rdb"
}
err := db.persister.GenerateRDB(rdbFilename)
if err != nil {
logger.Error(err)
}
}()
return protocol.MakeStatusReply("Background saving started")
}
// GetDBSize returns keys count and ttl key count
func (server *Server) GetDBSize(dbIndex int) (int, int) {
db := server.mustSelectDB(dbIndex)
return db.data.Len(), db.ttlMap.Len()
}
func (server *Server) startReplCron() {
go func(mdb *Server) {
ticker := time.Tick(time.Second * 10)
for range ticker {
mdb.slaveCron()
mdb.masterCron()
}
}(server)
}
// GetAvgTTL Calculate the average expiration time of keys
func (server *Server) GetAvgTTL(dbIndex, randomKeyCount int) int64 {
var ttlCount int64
db := server.mustSelectDB(dbIndex)
keys := db.data.RandomKeys(randomKeyCount)
for _, k := range keys {
t := time.Now()
rawExpireTime, ok := db.ttlMap.Get(k)
if !ok {
continue
}
expireTime, _ := rawExpireTime.(time.Time)
// if the key has already reached its expiration time during calculation, ignore it
if expireTime.Sub(t).Microseconds() > 0 {
ttlCount += expireTime.Sub(t).Microseconds()
}
}
return ttlCount / int64(len(keys))
}
func (server *Server) SetKeyInsertedCallback(cb database.KeyEventCallback) {
server.insertCallback = cb
for i := range server.dbSet {
db := server.mustSelectDB(i)
db.insertCallback = cb
}
}
func (server *Server) SetKeyDeletedCallback(cb database.KeyEventCallback) {
server.deleteCallback = cb
for i := range server.dbSet {
db := server.mustSelectDB(i)
db.deleteCallback = cb
}
}