Simplify queuedMessage and handler implementation

Author: izumin5210

context_helper.go

@@ -2,9 +2,13 @@ package subee
 
 import (
 	"context"
+	"time"
 )
 
-type loggerContextKey struct{}
+type (
+	loggerContextKey     struct{}
+	enqueuedAtContextKey struct{}
+)
 
 // GetLogger return Logger implementation set in the context.
 func GetLogger(ctx context.Context) Logger {
@@ -14,3 +18,11 @@ func GetLogger(ctx context.Context) Logger {
 func setLogger(ctx context.Context, l Logger) context.Context {
 	return context.WithValue(ctx, loggerContextKey{}, l)
 }
+
+func getEnqueuedAt(ctx context.Context) time.Time {
+	return ctx.Value(enqueuedAtContextKey{}).(time.Time)
+}
+
+func setEnqueuedAt(ctx context.Context, t time.Time) context.Context {
+	return context.WithValue(ctx, enqueuedAtContextKey{}, t)
+}

process.go

@@ -42,14 +42,8 @@ func (p *processImpl) startConsumingProcess(ctx context.Context) error {
 	consumer := chainConsumerInterceptors(p.Consumer, p.ConsumerInterceptors...)
 
 	err := p.subscribe(ctx, func(in Message) {
-		msg := &singleMessage{
-			Ctx:        p.createConsumingContext(),
-			Msg:        in,
-			EnqueuedAt: time.Now(),
-		}
-		p.handleMessage(msg, func(in queuedMessage) error {
-			m := in.(*singleMessage)
-			return errors.WithStack(consumer.Consume(m.Ctx, m.Msg))
+		p.handleMessage(p.createConsumingContext(), &singleMessage{Message: in}, func(ctx context.Context) error {
+			return errors.WithStack(consumer.Consume(ctx, in))
 		})
 	})
 	return errors.WithStack(err)
@@ -73,9 +67,9 @@ func (p *processImpl) startBatchConsumingProcess(ctx context.Context) error {
 		defer p.Logger.Print("Finish batch consuming process")
 
 		for m := range outCh {
-			p.handleMessage(m, func(in queuedMessage) error {
-				m := in.(*multiMessages)
-				return errors.WithStack(batchConsumer.BatchConsume(m.Ctx, m.Msgs))
+			msgs := m.Msgs
+			p.handleMessage(m.Ctx, m, func(ctx context.Context) error {
+				return errors.WithStack(batchConsumer.BatchConsume(ctx, msgs))
 			})
 		}
 	}()
@@ -93,10 +87,11 @@ func (p *processImpl) createConsumingContext() context.Context {
 	ctx := context.Background()
 	ctx = p.StatsHandler.TagProcess(ctx, &BeginTag{})
 	ctx = p.StatsHandler.TagProcess(ctx, &EnqueueTag{})
+	ctx = setEnqueuedAt(ctx, time.Now().UTC())
 	return ctx
 }
 
-func (p *processImpl) handleMessage(m queuedMessage, handle func(queuedMessage) error) {
+func (p *processImpl) handleMessage(ctx context.Context, m queuedMessage, handle func(context.Context) error) {
 	p.wg.Add(1)
 	go func() {
 		defer p.wg.Done()
@@ -107,16 +102,18 @@ func (p *processImpl) handleMessage(m queuedMessage, handle func(queuedMessage)
 			m.Ack()
 		}
 
-		p.StatsHandler.HandleProcess(m.Context(), &Dequeue{
-			BeginTime: m.GetEnqueuedAt(),
+		enqueuedAt := getEnqueuedAt(ctx)
+
+		p.StatsHandler.HandleProcess(ctx, &Dequeue{
+			BeginTime: enqueuedAt,
 			EndTime:   time.Now(),
 		})
 
 		beginTime := time.Now()
 
-		m.SetContext(p.StatsHandler.TagProcess(m.Context(), &ConsumeBeginTag{}))
+		ctx = p.StatsHandler.TagProcess(ctx, &ConsumeBeginTag{})
 
-		err := handle(m)
+		err := handle(ctx)
 
 		if !p.AckImmediately {
 			if err != nil {
@@ -126,15 +123,15 @@ func (p *processImpl) handleMessage(m queuedMessage, handle func(queuedMessage)
 			}
 		}
 
-		p.StatsHandler.HandleProcess(m.Context(), &ConsumeEnd{
+		p.StatsHandler.HandleProcess(ctx, &ConsumeEnd{
 			BeginTime: beginTime,
 			EndTime:   time.Now(),
 			Error:     err,
 		})
 
-		p.StatsHandler.HandleProcess(m.Context(), &End{
+		p.StatsHandler.HandleProcess(ctx, &End{
 			MsgCount:  m.Count(),
-			BeginTime: m.GetEnqueuedAt(),
+			BeginTime: enqueuedAt,
 			EndTime:   time.Now(),
 		})
 	}()

queue.go

@@ -8,33 +8,17 @@ import (
 type queuedMessage interface {
 	Acknowledger
 	Count() int
-	Context() context.Context
-	SetContext(ctx context.Context)
-	GetEnqueuedAt() time.Time
 }
 
 type singleMessage struct {
-	Ctx        context.Context
-	Msg        Message
-	EnqueuedAt time.Time
+	Message
 }
 
-func (s *singleMessage) Ack() { s.Msg.Ack() }
-
-func (s *singleMessage) Nack() { s.Msg.Nack() }
-
-func (s *singleMessage) Context() context.Context { return s.Ctx }
-
-func (s *singleMessage) SetContext(ctx context.Context) { s.Ctx = ctx }
-
 func (s *singleMessage) Count() int { return 1 }
 
-func (s *singleMessage) GetEnqueuedAt() time.Time { return s.EnqueuedAt }
-
 type multiMessages struct {
-	Ctx        context.Context
-	Msgs       []Message
-	EnqueuedAt time.Time
+	Ctx  context.Context
+	Msgs []Message
 }
 
 func (m *multiMessages) Ack() {
@@ -51,22 +35,16 @@ func (m *multiMessages) Nack() {
 
 func (m *multiMessages) Count() int { return len(m.Msgs) }
 
-func (m *multiMessages) Context() context.Context { return m.Ctx }
-
-func (m *multiMessages) SetContext(ctx context.Context) { m.Ctx = ctx }
-
-func (m *multiMessages) GetEnqueuedAt() time.Time { return m.EnqueuedAt }
-
 func createBufferedQueue(
 	createCtx func() context.Context,
 	chunkSize int,
 	flushInterval time.Duration,
 ) (
 	chan<- Message,
-	<-chan queuedMessage,
+	<-chan *multiMessages,
 ) {
 	inCh := make(chan Message, chunkSize)
-	outCh := make(chan queuedMessage)
+	outCh := make(chan *multiMessages)
 
 	go func() {
 		defer close(outCh)
@@ -76,9 +54,8 @@ func createBufferedQueue(
 
 			if len(msgs) > 0 {
 				outCh <- &multiMessages{
-					Ctx:        createCtx(),
-					Msgs:       msgs,
-					EnqueuedAt: time.Now(),
+					Ctx:  createCtx(),
+					Msgs: msgs,
 				}
 			}
 

queue_test.go

@@ -41,10 +41,6 @@ func TestCreateBufferedQueue(t *testing.T) {
 		if got, want := out.Count(), n; got != want {
 			t.Errorf("Item[%d] has %d messages, want %d", i, got, want)
 		}
-
-		if m, ok := out.(*multiMessages); !ok {
-			t.Errorf("Item[%d] is %T type, want *subee.multiMessages type", i, m)
-		}
 	}
 
 	_, ok := <-outCh