Repository: matrixorigin/matrixcube (go)

Issue:
Refactoring for the storage package
There are two major issues here - 

1. The Context interface returns a []Batch, which means write requests from multiple raft logs are handed over to the state machine in one go. With such added complexity, the actual value of such further batching is questionable, plus the current implementation of executeContext and state machine don't use it anyway. In fact,  the current applyContext implementation actually requires no more than one batch to be appended into it (only one rpc.RequestBatch in applyContext). 

2. We shouldn't batch any read operation. Reads should be concurrently issued to the state machine.   

Pull request:
Fix #257 refactored storage.Context
**What type of PR is this?**

- [ ] API-change
- [ ] BUG
- [ ] Improvement
- [ ] Documentation
- [ ] Feature
- [ ] Test and CI
- [x] Code Refactoring

**Which issue(s) this PR fixes:**

Fixes #257 

**What this PR does / why we need it:**

Not Available

**Special notes for your reviewer:**

Not Available

**Additional documentation (e.g. design docs, usage docs, etc.):**

Not Available


Patch:
diff --git a/raftstore/execute_context.go b/raftstore/execute_context.go
index b0fc876..a2ed1d2 100644
--- a/raftstore/execute_context.go
+++ b/raftstore/execute_context.go
@@ -8,98 +8,121 @@ import (
 	"github.com/matrixorigin/matrixcube/storage"
 )
 
