TinyKV Project 2B

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

// Server API for TinyKv service
type TinyKvServer interface {
  // KV commands with mvcc/txn supported.
  KvGet(context.Context, *kvrpcpb.GetRequest) (*kvrpcpb.GetResponse, error)
  KvScan(context.Context, *kvrpcpb.ScanRequest) (*kvrpcpb.ScanResponse, error)
  KvPrewrite(context.Context, *kvrpcpb.PrewriteRequest) (*kvrpcpb.PrewriteResponse, error)
  KvCommit(context.Context, *kvrpcpb.CommitRequest) (*kvrpcpb.CommitResponse, error)
  KvCheckTxnStatus(context.Context, *kvrpcpb.CheckTxnStatusRequest) (*kvrpcpb.CheckTxnStatusResponse, error)
  KvBatchRollback(context.Context, *kvrpcpb.BatchRollbackRequest) (*kvrpcpb.BatchRollbackResponse, error)
  KvResolveLock(context.Context, *kvrpcpb.ResolveLockRequest) (*kvrpcpb.ResolveLockResponse, error)
  // RawKV commands.
  RawGet(context.Context, *kvrpcpb.RawGetRequest) (*kvrpcpb.RawGetResponse, error)
  RawPut(context.Context, *kvrpcpb.RawPutRequest) (*kvrpcpb.RawPutResponse, error)
  RawDelete(context.Context, *kvrpcpb.RawDeleteRequest) (*kvrpcpb.RawDeleteResponse, error)
  RawScan(context.Context, *kvrpcpb.RawScanRequest) (*kvrpcpb.RawScanResponse, error)
  // Raft commands (tinykv <-> tinykv).
  Raft(TinyKv_RaftServer) error
  Snapshot(TinyKv_SnapshotServer) error
  // Coprocessor
  Coprocessor(context.Context, *coprocessor.Request) (*coprocessor.Response, error)
}

1
2
3
4
5
6
7
8
9

