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

     1  package aof
     2  
     3  import (
     4  	"io"
     5  	"os"
     6  	"strconv"
     7  
     8  	"github.com/hdt3213/godis/config"
     9  	"github.com/hdt3213/godis/lib/logger"
    10  	"github.com/hdt3213/godis/lib/utils"
    11  	"github.com/hdt3213/godis/redis/protocol"
    12  )
    13  
    14  func (persister *Persister) newRewriteHandler() *Persister {
    15  	h := &Persister{}
    16  	h.aofFilename = persister.aofFilename
    17  	h.db = persister.tmpDBMaker()
    18  	return h
    19  }
    20  
    21  // RewriteCtx holds context of an AOF rewriting procedure
    22  type RewriteCtx struct {
    23  	tmpFile  *os.File // tmpFile is the file handler of aof tmpFile
    24  	fileSize int64
    25  	dbIdx    int // selected db index when startRewrite
    26  }
    27  
    28  // Rewrite carries out AOF rewrite
    29  func (persister *Persister) Rewrite() error {
    30  	ctx, err := persister.StartRewrite()
    31  	if err != nil {
    32  		return err
    33  	}
    34  	err = persister.DoRewrite(ctx)
    35  	if err != nil {
    36  		return err
    37  	}
    38  
    39  	persister.FinishRewrite(ctx)
    40  	return nil
    41  }
    42  
    43  // DoRewrite actually rewrite aof file
    44  // makes DoRewrite public for testing only, please use Rewrite instead
    45  func (persister *Persister) DoRewrite(ctx *RewriteCtx) (err error) {
    46  	// start rewrite
    47  	if !config.Properties.AofUseRdbPreamble {
    48  		logger.Info("generate aof preamble")
    49  		err = persister.generateAof(ctx)
    50  	} else {
    51  		logger.Info("generate rdb preamble")
    52  		err = persister.generateRDB(ctx)
    53  	}
    54  	return err
    55  }
    56  
    57  // StartRewrite prepares rewrite procedure
    58  func (persister *Persister) StartRewrite() (*RewriteCtx, error) {
    59  	// pausing aof
    60  	persister.pausingAof.Lock()
    61  	defer persister.pausingAof.Unlock()
    62  
    63  	err := persister.aofFile.Sync()
    64  	if err != nil {
    65  		logger.Warn("fsync failed")
    66  		return nil, err
    67  	}
    68  
    69  	// get current aof file size
    70  	fileInfo, _ := os.Stat(persister.aofFilename)
    71  	filesize := fileInfo.Size()
    72  
    73  	// create tmp file
    74  	file, err := os.CreateTemp(config.GetTmpDir(), "*.aof")
    75  	if err != nil {
    76  		logger.Warn("tmp file create failed")
    77  		return nil, err
    78  	}
    79  	return &RewriteCtx{
    80  		tmpFile:  file,
    81  		fileSize: filesize,
    82  		dbIdx:    persister.currentDB,
    83  	}, nil
    84  }
    85  
    86  // FinishRewrite finish rewrite procedure
    87  func (persister *Persister) FinishRewrite(ctx *RewriteCtx) {
    88  	persister.pausingAof.Lock() // pausing aof
    89  	defer persister.pausingAof.Unlock()
    90  	tmpFile := ctx.tmpFile
    91  
    92  	// copy commands executed during rewriting to tmpFile
    93  	errOccurs := func() bool {
    94  		/* read write commands executed during rewriting */
    95  		src, err := os.Open(persister.aofFilename)
    96  		if err != nil {
    97  			logger.Error("open aofFilename failed: " + err.Error())
    98  			return true
    99  		}
   100  		defer func() {
   101  			_ = src.Close()
   102  			_ = tmpFile.Close()
   103  		}()
   104  
   105  		_, err = src.Seek(ctx.fileSize, 0)
   106  		if err != nil {
   107  			logger.Error("seek failed: " + err.Error())
   108  			return true
   109  		}
   110  		// sync tmpFile's db index with online aofFile
   111  		data := protocol.MakeMultiBulkReply(utils.ToCmdLine("SELECT", strconv.Itoa(ctx.dbIdx))).ToBytes()
   112  		_, err = tmpFile.Write(data)
   113  		if err != nil {
   114  			logger.Error("tmp file rewrite failed: " + err.Error())
   115  			return true
   116  		}
   117  		// copy data
   118  		_, err = io.Copy(tmpFile, src)
   119  		if err != nil {
   120  			logger.Error("copy aof filed failed: " + err.Error())
   121  			return true
   122  		}
   123  		return false
   124  	}()
   125  	if errOccurs {
   126  		return
   127  	}
   128  
   129  	// replace current aof file by tmp file
   130  	_ = persister.aofFile.Close()
   131  	if err := os.Rename(tmpFile.Name(), persister.aofFilename); err != nil {
   132  		logger.Warn(err)
   133  	}
   134  	// reopen aof file for further write
   135  	aofFile, err := os.OpenFile(persister.aofFilename, os.O_APPEND|os.O_CREATE|os.O_RDWR, 0600)
   136  	if err != nil {
   137  		panic(err)
   138  	}
   139  	persister.aofFile = aofFile
   140  
   141  	// write select command again to resume aof file selected db
   142  	// it should have the same db index with  persister.currentDB
   143  	data := protocol.MakeMultiBulkReply(utils.ToCmdLine("SELECT", strconv.Itoa(persister.currentDB))).ToBytes()
   144  	_, err = persister.aofFile.Write(data)
   145  	if err != nil {
   146  		panic(err)
   147  	}
   148  }