1 files changed, 57 insertions, 16 deletions

diff --git a/internal/mapr/client/aggregate.go b/internal/mapr/client/aggregate.go
index 2e9b61a..9989e8f 100644
--- a/internal/mapr/client/aggregate.go
+++ b/internal/mapr/client/aggregate.go
@@ -12,30 +12,46 @@ import (
 
 // Aggregate mapreduce data on the DTail client side.
 type Aggregate struct {
-	// This is the mapr query specified on the command line.
-	query *mapr.Query
 	// This represents aggregated data of a single remote server.
 	group *mapr.GroupSet
-	// This represents the merged aggregated data of all servers.
-	globalGroup *mapr.GlobalGroupSet
+	// Shared per-client session state.
+	session *SessionState
+	// The currently tracked shared generation.
+	generation uint64
 	// The server we aggregate the data for (logging and debugging purposes only)
 	server string
 }
 
 // NewAggregate create new client aggregator.
-func NewAggregate(server string, query *mapr.Query,
-	globalGroup *mapr.GlobalGroupSet) *Aggregate {
+func NewAggregate(server string, session *SessionState) *Aggregate {
+	generation := uint64(0)
+	if session != nil {
+		generation = session.Snapshot().Generation
+	}
 
 	return &Aggregate{
-		query:       query,
-		group:       mapr.NewGroupSet(),
-		globalGroup: globalGroup,
-		server:      server,
+		group:      mapr.NewGroupSet(),
+		session:    session,
+		generation: generation,
+		server:     server,
 	}
 }
 
 // Aggregate data from mapr log line into local (and global) group sets.
 func (a *Aggregate) Aggregate(message string) error {
+	if a.session == nil {
+		return fmt.Errorf("missing client mapreduce session state")
+	}
+
+	snapshot := a.session.Snapshot()
+	if snapshot.Query == nil || snapshot.GlobalGroup == nil {
+		return fmt.Errorf("missing client mapreduce query state")
+	}
+	if snapshot.Generation != a.generation {
+		a.group.InitSet()
+		a.generation = snapshot.Generation
+	}
+
 	parts := strings.Split(message, protocol.AggregateDelimiter)
 	if len(parts) < 4 {
 		return fmt.Errorf("aggregate message without any real data")
@@ -51,7 +67,7 @@ func (a *Aggregate) Aggregate(message string) error {
 	set := a.group.GetSet(groupKey)
 	var addedSamples bool
 
-	for _, sc := range a.query.Select {
+	for _, sc := range snapshot.Query.Select {
 		if val, ok := fields[sc.FieldStorage]; ok {
 			if err := set.Aggregate(sc.FieldStorage, sc.Operation, val, true); err != nil {
 				dlog.Client.Error(err)
@@ -65,9 +81,9 @@ func (a *Aggregate) Aggregate(message string) error {
 	}
 
 	// Merge data from group into global group.
-	isMerged, err := a.globalGroup.MergeNoblock(a.query, a.group)
+	isMerged, err := snapshot.GlobalGroup.MergeNoblock(snapshot.Query, a.group)
 	if err != nil {
-		panic(err)
+		return fmt.Errorf("unable to merge aggregate data for server %s: %w", a.server, err)
 	}
 	if isMerged {
 		// Re-init local group (make it empty again).
@@ -76,15 +92,40 @@ func (a *Aggregate) Aggregate(message string) error {
 	return nil
 }
 
+// Flush merges any pending per-server aggregate state into the shared global group.
+// The normal hot path uses MergeNoblock to avoid stalling on the global merge lock.
+// During shutdown we need a blocking flush so the last local batch is not lost.
+func (a *Aggregate) Flush() error {
+	if a.session == nil {
+		return fmt.Errorf("missing client mapreduce session state")
+	}
+
+	snapshot := a.session.Snapshot()
+	if snapshot.Query == nil || snapshot.GlobalGroup == nil {
+		return nil
+	}
+	if snapshot.Generation != a.generation {
+		a.group.InitSet()
+		a.generation = snapshot.Generation
+		return nil
+	}
+
+	if err := snapshot.GlobalGroup.Merge(snapshot.Query, a.group); err != nil {
+		return fmt.Errorf("unable to flush aggregate data for server %s: %w", a.server, err)
+	}
+	a.group.InitSet()
+	return nil
+}
+
 // Create a map of key-value pairs from a part list such as ["foo=bar",  "bar=baz"].
 func (a *Aggregate) makeFields(parts []string) map[string]string {
 	fields := make(map[string]string, len(parts))
 	for _, part := range parts {
-		kv := strings.SplitN(part, protocol.AggregateKVDelimiter, 2)
-		if len(kv) != 2 {
+		key, value, ok := strings.Cut(part, protocol.AggregateKVDelimiter)
+		if !ok {
 			continue
 		}
-		fields[kv[0]] = kv[1]
+		fields[key] = value
 	}
 	return fields
 }