task 4abe7505: reset dmap generation state

4 files changed, 274 insertions, 12 deletions

diff --git a/internal/mapr/client/aggregate.go b/internal/mapr/client/aggregate.go
index 2750643..8cbd339 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,7 +81,7 @@ 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 {
 		return fmt.Errorf("unable to merge aggregate data for server %s: %w", a.server, err)
 	}
diff --git a/internal/mapr/client/aggregate_test.go b/internal/mapr/client/aggregate_test.go
new file mode 100644
index 0000000..8ac94a1
--- /dev/null
+++ b/internal/mapr/client/aggregate_test.go
@@ -0,0 +1,70 @@
+package client
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/mimecast/dtail/internal/mapr"
+	"github.com/mimecast/dtail/internal/protocol"
+)
+
+func TestAggregateResetsPendingLocalStateOnGenerationChange(t *testing.T) {
+	query := mustSessionStateQuery(t, "select status,count(status) from stats group by status")
+	state := NewSessionState(query)
+	aggregate := NewAggregate("srv1", state)
+	countStorage := aggregateCountStorage(t, query)
+
+	oldSet := aggregate.group.GetSet("ERROR")
+	oldSet.Samples = 1
+	oldSet.FValues[countStorage] = 1
+
+	rawQuery := "select status,count(status) from warnings group by status"
+	if _, err := state.CommitQuery(rawQuery, 2); err != nil {
+		t.Fatalf("CommitQuery() error = %v", err)
+	}
+
+	snapshot := state.Snapshot()
+	message := strings.Join([]string{
+		"WARN",
+		"1",
+		aggregateCountStorage(t, snapshot.Query) + protocol.AggregateKVDelimiter + "1",
+		"",
+	}, protocol.AggregateDelimiter)
+
+	if err := aggregate.Aggregate(message); err != nil {
+		t.Fatalf("Aggregate() error = %v", err)
+	}
+
+	result, numRows, err := snapshot.GlobalGroup.Result(snapshot.Query, 10, nil)
+	if err != nil {
+		t.Fatalf("Result() error = %v", err)
+	}
+	if numRows != 1 {
+		t.Fatalf("numRows = %d, want 1", numRows)
+	}
+	if !strings.Contains(result, "1") {
+		t.Fatalf("expected one new-generation aggregate row, got %q", result)
+	}
+}
+
+func TestAggregateRejectsMalformedMessage(t *testing.T) {
+	query := mustSessionStateQuery(t, "select count(status) from stats group by status")
+	state := NewSessionState(query)
+	aggregate := NewAggregate("srv1", state)
+
+	if err := aggregate.Aggregate("broken"); err == nil {
+		t.Fatalf("expected Aggregate() to reject malformed messages")
+	}
+}
+
+func aggregateCountStorage(t *testing.T, query *mapr.Query) string {
+	t.Helper()
+
+	for _, selectCondition := range query.Select {
+		if selectCondition.Operation == mapr.Count {
+			return selectCondition.FieldStorage
+		}
+	}
+	t.Fatalf("query %q does not contain count() storage", query.RawQuery)
+	return ""
+}
diff --git a/internal/mapr/client/session_state.go b/internal/mapr/client/session_state.go
new file mode 100644
index 0000000..1983644
--- /dev/null
+++ b/internal/mapr/client/session_state.go
@@ -0,0 +1,95 @@
+package client
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/mimecast/dtail/internal/mapr"
+)
+
+// SessionSnapshot captures the current client-side mapreduce session state.
+type SessionSnapshot struct {
+	Generation  uint64
+	Query       *mapr.Query
+	GlobalGroup *mapr.GlobalGroupSet
+	LastResult  string
+}
+
+// SessionState keeps the mutable mapreduce query state shared by the client
+// reporter and per-server handlers.
+type SessionState struct {
+	mu         sync.RWMutex
+	generation uint64
+	query      *mapr.Query
+	global     *mapr.GlobalGroupSet
+	lastResult string
+	changedCh  chan struct{}
+}
+
+// NewSessionState returns a new shared mapreduce session state.
+func NewSessionState(query *mapr.Query) *SessionState {
+	return &SessionState{
+		query:     query,
+		global:    mapr.NewGlobalGroupSet(),
+		changedCh: make(chan struct{}, 1),
+	}
+}
+
+// Snapshot returns a point-in-time copy of the shared mapreduce state.
+func (s *SessionState) Snapshot() SessionSnapshot {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	return SessionSnapshot{
+		Generation:  s.generation,
+		Query:       s.query,
+		GlobalGroup: s.global,
+		LastResult:  s.lastResult,
+	}
+}
+
+// Changes returns a channel that is signaled whenever a new generation is committed.
+func (s *SessionState) Changes() <-chan struct{} {
+	return s.changedCh
+}
+
+// CommitQuery resets the shared aggregation state for a newly accepted query generation.
+func (s *SessionState) CommitQuery(rawQuery string, generation uint64) (*mapr.Query, error) {
+	query, err := mapr.NewQuery(rawQuery)
+	if err != nil {
+		return nil, fmt.Errorf("parse session query: %w", err)
+	}
+
+	s.mu.Lock()
+	s.generation = generation
+	s.query = query
+	s.global = mapr.NewGlobalGroupSet()
+	s.lastResult = ""
+	s.mu.Unlock()
+
+	s.notifyChange()
+	return query, nil
+}
+
+// CommitRenderedResult stores the last rendered result for the active generation.
+func (s *SessionState) CommitRenderedResult(generation uint64, result string) (changed bool, ok bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if s.generation != generation {
+		return false, false
+	}
+	if s.lastResult == result {
+		return false, true
+	}
+
+	s.lastResult = result
+	return true, true
+}
+
+func (s *SessionState) notifyChange() {
+	select {
+	case s.changedCh <- struct{}{}:
+	default:
+	}
+}
diff --git a/internal/mapr/client/session_state_test.go b/internal/mapr/client/session_state_test.go
new file mode 100644
index 0000000..f43ca70
--- /dev/null
+++ b/internal/mapr/client/session_state_test.go
@@ -0,0 +1,81 @@
+package client
+
+import (
+	"testing"
+
+	"github.com/mimecast/dtail/internal/mapr"
+)
+
+func TestSessionStateCommitQueryResetsGenerationAndResults(t *testing.T) {
+	query := mustSessionStateQuery(t, "select count(status) from stats group by status")
+	state := NewSessionState(query)
+
+	initial := state.Snapshot()
+	group := mapr.NewGroupSet()
+	set := group.GetSet("ERROR")
+	set.Samples = 1
+	set.FValues[query.Select[0].FieldStorage] = 1
+	if err := initial.GlobalGroup.Merge(query, group); err != nil {
+		t.Fatalf("Merge() error = %v", err)
+	}
+	if changed, ok := state.CommitRenderedResult(initial.Generation, "old-result"); !ok || !changed {
+		t.Fatalf("CommitRenderedResult() = changed:%v ok:%v, want changed and ok", changed, ok)
+	}
+
+	rawQuery := "select count(status) from warnings group by status"
+	updatedQuery, err := state.CommitQuery(rawQuery, 3)
+	if err != nil {
+		t.Fatalf("CommitQuery() error = %v", err)
+	}
+	if updatedQuery == nil || updatedQuery.RawQuery != rawQuery {
+		t.Fatalf("unexpected updated query: %#v", updatedQuery)
+	}
+
+	select {
+	case <-state.Changes():
+	default:
+		t.Fatalf("expected change notification after CommitQuery")
+	}
+
+	updated := state.Snapshot()
+	if updated.Generation != 3 {
+		t.Fatalf("generation = %d, want 3", updated.Generation)
+	}
+	if updated.Query == nil || updated.Query.RawQuery != rawQuery {
+		t.Fatalf("unexpected query after commit: %#v", updated.Query)
+	}
+	if !updated.GlobalGroup.IsEmpty() {
+		t.Fatalf("expected committed global group to be reset")
+	}
+	if updated.LastResult != "" {
+		t.Fatalf("last result = %q, want empty", updated.LastResult)
+	}
+}
+
+func TestSessionStateCommitQueryRejectsInvalidQuery(t *testing.T) {
+	query := mustSessionStateQuery(t, "select count(status) from stats group by status")
+	state := NewSessionState(query)
+	before := state.Snapshot()
+
+	if _, err := state.CommitQuery("select from", 5); err == nil {
+		t.Fatalf("expected CommitQuery() to reject invalid query")
+	}
+
+	after := state.Snapshot()
+	if after.Generation != before.Generation {
+		t.Fatalf("generation changed on invalid query: got %d want %d", after.Generation, before.Generation)
+	}
+	if after.Query == nil || after.Query.RawQuery != before.Query.RawQuery {
+		t.Fatalf("query changed on invalid query: before=%#v after=%#v", before.Query, after.Query)
+	}
+}
+
+func mustSessionStateQuery(t *testing.T, queryStr string) *mapr.Query {
+	t.Helper()
+
+	query, err := mapr.NewQuery(queryStr)
+	if err != nil {
+		t.Fatalf("NewQuery(%q) error = %v", queryStr, err)
+	}
+	return query
+}