Lab3A. 基于 Raft 实现容错的 kvDB

// 发起命令的复制
// idx 是命令复制成功后，其在各节点日志中的索引
// isLeader 表明当前节点是否为 leader
func (rf *Raft) Start(command interface{}) (idx int, term int, isLeader bool)

type Clerk struct {
    servers []*labrpc.ClientEnd // kv servers / raft peers
    leader  int                 // latest known leader
    cid     int64               // client id
    seq     int32               // latest request seq num
}

type PutAppendArgs struct {
    Cid   int64 // client id
    Seq   int32 // request sequential number
    Key   string
    Value string
    Op    string // Put/Append
}

func (ck *Clerk) PutAppend(key string, value string, op string) {
    curSeq := atomic.AddInt32(&ck.seq, 1)
    for {
        args := PutAppendArgs{
            Cid:   ck.cid,
            Seq:   curSeq,
            Key:   key,
            Value: value,
            Op:    op,
        }
        var reply PutAppendReply
        ok := ck.servers[ck.leader].Call("KVServer.PutAppend", &args, &reply)

        if !ok || reply.WrongLeader { // RPC call timeout, or ck.leader isn't current leader
            ck.changeLeader()
            continue // retry and re-use current sequential number
        }
        return
    }
}

type KVServer struct {
    mu      sync.Mutex
    me      int
    rf      *raft.Raft
    applyCh chan raft.ApplyMsg

    maxraftstate int               // snapshot if log grows this big
    db           map[string]string // kvDB
    cid2seq      map[int64]int32   // client id to max request sequential number
    agreeChs     map[int]chan Op   // command index to op channel
    killCh       chan struct{}     // kill KVServer
}

type Op struct {
    Cid   int64  // client id
    Seq   int32  // request sequence number
    Cmd   string // command type, Put/Append/Get
    Key   string
    Value string
}

func isSameOp(a, b Op) bool {
    return a.Cid == b.Cid && a.Seq == b.Seq && a.Cmd == b.Cmd && a.Key == b.Key && a.Value == b.Value
}

// wait agreement from Raft
func (kv *KVServer) waitAgree() {
    for {
        select {
        case <-kv.killCh:
            return
        case msg := <-kv.applyCh:
            op := msg.Command.(Op)
            kv.mu.Lock()
            maxSeq, ok := kv.cid2seq[op.Cid]
            if !ok || op.Seq > maxSeq { // only handle new request from specific client
                kv.cid2seq[op.Cid] = op.Seq
                switch op.Cmd {
                case "Put":
                    kv.db[op.Key] = op.Value
                case "Append":
                    kv.db[op.Key] += op.Value
                }
            }
            kv.mu.Unlock()

            kv.getAgreeCh(msg.CommandIndex) <- op
        }
    }
}

func (kv *KVServer) getAgreeCh(idx int) chan Op {
    kv.mu.Lock()
    defer kv.mu.Unlock()

    ch, ok := kv.agreeChs[idx]
    if !ok {
        ch = make(chan Op, 1) // never block this
        kv.agreeChs[idx] = ch
    }
    return ch
}

func (kv *KVServer) Get(args *GetArgs, reply *GetReply) {
    cmd := Op{Cmd: "Get", Key: args.Key,}
    idx, _, isLeader := kv.rf.Start(cmd)
    if !isLeader {
        reply.WrongLeader = true
        return
    }

    ch := kv.getAgreeCh(idx)
    var op Op
    select {
    case op = <-ch: // current leader can communicate with majority
        close(ch)
    case <-time.After(500 * time.Millisecond): // agreement may be failed, treat as timeout and client will retry
        reply.WrongLeader = true
        return
    }

  // if old leader has in net partition, may it's log may be overwrited, then reply value will be different
    if !isSameOp(cmd, op) {
        reply.WrongLeader = true
        return
    }
    
    kv.mu.Lock()
    reply.Value = kv.db[args.Key] // if key not exist, just return "" or return ErrNoKey
    kv.mu.Unlock()
}

func (kv *KVServer) PutAppend(args *PutAppendArgs, reply *PutAppendReply) {
    cmd := Op{
        Cid:   args.Cid,
        Seq:   args.Seq,
        Cmd:   args.Op,
        Key:   args.Key,
        Value: args.Value,
    }
    idx, _, isLeader := kv.rf.Start(cmd)
    if !isLeader {
        reply.WrongLeader = true
        return
    }

    ch := kv.getAgreeCh(idx) // sequence of PutAppend() and <-applyCh are uncertain
    var op Op
    select {
    case op = <-ch:
        close(ch)
    case <-time.After(500 * time.Millisecond):
        reply.WrongLeader = true
        return
    }

    if !isSameOp(cmd, op) {
        reply.WrongLeader = true
        return
    }
}