github.com/hdt3213/godis@v1.2.9/aof/rdb.go (about)

     1  package aof
     2  
     3  import (
     4  	"os"
     5  	"strconv"
     6  	"time"
     7  
     8  	"github.com/hdt3213/godis/config"
     9  	"github.com/hdt3213/godis/datastruct/dict"
    10  	List "github.com/hdt3213/godis/datastruct/list"
    11  	"github.com/hdt3213/godis/datastruct/set"
    12  	SortedSet "github.com/hdt3213/godis/datastruct/sortedset"
    13  	"github.com/hdt3213/godis/interface/database"
    14  	"github.com/hdt3213/godis/lib/logger"
    15  	rdb "github.com/hdt3213/rdb/encoder"
    16  	"github.com/hdt3213/rdb/model"
    17  )
    18  
    19  // todo: forbid concurrent rewrite
    20  
    21  // GenerateRDB generates rdb file from aof file
    22  func (persister *Persister) GenerateRDB(rdbFilename string) error {
    23  	ctx, err := persister.startGenerateRDB(nil, nil)
    24  	if err != nil {
    25  		return err
    26  	}
    27  	err = persister.generateRDB(ctx)
    28  	if err != nil {
    29  		return err
    30  	}
    31  	err = ctx.tmpFile.Close()
    32  	if err != nil {
    33  		return err
    34  	}
    35  	err = os.Rename(ctx.tmpFile.Name(), rdbFilename)
    36  	if err != nil {
    37  		return err
    38  	}
    39  	return nil
    40  }
    41  
    42  // GenerateRDBForReplication asynchronously generates rdb file from aof file and returns a channel to receive following data
    43  // parameter listener would receive following updates of rdb
    44  // parameter hook allows you to do something during aof pausing
    45  func (persister *Persister) GenerateRDBForReplication(rdbFilename string, listener Listener, hook func()) error {
    46  	ctx, err := persister.startGenerateRDB(listener, hook)
    47  	if err != nil {
    48  		return err
    49  	}
    50  
    51  	err = persister.generateRDB(ctx)
    52  	if err != nil {
    53  		return err
    54  	}
    55  	err = ctx.tmpFile.Close()
    56  	if err != nil {
    57  		return err
    58  	}
    59  	err = os.Rename(ctx.tmpFile.Name(), rdbFilename)
    60  	if err != nil {
    61  		return err
    62  	}
    63  	return nil
    64  }
    65  
    66  func (persister *Persister) startGenerateRDB(newListener Listener, hook func()) (*RewriteCtx, error) {
    67  	persister.pausingAof.Lock() // pausing aof
    68  	defer persister.pausingAof.Unlock()
    69  
    70  	err := persister.aofFile.Sync()
    71  	if err != nil {
    72  		logger.Warn("fsync failed")
    73  		return nil, err
    74  	}
    75  
    76  	// get current aof file size
    77  	fileInfo, _ := os.Stat(persister.aofFilename)
    78  	filesize := fileInfo.Size()
    79  	// create tmp file
    80  	file, err := os.CreateTemp(config.GetTmpDir(), "*.aof")
    81  	if err != nil {
    82  		logger.Warn("tmp file create failed")
    83  		return nil, err
    84  	}
    85  	if newListener != nil {
    86  		persister.listeners[newListener] = struct{}{}
    87  	}
    88  	if hook != nil {
    89  		hook()
    90  	}
    91  	return &RewriteCtx{
    92  		tmpFile:  file,
    93  		fileSize: filesize,
    94  	}, nil
    95  }
    96  
    97  // generateRDB generates rdb file from aof file
    98  func (persister *Persister) generateRDB(ctx *RewriteCtx) error {
    99  	// load aof tmpFile
   100  	tmpHandler := persister.newRewriteHandler()
   101  	tmpHandler.LoadAof(int(ctx.fileSize))
   102  
   103  	encoder := rdb.NewEncoder(ctx.tmpFile).EnableCompress()
   104  	err := encoder.WriteHeader()
   105  	if err != nil {
   106  		return err
   107  	}
   108  	auxMap := map[string]string{
   109  		"redis-ver":    "6.0.0",
   110  		"redis-bits":   "64",
   111  		"aof-preamble": "0",
   112  		"ctime":        strconv.FormatInt(time.Now().Unix(), 10),
   113  	}
   114  
   115  	// change aof preamble
   116  	if config.Properties.AofUseRdbPreamble {
   117  		auxMap["aof-preamble"] = "1"
   118  	}
   119  
   120  	for k, v := range auxMap {
   121  		err := encoder.WriteAux(k, v)
   122  		if err != nil {
   123  			return err
   124  		}
   125  	}
   126  
   127  	for i := 0; i < config.Properties.Databases; i++ {
   128  		keyCount, ttlCount := tmpHandler.db.GetDBSize(i)
   129  		if keyCount == 0 {
   130  			continue
   131  		}
   132  		err = encoder.WriteDBHeader(uint(i), uint64(keyCount), uint64(ttlCount))
   133  		if err != nil {
   134  			return err
   135  		}
   136  		// dump db
   137  		var err2 error
   138  		tmpHandler.db.ForEach(i, func(key string, entity *database.DataEntity, expiration *time.Time) bool {
   139  			var opts []interface{}
   140  			if expiration != nil {
   141  				opts = append(opts, rdb.WithTTL(uint64(expiration.UnixNano()/1e6)))
   142  			}
   143  			switch obj := entity.Data.(type) {
   144  			case []byte:
   145  				err = encoder.WriteStringObject(key, obj, opts...)
   146  			case List.List:
   147  				vals := make([][]byte, 0, obj.Len())
   148  				obj.ForEach(func(i int, v interface{}) bool {
   149  					bytes, _ := v.([]byte)
   150  					vals = append(vals, bytes)
   151  					return true
   152  				})
   153  				err = encoder.WriteListObject(key, vals, opts...)
   154  			case *set.Set:
   155  				vals := make([][]byte, 0, obj.Len())
   156  				obj.ForEach(func(m string) bool {
   157  					vals = append(vals, []byte(m))
   158  					return true
   159  				})
   160  				err = encoder.WriteSetObject(key, vals, opts...)
   161  			case dict.Dict:
   162  				hash := make(map[string][]byte)
   163  				obj.ForEach(func(key string, val interface{}) bool {
   164  					bytes, _ := val.([]byte)
   165  					hash[key] = bytes
   166  					return true
   167  				})
   168  				err = encoder.WriteHashMapObject(key, hash, opts...)
   169  			case *SortedSet.SortedSet:
   170  				var entries []*model.ZSetEntry
   171  				obj.ForEach(int64(0), obj.Len(), true, func(element *SortedSet.Element) bool {
   172  					entries = append(entries, &model.ZSetEntry{
   173  						Member: element.Member,
   174  						Score:  element.Score,
   175  					})
   176  					return true
   177  				})
   178  				err = encoder.WriteZSetObject(key, entries, opts...)
   179  			}
   180  			if err != nil {
   181  				err2 = err
   182  				return false
   183  			}
   184  			return true
   185  		})
   186  		if err2 != nil {
   187  			return err2
   188  		}
   189  	}
   190  	err = encoder.WriteEnd()
   191  	if err != nil {
   192  		return err
   193  	}
   194  	return nil
   195  }