func PipeReaderWriter(wg *sync.WaitGroup, r io.Reader, w io.Writer, nread, nwrite *AtomicInt64, total int64) { wg.Add(1) go func() { defer wg.Done() for total != 0 { p := make([]byte, 1024) if total > 0 && int64(len(p)) > total { p = p[:total] } if n, err := r.Read(p); err != nil { utils.Panic("read full error = '%s'", err) } else { p = p[:n] } delta := int64(len(p)) nread.Add(delta) for len(p) != 0 { n, err := w.Write(p) if err != nil { utils.Panic("write error = '%s'", err) } p = p[n:] } nwrite.Add(delta) if total > 0 { total -= delta } } }() }
func NewRdbLoader(wg *sync.WaitGroup, size int, reader *bufio.Reader, nread *AtomicInt64) *RdbLoader { wg.Add(1) pipe := make(chan *rdb.Entry, size) go func() { defer close(pipe) defer wg.Done() l := rdb.NewLoader(reader) if err := l.LoadHeader(); err != nil { utils.Panic("parse rdb header error = '%s'", err) } for { if entry, offset, err := l.LoadEntry(); err != nil { utils.Panic("parse rdb entry error = '%s'", err) } else { if entry != nil { nread.Set(offset) pipe <- entry } else { if err := l.LoadChecksum(); err != nil { utils.Panic("parse rdb checksum error = '%s'", err) } return } } } }() return &RdbLoader{pipe} }
func openReadFile(name string) (f *os.File, nsize int64) { var err error if f, err = os.Open(name); err != nil { utils.Panic("cannot open file-reader '%s', error = '%s'", name, err) } if fi, err := f.Stat(); err != nil { utils.Panic("cannot stat file-reader '%s', error = '%s'", name, err) } else { nsize = fi.Size() } return }
func openWriteFile(name string) *os.File { f, err := os.OpenFile(name, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666) if err != nil { utils.Panic("cannot open file-writer '%s', error = %s", name, err) } return f }
func openNetConn(target string) net.Conn { c, err := net.Dial("tcp", target) if err != nil { utils.Panic("cannot connect to '%s', error = '%s'", target, err) } return c }
func NewBufWriter(wg *sync.WaitGroup, size int, writer *bufio.Writer, nwrite *AtomicInt64) *BufWriter { wg.Add(1) pipe := make(chan string, size) go func() { defer wg.Done() for s := range pipe { if _, err := writer.WriteString(s); err != nil { utils.Panic("write error = '%s'", err) } if err := writer.Flush(); err != nil { utils.Panic("flush error = '%s'", err) } nwrite.Add(int64(len(s))) } }() return &BufWriter{pipe} }
func restoreRdbEntry(c redis.Conn, e *rdb.Entry) { var ttlms uint64 if e.ExpireAt != 0 { if now := uint64(time.Now().UnixNano() / int64(time.Millisecond)); now >= e.ExpireAt { ttlms = 1 } else { ttlms = e.ExpireAt - now } } s, err := redis.String(c.Do("slotsrestore", e.Key, ttlms, e.ValDump)) if err != nil { utils.Panic("restore command error = '%s'", err) } if s != "OK" { utils.Panic("restore command response = '%s', should be 'OK'", s) } }
func init() { usage := ` Usage: redis-port decode [--ncpu=N] [--input=INPUT] [--output=OUTPUT] redis-port restore [--ncpu=N] [--input=INPUT] --target=TARGET redis-port dump [--ncpu=N] --from=MASTER [--output=OUTPUT] redis-port sync [--ncpu=N] --from=MASTER --target=TARGET Options: -n N, --ncpu=N Set runtime.GOMAXPROCS to N. -i INPUT, --input=INPUT Set input file, default is stdin ('/dev/stdin'). -o OUTPUT, --output=OUTPUT Set output file, default is stdout ('/dev/stdout'). -f MASTER, --from=MASTER Set master redis. -t TARGET, --target=TARGET Set slave redis. ` var err error args, err = docopt.Parse(usage, nil, true, "", false) if err != nil { utils.Panic("parse usage error = '%s'", err) } ncpu = runtime.GOMAXPROCS(0) if s := strArg("--ncpu", false); len(s) != 0 { if n, err := strconv.ParseInt(s, 10, 64); err != nil { utils.Panic("parse --ncpu = %v, error = '%s'", s, err) } else if n <= 0 || n > 64 { utils.Panic("parse --ncpu = %d, only accept [1,64]", n) } else { ncpu = int(n) runtime.GOMAXPROCS(ncpu) } } if ncpu == 0 { ncpu = 1 } for _, s := range []string{"decode", "restore", "dump", "sync"} { if args[s].(bool) { code = s return } } utils.Panic("parse command code error") }
func strArg(name string, nonil bool) string { if v := args[name]; v != nil { if s, ok := v.(string); !ok { utils.Panic("argument `%s' is not a string", name) } else if len(s) != 0 { return s } } if nonil { invArg(name) } return "" }
func Dump(ncpu int, from, output string) { log.Printf("[ncpu=%d] dump from '%s' to '%s'\n", ncpu, from, output) fout := openWriteFile(output) defer fout.Close() master, wait := openSyncConn(from) defer master.Close() var nsize int64 for nsize == 0 { select { case nsize = <-wait: if nsize == 0 { log.Println("+") } case <-time.After(time.Second): log.Println("-") } } var nread, nwrite AtomicInt64 var wg sync.WaitGroup wg.Add(1) go func() { defer wg.Done() for { r, w := nread.Get(), nwrite.Get() p := 100 * r / nsize log.Printf("total = %d - %3d%%, read=%-14d write=%-14d\n", nsize, p, r, w) if nsize == r { log.Printf("done\n") return } time.Sleep(time.Second) } }() reader := bufio.NewReaderSize(master, 1024*1024*32) writer := bufio.NewWriterSize(fout, 1024*64) PipeReaderWriter(&wg, reader, writer, &nread, &nwrite, nsize) wg.Wait() if err := writer.Flush(); err != nil { utils.Panic("writer flush error = '%s'", err) } }
func openSyncConn(target string) (net.Conn, chan int64) { c := openNetConn(target) for cmd := []byte("sync\r\n"); len(cmd) != 0; { n, err := c.Write(cmd) if err != nil { utils.Panic("write sync command error = %s", err) } cmd = cmd[n:] } size := make(chan int64) go func() { var rsp string for { b := []byte{0} if _, err := c.Read(b); err != nil { utils.Panic("read sync response = '%s', error = %s", rsp, err) } if len(rsp) == 0 && b[0] == '\n' { size <- 0 continue } rsp += string(b) if strings.HasSuffix(rsp, "\r\n") { break } } if rsp[0] != '$' { utils.Panic("invalid sync response, rsp = '%s'", rsp) } n, err := strconv.Atoi(rsp[1 : len(rsp)-2]) if err != nil || n <= 0 { utils.Panic("invalid sync response = '%s', error = %s, n = %d", rsp, err, n) } size <- int64(n) }() return c, size }
func Restore(ncpu int, input, target string) { log.Printf("[ncpu=%d] restore from '%s' to '%s'\n", ncpu, input, target) fin, nsize := openReadFile(input) defer fin.Close() var nread, nrestore AtomicInt64 var wg sync.WaitGroup onTick := func() { r, s := nread.Get(), nrestore.Get() if nsize != 0 { p := 100 * r / nsize log.Printf("total = %d - %3d%%, read=%-14d restore=%-14d\n", nsize, p, r, s) } else { log.Printf("total = unknown - read=%-14d restore=%-14d\n", r, s) } } onClose := func() { onTick() log.Printf("done\n") } ticker := NewClockTicker(&wg, onTick, onClose) loader := NewRdbLoader(&wg, ncpu*32, bufio.NewReaderSize(fin, 1024*1024*32), &nread) for i, count := 0, AtomicInt64(ncpu); i < ncpu; i++ { wg.Add(1) go func() { defer wg.Done() defer func() { if count.Sub(1) == 0 { ticker.Close() } }() c := openRedisConn(target) defer c.Close() for e := range loader.Pipe() { if e.DB != 0 { utils.Panic("dbnum must be 0, but got %d", e.DB) } restoreRdbEntry(c, e) nrestore.Add(1) } }() } wg.Wait() }
func Decode(ncpu int, input, output string) { log.Printf("[ncpu=%d] decode from '%s' to '%s'\n", ncpu, input, output) fin, nsize := openReadFile(input) defer fin.Close() fout := openWriteFile(output) defer fout.Close() var nread, nwrite AtomicInt64 var wg sync.WaitGroup onTick := func() { r, w := nread.Get(), nwrite.Get() if nsize != 0 { p := 100 * r / nsize log.Printf("total = %d - %3d%%, read=%-14d write=%-14d\n", nsize, p, r, w) } else { log.Printf("total = unknown - read=%-14d write=%-14d\n", r, w) } } onClose := func() { onTick() log.Printf("done\n") } ticker := NewClockTicker(&wg, onTick, onClose) loader := NewRdbLoader(&wg, ncpu*32, bufio.NewReaderSize(fin, 1024*1024*32), &nread) writer := NewBufWriter(&wg, ncpu*32, bufio.NewWriterSize(fout, 128*1024), &nwrite) for i, count := 0, AtomicInt64(ncpu); i < ncpu; i++ { wg.Add(1) go func() { defer wg.Done() defer func() { if count.Sub(1) == 0 { writer.Close() ticker.Close() } }() toHexString := func(p []byte) string { var b bytes.Buffer b.WriteByte('{') for _, c := range p { switch { case c >= 'a' && c <= 'z': fallthrough case c >= 'A' && c <= 'Z': fallthrough case c >= '0' && c <= '9': b.WriteByte(c) default: b.WriteByte('.') } } b.WriteByte('|') b.WriteString(hex.EncodeToString(p)) b.WriteByte('}') return b.String() } for e := range loader.Pipe() { o, err := rdb.DecodeDump(e.ValDump) if err != nil { utils.Panic("decode error = '%s'", err) } var b bytes.Buffer key := toHexString(e.Key) switch obj := o.(type) { default: utils.Panic("unknown object %v", o) case rdb.String: val := toHexString(obj) fmt.Fprintf(&b, "db=%d type=%s expireat=%d key=%s value=%s\n", e.DB, "string", e.ExpireAt, key, val) case rdb.List: for _, x := range obj { ele := toHexString(x) fmt.Fprintf(&b, "db=%d type=%s expireat=%d key=%s element=%s\n", e.DB, "list", e.ExpireAt, key, ele) } case rdb.HashMap: for _, x := range obj { fld := toHexString(x.Field) mem := toHexString(x.Value) fmt.Fprintf(&b, "db=%d type=%s expireat=%d key=%s field=%s member=%s\n", e.DB, "hset", e.ExpireAt, key, fld, mem) } case rdb.Set: for _, x := range obj { mem := toHexString(x) fmt.Fprintf(&b, "db=%d type=%s expireat=%d key=%s member=%s\n", e.DB, "set", e.ExpireAt, key, mem) } case rdb.ZSet: for _, x := range obj { mem := toHexString(x.Member) fmt.Fprintf(&b, "db=%d type=%s expireat=%d key=%s member=%s score=%f\n", e.DB, "zset", e.ExpireAt, key, mem, x.Score) } } writer.Append(b.String()) } }() } wg.Wait() }
func invArg(name string) { utils.Panic("please specify argument `%s' correctly", name) }
func Sync(ncpu int, from, target string) { log.Printf("[ncpu=%d] sync from '%s' to '%s'\n", ncpu, from, target) master, wait := openSyncConn(from) defer master.Close() var nsize int64 for nsize == 0 { select { case nsize = <-wait: if nsize == 0 { log.Println("+") } case <-time.After(time.Second): log.Println("-") } } var nread, nrestore AtomicInt64 var wg sync.WaitGroup onTick := func() { r, s := nread.Get(), nrestore.Get() p := 100 * r / nsize log.Printf("sync: total = %d - %3d%%, read=%-14d restore=%-14d\n", nsize, p, r, s) } onClose := func() { onTick() log.Printf("sync: done\n") } ticker := NewClockTicker(&wg, onTick, onClose) reader := bufio.NewReaderSize(master, 1024*1024*32) loader := NewRdbLoader(&wg, ncpu*32, reader, &nread) for i, count := 0, AtomicInt64(ncpu); i < ncpu; i++ { wg.Add(1) go func() { defer wg.Done() defer func() { if count.Sub(1) == 0 { ticker.Close() } }() c := openRedisConn(target) defer c.Close() for e := range loader.Pipe() { if e.DB != 0 { utils.Panic("dbnum must b 0, but got %d", e.DB) } restoreRdbEntry(c, e) nrestore.Add(1) } }() } wg.Wait() slave := openNetConn(target) defer slave.Close() var nsend, nrecv, ndiscard AtomicInt64 PipeReaderWriter(&wg, reader, slave, &nsend, &ndiscard, -1) PipeReaderWriter(&wg, slave, ioutil.Discard, &nrecv, &ndiscard, -1) wg.Add(1) go func() { defer wg.Done() for { time.Sleep(time.Second) s, r := nsend.Reset(), nrecv.Reset() log.Printf("pipe: send=%-14d recv=%-14d\n", s, r) } }() wg.Wait() }