diff --git a/akita.code-workspace b/akita.code-workspace
new file mode 100644
index 00000000..abcb2ae7
--- /dev/null
+++ b/akita.code-workspace
@@ -0,0 +1,11 @@
+{
+	"folders": [
+		{
+			"path": "."
+		},
+		{
+			"path": "../../rewrite_mem"
+		}
+	],
+	"settings": {}
+}
\ No newline at end of file
diff --git a/mem/idealmemcontroller/111.txt b/mem/idealmemcontroller/111.txt
new file mode 100644
index 00000000..5af30be3
--- /dev/null
+++ b/mem/idealmemcontroller/111.txt
@@ -0,0 +1,14 @@
+package idealmemcontroller
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+//go:generate mockgen -destination "mock_sim_test.go" -package $GOPACKAGE -write_package_comment=false github.com/sarchlab/akita/v4/sim Port,Connection,Engine
+func TestIdealmemcontroller(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Idealmemcontroller Suite")
+}
diff --git a/mem/idealmemcontroller/builder.go b/mem/idealmemcontroller/builder.go
index a5adf841..d65343fa 100644
--- a/mem/idealmemcontroller/builder.go
+++ b/mem/idealmemcontroller/builder.go
@@ -90,6 +90,7 @@ func (b Builder) Build(
 	c := &Comp{
 		Latency: b.latency,
 		width:   b.width,
+		state:   "enable",
 	}
 
 	c.TickingComponent = sim.NewTickingComponent(name, b.engine, b.freq, c)
@@ -102,11 +103,15 @@ func (b Builder) Build(
 		c.Storage = b.storage
 	}
 
+	ctrlMiddleware := &ctrlMiddleware{Comp: c}
+	c.AddMiddleware(ctrlMiddleware)
+	funcMiddleware := &funcMiddleware{Comp: c}
+	c.AddMiddleware(funcMiddleware)
+
 	c.topPort = sim.NewLimitNumMsgPort(c, b.topBufSize, name+".TopPort")
 	c.AddPort("Top", c.topPort)
-
-	middleware := &middleware{Comp: c}
-	c.AddMiddleware(middleware)
+	c.ctrlPort = sim.NewLimitNumMsgPort(c, b.topBufSize, name+".CtrlPort")
+	c.AddPort("Control", c.ctrlPort)
 
 	return c
 }
