From e65fa58ab56e3198e46e298090f396c36ec111d0 Mon Sep 17 00:00:00 2001
From: Alex Kaminskii <alex@openreplay.com>
Date: Wed, 11 May 2022 18:51:55 +0200
Subject: [PATCH] refactor(backend-internal): dry builder

---
 backend/internal/builder/builder.go    |  4 +--
 backend/internal/builder/builderMap.go | 45 +++++++++++++++-----------
 2 files changed, 29 insertions(+), 20 deletions(-)

diff --git a/backend/internal/builder/builder.go b/backend/internal/builder/builder.go
index 4916764cd..ff3d91e1b 100644
--- a/backend/internal/builder/builder.go
+++ b/backend/internal/builder/builder.go
@@ -18,7 +18,7 @@ func NewBuilder(handlers ...handlers.MessageProcessor) *builder {
 	}
 }
 
-func (b *builder) iterateReadyMessage(iter func(msg Message)) {
+func (b *builder) iterateReadyMessages(iter func(msg Message)) {
 	for _, readyMsg := range b.readyMsgs {
 		iter(readyMsg)
 	}
@@ -44,10 +44,10 @@ func (b *builder) handleMessage(message Message, messageID uint64) {
 		return
 	}
 
-	b.checkSessionEnd(message)
 	for _, p := range b.processors {
 		if rm := p.Handle(message, messageID, b.timestamp); rm != nil {
 			b.readyMsgs = append(b.readyMsgs, rm)
 		}
 	}
+	b.checkSessionEnd(message)
 }
diff --git a/backend/internal/builder/builderMap.go b/backend/internal/builder/builderMap.go
index b393bcd28..f6d81b995 100644
--- a/backend/internal/builder/builderMap.go
+++ b/backend/internal/builder/builderMap.go
@@ -20,7 +20,7 @@ func NewBuilderMap(handlers ...handlers.MessageProcessor) *builderMap {
 func (m *builderMap) GetBuilder(sessionID uint64) *builder {
 	b := m.sessions[sessionID]
 	if b == nil {
-		b = NewBuilder(m.handlers...)
+		b = NewBuilder(m.handlers...) // Should create new instances
 		m.sessions[sessionID] = b
 	}
 	return b
@@ -31,21 +31,29 @@ func (m *builderMap) HandleMessage(sessionID uint64, msg Message, messageID uint
 	b.handleMessage(msg, messageID)
 }
 
-func (m *builderMap) IterateReadyMessages(iter func(sessionID uint64, msg Message)) {
-	for sessionID, b := range m.sessions {
-		if b.ended {
-			for _, p := range b.processors {
-				if rm := p.Build(); rm != nil {
-					b.readyMsgs = append(b.readyMsgs, rm)
-				}
+func (m *builderMap) iterateSessionReadyMessages(sessionID uint64, b *builder, iter func(msg Message)) {
+	if b.ended {
+		for _, p := range b.processors {
+			if rm := p.Build(); rm != nil {
+				b.readyMsgs = append(b.readyMsgs, rm)
 			}
 		}
-		b.iterateReadyMessage(func(msg Message) {
-			iter(sessionID, msg)
-		})
-		if b.ended {
-			delete(m.sessions, sessionID)
-		}
+	}
+	b.iterateReadyMessage(iter)
+	if b.ended {
+		delete(m.sessions, sessionID)
+	}
+}
+
+func (m *builderMap) IterateReadyMessages(iter func(sessionID uint64, msg Message)) {
+	for sessionID, session := range m.sessions {
+		m.iterateSessionReadyMessages(
+			sessionID,
+			session,
+			func(msg Message) {
+				iter(sessionID, msg)
+			},
+		)
 	}
 }
 
@@ -54,8 +62,9 @@ func (m *builderMap) IterateSessionReadyMessages(sessionID uint64, iter func(msg
 	if !ok {
 		return
 	}
-	session.iterateReadyMessage(iter)
-	if session.ended {
-		delete(m.sessions, sessionID)
-	}
+	m.iterateSessionReadyMessages(
+		sessionID,
+		session,
+		inter,
+	)
 }