-type executeContext struct {
+type writeContext struct {
 	shard        Shard
 	wb           storage.Resetable
 	buf          *buf.ByteBuf
-	batches      []batch
-	requests     []storage.Batch
+	batch        storage.Batch
 	responses    [][]byte
 	writtenBytes uint64
 	diffBytes    int64
-	readBytes    uint64
 }
 
-func newExecuteContext(base storage.BaseStorage) *executeContext {
-	return &executeContext{
+var _ storage.WriteContext = (*writeContext)(nil)
+
+func newWriteContext(base storage.BaseStorage) *writeContext {
+	return &writeContext{
 		buf: buf.NewByteBuf(128),
 		wb:  base.NewWriteBatch(),
 	}
 }
 
-func (ctx *executeContext) close() {
+func (ctx *writeContext) close() {
 	ctx.buf.Release()
 }
 
-func (ctx *executeContext) appendRequestBatch(req rpc.RequestBatch) {
-	ctx.appendBatch(batch{requestBatch: req})
-}
-
-func (ctx *executeContext) appendBatch(c batch) {
-	b := storage.Batch{}
-	for _, req := range c.requestBatch.Requests {
-		b.Requests = append(b.Requests, storage.Request{
-			CmdType: req.CustemType,
-			Key:     req.Key,
-			Cmd:     req.Cmd,
+func (ctx *writeContext) setRequestBatch(batch rpc.RequestBatch) {
+	ctx.batch = storage.Batch{}
+	for _, r := range batch.Requests {
+		ctx.batch.Requests = append(ctx.batch.Requests, storage.Request{
+			CmdType: r.CustemType,
+			Key:     r.Key,
+			Cmd:     r.Cmd,
 		})
 	}
-
-	ctx.requests = append(ctx.requests, b)
-	ctx.batches = append(ctx.batches, c)
 }
 
-func (ctx *executeContext) hasRequest() bool {
-	return len(ctx.batches) > 0
+func (ctx *writeContext) hasRequest() bool {
+	return len(ctx.batch.Requests) > 0
 }
 
-func (ctx *executeContext) WriteBatch() storage.Resetable {
+func (ctx *writeContext) WriteBatch() storage.Resetable {
 	return ctx.wb
 }
 
-func (ctx *executeContext) ByteBuf() *buf.ByteBuf {
+func (ctx *writeContext) ByteBuf() *buf.ByteBuf {
 	return ctx.buf
 }
 
-func (ctx *executeContext) Shard() Shard {
+func (ctx *writeContext) Shard() Shard {
 	return ctx.shard
 }
 
-func (ctx *executeContext) Batches() []storage.Batch {
-	return ctx.requests
+func (ctx *writeContext) Batch() storage.Batch {
+	return ctx.batch
 }
 
-func (ctx *executeContext) AppendResponse(resp []byte) {
+func (ctx *writeContext) AppendResponse(resp []byte) {
 	ctx.responses = append(ctx.responses, resp)
 }
 
-func (ctx *executeContext) SetWrittenBytes(value uint64) {
+func (ctx *writeContext) SetWrittenBytes(value uint64) {
 	ctx.writtenBytes = value
 }
 
-func (ctx *executeContext) SetReadBytes(value uint64) {
-	ctx.readBytes = value
-}
-
-func (ctx *executeContext) SetDiffBytes(value int64) {
+func (ctx *writeContext) SetDiffBytes(value int64) {
 	ctx.diffBytes = value
 }
 
-func (ctx *executeContext) reset(shard Shard) {
+func (ctx *writeContext) reset(shard Shard) {
 	ctx.buf.Clear()
 	ctx.shard = shard
-	ctx.batches = ctx.batches[:0]
-	ctx.requests = ctx.requests[:0]
+	ctx.batch = storage.Batch{}
 	ctx.responses = ctx.responses[:0]
 	ctx.writtenBytes = 0
 	ctx.diffBytes = 0
+}
+
+type readContext struct {
+	shard     Shard
+	buf       *buf.ByteBuf
+	request   storage.Request
+	readBytes uint64
+}
+
+var _ storage.ReadContext = (*readContext)(nil)
+
+func newReadContext() *readContext {
+	return &readContext{
+		buf: buf.NewByteBuf(128),
+	}
+}
+
+func (ctx *readContext) ByteBuf() *buf.ByteBuf {
+	return ctx.buf
+}
+
+func (ctx *readContext) Shard() Shard {
+	return ctx.shard
+}
+
+func (ctx *readContext) Request() storage.Request {
+	return ctx.request
+}
+
+func (ctx *readContext) SetReadBytes(value uint64) {
+	ctx.readBytes = value
+}
+
+func (ctx *readContext) reset(shard Shard, req storage.Request) {
+	ctx.shard = shard
+	ctx.request = req
+	ctx.buf.Clear()
 	ctx.readBytes = 0
 }
 
-// TODO: this implies that we can't have more than one batch in the
-// executeContext
 type applyContext struct {
-	writeCtx    *executeContext
+	writeCtx    *writeContext
 	req         rpc.RequestBatch
 	entry       raftpb.Entry
 	adminResult *adminResult
@@ -108,7 +131,7 @@ type applyContext struct {
 
 func newApplyContext(base storage.BaseStorage) *applyContext {
 	return &applyContext{
-		writeCtx: newExecuteContext(base),
+		writeCtx: newWriteContext(base),
 		req:      rpc.RequestBatch{},
 	}
 }
diff --git a/raftstore/replica.go b/raftstore/replica.go
index a5c33ae..161e297 100644
--- a/raftstore/replica.go
+++ b/raftstore/replica.go
@@ -84,8 +84,6 @@ type replica struct {
 
 	metrics  localMetrics
 	stopOnce sync.Once
-
-	readCtx *executeContext
 }
 
 // createReplica called in:
@@ -155,9 +153,7 @@ func (pr *replica) start() {
 		}
 	}
 
-	dataStorage := pr.store.DataStorageByGroup(shard.Group)
 	pr.batch = newProposeBatch(pr.logger, uint64(pr.store.cfg.Raft.MaxEntryBytes), shard.ID, pr.replica)
-	pr.readCtx = newExecuteContext(dataStorage)
 	pr.ticks = task.New(32)
 	pr.messages = task.New(32)
 	pr.requests = task.New(32)
diff --git a/raftstore/replica_event_loop.go b/raftstore/replica_event_loop.go
index eaaa46f..4007f1e 100644
--- a/raftstore/replica_event_loop.go
+++ b/raftstore/replica_event_loop.go
@@ -117,8 +117,8 @@ func (pr *replica) onStop() {
 		}
 
 		// resp all pending requests in batch and queue
-		for _, c := range pr.pendingReads.reads {
-			for _, req := range c.req.Requests {
+		for _, rr := range pr.pendingReads.reads {
+			for _, req := range rr.batch.Requests {
 				req.Key = keys.DecodeDataKey(req.Key)
 				respStoreNotMatch(errStoreNotMatch, req, pr.store.shardsProxy.OnResponse)
 			}
diff --git a/raftstore/replica_event_read_index.go b/raftstore/replica_event_read_index.go
index 8547811..d4bb755 100644
--- a/raftstore/replica_event_read_index.go
+++ b/raftstore/replica_event_read_index.go
@@ -15,13 +15,15 @@ package raftstore
 
 import (
 	"github.com/fagongzi/util/protoc"
-	"github.com/matrixorigin/matrixcube/pb/rpc"
 	"go.etcd.io/etcd/raft/v3"
 	"go.uber.org/zap"
+
+	"github.com/matrixorigin/matrixcube/pb/rpc"
+	"github.com/matrixorigin/matrixcube/storage"
 )
 
 type readyRead struct {
-	req   rpc.RequestBatch
+	batch rpc.RequestBatch
 	index uint64
 }
 
@@ -35,10 +37,10 @@ func (q *readIndexQueue) reset() {
 }
 
 func (q *readIndexQueue) ready(state raft.ReadState) {
-	var req rpc.RequestBatch
-	protoc.MustUnmarshal(&req, state.RequestCtx)
+	var batch rpc.RequestBatch
+	protoc.MustUnmarshal(&batch, state.RequestCtx)
 	q.reads = append(q.reads, readyRead{
-		req:   req,
+		batch: batch,
 		index: state.Index,
 	})
 }
@@ -49,43 +51,44 @@ func (q *readIndexQueue) process(appliedIndex uint64, pr *replica) {
 	}
 
 	newReady := q.reads[:0] // avoid alloc new slice
-	pr.readCtx.reset(pr.getShard())
-	for _, r := range q.reads {
-		if r.index > 0 && r.index <= appliedIndex {
-			c := batch{
-				requestBatch: r.req,
-				cb:           pr.store.shardsProxy.OnResponse,
-				tp:           read,
-			}
-			pr.readCtx.appendBatch(c)
-		} else {
-			newReady = append(newReady, r)
-		}
-	}
+	shard := pr.getShard()
+	readCtx := newReadContext()
+	ds := pr.store.DataStorageByGroup(shard.Group)
 
 	// TODO (lni):
 	// multiple read requests issued to the data storage, but we won't get
 	// anything back until the completion of the slowest one.
-	q.reads = newReady
-	if pr.readCtx.hasRequest() {
-		ds := pr.store.DataStorageByGroup(pr.getShard().Group)
-		if err := ds.Read(pr.readCtx); err != nil {
-			pr.logger.Fatal("fail to exec read batch",
-				zap.Error(err))
-		}
-
-		pr.readBytes += pr.readCtx.readBytes
-		pr.readKeys += uint64(len(pr.readCtx.batches))
-		idx := 0
-		for _, c := range pr.readCtx.batches {
+	for _, ready := range q.reads {
+		if ready.index > 0 && ready.index <= appliedIndex {
 			resp := rpc.ResponseBatch{}
-			for i := 0; i < len(c.requestBatch.Requests); i++ {
-				r := rpc.Response{}
-				r.Value = pr.readCtx.responses[idx]
+			for _, req := range ready.batch.Requests {
+				rr := storage.Request{
+					// FIXME: fix the typo below in the field name
+					CmdType: req.CustemType,
+					Key:     req.Key,
+					Cmd:     req.Cmd,
+				}
+				readCtx.reset(shard, rr)
+				v, err := ds.Read(readCtx)
+				if err != nil {
+					// FIXME: some read failures should be tolerated.
+					pr.logger.Fatal("fail to exec read batch",
+						zap.Error(err))
+				}
+				pr.readBytes += readCtx.readBytes
+				pr.readKeys += 1
+				r := rpc.Response{Value: v}
 				resp.Responses = append(resp.Responses, r)
-				idx++
 			}
-			c.resp(resp)
+			// FIXME: it is strange to require a batch{} to be created/used to send
+			// responses. why sending responses is coupled with the batch type?
+			// FIXME: input parameter logger is missing
+			b := newBatch(nil, ready.batch, pr.store.shardsProxy.OnResponse, 0, 0)
+			b.resp(resp)
+		} else {
+			newReady = append(newReady, ready)
 		}
 	}
+
+	q.reads = newReady
 }
diff --git a/raftstore/replica_state_machine_exec.go b/raftstore/replica_state_machine_exec.go
index 57ab6cc..0764312 100644
--- a/raftstore/replica_state_machine_exec.go
+++ b/raftstore/replica_state_machine_exec.go
@@ -409,7 +409,7 @@ func (d *stateMachine) doExecSplit(ctx *applyContext) (rpc.ResponseBatch, error)
 
 func (d *stateMachine) execWriteRequest(ctx *applyContext) rpc.ResponseBatch {
 	ctx.writeCtx.reset(d.getShard())
-	ctx.writeCtx.appendRequestBatch(ctx.req)
+	ctx.writeCtx.setRequestBatch(ctx.req)
 	for _, req := range ctx.req.Requests {
 		if ce := d.logger.Check(zapcore.DebugLevel, "begin to execute write"); ce != nil {
 			ce.Write(log.HexField("id", req.ID))
diff --git a/storage/executor/simple/kv_executor.go b/storage/executor/simple/kv_executor.go
index b8386c7..aafc425 100644
--- a/storage/executor/simple/kv_executor.go
+++ b/storage/executor/simple/kv_executor.go
@@ -43,22 +43,20 @@ func NewSimpleKVExecutor(kv storage.KVStorage) storage.Executor {
 	return &simpleKVExecutor{kv: kv}
 }
 
-func (ce *simpleKVExecutor) UpdateWriteBatch(ctx storage.Context) error {
+func (ce *simpleKVExecutor) UpdateWriteBatch(ctx storage.WriteContext) error {
 	writtenBytes := uint64(0)
 	r := ctx.WriteBatch()
 	wb := r.(util.WriteBatch)
-	batches := ctx.Batches()
-	for i := range batches {
-		requests := batches[i].Requests
-		for j := range requests {
-			switch requests[j].CmdType {
-			case setCmd:
-				wb.Set(requests[j].Key, requests[j].Cmd)
-				writtenBytes += uint64(len(requests[j].Key) + len(requests[j].Cmd))
-				ctx.AppendResponse(OK)
-			default:
-				panic(fmt.Errorf("invalid write cmd %d", requests[j].CmdType))
-			}
+	batch := ctx.Batch()
+	requests := batch.Requests
+	for j := range requests {
+		switch requests[j].CmdType {
+		case setCmd:
+			wb.Set(requests[j].Key, requests[j].Cmd)
+			writtenBytes += uint64(len(requests[j].Key) + len(requests[j].Cmd))
+			ctx.AppendResponse(OK)
+		default:
+			panic(fmt.Errorf("invalid write cmd %d", requests[j].CmdType))
 		}
 	}
 
@@ -73,29 +71,20 @@ func (ce *simpleKVExecutor) ApplyWriteBatch(r storage.Resetable) error {
 	return ce.kv.Write(wb, false)
 }
 
-func (ce *simpleKVExecutor) Read(ctx storage.Context) error {
-	readBytes := uint64(0)
-	batches := ctx.Batches()
-	for i := range batches {
-		requests := batches[i].Requests
-		for j := range requests {
-			switch requests[j].CmdType {
-			case getCmd:
-				v, err := ce.kv.Get(requests[j].Key)
-				if err != nil {
-					return err
-				}
-
-				readBytes += uint64(len(v))
-				ctx.AppendResponse(v)
-			default:
-				panic(fmt.Errorf("invalid read cmd %d", requests[j].CmdType))
-			}
+func (ce *simpleKVExecutor) Read(ctx storage.ReadContext) ([]byte, error) {
+	request := ctx.Request()
+	switch request.CmdType {
+	case getCmd:
+		v, err := ce.kv.Get(request.Key)
+		if err != nil {
+			return nil, err
 		}
+		ctx.SetReadBytes(uint64(len(v)))
+		return v, nil
+	default:
+		panic(fmt.Errorf("invalid read cmd %d", request.CmdType))
 	}
-
-	ctx.SetReadBytes(readBytes)
-	return nil
+	panic("not suppose to reach here")
 }
 
 // NewWriteRequest return write request
diff --git a/storage/kv/kv_data_storage.go b/storage/kv/kv_data_storage.go
index 2e31b11..aa48a7e 100644
--- a/storage/kv/kv_data_storage.go
+++ b/storage/kv/kv_data_storage.go
@@ -91,23 +91,25 @@ func (kv *kvDataStorage) NewWriteBatch() storage.Resetable {
 	return kv.base.NewWriteBatch()
 }
 
-func (kv *kvDataStorage) Write(ctx storage.Context) error {
+func (kv *kvDataStorage) Write(ctx storage.WriteContext) error {
 	if err := kv.executor.UpdateWriteBatch(ctx); err != nil {
 		return err
 	}
 	r := ctx.WriteBatch()
 	defer r.Reset()
-	batches := ctx.Batches()
-	lastLogIndex := batches[len(batches)-1].Index
-	kv.setAppliedIndexToWriteBatch(ctx, lastLogIndex)
-	kv.updateAppliedIndex(ctx.Shard().ID, lastLogIndex)
+	batch := ctx.Batch()
+	if batch.Index == 0 {
+		panic("empty batch?")
+	}
+	kv.setAppliedIndexToWriteBatch(ctx, batch.Index)
+	kv.updateAppliedIndex(ctx.Shard().ID, batch.Index)
 	if err := kv.executor.ApplyWriteBatch(r); err != nil {
 		return err
 	}
 	return kv.trySync()
 }
 
-func (kv *kvDataStorage) setAppliedIndexToWriteBatch(ctx storage.Context, index uint64) {
+func (kv *kvDataStorage) setAppliedIndexToWriteBatch(ctx storage.WriteContext, index uint64) {
 	r := ctx.WriteBatch()
 	wb := r.(util.WriteBatch)
 	ctx.ByteBuf().MarkWrite()
@@ -117,7 +119,7 @@ func (kv *kvDataStorage) setAppliedIndexToWriteBatch(ctx storage.Context, index
 	wb.Set(key, val)
 }
 
-func (kv *kvDataStorage) Read(ctx storage.Context) error {
+func (kv *kvDataStorage) Read(ctx storage.ReadContext) ([]byte, error) {
 	return kv.executor.Read(ctx)
 }
 
diff --git a/storage/storage_data.go b/storage/storage_data.go
index 3e9c715..fb070cc 100644
--- a/storage/storage_data.go
+++ b/storage/storage_data.go
@@ -47,17 +47,17 @@ type Resetable interface {
 
 // Executor is used to execute read/write requests.
 type Executor interface {
-	// UpdateWriteBatch applies write requests into the provided `Context`. No
-	// writes is allowed to be written to the actual underlying data storage.
-	UpdateWriteBatch(Context) error
+	// UpdateWriteBatch applies write requests into the provided `WriteContext`.
+	// No writes is allowed to be written to the actual underlying data storage.
+	UpdateWriteBatch(WriteContext) error
 	// ApplyWriteBatch atomically applies the write batch into the underlying
 	// data storage.
 	ApplyWriteBatch(wb Resetable) error
-	// Read execute read requests. The `Context` holds all the requests involved
-	// in this execution. The implementation should call the `SetReadBytes` method
-	// of `Context` to report the statistical changes involved in this execution
-	// before returning.
-	Read(Context) error
+	// Read executes the read request and returns the result. The `ReadContext`
+	// holds the read request to be invoked in this execution. The implementation
+	// should call the `SetReadBytes` method of `Context` to report the
+	// statistical changes involved in this execution before returning.
+	Read(ReadContext) ([]byte, error)
 }
 
 // BaseStorage is the interface to be implemented by all storage types.
@@ -97,20 +97,23 @@ type BaseStorage interface {
 // GetPersistentLogIndex() + 1.
 type DataStorage interface {
 	BaseStorage
-	// Write applies write requests into the underlying data storage. The `Context`
-	// holds all requests involved and packs as many requests from multiple Raft
-	// logs together as possible. The implementation must ensure that the content
-	// of each LogRequest instance provided by the Context is atomically applied
-	// into the underlying storage. The implementation should call the
-	// `SetWrittenBytes`, `SetReadBytes`, `SetDiffBytes` methods of `Context` to
-	// report the statistical changes involved in applying the specified `Context`
-	// before returning.
-	Write(Context) error
-	// Read execute read requests. The `Context` holds all the requests involved
-	// in this execution. The implementation should call the `SetReadBytes` method
-	// of `Context` to report the statistical changes involved in this execution
-	// before returning.
-	Read(Context) error
+	// Write applies write requests into the underlying data storage. The
+	// `WriteContext` holds all requests involved and packs as many requests from
+	// multiple Raft logs together as possible. The implementation must ensure
+	// that the content of each LogRequest instance provided by the WriteContext
+	// is atomically applied into the underlying storage. The implementation
+	// should call the `SetWrittenBytes` and `SetDiffBytes` methods of the
+	// `WriteContext` to report the statistical changes involved in applying
+	// the specified `WriteContext` before returning.
+	Write(WriteContext) error
+	// TODO: refactor this method again to consider what is the best approach
+	// to avoid extra allocation.
+
+	// Read execute read requests and returns the read result. The `ReadContext`
+	// holds the read request to invoked. The implementation should call the
+	// `SetReadBytes` method of `ReadContext` to report the statistical changes
+	// involved in this execution before returning.
+	Read(ReadContext) ([]byte, error)
 	// GetInitialStates returns the most recent shard states of all shards known
 	// to the DataStorage instance that are consistent with their related table
 	// shards data. The shard metadata is last changed by the raft log identified
@@ -144,9 +147,9 @@ type ShardMetadata struct {
 	Metadata []byte
 }
 
-// Context contains the details of write and read requests that need to be
-// handled by the data storage.
-type Context interface {
+// WriteContext contains the details of write requests to be handled by the
+// data storage.
+type WriteContext interface {
 	// ByteBuf returns the bytebuf that can be used to avoid memory allocation.
 	ByteBuf() *buf.ByteBuf
 	// WriteBatch returns a write batch which will be used to hold a sequence of
@@ -157,9 +160,8 @@ type Context interface {
 	WriteBatch() Resetable
 	// Shard returns the current shard details.
 	Shard() meta.Shard
-	// Batches returns a list of Batch instance, each representing requests from a
-	// single Raft log.
-	Batches() []Batch
+	// Batch returns the Batch instance transformed from a single Raft log.
+	Batch() Batch
 	// AppendResponse is used for appending responses once each request is handled.
 	AppendResponse([]byte)
 	// SetWrittenBytes set the number of bytes written to storage for all requests
@@ -167,10 +169,6 @@ type Context interface {
 	// contributes to the scheduler's auto-rebalancing feature.
 	// This method must be called before `Read` or `Write` returns.
 	SetWrittenBytes(uint64)
-	// SetReadBytes set the number of bytes read from storage for all requests in
-	// the current context. This is an approximation value that contributes to the
-	// scheduler's auto-rebalancing feature.
-	SetReadBytes(uint64)
 	// SetDiffBytes set the diff of the bytes stored in storage after Write is
 	// executed. This is an approximation value used to modify the approximate
 	// amount of data in the `Shard` which is used for triggering the auto-split
@@ -178,6 +176,19 @@ type Context interface {
 	SetDiffBytes(int64)
 }
 
+type ReadContext interface {
+	// ByteBuf returns the bytebuf that can be used to avoid memory allocation.
+	ByteBuf() *buf.ByteBuf
+	// Shard returns the current shard details.
+	Shard() meta.Shard
+	// Requeset returns the read request to be processed on the storage engine.
+	Request() Request
+	// SetReadBytes set the number of bytes read from storage for all requests in
+	// the current context. This is an approximation value that contributes to the
+	// scheduler's auto-rebalancing feature.
+	SetReadBytes(uint64)
+}
+
 // Batch contains a list of requests. For write batches, all requests are from
 // the same raft log specified by the Index value. They must be atomically
 // applied into the data storage together with the Index value itself. For
@@ -200,41 +211,41 @@ type Request struct {
 	Cmd []byte
 }
 
-// SimpleContext is a simple Context implementation used for testing.
-type SimpleContext struct {
+// SimpleWriteContext is a simple WriteContext implementation used for testing.
+type SimpleWriteContext struct {
 	buf          *buf.ByteBuf
 	shard        meta.Shard
 	wb           Resetable
-	requests     []Batch
+	batch        Batch
 	responses    [][]byte
 	writtenBytes uint64
 	readBytes    uint64
 	diffBytes    int64
 }
 
-var _ Context = (*SimpleContext)(nil)
+var _ WriteContext = (*SimpleWriteContext)(nil)
 
-// NewSimpleContext returns a testing context.
-func NewSimpleContext(shard uint64,
-	base KVStorage, requests []Batch) *SimpleContext {
-	c := &SimpleContext{
-		wb:       base.NewWriteBatch(),
-		buf:      buf.NewByteBuf(32),
-		requests: requests,
+// NewSimpleWriteContext returns a testing context.
+func NewSimpleWriteContext(shardID uint64,
+	base KVStorage, batch Batch) *SimpleWriteContext {
+	c := &SimpleWriteContext{
+		wb:    base.NewWriteBatch(),
+		buf:   buf.NewByteBuf(32),
+		batch: batch,
 	}
-	c.shard.ID = shard
+	c.shard.ID = shardID
 	return c
 }
 
-func (ctx *SimpleContext) ByteBuf() *buf.ByteBuf        { return ctx.buf }
-func (ctx *SimpleContext) WriteBatch() Resetable        { return ctx.wb }
-func (ctx *SimpleContext) Shard() meta.Shard            { return ctx.shard }
-func (ctx *SimpleContext) Batches() []Batch             { return ctx.requests }
-func (ctx *SimpleContext) AppendResponse(value []byte)  { ctx.responses = append(ctx.responses, value) }
-func (ctx *SimpleContext) SetWrittenBytes(value uint64) { ctx.writtenBytes = value }
-func (ctx *SimpleContext) SetReadBytes(value uint64)    { ctx.readBytes = value }
-func (ctx *SimpleContext) SetDiffBytes(value int64)     { ctx.diffBytes = value }
-func (ctx *SimpleContext) GetWrittenBytes() uint64      { return ctx.writtenBytes }
-func (ctx *SimpleContext) GetReadBytes() uint64         { return ctx.readBytes }
-func (ctx *SimpleContext) GetDiffBytes() int64          { return ctx.diffBytes }
-func (ctx *SimpleContext) Responses() [][]byte          { return ctx.responses }
+func (ctx *SimpleWriteContext) ByteBuf() *buf.ByteBuf { return ctx.buf }
+func (ctx *SimpleWriteContext) WriteBatch() Resetable { return ctx.wb }
+func (ctx *SimpleWriteContext) Shard() meta.Shard     { return ctx.shard }
+func (ctx *SimpleWriteContext) Batch() Batch          { return ctx.batch }
+func (ctx *SimpleWriteContext) AppendResponse(value []byte) {
+	ctx.responses = append(ctx.responses, value)
+}
+func (ctx *SimpleWriteContext) SetWrittenBytes(value uint64) { ctx.writtenBytes = value }
+func (ctx *SimpleWriteContext) SetDiffBytes(value int64)     { ctx.diffBytes = value }
+func (ctx *SimpleWriteContext) GetWrittenBytes() uint64      { return ctx.writtenBytes }
+func (ctx *SimpleWriteContext) GetDiffBytes() int64          { return ctx.diffBytes }
+func (ctx *SimpleWriteContext) Responses() [][]byte          { return ctx.responses }


Test patch:
diff --git a/raftstore/execute_context_test.go b/raftstore/execute_context_test.go
index 800e526..b7750ac 100644
--- a/raftstore/execute_context_test.go
+++ b/raftstore/execute_context_test.go
@@ -23,7 +23,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestExecuteContextCanBeAppendedAndReset(t *testing.T) {
+func TestWriteContextCanBeSetAndReset(t *testing.T) {
 	cases := []struct {
 		batch rpc.RequestBatch
 	}{
@@ -39,22 +39,21 @@ func TestExecuteContextCanBeAppendedAndReset(t *testing.T) {
 	}
 
 	fs := vfs.GetTestFS()
-	ctx := newExecuteContext(mem.NewStorage(fs))
+	ctx := newWriteContext(mem.NewStorage(fs))
 	assert.False(t, ctx.hasRequest())
 	for i, c := range cases {
-		ctx.appendRequestBatch(c.batch)
+		ctx.setRequestBatch(c.batch)
 		assert.True(t, ctx.hasRequest())
-		assert.Equal(t, len(ctx.batches[i].requestBatch.Requests), len(c.batch.Requests), "index %d", i)
-		assert.Equal(t, ctx.batches[i].requestBatch, c.batch, "index %d", i)
-
-		for j, req := range ctx.batches[i].requestBatch.Requests {
-			assert.Equal(t, req, c.batch.Requests[j], "index %d, %d request", i, j)
+		assert.Equal(t, len(c.batch.Requests), len(ctx.batch.Requests), "index %d", i)
+		for idx := range c.batch.Requests {
+			assert.Equal(t, c.batch.Requests[idx].CustemType, ctx.batch.Requests[idx].CmdType)
+			assert.Equal(t, c.batch.Requests[idx].Key, ctx.batch.Requests[idx].Key)
+			assert.Equal(t, c.batch.Requests[idx].Cmd, ctx.batch.Requests[idx].Cmd)
 		}
 	}
 
 	ctx.reset(Shard{})
-	assert.Empty(t, ctx.requests)
-	assert.Empty(t, ctx.batches)
+	assert.Empty(t, ctx.batch)
 	assert.Empty(t, ctx.responses)
 	assert.Equal(t, Shard{}, ctx.shard)
 }
diff --git a/storage/kv/kv_data_storage_test.go b/storage/kv/kv_data_storage_test.go
index f77f3cf..eeabc11 100644
--- a/storage/kv/kv_data_storage_test.go
+++ b/storage/kv/kv_data_storage_test.go
@@ -158,7 +158,7 @@ func TestGetPersistentLogIndex(t *testing.T) {
 		{
 			sample:       2,
 			requests:     2,
-			appliedIndex: 1,
+			appliedIndex: 2,
 		},
 	}
 
@@ -172,12 +172,12 @@ func TestGetPersistentLogIndex(t *testing.T) {
 			defer s.Close()
 
 			s.(*kvDataStorage).mu.loaded = true
-			for i := 0; i < c.requests; i++ {
+			for i := 1; i <= c.requests; i++ {
 				var batch storage.Batch
 				batch.Index = uint64(i)
 				k := []byte(fmt.Sprintf("%d", i))
 				batch.Requests = append(batch.Requests, simple.NewWriteRequest(k, k))
-				ctx := storage.NewSimpleContext(0, base, []storage.Batch{batch})
+				ctx := storage.NewSimpleWriteContext(0, base, batch)
 				assert.NoError(t, s.Write(ctx), "index %d", i)
 			}
 
@@ -188,43 +188,11 @@ func TestGetPersistentLogIndex(t *testing.T) {
 			kv := s.(*kvDataStorage)
 			v, err := kv.base.Get(keys.GetAppliedIndexKey(0))
 			assert.NoError(t, err)
-			assert.Equal(t, uint64(c.requests)-1, buf.Byte2UInt64(v))
+			assert.Equal(t, uint64(c.requests), buf.Byte2UInt64(v))
 		}()
 	}
 }
 
-/*
-func TestMemFS(t *testing.T) {
-	memfs := vfs.NewMemFS()
-	defer vfs.ReportLeakedFD(memfs, t)
-	opts := &cpebble.Options{
-		FS: vfs.NewPebbleFS(memfs),
-	}
-	func() {
-		memfs.MkdirAll("/test-data", 0755)
-		dir, err := memfs.OpenDir("/")
-		assert.NoError(t, err)
-		dir.Sync()
-		base, err := pebble.NewStorage("test-data", opts)
-		assert.NoError(t, err)
-		defer func() {
-			memfs.(*pvfs.MemFS).SetIgnoreSyncs(true)
-			base.Close()
-		}()
-		base.Set([]byte("key-1"), []byte("value-1"), false)
-		base.Sync()
-		base.Set([]byte("key-2"), []byte("value-2"), false)
-	}()
-	memfs.(*pvfs.MemFS).ResetToSyncedState()
-	memfs.(*pvfs.MemFS).SetIgnoreSyncs(false)
-	base, err := pebble.NewStorage("test-data", opts)
-	assert.NoError(t, err)
-	defer base.Close()
-	v, err := base.Get([]byte("key-1"))
-	assert.NoError(t, err)
-	assert.True(t, len(v) > 0)
-}*/
-
 func TestKVDataStorageRestartWithNotSyncedDataLost(t *testing.T) {
 	for _, sample := range []uint64{10, 11} {
 		memfs := vfs.NewMemFS()
@@ -253,7 +221,7 @@ func TestKVDataStorageRestartWithNotSyncedDataLost(t *testing.T) {
 				batch.Index = index
 				k := []byte(fmt.Sprintf("%d", 100))
 				batch.Requests = append(batch.Requests, simple.NewWriteRequest(k, k))
-				ctx := storage.NewSimpleContext(shardID, base, []storage.Batch{batch})
+				ctx := storage.NewSimpleWriteContext(shardID, base, batch)
 				assert.NoError(t, s.Write(ctx))
 			}