type Server struct {
  storage storage.Storage

  // (Used in 4B)
  Latches *latches.Latches		// 锁集合，实现多版本并发控制

  // coprocessor API handler, out of course scope
  copHandler *coprocessor.CopHandler
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

type RaftStorage struct {
  engines *engine_util.Engines
  config  *config.Config

  node          *raftstore.Node
  snapManager   *snap.SnapManager
  raftRouter    *raftstore.RaftstoreRouter
  raftSystem    *raftstore.Raftstore
  resolveWorker *worker.Worker
  snapWorker    *worker.Worker

  wg sync.WaitGroup
}

type RaftstoreRouter struct {
  router *router
}

type router struct {
  peers       sync.Map // regionID -> peerState
  peerSender  chan message.Msg
  storeSender chan<- message.Msg
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

type raftWorker struct {
  pr *router

  // receiver of messages should sent to raft, including:
  // * raft command from `raftStorage`
  // * raft inner messages from other peers sent by network
  raftCh chan message.Msg
  ctx    *GlobalContext

  closeCh <-chan struct{}
}

func newRaftWorker(ctx *GlobalContext, pm *router) *raftWorker {
  return &raftWorker{
    raftCh: pm.peerSender,
    ctx:    ctx,
    pr:     pm,
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

func (rw *raftWorker) run(closeCh <-chan struct{}, wg *sync.WaitGroup) {
  defer wg.Done()
  var msgs []message.Msg
  for {
    msgs = msgs[:0]
    select {
    case <-closeCh:
      return
    case msg := <-rw.raftCh:
      msgs = append(msgs, msg)
    }
    pending := len(rw.raftCh)
    for i := 0; i < pending; i++ {
      msgs = append(msgs, <-rw.raftCh)
    }
    peerStateMap := make(map[uint64]*peerState)
    for _, msg := range msgs {
      peerState := rw.getPeerState(peerStateMap, msg.RegionID)
      if peerState == nil {
        continue
      }
      newPeerMsgHandler(peerState.peer, rw.ctx).HandleMsg(msg)
    }
    for _, peerState := range peerStateMap {
      newPeerMsgHandler(peerState.peer, rw.ctx).HandleRaftReady()
    }
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

func (d *peerMsgHandler) HandleMsg(msg message.Msg) {
  switch msg.Type {
  case message.MsgTypeRaftMessage:
    raftMsg := msg.Data.(*rspb.RaftMessage)
    if err := d.onRaftMsg(raftMsg); err != nil {
      log.Errorf("%s handle raft message error %v", d.Tag, err)
    }
  case message.MsgTypeRaftCmd:
    raftCMD := msg.Data.(*message.MsgRaftCmd)
    d.proposeRaftCommand(raftCMD.Request, raftCMD.Callback)
  case message.MsgTypeTick:
    d.onTick()
  case message.MsgTypeSplitRegion:
    split := msg.Data.(*message.MsgSplitRegion)
    log.Infof("%s on split with %v", d.Tag, split.SplitKey)
    d.onPrepareSplitRegion(split.RegionEpoch, split.SplitKey, split.Callback)
  case message.MsgTypeRegionApproximateSize:
    d.onApproximateRegionSize(msg.Data.(uint64))
  case message.MsgTypeGcSnap:
    gcSnap := msg.Data.(*message.MsgGCSnap)
    d.onGCSnap(gcSnap.Snaps)
  case message.MsgTypeStart:
    d.startTicker()
  }
}

1
2
3
4
5
6
7
8

func (d *peerMsgHandler) proposeRaftCommand(msg *raft_cmdpb.RaftCmdRequest, cb *message.Callback) {
  err := d.preProposeRaftCommand(msg)
  if err != nil {
    cb.Done(ErrResp(err))
    return
  }
  // Your Code Here (2B).
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

func (rw *raftWorker) run(closeCh <-chan struct{}, wg *sync.WaitGroup) {
  defer wg.Done()
  var msgs []message.Msg
  for {
    msgs = msgs[:0]
    select {
    case <-closeCh:
      return
    case msg := <-rw.raftCh:
      msgs = append(msgs, msg)
    }
    pending := len(rw.raftCh)
    for i := 0; i < pending; i++ {
      msgs = append(msgs, <-rw.raftCh)
    }
    peerStateMap := make(map[uint64]*peerState)
    for _, msg := range msgs {
      peerState := rw.getPeerState(peerStateMap, msg.RegionID)
      if peerState == nil {
        continue
      }
      newPeerMsgHandler(peerState.peer, rw.ctx).HandleMsg(msg)
    }
    for _, peerState := range peerStateMap {
      newPeerMsgHandler(peerState.peer, rw.ctx).HandleRaftReady()
    }
  }
}

1
2
3
4
5
6

// Append the given entries to the raft log and update ps.raftState also delete log entries that will
// never be committed
func (ps *PeerStorage) Append(entries []eraftpb.Entry, raftWB *engine_util.WriteBatch) error {
  // Your Code Here (2B).
  return nil
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

func (ps *PeerStorage) SaveReadyState(ready *raft.Ready) (*ApplySnapResult, error) {
  // Hint: you may call `Append()` and `ApplySnapshot()` in this function
  // Your Code Here (2B/2C).

  // 1. append
  wb := &engine_util.WriteBatch{}
  ps.Append(ready.Entries, wb)
  // 2. update hardstate and save raftState
  if !raft.IsEmptyHardState(ready.HardState) {
    ps.raftState.HardState = &ready.HardState
  }
  wb.SetMeta(meta.RaftStateKey(ps.region.Id), ps.raftState)
  // 3. write to DB
  ps.Engines.WriteRaft(wb)

  return nil, nil
}

1
2
3
4
5

if !d.IsLeader() {
  leader := d.getPeerFromCache(d.LeaderId())
  cb.Done(ErrResp(&util.ErrNotLeader{RegionId: d.regionId, Leader: leader}))
  return
}

1
2
3
4
5
6
7
8
9

func (d *peerMsgHandler) proposeRaftCommand(msg *raft_cmdpb.RaftCmdRequest, cb *message.Callback) {
  err := d.preProposeRaftCommand(msg)
  if err != nil {
    cb.Done(ErrResp(err))
    return
  }
  // Your Code Here (2B).
  return
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

// Record the callback of the proposals
// (Used in 2B)
proposals []*proposal

type proposal struct {
  // index + term for unique identification
  index uint64
  term  uint64
  cb    *message.Callback
}

1
2
3
4
5

proposal := &proposal{
  index: d.nextProposalIndex(),
  term:  d.Term(),
  cb:    cb,
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

  proposal := &proposal{
    index: d.nextProposalIndex(),
    term:  d.Term(),
    cb:    cb,
  }
  data, marErr := msg.Marshal()
  if marErr != nil {
    cb.Done(ErrResp(marErr))
    return
  }
  perr := d.RaftGroup.Propose(data)
  if perr != nil {
    cb.Done(ErrResp(perr))
    return
  }
  // log.DPrintf("kv/raftstore/peer_msg_handler.go: proposal: %v", proposal)
  d.proposals = append(d.proposals, proposal)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

for {
  select {
  case <-s.Ticker:
    Node.Tick()
  default:
    if Node.HasReady() {
      rd := Node.Ready()
      saveToStorage(rd.State, rd.Entries, rd.Snapshot)
      send(rd.Messages)
      for _, entry := range rd.CommittedEntries {
        process(entry)
      }
      s.Node.Advance(rd)
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

// if return true, d.proposals[0] is the target proposal
func (d *peerMsgHandler) clearStaleAndGetTargetProposal(entry *eraftpb.Entry) bool {
  d.clearStaleProposals(entry)
  if len(d.proposals) > 0 && d.proposals[0].index == entry.Index {
    p := d.proposals[0]
    if p.term != entry.Term {
      NotifyStaleReq(entry.Term, p.cb)
      d.proposals = d.proposals[1:]
      return false
    } else {
      return true
    }
  } else {
    return false
  }
}

func (d *peerMsgHandler) clearStaleProposals(entry *eraftpb.Entry) {
  var i int
  for i = 0; i < len(d.proposals) && d.proposals[i].index < entry.Index; i++ {
    d.proposals[i].cb.Done(ErrResp(&util.ErrStaleCommand{}))
  }
  d.proposals = d.proposals[i:]
}

1
2
3
4
5
6
7
8

type RaftCmdResponse struct {
  Header               *RaftResponseHeader `protobuf:"bytes,1,opt,name=header" json:"header,omitempty"`
  Responses            []*Response         `protobuf:"bytes,2,rep,name=responses" json:"responses,omitempty"`
  AdminResponse        *AdminResponse      `protobuf:"bytes,3,opt,name=admin_response,json=adminResponse" json:"admin_response,omitempty"`
  XXX_NoUnkeyedLiteral struct{}            `json:"-"`
  XXX_unrecognized     []byte              `json:"-"`
  XXX_sizecache        int32               `json:"-"`
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

var p *proposal
resp := &raft_cmdpb.RaftCmdResponse{
  Header:    &raft_cmdpb.RaftResponseHeader{},
  Responses: []*raft_cmdpb.Response{},
}
for _, req := range cmd.Requests {
  switch req.CmdType {
  case raft_cmdpb.CmdType_Put:
    wb.SetCF(req.Put.Cf, req.Put.Key, req.Put.Value)

  case raft_cmdpb.CmdType_Delete:
    wb.DeleteCF(req.Delete.Cf, req.Delete.Key)

  case raft_cmdpb.CmdType_Get:
  case raft_cmdpb.CmdType_Snap:
  }

  if d.clearStaleAndGetTargetProposal(entry) {
    p = d.proposals[0]
    switch req.CmdType {
    case raft_cmdpb.CmdType_Put:
      resp.Responses = append(resp.Responses, &raft_cmdpb.Response{CmdType: raft_cmdpb.CmdType_Put, Put: &raft_cmdpb.PutResponse{}})
    case raft_cmdpb.CmdType_Delete:
      resp.Responses = append(resp.Responses, &raft_cmdpb.Response{CmdType: raft_cmdpb.CmdType_Delete, Delete: &raft_cmdpb.DeleteResponse{}})
    case raft_cmdpb.CmdType_Get:
      val, err := engine_util.GetCF(d.peerStorage.Engines.Kv, req.Get.Cf, req.Get.Key)
      if err != nil {
        panic(err)
      }
      resp.Responses = append(resp.Responses, &raft_cmdpb.Response{CmdType: raft_cmdpb.CmdType_Get, Get: &raft_cmdpb.GetResponse{Value: val}})
    case raft_cmdpb.CmdType_Snap:
      resp.Responses = append(resp.Responses, &raft_cmdpb.Response{CmdType: raft_cmdpb.CmdType_Snap, Snap: &raft_cmdpb.SnapResponse{Region: d.Region()}})
      p.cb.Txn = d.peerStorage.Engines.Kv.NewTransaction(false)
    default:
      resp.Responses = append(resp.Responses, &raft_cmdpb.Response{CmdType: raft_cmdpb.CmdType_Invalid})
    }
  }
}
if p != nil {
  p.cb.Done(resp)
  d.proposals = d.proposals[1:]
}