fix(edge): delete group stats on emit not collect

isabella232 · Apr 15, 2019 · 377dde3 · 377dde3
1 parent 919c3f3
commit 377dde3
Showing 1 changed file with 9 additions and 10 deletions.
diff --git a/edge/stats.go b/edge/stats.go
@@ -167,11 +167,8 @@ func (e *batchStatsEdge) Collect(m Message) error {
 		e.collected.Add(1)
 		begin := b.Begin()
 		e.incCollected(begin.GroupID(), begin.GroupInfo, int64(len(b.Points())))
-	case DeleteGroupMessage:
-		e.deleteGroup(b.GroupID())
 	default:
 		// Do not count other messages
-		// TODO(nathanielc): How should we count other messages?
 	}
 	return nil
 }
@@ -197,9 +194,10 @@ func (e *batchStatsEdge) Emit() (m Message, ok bool) {
 			e.emitted.Add(1)
 			begin := b.Begin()
 			e.incEmitted(begin.GroupID(), begin.GroupInfo, int64(len(b.Points())))
+		case DeleteGroupMessage:
+			e.deleteGroup(b.GroupID())
 		default:
 			// Do not count other messages
-			// TODO(nathanielc): How should we count other messages?
 		}
 	}
 	return
@@ -221,21 +219,22 @@ func (e *streamStatsEdge) Collect(m Message) error {
 	case PointMessage:
 		e.collected.Add(1)
 		e.incCollected(m.GroupID(), m.GroupInfo, 1)
-	case DeleteGroupMessage:
-		e.deleteGroup(m.GroupID())
 	default:
 		// Do not count other messages
-		// TODO(nathanielc): How should we count other messages?
 	}
 	return nil
 }
 
 func (e *streamStatsEdge) Emit() (m Message, ok bool) {
 	m, ok = e.edge.Emit()
-	if ok && m.Type() == Point {
+	switch m := m.(type) {
+	case PointMessage:
 		e.emitted.Add(1)
-		p := m.(GroupInfoer)
-		e.incEmitted(p.GroupID(), p.GroupInfo, 1)
+		e.incEmitted(m.GroupID(), m.GroupInfo, 1)
+	case DeleteGroupMessage:
+		e.deleteGroup(m.GroupID())
+	default:
+		// Do not count other messages
 	}
 	return
 }