feat: DTail fork — server/client feature development

Squashed development of the snonux/dtail fork's product code (internal/, cmd/)
since diverging from mimecast/dtail. Major areas:

- Read/output path: the former "turbo" channel-less path is now the single,
  default server-side read/output path for cat/grep/tail and MapReduce; the old
  channel-based path and its config/env toggles were removed.
- MapReduce: single aggregate implementation (server + serverless) fed directly
  by a processor pipeline, with input-exhausted finalization via the shutdown
  coordinator; high-concurrency and data-race fixes.
- Journal source reads (journal:unit.service) via journalctl, Linux-gated behind
  a journal-v1 capability.
- Auth-key fast reconnect: in-memory per-user public-key cache with TTL/max-keys,
  registered over an authenticated session (AUTHKEY), checked before
  authorized_keys.
- Interactive query reload (--interactive-query) with SESSION START/UPDATE
  generation boundaries and capability negotiation.
- Client-side deadlines: --timeout / --shutdownAfter as context deadlines;
  follow shutdown handling.
- Client logging: diagnostics-only daily log by default, opt-in payload tee via
  --log-payload.
- Numerous correctness fixes (buffer-pool double-recycle races, EOF-sentinel
  leaks, glob-expansion cap, TOCTOU in CSV parsing) with accompanying unit tests.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

5 files changed, 1252 insertions, 185 deletions

diff --git a/internal/mapr/server/aggregate.go b/internal/mapr/server/aggregate.go
index 4f14751..26d1211 100644
--- a/internal/mapr/server/aggregate.go
+++ b/internal/mapr/server/aggregate.go
@@ -1,37 +1,102 @@
 package server
 
 import (
+	"bytes"
 	"context"
 	"strings"
+	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/mimecast/dtail/internal"
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/dlog"
-	"github.com/mimecast/dtail/internal/io/line"
+	"github.com/mimecast/dtail/internal/io/pool"
 	"github.com/mimecast/dtail/internal/mapr"
 	"github.com/mimecast/dtail/internal/mapr/logformat"
-	"github.com/mimecast/dtail/internal/protocol"
 )
 
-// Aggregate is for aggregating mapreduce data on the DTail server side.
+// Aggregate is a high-performance aggregator for MapReduce operations.
+// It processes lines directly without channels for maximum throughput.
 type Aggregate struct {
 	done *internal.Done
-	// NextLinesCh can be used to use a new line ch.
-	NextLinesCh chan chan *line.Line
-	linesCh     chan *line.Line
+	// inputFinished is signaled via FinishInput once all one-shot input file
+	// reads (cat/grep style) feeding this aggregate have drained. Start then
+	// emits a final serialization and returns instead of blocking until
+	// session teardown. Follow-mode (tail) inputs never signal it, so
+	// interval-based streaming aggregation keeps running.
+	inputFinished *internal.Done
 	// Hostname of the current server (used to populate $hostname field).
 	hostname string
-	// Signals to serialize data.
-	serialize chan struct{}
 	// The mapr query
 	query *mapr.Query
 	// The mapr log format parser
 	parser logformat.Parser
+	// Group sets are swapped out during serialization to avoid clone-heavy flushes.
+	groupMu   sync.Mutex
+	groupSets map[string]*mapr.AggregateSet
+	// serializeMu ensures only one serialization runs at a time.
+	serializeMu sync.Mutex
+	// Batch processing
+	batchMu   sync.Mutex
+	batch     []rawLine
+	batchSize int
+	// Periodic serialization.
+	// serializeTicker is published once by Start (before the serializationLoop
+	// goroutine is launched) and read from two other places: serializationLoop's
+	// select (ordered after the store by the go statement) and
+	// stopSerializeTicker, which is reachable from Shutdown/Abort on the external
+	// teardown goroutine with no other happens-before edge to Start's write. An
+	// atomic.Pointer gives that publish-once/read-many access a lock-free
+	// happens-before guarantee without coupling the ticker to serializeMu:
+	// guarding it with serializeMu would make Abort's stop block behind an
+	// in-flight doSerialize, violating Abort's immediate, non-blocking preemption
+	// contract.
+	serializeTicker atomic.Pointer[time.Ticker]
+	serialize       chan struct{}
+	// maprMessages is the output channel for serialized results. It is
+	// published once by Start and read by doSerialize; both accesses are
+	// guarded by serializeMu so the write in Start happens-before any read in
+	// doSerialize, even when doSerialize runs from a different goroutine via
+	// Shutdown.
+	maprMessages chan<- string
+	// Stats
+	linesProcessed atomic.Uint64
+	errors         atomic.Uint64
+	filesProcessed atomic.Uint64
+	// Synchronization for clean shutdown.
+	processorsWg sync.WaitGroup
+	// Track active file processors
+	activeProcessors atomic.Int32
+	startOnce        sync.Once
+	started          chan struct{}
 }
 
