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 }