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()
}