-// NewAggregate return a new server side aggregator.
-func NewAggregate(queryStr string) (*Aggregate, error) {
+type rawLine struct {
+	content  *bytes.Buffer
+	sourceID string
+}
+
+func (a *Aggregate) stopping() bool {
+	select {
+	case <-a.done.Done():
+		return true
+	default:
+		return false
+	}
+}
+
+func (a *Aggregate) stopSerializeTicker() {
+	// Load is safe from the external Shutdown/Abort teardown goroutine: Start
+	// publishes the ticker with an atomic Store, so a nil load simply means
+	// Start has not created it yet and there is nothing to stop.
+	if ticker := a.serializeTicker.Load(); ticker != nil {
+		ticker.Stop()
+	}
+}
+
+// NewAggregate returns a new aggregator.
+func NewAggregate(queryStr string, defaultLogFormat string) (*Aggregate, error) {
 	query, err := mapr.NewQuery(queryStr)
 	if err != nil {
 		return nil, err
@@ -43,18 +108,12 @@ func NewAggregate(queryStr string) (*Aggregate, error) {
 	}
 	s := strings.Split(fqdn, ".")
 
-	var parserName string
-	switch query.LogFormat {
-	case "":
-		parserName = config.Server.MapreduceLogFormat
-		if query.Table == "" {
-			parserName = "generic"
-		}
-	default:
-		parserName = query.LogFormat
-	}
+	parserName := resolveParserName(query, defaultLogFormat)
 
-	dlog.Server.Info("Creating log format parser", parserName)
+	dlog.Server.Info("Creating log format parser",
+		"parserName", parserName,
+		"queryTable", query.Table,
+		"queryLogFormat", query.LogFormat)
 	logParser, err := logformat.NewParser(parserName, query)
 	if err != nil {
 		dlog.Server.Error("Could not create log format parser. Falling back to 'generic'", err)
@@ -64,236 +123,427 @@ func NewAggregate(queryStr string) (*Aggregate, error) {
 	}
 
 	return &Aggregate{
-		done:        internal.NewDone(),
-		NextLinesCh: make(chan chan *line.Line, 100),
-		serialize:   make(chan struct{}),
-		hostname:    s[0],
-		query:       query,
-		parser:      logParser,
+		done:          internal.NewDone(),
+		inputFinished: internal.NewDone(),
+		serialize:     make(chan struct{}, 1), // Buffered to avoid blocking
+		hostname:      s[0],
+		query:         query,
+		parser:        logParser,
+		groupSets:     make(map[string]*mapr.AggregateSet),
+		batchSize:     100, // Process 100 lines at a time
+		batch:         make([]rawLine, 0, 100),
+		started:       make(chan struct{}),
 	}, nil
 }
 
+// countGroups returns the current number of groups in the aggregation.
+func (a *Aggregate) countGroups() int {
+	a.groupMu.Lock()
+	defer a.groupMu.Unlock()
+	return len(a.groupSets)
+}
+
 // Shutdown the aggregation engine.
 func (a *Aggregate) Shutdown() {
 	a.done.Shutdown()
-}
-
-// Start an aggregation.
-func (a *Aggregate) Start(ctx context.Context, maprMessages chan<- string) {
-	myCtx, cancel := context.WithCancel(ctx)
+	a.stopSerializeTicker()
+	a.processorsWg.Wait()
+	a.processBatchAndWait()
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
+	a.doSerialize(ctx)
+}
 
-	go func() {
-		select {
-		case <-myCtx.Done():
-			a.done.Shutdown()
-		case <-a.done.Done():
-			cancel()
-		}
-	}()
+// Abort stops background processing without waiting for final serialization.
+// Session generation replacement uses this to preempt old query work immediately.
+func (a *Aggregate) Abort() {
+	a.done.Shutdown()
+	a.stopSerializeTicker()
+}
 
-	fieldsCh := a.fieldsFromLines(myCtx)
-	// Add fields (e.g. via 'set' clause)
-	if len(a.query.Set) > 0 {
-		fieldsCh = a.setAdditionalFields(myCtx, fieldsCh)
-	}
-	// Periodically pre-aggregate data every a.query.Interval seconds.
-	go a.aggregateTimer(myCtx)
-	a.aggregateAndSerialize(myCtx, fieldsCh, maprMessages)
+// FinishInput signals that all one-shot input (cat/grep style file reads)
+// feeding this aggregate has been fully consumed and no further processors
+// will register. Start reacts by emitting a final serialization and
+// returning, which lets a server-side map command complete once its input
+// line channels are exhausted.
+// Follow-mode (tail) inputs must never call this so that interval-based
+// streaming aggregation keeps running. Safe to call multiple times.
+func (a *Aggregate) FinishInput() {
+	a.inputFinished.Shutdown()
 }
 
-func (a *Aggregate) aggregateTimer(ctx context.Context) {
-	for {
-		select {
-		case <-time.After(a.query.Interval):
-			a.Serialize(ctx)
-		case <-ctx.Done():
-			return
-		}
+// Start the aggregation. It blocks until the context is canceled, the
+// aggregate is shut down, or — for one-shot inputs — FinishInput signals that
+// all input has been consumed, in which case all remaining data is flushed
+// and serialized before returning.
+func (a *Aggregate) Start(ctx context.Context, maprMessages chan<- string) {
+	// Publish the output channel under serializeMu. doSerialize reads
+	// a.maprMessages while holding serializeMu (see line ~355), and it can be
+	// invoked from a different goroutine than this one — Shutdown() is called
+	// from the handler teardown path (baseHandler.Shutdown) concurrently with
+	// this Start goroutine. Writing under the same lock the reader holds
+	// establishes a happens-before edge, so the read is never torn or stale.
+	// The internal serializationLoop reader is already ordered by the go
+	// statement below, but the external Shutdown reader needs this lock.
+	a.serializeMu.Lock()
+	a.maprMessages = maprMessages
+	a.serializeMu.Unlock()
+	interval := a.query.Interval
+	if interval <= 0 {
+		interval = time.Second
 	}
-}
+	// Publish the ticker before launching serializationLoop below. The store
+	// happens-before that goroutine's Load via the go statement, and any later
+	// stopSerializeTicker on the teardown goroutine observes it through the
+	// atomic (see the serializeTicker field comment).
+	a.serializeTicker.Store(time.NewTicker(interval))
+	a.startOnce.Do(func() {
+		if a.started != nil {
+			close(a.started)
+		}
+	})
+	defer a.stopSerializeTicker()
 
-func (a *Aggregate) nextLine() (line *line.Line, ok bool, noMoreChannels bool) {
-	dlog.Server.Trace("nextLine.enter", line, ok, noMoreChannels)
+	loopDone := make(chan struct{})
+	go func() {
+		defer close(loopDone)
+		a.serializationLoop(ctx)
+	}()
 
 	select {
-	case line, ok = <-a.linesCh:
-		if !ok {
-			// Channel is closed, go to next channel.
-			select {
-			case a.linesCh = <-a.NextLinesCh:
-			default:
-				noMoreChannels = true
-			}
-		}
-	default:
-		// No new line from current lines channel. Try next one.
-		select {
-		case newLinesCh := <-a.NextLinesCh:
-			oldLinesCh := a.linesCh
-			go func() { a.NextLinesCh <- oldLinesCh }()
-			a.linesCh = newLinesCh
-		default:
-			// No new lines channel found.
-		}
+	case <-ctx.Done():
+	case <-a.done.Done():
+	case <-a.inputFinished.Done():
+		// All one-shot input is consumed: emit the final result and stop.
+		// Shutdown waits for the processors, drains the batch and performs
+		// the final serialization. Without this path a server-mode dmap
+		// command would block here until session teardown while keeping the
+		// session's active-command count nonzero — a circular wait that hung
+		// the client forever even though all results had been transmitted.
+		a.Shutdown()
 	}
-	dlog.Server.Trace("nextLine.exit", line, ok, noMoreChannels)
-	return
-}
 
-func (a *Aggregate) fieldsFromLines(ctx context.Context) <-chan map[string]string {
-	fieldsCh := make(chan map[string]string)
+	// Stop the serialization loop and wait for it to exit before returning,
+	// so no serialization can send on maprMessages once the caller closes its
+	// side of the channel right after Start returns.
+	a.done.Shutdown()
+	<-loopDone
+}
 
-	go func() {
-		defer close(fieldsCh)
+// ProcessLineDirect processes a line directly without channels.
+// This is called from the AggregateProcessor.
+func (a *Aggregate) ProcessLineDirect(lineContent *bytes.Buffer, sourceID string) error {
+	if a.stopping() {
+		pool.RecycleBytesBuffer(lineContent)
+		return nil
+	}
 
-		// Gather first lines channel (first input file)
-		select {
-		case a.linesCh = <-a.NextLinesCh:
-		case <-ctx.Done():
-			return
-		}
+	a.linesProcessed.Add(1)
 
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			default:
-			}
+	// Add to batch
+	a.batchMu.Lock()
+	a.batch = append(a.batch, rawLine{content: lineContent, sourceID: sourceID})
+	shouldProcess := len(a.batch) >= a.batchSize
+	a.batchMu.Unlock()
 
-			// Gather first lines channel (first input file)
-			line, ok, noMoreChannels := a.nextLine()
-			if !ok {
-				if noMoreChannels {
-					return
-				}
-				time.Sleep(time.Millisecond * 100)
-				continue
-			}
+	if shouldProcess {
+		a.processBatch()
+	}
 
-			if err := a.fieldFromLine(ctx, line, fieldsCh); err != nil {
-				dlog.Server.Error(err)
-			}
-		}
-	}()
+	return nil
+}
 
-	return fieldsCh
+// processBatch processes a full batch immediately.
+func (a *Aggregate) processBatch() {
+	a.processRawBatch(a.takeBatch())
 }
 
-func (a *Aggregate) fieldFromLine(ctx context.Context, line *line.Line,
-	fieldsCh chan<- map[string]string) error {
+// processBatchAndWait processes a batch of lines synchronously and waits for completion.
+// This is used when flushing to ensure all data is processed before continuing.
+func (a *Aggregate) processBatchAndWait() {
+	a.processRawBatch(a.takeBatch())
+}
 
-	maprLine := strings.TrimSpace(line.Content.String())
+func (a *Aggregate) takeBatch() []rawLine {
+	a.batchMu.Lock()
+	if len(a.batch) == 0 {
+		a.batchMu.Unlock()
+		return nil
+	}
+	batch := a.batch
+	a.batch = make([]rawLine, 0, a.batchSize)
+	a.batchMu.Unlock()
+	return batch
+}
 
-	// after recycling it, don't use line object anymore!!!
-	line.Recycle()
-	fields, err := a.parser.MakeFields(maprLine)
+func (a *Aggregate) processRawBatch(batch []rawLine) {
+	for i := range batch {
+		if err := a.processLine(batch[i].content, batch[i].sourceID); err != nil {
+			a.errors.Add(1)
+			dlog.Server.Error("Error processing line:", err, "lineIndex", i)
+		}
+		if batch[i].content != nil {
+			pool.RecycleBytesBuffer(batch[i].content)
+		}
+	}
+}
 
+// processLine processes a single line and aggregates it.
+func (a *Aggregate) processLine(lineContent *bytes.Buffer, sourceID string) error {
+	maprLine := strings.TrimSpace(lineContent.String())
+	parsedFields, err := a.parser.MakeFields(maprLine, sourceID)
 	if err != nil {
-		// Should fields be ignored anyway?
 		if err != logformat.ErrIgnoreFields {
 			return err
 		}
 		return nil
 	}
-	if !a.query.WhereClause(fields) {
+
+	// Apply where clause
+	if !a.query.WhereClause(parsedFields) {
 		return nil
 	}
 
-	select {
-	case fieldsCh <- fields:
-	case <-ctx.Done():
+	// Apply set clause if needed
+	if len(a.query.Set) > 0 {
+		if err := a.query.SetClause(parsedFields); err != nil {
+			return err
+		}
 	}
 
+	// Aggregate the fields
+	a.aggregate(parsedFields)
 	return nil
 }
 
-func (a *Aggregate) setAdditionalFields(ctx context.Context,
-	fieldsCh <-chan map[string]string) <-chan map[string]string {
+// aggregate adds fields to the appropriate group. The set is only created (or
+// looked up) after at least one select field matches, preventing empty sets with
+// Samples==0 from entering the map and causing 0/0 = NaN on the client for Avg.
+func (a *Aggregate) aggregate(fields map[string]string) {
+	groupKey := buildGroupKey(a.query.GroupBy, fields)
+	a.groupMu.Lock()
 
-	newFieldsCh := make(chan map[string]string)
-	go func() {
-		defer close(newFieldsCh)
-		for {
-			fields, ok := <-fieldsCh
-			if !ok {
-				return
-			}
-			if err := a.query.SetClause(fields); err != nil {
-				dlog.Server.Error(err)
-			}
+	var set *mapr.AggregateSet
+	var addedSample bool
 
-			select {
-			case newFieldsCh <- fields:
-			case <-ctx.Done():
+	for _, sc := range a.query.Select {
+		val, ok := fields[sc.Field]
+		if !ok {
+			continue
+		}
+		// Lazily look up or allocate the aggregate set on the first matching
+		// field so that lines with no matching fields never create empty entries.
+		if set == nil {
+			set, ok = a.groupSets[groupKey]
+			if !ok {
+				set = mapr.NewAggregateSet()
+				a.groupSets[groupKey] = set
 			}
 		}
-	}()
-	return newFieldsCh
+		if err := set.Aggregate(sc.FieldStorage, sc.Operation, val, false); err != nil {
+			dlog.Server.Error("Aggregate aggregation error", err, "field", sc.Field, "operation", sc.Operation)
+			continue
+		}
+		addedSample = true
+	}
+	if addedSample {
+		set.Samples++
+	}
+	a.groupMu.Unlock()
 }
 
-func (a *Aggregate) aggregateAndSerialize(ctx context.Context,
-	fieldsCh <-chan map[string]string, maprMessages chan<- string) {
-
-	group := mapr.NewGroupSet()
-	serialize := func() {
-		dlog.Server.Info("Serializing mapreduce result")
-		group.Serialize(ctx, maprMessages)
-		group = mapr.NewGroupSet()
-	}
+// serializationLoop handles periodic serialization.
+func (a *Aggregate) serializationLoop(ctx context.Context) {
+	// Start stores serializeTicker before launching this goroutine, so the load
+	// is ordered-after that store and never nil here. The ticker pointer is
+	// never replaced after Start, so loading it once is sufficient.
+	ticker := a.serializeTicker.Load()
 	for {
 		select {
-		case fields, ok := <-fieldsCh:
-			if !ok {
-				serialize()
-				return
-			}
-			a.aggregate(group, fields)
-		case <-a.serialize:
-			serialize()
 		case <-ctx.Done():
 			return
+		case <-a.done.Done():
+			return
+		case <-ticker.C:
+			a.Serialize(ctx)
+		case <-a.serialize:
+			a.doSerialize(ctx)
 		}
 	}
 }
 
-func (a *Aggregate) aggregate(group *mapr.GroupSet, fields map[string]string) {
-	var sb strings.Builder
-	for i, field := range a.query.GroupBy {
-		if i > 0 {
-			sb.WriteString(protocol.AggregateGroupKeyCombinator)
-		}
-		if val, ok := fields[field]; ok {
-			sb.WriteString(val)
+// Serialize triggers serialization of all aggregated data.
+func (a *Aggregate) Serialize(ctx context.Context) {
+	select {
+	case a.serialize <- struct{}{}:
+	case <-time.After(time.Minute):
+		dlog.Server.Warn("Starting to serialize mapreduce data takes over a minute")
+	case <-ctx.Done():
+	}
+}
+
+// doSerialize performs the actual serialization.
+func (a *Aggregate) doSerialize(ctx context.Context) {
+	a.serializeMu.Lock()
+	defer a.serializeMu.Unlock()
+
+	a.processBatchAndWait()
+	if a.maprMessages == nil {
+		dlog.Server.Error("Aggregate maprMessages channel is nil")
+		return
+	}
+
+	snapshot := a.swapGroupSets()
+	if len(snapshot) == 0 {
+		return
+	}
+
+	group := mapr.NewGroupSet()
+	for groupKey, aggregateSet := range snapshot {
+		groupSet := group.GetSet(groupKey)
+		*groupSet = *aggregateSet
+	}
+
+	serializeCtx := ctx
+	if _, ok := ctx.Deadline(); ok {
+		var cancel context.CancelFunc
+		serializeCtx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+	}
+	remaining := group.Serialize(serializeCtx, a.maprMessages)
+	if len(remaining) > 0 {
+		a.mergeRemainingLocked(remaining)
+	}
+}
+
+// mergeRemainingLocked re-inserts aggregate sets that could not be sent during
+// serialization back into the live groupSets map. Without this path the
+// snapshot taken by swapGroupSets would be silently discarded on ctx
+// cancellation and the next Serialize would not be able to retry the data.
+//
+// Concurrent ProcessLine calls may have already added new samples for the same
+// group keys while the serialize was in flight. In that case we must preserve
+// the newer live state for overwrite-style aggregations such as last() and
+// len(), while still adding numeric contributions from the canceled snapshot.
+func (a *Aggregate) mergeRemainingLocked(remaining map[string]*mapr.AggregateSet) {
+	a.groupMu.Lock()
+	defer a.groupMu.Unlock()
+	for key, set := range remaining {
+		existing, ok := a.groupSets[key]
+		if !ok {
+			a.groupSets[key] = set
+			continue
 		}
+		mergeCancelledSnapshot(a.query, existing, set)
 	}
-	groupKey := sb.String()
-	set := group.GetSet(groupKey)
+	dlog.Server.Warn("Aggregate serialize interrupted; re-merged unsent groups",
+		"remaining", len(remaining))
+}
 
-	var addedSample bool
-	for _, sc := range a.query.Select {
-		if val, ok := fields[sc.Field]; ok {
-			if err := set.Aggregate(sc.FieldStorage, sc.Operation, val, false); err != nil {
-				dlog.Server.Error(err)
+func mergeCancelledSnapshot(query *mapr.Query, live, snapshot *mapr.AggregateSet) {
+	live.Samples += snapshot.Samples
+	for _, sc := range query.Select {
+		storage := sc.FieldStorage
+		switch sc.Operation {
+		case mapr.Count, mapr.Sum, mapr.Avg, mapr.Percentage, mapr.Percentile:
+			live.FValues[storage] += snapshot.FValues[storage]
+		case mapr.Min:
+			liveValue, ok := live.FValues[storage]
+			if !ok {
+				live.FValues[storage] = snapshot.FValues[storage]
 				continue
 			}
-			addedSample = true
+			if snapshotValue := snapshot.FValues[storage]; snapshotValue < liveValue {
+				live.FValues[storage] = snapshotValue
+			}
+		case mapr.Max:
+			liveValue, ok := live.FValues[storage]
+			if !ok {
+				live.FValues[storage] = snapshot.FValues[storage]
+				continue
+			}
+			if snapshotValue := snapshot.FValues[storage]; snapshotValue > liveValue {
+				live.FValues[storage] = snapshotValue
+			}
+		case mapr.Last:
+			if _, ok := live.SValues[storage]; !ok {
+				if snapshotValue, ok := snapshot.SValues[storage]; ok {
+					live.SValues[storage] = snapshotValue
+				}
+			}
+		case mapr.Len:
+			if _, ok := live.SValues[storage]; !ok {
+				if snapshotValue, ok := snapshot.SValues[storage]; ok {
+					live.SValues[storage] = snapshotValue
+					live.FValues[storage] = snapshot.FValues[storage]
+				}
+			}
+		default:
+			dlog.Server.Error("Aggregate re-merge encountered unsupported aggregation",
+				"operation", sc.Operation, "storage", storage)
 		}
 	}
+}
 
-	if addedSample {
-		set.Samples++
-		return
+func (a *Aggregate) swapGroupSets() map[string]*mapr.AggregateSet {
+	a.groupMu.Lock()
+	defer a.groupMu.Unlock()
+
+	if len(a.groupSets) == 0 {
+		return nil
 	}
-	dlog.Server.Trace("Aggregated data locally without adding new samples")
+
+	snapshot := a.groupSets
+	a.groupSets = make(map[string]*mapr.AggregateSet, len(snapshot))
+	return snapshot
 }
 
-// Serialize all the aggregated data.
-func (a *Aggregate) Serialize(ctx context.Context) {
-	select {
-	case a.serialize <- struct{}{}:
-	case <-time.After(time.Minute):
-		dlog.Server.Warn("Starting to serialize mapredice data takes over a minute")
-	case <-ctx.Done():
+// AggregateProcessor implements the line processor interface for aggregation.
+type AggregateProcessor struct {
+	aggregate *Aggregate
+	globID    string
+	flushOnce sync.Once
+	closeOnce sync.Once
+}
+
+// NewAggregateProcessor creates a new aggregate processor.
+func NewAggregateProcessor(aggregate *Aggregate, globID string) *AggregateProcessor {
+	aggregate.processorsWg.Add(1)
+	aggregate.activeProcessors.Add(1)
+	return &AggregateProcessor{
+		aggregate: aggregate,
+		globID:    globID,
 	}
 }
+
+// ProcessLine processes a line directly to the aggregate.
+func (p *AggregateProcessor) ProcessLine(lineContent *bytes.Buffer, _ uint64, sourceID string) error {
+	if p.aggregate.stopping() {
+		pool.RecycleBytesBuffer(lineContent)
+		return nil
+	}
+	return p.aggregate.ProcessLineDirect(lineContent, sourceID)
+}
+
+// Flush ensures all buffered data is processed.
+func (p *AggregateProcessor) Flush() error {
+	if p.aggregate.stopping() {
+		return nil
+	}
+
+	p.flushOnce.Do(func() {
+		p.aggregate.processBatchAndWait()
+		p.aggregate.filesProcessed.Add(1)
+	})
+	return nil
+}
+
+// Close flushes any remaining data.
+func (p *AggregateProcessor) Close() error {
+	err := p.Flush()
+	p.closeOnce.Do(func() {
+		p.aggregate.activeProcessors.Add(-1)
+		p.aggregate.processorsWg.Done()
+	})
+	return err
+}
diff --git a/internal/mapr/server/aggregate_test.go b/internal/mapr/server/aggregate_test.go
new file mode 100644
index 0000000..1ef060f
--- /dev/null
+++ b/internal/mapr/server/aggregate_test.go
@@ -0,0 +1,714 @@
+package server
+
+import (
+	"bytes"
+	"context"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/mimecast/dtail/internal"
+	"github.com/mimecast/dtail/internal/config"
+	"github.com/mimecast/dtail/internal/io/dlog"
+	"github.com/mimecast/dtail/internal/mapr"
+	"github.com/mimecast/dtail/internal/source"
+)
+
+// ensureTestServerConfig initialises the minimum globals required by
+// aggregate tests. Safe to call from multiple tests; it is idempotent.
+func ensureTestServerConfig(t *testing.T) {
+	t.Helper()
+	if config.Common == nil {
+		config.Common = &config.CommonConfig{
+			Logger:   "none",
+			LogLevel: "error",
+		}
+	}
+	if config.Server == nil {
+		config.Server = &config.ServerConfig{
+			MapreduceLogFormat: "default",
+		}
+	}
+	// dlog.Server.Error touches config.Client (TermColorsEnable) when it logs,
+	// e.g. the nil-maprMessages branch in doSerialize. Provide a minimal client
+	// config so those log calls do not nil-panic under test.
+	if config.Client == nil {
+		config.Client = &config.ClientConfig{TermColorsEnable: false}
+	}
+	if dlog.Server == nil {
+		ctx, cancel := context.WithCancel(context.Background())
+		t.Cleanup(cancel)
+		var wg sync.WaitGroup
+		wg.Add(1)
+		dlog.Start(ctx, &wg, source.Server)
+	}
+}
+
+// TestAggregateDoSerializeReMergesOnCtxCancel verifies that when a
+// serialize is cancelled after the live map has already advanced, the
+// canceled snapshot is merged back without overwriting newer overwrite-style
+// values. This guards against stale last()/len() values clobbering more recent
+// updates that arrived after swapGroupSets.
+func TestAggregateDoSerializeReMergesOnCtxCancel(t *testing.T) {
+	ensureTestServerConfig(t)
+
+	queryStr := `from STATS select count($time),last($message),len($message) from - group by $service`
+	agg, err := NewAggregate(queryStr, config.Server.MapreduceLogFormat)
+	if err != nil {
+		t.Fatalf("NewAggregate failed: %v", err)
+	}
+
+	countStorage := agg.query.Select[0].FieldStorage
+	lastStorage := agg.query.Select[1].FieldStorage
+	lenStorage := agg.query.Select[2].FieldStorage
+
+	agg.groupMu.Lock()
+	agg.groupSets["svc"] = &mapr.AggregateSet{
+		Samples: 1,
+		FValues: map[string]float64{
+			countStorage: 1,
+			lenStorage:   float64(len("old-len")),
+		},
+		SValues: map[string]string{
+			lastStorage: "old-last",
+			lenStorage:  "old-len",
+		},
+	}
+	agg.groupMu.Unlock()
+
+	if got := agg.countGroups(); got != 1 {
+		t.Fatalf("precondition: expected 1 group, got %d", got)
+	}
+
+	// Block the first send so doSerialize captures a snapshot and then waits
+	// in AggregateSet.Serialize. While it is blocked we advance the live state
+	// for the same group, then cancel the serialize context.
+	messages := make(chan string)
+	agg.maprMessages = messages
+
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+	go func() {
+		agg.doSerialize(ctx)
+		close(done)
+	}()
+
+	deadline := time.After(2 * time.Second)
+	for {
+		if got := agg.countGroups(); got == 0 {
+			break
+		}
+		select {
+		case <-deadline:
+			t.Fatal("timed out waiting for aggregate to swap live state")
+		case <-time.After(5 * time.Millisecond):
+		}
+	}
+
+	agg.groupMu.Lock()
+	agg.groupSets["svc"] = &mapr.AggregateSet{
+		Samples: 2,
+		FValues: map[string]float64{
+			countStorage: 2,
+			lenStorage:   float64(len("new-len")),
+		},
+		SValues: map[string]string{
+			lastStorage: "new-last",
+			lenStorage:  "new-len",
+		},
+	}
+	agg.groupMu.Unlock()
+
+	cancel()
+
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("doSerialize did not return after ctx cancel")
+	}
+
+	agg.groupMu.Lock()
+	set, ok := agg.groupSets["svc"]
+	agg.groupMu.Unlock()
+	if !ok {
+		t.Fatal("expected svc group to be re-merged after ctx cancel")
+	}
+	if got := set.Samples; got != 3 {
+		t.Fatalf("expected merged samples to be 3, got %d", got)
+	}
+	if got := set.FValues[countStorage]; got != 3 {
+		t.Fatalf("expected merged count to be 3, got %v", got)
+	}
+	if got := set.SValues[lastStorage]; got != "new-last" {
+		t.Fatalf("expected latest last() value to survive cancel, got %q", got)
+	}
+	if got := set.SValues[lenStorage]; got != "new-len" {
+		t.Fatalf("expected latest len() string value to survive cancel, got %q", got)
+	}
+	if got := set.FValues[lenStorage]; got != float64(len("new-len")) {
+		t.Fatalf("expected latest len() numeric value to survive cancel, got %v", got)
+	}
+}
+
+// TestAggregateProducesResults verifies the aggregate processes all
+// input lines and produces serialized results. It was formerly a
+// two-aggregator comparison that also exercised the regular channel-based
+// server.Aggregate; that regular aggregate was deleted once this aggregate
+// became the only aggregate path (task hv0), so only this subtest remains.
+func TestAggregateProducesResults(t *testing.T) {
+	// Initialize minimal config and logging
+	if config.Common == nil {
+		config.Common = &config.CommonConfig{
+			Logger:   "none",
+			LogLevel: "error",
+		}
+	}
+	if config.Server == nil {
+		config.Server = &config.ServerConfig{
+			MapreduceLogFormat: "default",
+		}
+	}
+	if dlog.Server == nil {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+		var wg sync.WaitGroup
+		wg.Add(1)
+		dlog.Start(ctx, &wg, source.Server)
+	}
+
+	// Test query
+	queryStr := `from STATS select count($time),$time,avg($goroutines) from - group by $time order by $time`
+
+	// Test data - DTail MapReduce format
+	testLines := []string{
+		"INFO|1002-071143|1|stats.go:56|8|15|7|0.21|471h0m21s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=1",
+		"INFO|1002-071143|1|stats.go:56|8|16|7|0.21|471h0m21s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=1",
+		"INFO|1002-071143|1|stats.go:56|8|17|7|0.21|471h0m21s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=1",
+		"INFO|1002-071147|1|stats.go:56|8|10|7|0.21|471h0m21s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=1",
+		"INFO|1002-071147|1|stats.go:56|8|11|7|0.21|471h0m21s|MAPREDUCE:STATS|currentConnections=0|lifetimeConnections=1",
+	}
+
+	t.Run("Aggregate", func(t *testing.T) {
+		// Create aggregate
+		agg, err := NewAggregate(queryStr, config.Server.MapreduceLogFormat)
+		if err != nil {
+			t.Fatalf("Failed to create aggregate: %v", err)
+		}
+
+		// Channel to collect messages
+		messages := make(chan string, 100)
+		// Use a cancellable context
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		startDone := make(chan struct{})
+		go func() {
+			defer close(startDone)
+			agg.Start(ctx, messages)
+		}()
+		waitForAggregateStart(t, agg)
+
+		// Process lines
+		processor := NewAggregateProcessor(agg, "test")
+		for i, line := range testLines {
+			buf := bytes.NewBufferString(line)
+			err := processor.ProcessLine(buf, uint64(i+1), "test")
+			if err != nil {
+				t.Errorf("Failed to process line %d: %v", i+1, err)
+			}
+		}
+
+		// Flush to ensure all data is processed
+		err = processor.Flush()
+		if err != nil {
+			t.Errorf("Failed to flush: %v", err)
+		}
+
+		// Close the processor to decrement activeProcessors
+		err = processor.Close()
+		if err != nil {
+			t.Errorf("Failed to close processor: %v", err)
+		}
+
+		// Shutdown and get results
+		agg.Shutdown()
+
+		// Cancel context to stop background goroutines
+		cancel()
+		<-startDone
+
+		// Collect results with timeout
+		done := make(chan struct{})
+		var results []string
+		go func() {
+			for msg := range messages {
+				results = append(results, msg)
+			}
+			close(done)
+		}()
+
+		// Wait a bit for serialization
+		time.Sleep(200 * time.Millisecond)
+		close(messages)
+
+		// Wait for collection to complete with timeout
+		select {
+		case <-done:
+			// Good, collected all messages
+		case <-time.After(2 * time.Second):
+			t.Error("Timeout collecting messages")
+		}
+
+		t.Logf("Aggregate processed %d lines", agg.linesProcessed.Load())
+		t.Logf("Aggregate results: %d messages", len(results))
+		for _, r := range results {
+			t.Logf("Result: %s", r)
+		}
+
+		// Verify we got results
+		if len(results) == 0 {
+			t.Error("Aggregate produced no results")
+		}
+
+		// Check line count
+		if agg.linesProcessed.Load() != uint64(len(testLines)) {
+			t.Errorf("Expected %d lines processed, got %d", len(testLines), agg.linesProcessed.Load())
+		}
+	})
+}
+
+// TestAggregateConcurrency tests aggregate with concurrent file processing
+func TestAggregateConcurrency(t *testing.T) {
+	// Initialize minimal config and logging
+	if config.Common == nil {
+		config.Common = &config.CommonConfig{
+			Logger:   "none",
+			LogLevel: "error",
+		}
+	}
+	if config.Server == nil {
+		config.Server = &config.ServerConfig{
+			MapreduceLogFormat: "default",
+		}
+	}
+	if dlog.Server == nil {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+		var wg sync.WaitGroup
+		wg.Add(1)
+		dlog.Start(ctx, &wg, source.Server)
+	}
+
+	queryStr := `from STATS select count($time),$time from - group by $time`
+
+	// Create aggregate
+	agg, err := NewAggregate(queryStr, config.Server.MapreduceLogFormat)
+	if err != nil {
+		t.Fatalf("Failed to create aggregate: %v", err)