diff --git a/mem/idealmemcontroller/comp.go b/mem/idealmemcontroller/comp.go
index 707895af..cfea0d47 100644
--- a/mem/idealmemcontroller/comp.go
+++ b/mem/idealmemcontroller/comp.go
@@ -1,14 +1,9 @@
 package idealmemcontroller
 
 import (
-	"log"
-	"reflect"
-
 	"github.com/sarchlab/akita/v4/mem/mem"
 
 	"github.com/sarchlab/akita/v4/sim"
-
-	"github.com/sarchlab/akita/v4/tracing"
 )
 
 type readRespondEvent struct {
@@ -41,159 +36,19 @@ type Comp struct {
 	sim.MiddlewareHolder
 
 	topPort          sim.Port
+	ctrlPort         sim.Port
 	Storage          *mem.Storage
 	Latency          int
 	addressConverter mem.AddressConverter
 
-	width int
-}
-
-func (c *Comp) Tick() bool {
-	return c.MiddlewareHolder.Tick()
-}
-
-// Handle defines how the Comp handles event
-func (c *Comp) Handle(e sim.Event) error {
-	switch e := e.(type) {
-	case *readRespondEvent:
-		return c.handleReadRespondEvent(e)
-	case *writeRespondEvent:
-		return c.handleWriteRespondEvent(e)
-	case sim.TickEvent:
-		return c.TickingComponent.Handle(e)
-	default:
-		log.Panicf("cannot handle event of %s", reflect.TypeOf(e))
-	}
-
-	return nil
-}
-
-type middleware struct {
-	*Comp
-}
-
-// Tick updates ideal memory controller state.
-func (m *middleware) Tick() bool {
-	msg := m.topPort.RetrieveIncoming()
-	if msg == nil {
-		return false
-	}
-
-	tracing.TraceReqReceive(msg, m.Comp)
-
-	switch msg := msg.(type) {
-	case *mem.ReadReq:
-		m.handleReadReq(msg)
-		return true
-	case *mem.WriteReq:
-		m.handleWriteReq(msg)
-		return true
-	default:
-		log.Panicf("cannot handle request of type %s", reflect.TypeOf(msg))
-	}
-
-	return false
-}
-
-func (m *middleware) handleReadReq(req *mem.ReadReq) {
-	now := m.CurrentTime()
-	timeToSchedule := m.Freq.NCyclesLater(m.Latency, now)
-	respondEvent := newReadRespondEvent(timeToSchedule, m.Comp, req)
-	m.Engine.Schedule(respondEvent)
-}
-
-func (m *middleware) handleWriteReq(req *mem.WriteReq) {
-	now := m.CurrentTime()
-	timeToSchedule := m.Freq.NCyclesLater(m.Latency, now)
-	respondEvent := newWriteRespondEvent(timeToSchedule, m.Comp, req)
-	m.Engine.Schedule(respondEvent)
-}
-
-func (c *Comp) handleReadRespondEvent(e *readRespondEvent) error {
-	now := e.Time()
-	req := e.req
-
-	addr := req.Address
-	if c.addressConverter != nil {
-		addr = c.addressConverter.ConvertExternalToInternal(addr)
-	}
-
-	data, err := c.Storage.Read(addr, req.AccessByteSize)
-	if err != nil {
-		log.Panic(err)
-	}
+	respondReq *mem.ControlMsg
+	width      int
 
-	rsp := mem.DataReadyRspBuilder{}.
-		WithSrc(c.topPort).
-		WithDst(req.Src).
-		WithRspTo(req.ID).
-		WithData(data).
-		Build()
+	state string
 
-	networkErr := c.topPort.Send(rsp)
-
-	if networkErr != nil {
-		retry := newReadRespondEvent(c.Freq.NextTick(now), c, req)
-		c.Engine.Schedule(retry)
-		return nil
-	}
-
-	tracing.TraceReqComplete(req, c)
-	c.TickLater()
-
-	return nil
+	inflightbuffer []sim.Msg
 }
 
-func (c *Comp) handleWriteRespondEvent(e *writeRespondEvent) error {
-	now := e.Time()
-	req := e.req
-
-	rsp := mem.WriteDoneRspBuilder{}.
-		WithSrc(c.topPort).
-		WithDst(req.Src).
-		WithRspTo(req.ID).
-		Build()
-
-	networkErr := c.topPort.Send(rsp)
-	if networkErr != nil {
-		retry := newWriteRespondEvent(c.Freq.NextTick(now), c, req)
-		c.Engine.Schedule(retry)
-		return nil
-	}
-
-	addr := req.Address
-
-	if c.addressConverter != nil {
-		addr = c.addressConverter.ConvertExternalToInternal(addr)
-	}
-
-	if req.DirtyMask == nil {
-		err := c.Storage.Write(addr, req.Data)
-		if err != nil {
-			log.Panic(err)
-		}
-	} else {
-		data, err := c.Storage.Read(addr, uint64(len(req.Data)))
-		if err != nil {
-			panic(err)
-		}
-		for i := 0; i < len(req.Data); i++ {
-			if req.DirtyMask[i] {
-				data[i] = req.Data[i]
-			}
-		}
-		err = c.Storage.Write(addr, data)
-		if err != nil {
-			panic(err)
-		}
-	}
-
-	tracing.TraceReqComplete(req, c)
-	c.TickLater()
-
-	return nil
-}
-
-func (c *Comp) CurrentTime() sim.VTimeInSec {
-	return c.Engine.CurrentTime()
+func (c *Comp) Tick() bool {
+	return c.MiddlewareHolder.Tick()
 }
diff --git a/mem/idealmemcontroller/ctrlMiddleware_internal_test.go b/mem/idealmemcontroller/ctrlMiddleware_internal_test.go
new file mode 100644
index 00000000..14bcdcd2
--- /dev/null
+++ b/mem/idealmemcontroller/ctrlMiddleware_internal_test.go
@@ -0,0 +1,119 @@
+package idealmemcontroller
+
+import (
+	"github.com/golang/mock/gomock"
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+	"github.com/sarchlab/akita/v4/mem/mem"
+	"github.com/sarchlab/akita/v4/sim"
+)
+
+var _ = FDescribe("CtrlMiddleware", func() {
+	var (
+		mockCtrl       *gomock.Controller
+		comp           *Comp
+		engine         *MockEngine
+		remoteCtrlPort *MockPort
+		ctrlPort       *MockPort
+		ctrlMW         *ctrlMiddleware
+	)
+
+	BeforeEach(func() {
+		mockCtrl = gomock.NewController(GinkgoT())
+
+		engine = NewMockEngine(mockCtrl)
+		ctrlPort = NewMockPort(mockCtrl)
+		remoteCtrlPort = NewMockPort(mockCtrl)
+
+		comp = MakeBuilder().
+			WithEngine(engine).
+			WithNewStorage(1 * mem.MB).
+			Build("MemCtrl")
+		comp.Freq = 1000 * sim.MHz
+		comp.Latency = 10
+		comp.ctrlPort = ctrlPort
+
+		ctrlMW = &ctrlMiddleware{
+			Comp: comp,
+		}
+	})
+
+	AfterEach(func() {
+		mockCtrl.Finish()
+	})
+
+	It("should do nothing if no ctrl message", func() {
+		ctrlPort.EXPECT().RetrieveIncoming().Return(nil)
+
+		madeProgress := ctrlMW.Tick()
+
+		Expect(madeProgress).To(BeFalse())
+	})
+
+	It("should handle enable message", func() {
+		comp.state = "paused"
+
+		ctrlMsg := mem.ControlMsgBuilder{}.
+			WithSrc(remoteCtrlPort).
+			WithDst(ctrlPort).
+			WithCtrlInfo(true, false, false, false).
+			Build()
+		ctrlPort.EXPECT().RetrieveIncoming().Return(ctrlMsg)
+		ctrlPort.EXPECT().
+			Send(gomock.Any()).
+			Do(func(msg *sim.GeneralRsp) {
+				Expect(msg.Src).To(Equal(ctrlPort))
+				Expect(msg.Dst).To(Equal(remoteCtrlPort))
+				Expect(msg.OriginalReq).To(Equal(ctrlMsg))
+			}).
+			Return(nil).
+			AnyTimes()
+
+		madeProgress := ctrlMW.Tick()
+
+		Expect(madeProgress).To(BeTrue())
+		Expect(comp.state).To(Equal("enable"))
+	})
+
+	It("should handle pause message", func() {
+		comp.state = "enable"
+
+		ctrlMsg := mem.ControlMsgBuilder{}.
+			WithSrc(remoteCtrlPort).
+			WithDst(ctrlPort).
+			WithCtrlInfo(false, false, false, false).
+			Build()
+		ctrlPort.EXPECT().RetrieveIncoming().Return(ctrlMsg)
+		ctrlPort.EXPECT().
+			Send(gomock.Any()).
+			Do(func(msg *sim.GeneralRsp) {
+				Expect(msg.Src).To(Equal(ctrlPort))
+				Expect(msg.Dst).To(Equal(remoteCtrlPort))
+				Expect(msg.OriginalReq).To(Equal(ctrlMsg))
+			}).
+			Return(nil).
+			AnyTimes()
+
+		madeProgress := ctrlMW.Tick()
+
+		Expect(madeProgress).To(BeTrue())
+		Expect(comp.state).To(Equal("pause"))
+	})
+
+	It("should handle drain message", func() {
+		comp.state = "enable"
+
+		ctrlMsg := mem.ControlMsgBuilder{}.
+			WithSrc(remoteCtrlPort).
+			WithDst(ctrlPort).
+			WithCtrlInfo(false, true, false, false).
+			Build()
+		ctrlPort.EXPECT().RetrieveIncoming().Return(ctrlMsg)
+		madeProgress := ctrlMW.Tick()
+
+		Expect(madeProgress).To(BeTrue())
+		Expect(comp.state).To(Equal("drain"))
+		Expect(comp.respondReq).To(Equal(ctrlMsg))
+	})
+
+})
diff --git a/mem/idealmemcontroller/ctrlmiddleware.go b/mem/idealmemcontroller/ctrlmiddleware.go
new file mode 100644
index 00000000..1f2b7f84
--- /dev/null
+++ b/mem/idealmemcontroller/ctrlmiddleware.go
@@ -0,0 +1,83 @@
+package idealmemcontroller
+
+import "github.com/sarchlab/akita/v4/mem/mem"
+
+type ctrlMiddleware struct {
+	*Comp
+}
+
+func (m *ctrlMiddleware) Tick() (madeProgress bool) {
+	msg := m.ctrlPort.RetrieveIncoming()
+	if msg == nil {
+		return false
+	}
+
+	ctrlMsg := msg.(*mem.ControlMsg)
+
+	m.ctrlMsgMustBeValid(ctrlMsg)
+
+	madeProgress = m.handleEnable(ctrlMsg) || madeProgress
+	madeProgress = m.handlePause(ctrlMsg) || madeProgress
+	madeProgress = m.handleDrain(ctrlMsg) || madeProgress
+
+	return madeProgress
+}
+
+func (m *ctrlMiddleware) handleEnable(ctrlMsg *mem.ControlMsg) bool {
+	if ctrlMsg.Enable {
+		m.state = "enable"
+		rsp := ctrlMsg.GenerateRsp()
+		m.ctrlPort.Send(rsp)
+		return true
+	}
+	return false
+}
+
+func (m *ctrlMiddleware) handlePause(ctrlMsg *mem.ControlMsg) bool {
+	if !ctrlMsg.Enable && !ctrlMsg.Drain {
+		m.state = "pause"
+		rsp := ctrlMsg.GenerateRsp()
+		m.ctrlPort.Send(rsp)
+		return true
+	}
+
+	return false
+}
+
+func (m *ctrlMiddleware) handleDrain(ctrlMsg *mem.ControlMsg) bool {
+	if !ctrlMsg.Enable && ctrlMsg.Drain {
+		m.state = "drain"
+		m.respondReq = ctrlMsg
+		return true
+	}
+
+	return false
+}
+
+func (m *ctrlMiddleware) ctrlMsgMustBeValid(ctrlMsg *mem.ControlMsg) {
+	if ctrlMsg.Enable {
+		if ctrlMsg.Drain {
+			panic("Enable and Drain should not be set at the same time")
+		}
+
+		if ctrlMsg.Invalid {
+			panic("Enable and Invalid should not be set at the same time")
+		}
+
+		if ctrlMsg.Flush {
+			panic("Enable and Flush should not be set at the same time")
+		}
+	}
+
+	if !ctrlMsg.Enable {
+		if ctrlMsg.Drain {
+			if ctrlMsg.Invalid {
+				panic("Drain and Invalid should not be set at the same time")
+			}
+
+			if ctrlMsg.Flush {
+				panic("Drain and Flush should not be set at the same time")
+			}
+		}
+	}
+}
diff --git a/mem/idealmemcontroller/funcMiddleware.go b/mem/idealmemcontroller/funcMiddleware.go
new file mode 100644
index 00000000..f89d8e59
--- /dev/null
+++ b/mem/idealmemcontroller/funcMiddleware.go
@@ -0,0 +1,135 @@
+package idealmemcontroller
+
+import (
+	"log"
+	"reflect"
+
+	"github.com/sarchlab/akita/v4/mem/mem"
+	"github.com/sarchlab/akita/v4/sim"
+	"github.com/sarchlab/akita/v4/tracing"
+)
+
+type funcMiddleware struct {
+	*Comp
+}
+
+func (m *funcMiddleware) Tick() bool {
+	madeProgress := false
+
+	madeProgress = m.tickInflightBuffer() || madeProgress
+	madeProgress = m.checkAndExecuteState() || madeProgress
+
+	return madeProgress
+}
+
+func (m *funcMiddleware) checkAndExecuteState() bool {
+	madeProgress := false
+
+	switch state := m.state; state {
+	case "enable":
+		madeProgress = m.handleInflightMemReqs()
+	case "pause":
+		madeProgress = true
+	case "drain":
+		madeProgress = m.handleDrainReq()
+	}
+
+	return madeProgress
+}
+
+func (m *funcMiddleware) handleDrainReq() bool {
+	madeProgress := false
+	for len(m.inflightbuffer) != 0 {
+		madeProgress = m.handleMemReqs()
+	}
+	if !m.setState("pause", m.respondReq) {
+		return true
+	}
+	return madeProgress
+}
+
+func (m *funcMiddleware) tickInflightBuffer() bool {
+	if m.state == "pause" {
+		return false
+	}
+
+	for i := 0; i < m.width; i++ {
+		msg := m.topPort.RetrieveIncoming()
+		if msg == nil {
+			return false
+		}
+
+		m.inflightbuffer = append(m.inflightbuffer, msg)
+	}
+	return true
+}
+
+// updateMemCtrl updates ideal memory controller state.
+func (m *funcMiddleware) handleInflightMemReqs() bool {
+	madeProgress := false
+	for i := 0; i < m.width; i++ {
+		madeProgress = m.handleMemReqs()
+	}
+
+	return madeProgress
+}
+
+func (m *funcMiddleware) handleMemReqs() bool {
+	if len(m.inflightbuffer) == 0 {
+		return false
+	}
+
+	msg := m.inflightbuffer[0]
+	m.inflightbuffer = m.inflightbuffer[1:]
+
+	tracing.TraceReqReceive(msg, m)
+
+	switch msg := msg.(type) {
+	case *mem.ReadReq:
+		m.handleReadReq(msg)
+		return true
+	case *mem.WriteReq:
+		m.handleWriteReq(msg)
+		return true
+	default:
+		log.Panicf("cannot handle request of type %s", reflect.TypeOf(msg))
+	}
+	return false
+}
+
+func (m *funcMiddleware) handleReadReq(req *mem.ReadReq) {
+	now := m.CurrentTime()
+	timeToSchedule := m.Freq.NCyclesLater(m.Latency, now)
+	respondEvent := newReadRespondEvent(timeToSchedule, m, req)
+	m.Engine.Schedule(respondEvent)
+}
+
+func (m *funcMiddleware) handleWriteReq(req *mem.WriteReq) {
+	now := m.CurrentTime()
+	timeToSchedule := m.Freq.NCyclesLater(m.Latency, now)
+	respondEvent := newWriteRespondEvent(timeToSchedule, m, req)
+	m.Engine.Schedule(respondEvent)
+}
+
+func (m *funcMiddleware) CurrentTime() sim.VTimeInSec {
+	return m.Engine.CurrentTime()
+}
+
+func (m *funcMiddleware) setState(state string, rspMessage *mem.ControlMsg) bool {
+	ctrlRsp := sim.GeneralRspBuilder{}.
+		WithSrc(m.ctrlPort).
+		WithDst(rspMessage.Src).
+		WithOriginalReq(rspMessage).
+		Build()
+
+	err := m.ctrlPort.Send(ctrlRsp)
+
+	if err != nil {
+		return false
+	}
+
+	m.ctrlPort.RetrieveIncoming()
+	m.state = state
+
+	return true
+}
diff --git a/mem/idealmemcontroller/funcMiddleware_internal_test.go b/mem/idealmemcontroller/funcMiddleware_internal_test.go
new file mode 100644
index 00000000..9cf67c9f
--- /dev/null
+++ b/mem/idealmemcontroller/funcMiddleware_internal_test.go
@@ -0,0 +1,88 @@
+package idealmemcontroller
+
+import (
+	"github.com/golang/mock/gomock"
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+	"github.com/sarchlab/akita/v4/mem/mem"
+	"github.com/sarchlab/akita/v4/sim"
+)
+
+var _ = Describe("FuncMiddleware", func() {
+	var (
+		mockCtrl *gomock.Controller
+		comp     *Comp
+		engine   *MockEngine
+		ctrlPort *MockPort
+		topPort  *MockPort
+		funcMW   *funcMiddleware
+	)
+
+	BeforeEach(func() {
+		mockCtrl = gomock.NewController(GinkgoT())
+
+		engine = NewMockEngine(mockCtrl)
+		topPort = NewMockPort(mockCtrl)
+
+		comp = MakeBuilder().
+			WithEngine(engine).
+			WithNewStorage(1 * mem.MB).
+			WithWidth(1).
+			Build("MemCtrl")
+		comp.Freq = 1000 * sim.MHz
+		comp.Latency = 10
+		comp.topPort = topPort
+		comp.ctrlPort = ctrlPort
+
+		funcMW = &funcMiddleware{
+			Comp: comp,
+		}
+	})
+
+	AfterEach(func() {
+		mockCtrl.Finish()
+	})
+
+	FIt("should do nothing if no memory read request or memory store request",
+		func() {
+			topPort.EXPECT().RetrieveIncoming().Return(nil)
+
+			madeProgress := funcMW.Tick()
+
+			Expect(madeProgress).To(BeFalse())
+		})
+
+	FIt("should handle memory read request", func() {
+		readReq := mem.ReadReqBuilder{}.
+			WithDst(funcMW.topPort).
+			WithAddress(0).
+			WithByteSize(4).
+			Build()
+		topPort.EXPECT().RetrieveIncoming().Return(readReq)
+		engine.EXPECT().CurrentTime().Return(sim.VTimeInSec(10))
+
+		engine.EXPECT().
+			Schedule(gomock.AssignableToTypeOf(&readRespondEvent{}))
+
+		madeProgress := funcMW.Tick()
+
+		Expect(madeProgress).To(BeTrue())
+	})
+
+	FIt("should process write request", func() {
+		writeReq := mem.WriteReqBuilder{}.
+			WithDst(funcMW.topPort).
+			WithAddress(0).
+			WithData([]byte{0, 1, 2, 3}).
+			WithDirtyMask([]bool{false, false, true, false}).
+			Build()
+		topPort.EXPECT().RetrieveIncoming().Return(writeReq)
+		engine.EXPECT().CurrentTime().Return(sim.VTimeInSec(10))
+
+		engine.EXPECT().
+			Schedule(gomock.AssignableToTypeOf(&writeRespondEvent{}))
+
+		madeProgress := funcMW.Tick()
+		Expect(madeProgress).To(BeTrue())
+	})
+})
diff --git a/mem/idealmemcontroller/idealmemcontroller_suite_test.go b/mem/idealmemcontroller/idealmemcontroller_suite_test.go
index 5af30be3..3e1c2022 100644
--- a/mem/idealmemcontroller/idealmemcontroller_suite_test.go
+++ b/mem/idealmemcontroller/idealmemcontroller_suite_test.go
@@ -1,4 +1,4 @@
-package idealmemcontroller
+package idealmemcontroller_test
 
 import (
 	"testing"
@@ -7,7 +7,6 @@ import (
 	. "github.com/onsi/gomega"
 )
 
-//go:generate mockgen -destination "mock_sim_test.go" -package $GOPACKAGE -write_package_comment=false github.com/sarchlab/akita/v4/sim Port,Connection,Engine
 func TestIdealmemcontroller(t *testing.T) {
 	RegisterFailHandler(Fail)
 	RunSpecs(t, "Idealmemcontroller Suite")
diff --git a/mem/idealmemcontroller/idealmemcontroller_test.go b/mem/idealmemcontroller/idealmemcontroller_test.go
index c92f57a6..fc9031e0 100644
--- a/mem/idealmemcontroller/idealmemcontroller_test.go
+++ b/mem/idealmemcontroller/idealmemcontroller_test.go
@@ -17,6 +17,7 @@ var _ = Describe("Ideal Memory Controller", func() {
 		engine        *MockEngine
 		memController *Comp
 		port          *MockPort
+		ctrlPort      *MockPort
 	)
 
 	BeforeEach(func() {
@@ -24,6 +25,7 @@ var _ = Describe("Ideal Memory Controller", func() {
 
 		engine = NewMockEngine(mockCtrl)
 		port = NewMockPort(mockCtrl)
+		ctrlPort = NewMockPort(mockCtrl)
 
 		memController = MakeBuilder().
 			WithEngine(engine).
@@ -32,6 +34,7 @@ var _ = Describe("Ideal Memory Controller", func() {
 		memController.Freq = 1000 * sim.MHz
 		memController.Latency = 10
 		memController.topPort = port
+		memController.ctrlPort = ctrlPort
 	})
 
 	AfterEach(func() {
@@ -45,6 +48,7 @@ var _ = Describe("Ideal Memory Controller", func() {
 			WithByteSize(4).
 			Build()
 		port.EXPECT().RetrieveIncoming().Return(readReq)
+		ctrlPort.EXPECT().PeekIncoming().Return(nil)
 		engine.EXPECT().CurrentTime().Return(sim.VTimeInSec(10))
 
 		engine.EXPECT().
@@ -63,6 +67,7 @@ var _ = Describe("Ideal Memory Controller", func() {
 			WithDirtyMask([]bool{false, false, true, false}).
 			Build()
 		port.EXPECT().RetrieveIncoming().Return(writeReq)
+		ctrlPort.EXPECT().PeekIncoming().Return(nil)
 		engine.EXPECT().CurrentTime().Return(sim.VTimeInSec(10))
 
 		engine.EXPECT().
diff --git a/mem/mem/protocol.go b/mem/mem/protocol.go
index 6aea00f9..30376523 100644
--- a/mem/mem/protocol.go
+++ b/mem/mem/protocol.go
@@ -399,17 +399,21 @@ func (b WriteDoneRspBuilder) Build() *WriteDoneRsp {
 }
 
 // ControlMsg is the commonly used message type for controlling the components
-// on the memory hierarchy. It is also used for resonpding the original
-// requester with the Done field.
+// on the memory hierarchy. It is also used for responding the original
+// requester with the Done field. Drain is used to process all the requests in
+// the queue when drain happens the component will not accept any new requests
+// Enable enables the component work. If the enable = false, it will not process
+// any requests.
 type ControlMsg struct {
 	sim.MsgMeta
 
-	DiscardTransations bool
-	Restart            bool
-	NotifyDone         bool
+	Enable  bool
+	Drain   bool
+	Flush   bool
+	Invalid bool
 }
 
-// Meta returns the meta data assocated with the ControlMsg.
+// Meta returns the meta data associated with the ControlMsg.
 func (m *ControlMsg) Meta() *sim.MsgMeta {
 	return &m.MsgMeta
 }
@@ -422,12 +426,24 @@ func (m *ControlMsg) Clone() sim.Msg {
 	return &cloneMsg
 }
 
+// GenerateRsp generate GeneralRsp to ControlMsg
+func (m *ControlMsg) GenerateRsp() sim.Rsp {
+	rsp := sim.GeneralRspBuilder{}.
+		WithSrc(m.Dst).
+		WithDst(m.Src).
+		WithOriginalReq(m).
+		Build()
+
+	return rsp
+}
+
 // A ControlMsgBuilder can build control messages.
 type ControlMsgBuilder struct {
-	src, dst            sim.Port
-	discardTransactions bool
-	restart             bool
-	notifyDone          bool
+	src, dst sim.Port
+	Enable   bool
+	Drain    bool
+	Flush    bool
+	Invalid  bool
 }
 
 // WithSrc sets the source of the request to build.
@@ -442,22 +458,16 @@ func (b ControlMsgBuilder) WithDst(dst sim.Port) ControlMsgBuilder {
 	return b
 }
 
-// ToDiscardTransactions sets the discard transactions bit of the control
-// messages to 1.
-func (b ControlMsgBuilder) ToDiscardTransactions() ControlMsgBuilder {
-	b.discardTransactions = true
-	return b
-}
-
-// ToRestart sets the restart bit of the control messages to 1.
-func (b ControlMsgBuilder) ToRestart() ControlMsgBuilder {
-	b.restart = true
-	return b
-}
-
-// ToNotifyDone sets the "notify done" bit of the control messages to 1.
-func (b ControlMsgBuilder) ToNotifyDone() ControlMsgBuilder {
-	b.notifyDone = true
+func (b ControlMsgBuilder) WithCtrlInfo(
+	enableFlag bool,
+	drainFlag bool,
+	flushFlag bool,
+	invalidFlag bool,
+) ControlMsgBuilder {
+	b.Enable = enableFlag
+	b.Drain = drainFlag
+	b.Flush = flushFlag
+	b.Invalid = invalidFlag
 	return b
 }
 
@@ -469,9 +479,10 @@ func (b ControlMsgBuilder) Build() *ControlMsg {
 	m.Dst = b.dst
 	m.TrafficBytes = controlMsgByteOverhead
 
-	m.DiscardTransations = b.discardTransactions
-	m.Restart = b.restart
-	m.NotifyDone = b.notifyDone
+	m.Enable = b.Enable
+	m.Drain = b.Drain
+	m.Flush = b.Flush
+	m.Invalid = b.Invalid
 
 	return m
 }