func (ms *master) convertMigPkg(pkg []byte, head *proto.PkgHead) ([]byte, error) {
_, err := head.Decode(pkg)
if err != nil {
return nil, err
}
if !ms.migration {
return pkg, nil
}
switch head.Cmd {
case proto.CmdIncr:
fallthrough
case proto.CmdDel:
fallthrough
case proto.CmdSet:
var p proto.PkgOneOp
_, err = p.Decode(pkg)
if err != nil {
return nil, err
}
if ms.slotId == ctrl.GetSlotId(p.DbId, p.TableId, p.RowKey) {
return pkg, nil
} else {
return nil, nil
}
case proto.CmdSync:
fallthrough
case proto.CmdMIncr:
fallthrough
case proto.CmdMDel:
fallthrough
case proto.CmdMSet:
var p proto.PkgMultiOp
_, err = p.Decode(pkg)
if err != nil {
return nil, err
}
var kvs []proto.KeyValue
for i := 0; i < len(p.Kvs); i++ {
if ms.slotId == ctrl.GetSlotId(p.DbId, p.Kvs[i].TableId, p.Kvs[i].RowKey) {
kvs = append(kvs, p.Kvs[i])
}
}
if len(kvs) == 0 {
return nil, nil
} else {
p.Kvs = kvs
pkg = make([]byte, p.Length())
_, err = p.Encode(pkg)
if err != nil {
return nil, err
}
return pkg, nil
}
}
return nil, nil
}
func TestTableDel(t *testing.T) {
var in proto.PkgOneOp
in.Cmd = proto.CmdSet
in.DbId = 1
in.Seq = 10
in.KeyValue = getTestKV(2, []byte("row1"), []byte("col1"), nil, 0, 0)
myDel(in, testAuth, getTestWA(), true, t)
}
// Cache authorize result. When authorizing again, return directly.
func (c *Client) cachAuth(pkg []byte) {
var one proto.PkgOneOp
_, err := one.Decode(pkg)
if err == nil && one.ErrCode == 0 {
c.mtx.Lock()
if c.authBM == nil {
c.authBM = util.NewBitMap(256 / 8)
}
c.authBM.Set(uint(one.DbId))
c.mtx.Unlock()
}
}
func TestTableSetCas(t *testing.T) {
var in proto.PkgOneOp
in.Cmd = proto.CmdSet
in.DbId = 1
in.Seq = 10
in.KeyValue = getTestKV(2, []byte("row1"), []byte("col1"), []byte("v1"), 30, 600)
out := mySet(in, testAuth, getTestWA(), false, t)
if out.ErrCode != table.EcCasNotMatch {
t.Fatalf("Should fail with EcCasNotMatch")
}
}
func (ms *master) syncStatus(key string, lastSeq uint64) {
var p proto.PkgOneOp
p.Cmd = proto.CmdSyncSt
p.DbId = proto.AdminDbId
p.Seq = lastSeq
p.RowKey = []byte(key)
var pkg = make([]byte, p.Length())
p.Encode(pkg)
ms.cli.AddResp(pkg)
}
func (slv *slave) SendAuthToMaster() error {
slv.mtx.Lock()
var cli = slv.cli
slv.mtx.Unlock()
if cli == nil {
return nil
}
if len(slv.adminPwd) == 0 {
return nil
}
var p proto.PkgOneOp
p.DbId = proto.AdminDbId
p.Cmd = proto.CmdAuth
p.RowKey = []byte(slv.adminPwd)
var pkg = make([]byte, p.Length())
_, err := p.Encode(pkg)
if err != nil {
return err
}
cli.AddResp(pkg)
return nil
}
func TestTableIncr(t *testing.T) {
var in proto.PkgOneOp
in.Cmd = proto.CmdGet
in.DbId = 1
in.Seq = 10
in.KeyValue = getTestKV(2, []byte("row1"), []byte("col1"), nil, -21, 0)
out := myIncr(in, testAuth, getTestWA(), true, t)
if len(out.Value) != 0 {
t.Fatalf("Value mismatch: %q", out.Value)
}
if out.Score != -21 {
t.Fatalf("Score mismatch")
}
}
func TestTableGet(t *testing.T) {
var in proto.PkgOneOp
in.Cmd = proto.CmdGet
in.DbId = 1
in.Seq = 10
in.KeyValue = getTestKV(2, []byte("row1"), []byte("col1"), nil, 0, 0)
out := myGet(in, testAuth, getTestWA(), t)
if bytes.Compare(out.Value, []byte("v1")) != 0 {
t.Fatalf("Value mismatch: %q", out.Value)
}
if out.Score != 30 {
t.Fatalf("Score mismatch")
}
}
func (tbl *Table) Auth(req *PkgArgs, au Authorize) []byte {
var in proto.PkgOneOp
_, err := in.Decode(req.Pkg)
if err != nil {
in.CtrlFlag &^= 0xFF // Clear all ctrl flags
in.SetErrCode(table.EcDecodeFail)
return replyHandle(&in)
}
in.ErrCode = 0
var authDB uint8
var already bool
if au.IsAuth(proto.AdminDbId) {
authDB = proto.AdminDbId
already = true
} else if in.DbId != proto.AdminDbId && au.IsAuth(in.DbId) {
authDB = in.DbId
already = true
}
if already {
return replyHandle(&in)
}
password := string(in.RowKey)
tbl.mtx.Lock()
// Admin password
if tbl.authPwd == nil || tbl.authPwd[proto.AdminDbId] == password {
authDB = proto.AdminDbId
} else {
// Selected DB password
if len(password) > 0 && tbl.authPwd[in.DbId] == password {
authDB = in.DbId
} else {
in.SetErrCode(table.EcAuthFailed)
}
}
tbl.mtx.Unlock()
// Success
if in.ErrCode == 0 {
in.DbId = authDB
au.SetAuth(authDB)
}
return replyHandle(&in)
}
func myGet(in proto.PkgOneOp, au Authorize, wa *WriteAccess,
t *testing.T) proto.PkgOneOp {
var pkg = make([]byte, in.Length())
_, err := in.Encode(pkg)
if err != nil {
t.Fatalf("Encode failed: ", err)
}
pkg = testTbl.Get(&PkgArgs{in.Cmd, in.DbId, in.Seq, pkg}, au, wa)
var out proto.PkgOneOp
_, err = out.Decode(pkg)
if err != nil {
t.Fatalf("Decode failed: ", err)
}
if out.ErrCode < 0 {
t.Fatalf("Failed with ErrCode %d", out.ErrCode)
}
if out.DbId != in.DbId || out.Seq != in.Seq || out.TableId != in.TableId {
t.Fatalf("DbId/Seq/TableId mismatch")
}
if bytes.Compare(out.RowKey, in.RowKey) != 0 ||
bytes.Compare(out.ColKey, in.ColKey) != 0 {
t.Fatalf("RowKey/ColKey mismatch")
}
if in.Cas == 0 && out.Cas != 0 {
t.Fatalf("Invalid default Cas value")
}
return out
}
func (srv *Server) replyOneOp(req *Request, errCode int8) {
var out proto.PkgOneOp
out.Cmd = req.Cmd
out.DbId = req.DbId
out.Seq = req.Seq
out.ErrCode = errCode
if out.ErrCode != 0 {
out.CtrlFlag |= proto.CtrlErrCode
}
var pkg = make([]byte, out.Length())
_, err := out.Encode(pkg)
if err != nil {
log.Fatalf("Encode failed: %s\n", err)
}
srv.sendResp(false, req, pkg)
}
func (srv *Server) syncStatus(req *Request) {
var cliType uint32 = ClientTypeNormal
if req.Cli != nil {
cliType = req.Cli.ClientType()
}
switch cliType {
case ClientTypeSlave:
var in proto.PkgOneOp
_, err := in.Decode(req.Pkg)
if err != nil || len(in.RowKey) == 0 {
return
}
rowKey := string(in.RowKey)
switch rowKey {
case store.KeyFullSyncEnd:
srv.mc.SetStatus(ctrl.SlaveIncrSync)
log.Printf("Switch sync status to SlaveIncrSync\n")
case store.KeyIncrSyncEnd:
var st = srv.mc.Status()
srv.mc.SetStatus(ctrl.SlaveReady)
var now = time.Now()
srv.rwMtx.Lock()
var lastTime = srv.readyTime
srv.readyTime = now
srv.rwMtx.Unlock()
if st != ctrl.SlaveReady || now.Sub(lastTime).Seconds() > 120 {
log.Printf("Switch sync status to SlaveReady")
}
case store.KeySyncLogMissing:
srv.mc.SetStatus(ctrl.SlaveNeedClear)
lastSeq, _ := srv.bin.GetMasterSeq()
// Any better solution?
log.Fatalf("Slave lastSeq %d is out of sync, please clear old data! "+
"(Restart may fix this issue)", lastSeq)
}
if req.Seq > 0 {
in.RowKey = nil // Set it as an empty OP
req.Pkg = make([]byte, in.Length())
in.Encode(req.Pkg)
srv.sendResp(true, req, nil)
}
case ClientTypeNormal:
log.Printf("User cannot send SYNCST command\n")
case ClientTypeMaster:
log.Printf("Slave SYNCST failed: [%d, %d]\n", req.DbId, req.Seq)
req.Cli.Close()
}
}
func (srv *Server) auth(req *Request) {
var cliType uint32 = ClientTypeNormal
if req.Cli != nil {
cliType = req.Cli.ClientType()
}
switch cliType {
case ClientTypeNormal:
var pkg = srv.tbl.Auth(&req.PkgArgs, req.Cli)
srv.sendResp(false, req, pkg)
case ClientTypeSlave:
var in proto.PkgOneOp
_, err := in.Decode(req.Pkg)
if err != nil {
log.Printf("Decode failed for auth reply(%s), close slave!\n", err)
} else if in.ErrCode != 0 {
log.Printf("Auth failed (%d), close slave!\n", in.ErrCode)
}
if err != nil || in.ErrCode != 0 {
if req.Slv != nil {
req.Slv.Close()
} else {
req.Cli.Close()
}
return
}
// Slave auth succeed
if req.Slv != nil {
err = req.Slv.SendSlaveOfToMaster()
if err != nil {
log.Printf("SendSlaveOfToMaster failed(%s), close slave!\n", err)
req.Slv.Close()
}
}
case ClientTypeMaster:
log.Printf("Invalid client type %d for AUTH cmd, close now!\n", cliType)
req.Cli.Close()
}
}
func checkOneOp(in *proto.PkgOneOp, req *PkgArgs, au Authorize) bool {
n, err := in.Decode(req.Pkg)
if err != nil || n != len(req.Pkg) {
in.ErrCode = table.EcDecodeFail
}
if in.ErrCode == 0 && in.DbId == proto.AdminDbId {
in.ErrCode = table.EcInvDbId
}
if in.ErrCode == 0 && !au.IsAuth(in.DbId) {
in.ErrCode = table.EcNoPrivilege
}
if in.ErrCode != 0 {
in.CtrlFlag &^= 0xFF // Clear all ctrl flags
in.CtrlFlag |= proto.CtrlErrCode
}
return in.ErrCode == 0
}
func mySet(in proto.PkgOneOp, au Authorize, wa *WriteAccess, expected bool,
t *testing.T) proto.PkgOneOp {
var pkg = make([]byte, in.Length())
_, err := in.Encode(pkg)
if err != nil {
t.Fatalf("Encode failed: ", err)
}
pkg, ok := testTbl.Set(&PkgArgs{in.Cmd, in.DbId, in.Seq, pkg}, au, wa)
if ok != expected {
if expected {
t.Fatalf("Set failed")
} else {
t.Fatalf("Set should fail")
}
}
var out proto.PkgOneOp
_, err = out.Decode(pkg)
if err != nil {
t.Fatalf("Decode failed: ", err)
}
if expected {
if out.ErrCode != 0 {
t.Fatalf("Failed with ErrCode %d", out.ErrCode)
}
}
if len(out.Value) != 0 || out.Score != 0 || out.Cas != 0 {
t.Fatalf("Invalid default values")
}
if out.Seq != in.Seq || out.DbId != in.DbId || out.TableId != in.TableId {
t.Fatalf("Seq/DbId/TableId mismatch")
}
if bytes.Compare(out.RowKey, in.RowKey) != 0 ||
bytes.Compare(out.ColKey, in.ColKey) != 0 {
t.Fatalf("RowKey/ColKey mismatch")
}
return out
}
func TestTableZopSetGet(t *testing.T) {
var in proto.PkgOneOp
in.PkgFlag |= proto.FlagZop
in.Cmd = proto.CmdSet
in.DbId = 1
in.Seq = 10
in.KeyValue = getTestKV(2, []byte("row1"), []byte("col1"), []byte("v1"), 30, 0)
// ZSET
mySet(in, testAuth, getTestWA(), true, t)
// ZGET
in.Cmd = proto.CmdGet
out := myGet(in, testAuth, getTestWA(), t)
if bytes.Compare(out.Value, []byte("v1")) != 0 {
t.Fatalf("Value mismatch: %q", out.Value)
}
if out.Score != 30 {
t.Fatalf("Score mismatch")
}
}
// Get, Set, Del, Incr, ZGet, ZSet, ZDel, ZIncr
func (c *Context) goOneOp(zop bool, cmd, tableId uint8,
rowKey, colKey, value []byte, score int64, cas uint32,
done chan *Call) (*Call, error) {
call := c.cli.newCall(cmd, done)
if call.err != nil {
return call, call.err
}
var p proto.PkgOneOp
p.Seq = call.seq
p.DbId = c.dbId
p.Cmd = call.cmd
p.TableId = tableId
p.RowKey = rowKey
p.ColKey = colKey
p.SetCas(cas)
p.SetScore(score)
p.SetValue(value)
// ZGet, ZSet, ZDel, ZIncr
if zop {
p.PkgFlag |= proto.FlagZop
}
var pkgLen = p.Length()
if pkgLen > proto.MaxPkgLen {
c.cli.errCall(call, ErrInvPkgLen)
return call, call.err
}
call.pkg = make([]byte, pkgLen)
_, err := p.Encode(call.pkg)
if err != nil {
c.cli.errCall(call, err)
return call, err
}
// put request pkg to sending channel
c.cli.sending <- call
return call, nil
}
func TestTableGetCas(t *testing.T) {
var in proto.PkgOneOp
in.Cmd = proto.CmdGet
in.DbId = 1
in.Seq = 10
in.KeyValue = getTestKV(2, []byte("row1"), []byte("col1"), []byte("v1"), 0, 2)
out := myGet(in, testAuth, getTestWA(), t)
if bytes.Compare(out.Value, []byte("v1")) != 0 {
t.Fatalf("Value mismatch: %q", out.Value)
}
if out.Score != 30 {
t.Fatalf("Score mismatch")
}
if out.Cas == 0 {
t.Fatalf("Should return new cas")
}
// Set
in.Cmd = proto.CmdSet
in.SetValue(append(out.Value, []byte("-cas")...))
in.SetScore(32)
in.SetCas(out.Cas)
mySet(in, testAuth, getTestWA(), true, t)
// Set again should fail
mySet(in, testAuth, getTestWA(), false, t)
// Get
in.Cmd = proto.CmdGet
in.Cas = 0
in.CtrlFlag &^= 0xFF
out = myGet(in, testAuth, getTestWA(), t)
if bytes.Compare(out.Value, []byte("v1-cas")) != 0 {
t.Fatalf("Value mismatch: %q", out.Value)
}
if out.Score != 32 {
t.Fatalf("Score mismatch")
}
}
// Get call reply. The real reply types are:
// Auth/Ping/(Z)Set/(Z)Del: nil;
// (Z)Get: GetReply;
// (Z)Incr: IncrReply;
// (Z)MGet: []GetReply;
// (Z)MSet: []SetReply;
// (Z)MDel: []DelReply;
// (Z)MIncr: []IncrReply;
// (Z)Scan: ScanReply;
// Dump: DumpReply;
func (call *Call) Reply() (interface{}, error) {
if call.err != nil {
return nil, call.err
}
if !call.ready {
return nil, ErrCallNotReady
}
if proto.CmdAuth == call.cmd ||
proto.CmdPing == call.cmd ||
proto.CmdIncr == call.cmd ||
proto.CmdDel == call.cmd ||
proto.CmdSet == call.cmd ||
proto.CmdGet == call.cmd {
var p proto.PkgOneOp
_, err := p.Decode(call.pkg)
if err != nil {
call.err = err
return nil, call.err
}
if p.ErrCode < 0 {
return nil, getErr(p.ErrCode)
}
switch call.cmd {
case proto.CmdAuth:
return nil, nil
case proto.CmdPing:
return nil, nil
case proto.CmdIncr:
return IncrReply{p.ErrCode, p.TableId, copyBytes(p.RowKey),
copyBytes(p.ColKey), copyBytes(p.Value), p.Score}, nil
case proto.CmdDel:
return nil, nil
case proto.CmdSet:
return nil, nil
case proto.CmdGet:
return GetReply{p.ErrCode, p.TableId, copyBytes(p.RowKey),
copyBytes(p.ColKey), copyBytes(p.Value), p.Score, p.Cas}, nil
}
}
if proto.CmdMIncr == call.cmd ||
proto.CmdMDel == call.cmd ||
proto.CmdMSet == call.cmd ||
proto.CmdMGet == call.cmd {
var p proto.PkgMultiOp
_, err := p.Decode(call.pkg)
if err != nil {
call.err = err
return nil, call.err
}
if p.ErrCode < 0 {
return nil, getErr(p.ErrCode)
}
switch call.cmd {
case proto.CmdMIncr:
var r = make([]IncrReply, len(p.Kvs))
for i := 0; i < len(r); i++ {
r[i] = IncrReply{p.Kvs[i].ErrCode, p.Kvs[i].TableId,
copyBytes(p.Kvs[i].RowKey), copyBytes(p.Kvs[i].ColKey),
copyBytes(p.Kvs[i].Value), p.Kvs[i].Score}
}
return r, nil
case proto.CmdMDel:
var r = make([]DelReply, len(p.Kvs))
for i := 0; i < len(r); i++ {
r[i] = DelReply{p.Kvs[i].ErrCode, p.Kvs[i].TableId,
copyBytes(p.Kvs[i].RowKey), copyBytes(p.Kvs[i].ColKey)}
}
return r, nil
case proto.CmdMSet:
var r = make([]SetReply, len(p.Kvs))
for i := 0; i < len(r); i++ {
r[i] = SetReply{p.Kvs[i].ErrCode, p.Kvs[i].TableId,
copyBytes(p.Kvs[i].RowKey), copyBytes(p.Kvs[i].ColKey)}
}
return r, nil
case proto.CmdMGet:
var r = make([]GetReply, len(p.Kvs))
for i := 0; i < len(r); i++ {
r[i] = GetReply{p.Kvs[i].ErrCode, p.Kvs[i].TableId,
copyBytes(p.Kvs[i].RowKey), copyBytes(p.Kvs[i].ColKey),
copyBytes(p.Kvs[i].Value), p.Kvs[i].Score, p.Kvs[i].Cas}
}
return r, nil
}
}
switch call.cmd {
case proto.CmdScan:
var p proto.PkgScanResp
_, err := p.Decode(call.pkg)
if err != nil {
call.err = err
return nil, call.err
}
if p.ErrCode < 0 {
return nil, getErr(p.ErrCode)
}
var r ScanReply
r.ctx = call.ctx.(scanContext)
r.TableId = r.ctx.tableId
r.RowKey = r.ctx.rowKey
r.End = (p.PkgFlag&proto.FlagScanEnd != 0)
r.Kvs = make([]ScanKV, len(p.Kvs))
for i := 0; i < len(p.Kvs); i++ {
r.Kvs[i] = ScanKV{copyBytes(p.Kvs[i].ColKey),
copyBytes(p.Kvs[i].Value), p.Kvs[i].Score}
}
return r, nil
case proto.CmdDump:
var p proto.PkgDumpResp
_, err := p.Decode(call.pkg)
if err != nil {
call.err = err
return nil, call.err
}
if p.ErrCode < 0 {
return nil, getErr(p.ErrCode)
}
var r DumpReply
r.ctx = call.ctx.(dumpContext)
r.ctx.lastSlotId = p.LastSlotId
r.ctx.slotStart = (p.PkgFlag&proto.FlagDumpSlotStart != 0)
r.End = (p.PkgFlag&proto.FlagDumpEnd != 0)
r.Kvs = make([]DumpKV, len(p.Kvs))
for i := 0; i < len(p.Kvs); i++ {
r.Kvs[i] = DumpKV{p.Kvs[i].TableId, p.Kvs[i].ColSpace,
copyBytes(p.Kvs[i].RowKey), copyBytes(p.Kvs[i].ColKey),
copyBytes(p.Kvs[i].Value), p.Kvs[i].Score}
}
return r, nil
}
switch call.cmd {
case proto.CmdSlaveOf:
return call.replyInnerCtrl(&ctrl.PkgSlaveOf{})
case proto.CmdMigrate:
return call.replyInnerCtrl(&ctrl.PkgMigrate{})
case proto.CmdSlaveSt:
return call.replyInnerCtrl(&ctrl.PkgSlaveStatus{})
case proto.CmdDelSlot:
return call.replyInnerCtrl(&ctrl.PkgDelSlot{})
}
return nil, ErrUnknownCmd
}