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

     1  package database
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"sync/atomic"
     7  
     8  	"github.com/hdt3213/godis/aof"
     9  	"github.com/hdt3213/godis/config"
    10  	"github.com/hdt3213/godis/datastruct/dict"
    11  	List "github.com/hdt3213/godis/datastruct/list"
    12  	HashSet "github.com/hdt3213/godis/datastruct/set"
    13  	SortedSet "github.com/hdt3213/godis/datastruct/sortedset"
    14  	"github.com/hdt3213/godis/interface/database"
    15  	"github.com/hdt3213/rdb/core"
    16  	rdb "github.com/hdt3213/rdb/parser"
    17  )
    18  
    19  // loadRdbFile loads rdb file from disk
    20  func (server *Server) loadRdbFile() error {
    21  	rdbFile, err := os.Open(config.Properties.RDBFilename)
    22  	if err != nil {
    23  		return fmt.Errorf("open rdb file failed " + err.Error())
    24  	}
    25  	defer func() {
    26  		_ = rdbFile.Close()
    27  	}()
    28  	decoder := rdb.NewDecoder(rdbFile)
    29  	err = server.LoadRDB(decoder)
    30  	if err != nil {
    31  		return fmt.Errorf("dump rdb file failed " + err.Error())
    32  	}
    33  	return nil
    34  }
    35  
    36  // LoadRDB real implementation of loading rdb file
    37  func (server *Server) LoadRDB(dec *core.Decoder) error {
    38  	return dec.Parse(func(o rdb.RedisObject) bool {
    39  		db := server.mustSelectDB(o.GetDBIndex())
    40  		var entity *database.DataEntity
    41  		switch o.GetType() {
    42  		case rdb.StringType:
    43  			str := o.(*rdb.StringObject)
    44  			entity = &database.DataEntity{
    45  				Data: str.Value,
    46  			}
    47  		case rdb.ListType:
    48  			listObj := o.(*rdb.ListObject)
    49  			list := List.NewQuickList()
    50  			for _, v := range listObj.Values {
    51  				list.Add(v)
    52  			}
    53  			entity = &database.DataEntity{
    54  				Data: list,
    55  			}
    56  		case rdb.HashType:
    57  			hashObj := o.(*rdb.HashObject)
    58  			hash := dict.MakeSimple()
    59  			for k, v := range hashObj.Hash {
    60  				hash.Put(k, v)
    61  			}
    62  			entity = &database.DataEntity{
    63  				Data: hash,
    64  			}
    65  		case rdb.SetType:
    66  			setObj := o.(*rdb.SetObject)
    67  			set := HashSet.Make()
    68  			for _, mem := range setObj.Members {
    69  				set.Add(string(mem))
    70  			}
    71  			entity = &database.DataEntity{
    72  				Data: set,
    73  			}
    74  		case rdb.ZSetType:
    75  			zsetObj := o.(*rdb.ZSetObject)
    76  			zSet := SortedSet.Make()
    77  			for _, e := range zsetObj.Entries {
    78  				zSet.Add(e.Member, e.Score)
    79  			}
    80  			entity = &database.DataEntity{
    81  				Data: zSet,
    82  			}
    83  		}
    84  		if entity != nil {
    85  			db.PutEntity(o.GetKey(), entity)
    86  			if o.GetExpiration() != nil {
    87  				db.Expire(o.GetKey(), *o.GetExpiration())
    88  			}
    89  			// add to aof
    90  			db.addAof(aof.EntityToCmd(o.GetKey(), entity).Args)
    91  		}
    92  		return true
    93  	})
    94  }
    95  
    96  func NewPersister(db database.DBEngine, filename string, load bool, fsync string) (*aof.Persister, error) {
    97  	return aof.NewPersister(db, filename, load, fsync, func() database.DBEngine {
    98  		return MakeAuxiliaryServer()
    99  	})
   100  }
   101  
   102  func (server *Server) AddAof(dbIndex int, cmdLine CmdLine) {
   103  	if server.persister != nil {
   104  		server.persister.SaveCmdLine(dbIndex, cmdLine)
   105  	}
   106  }
   107  
   108  func (server *Server) bindPersister(aofHandler *aof.Persister) {
   109  	server.persister = aofHandler
   110  	// bind SaveCmdLine
   111  	for _, db := range server.dbSet {
   112  		singleDB := db.Load().(*DB)
   113  		singleDB.addAof = func(line CmdLine) {
   114  			if config.Properties.AppendOnly { // config may be changed during runtime
   115  				server.persister.SaveCmdLine(singleDB.index, line)
   116  			}
   117  		}
   118  	}
   119  }
   120  
   121  // MakeAuxiliaryServer create a Server only with basic capabilities for aof rewrite and other usages
   122  func MakeAuxiliaryServer() *Server {
   123  	mdb := &Server{}
   124  	mdb.dbSet = make([]*atomic.Value, config.Properties.Databases)
   125  	for i := range mdb.dbSet {
   126  		holder := &atomic.Value{}
   127  		holder.Store(makeBasicDB())
   128  		mdb.dbSet[i] = holder
   129  	}
   130  	return mdb
   131  }