Add initial implementation of aggregator

aggregator.go

@@ -0,0 +1,129 @@
+// Copyright 2022 Kentaro Hibino. All rights reserved.
+// Use of this source code is governed by a MIT license
+// that can be found in the LICENSE file.
+
+package asynq
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"github.com/hibiken/asynq/internal/base"
+	"github.com/hibiken/asynq/internal/log"
+)
+
+// An aggregator is responsible for checking groups and aggregate into one task
+// if any of the grouping condition is met.
+type aggregator struct {
+	logger *log.Logger
+	broker base.Broker
+	client *Client
+
+	// channel to communicate back to the long running "aggregator" goroutine.
+	done chan struct{}
+
+	// list of queue names to check and aggregate.
+	queues []string
+
+	// Group configurations
+	gracePeriod time.Duration
+	maxDelay    time.Duration
+	maxSize     int
+
+	// Aggregation function
+	aggregateFunc func(gname string, tasks []*Task) *Task
+}
+
+type aggregatorParams struct {
+	logger        *log.Logger
+	broker        base.Broker
+	queues        []string
+	gracePeriod   time.Duration
+	maxDelay      time.Duration
+	maxSize       int
+	aggregateFunc func(gname string, msgs []*Task) *Task
+}
+
+func newAggregator(params aggregatorParams) *aggregator {
+	return &aggregator{
+		logger:        params.logger,
+		broker:        params.broker,
+		client:        &Client{broker: params.broker},
+		done:          make(chan struct{}),
+		queues:        params.queues,
+		gracePeriod:   params.gracePeriod,
+		maxDelay:      params.maxDelay,
+		maxSize:       params.maxSize,
+		aggregateFunc: params.aggregateFunc,
+	}
+}
+
+func (a *aggregator) shutdown() {
+	a.logger.Debug("Aggregator shutting down...")
+	// Signal the aggregator goroutine to stop.
+	a.done <- struct{}{}
+}
+
+func (a *aggregator) start(wg *sync.WaitGroup) {
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		timer := time.NewTimer(a.gracePeriod)
+		for {
+			select {
+			case <-a.done:
+				a.logger.Debug("Aggregator done")
+				return
+			case <-timer.C:
+				a.exec()
+				timer.Reset(a.gracePeriod)
+			}
+		}
+	}()
+}
+
+func (a *aggregator) exec() {
+	for _, qname := range a.queues {
+		groups, err := a.broker.ListGroups(qname)
+		if err != nil {
+			a.logger.Errorf("Failed to list groups in queue: %q", qname)
+			continue
+		}
+		for _, gname := range groups {
+			aggregationSetID, err := a.broker.AggregationCheck(qname, gname)
+			if err != nil {
+				a.logger.Errorf("Failed to run aggregation check: queue=%q group=%q", qname, gname)
+				continue
+			}
+			if aggregationSetID == "" {
+				a.logger.Debugf("No aggregation needed at this time: queue=%q group=%q", qname, gname)
+				continue
+			}
+
+			// Aggregate and enqueue.
+			msgs, deadline, err := a.broker.ReadAggregationSet(qname, gname, aggregationSetID)
+			if err != nil {
+				a.logger.Errorf("Failed to read aggregation set: queue=%q, group=%q, setID=%q",
+					qname, gname, aggregationSetID)
+				continue
+			}
+			tasks := make([]*Task, len(msgs))
+			for i, m := range msgs {
+				tasks[i] = NewTask(m.Type, m.Payload)
+			}
+			aggregatedTask := a.aggregateFunc(gname, tasks)
+			ctx, cancel := context.WithDeadline(context.Background(), deadline)
+			if _, err := a.client.EnqueueContext(ctx, aggregatedTask); err != nil {
+				a.logger.Errorf("Failed to enqueue aggregated task: %v", err)
+				cancel()
+				continue
+			}
+			if err := a.broker.DeleteAggregationSet(ctx, qname, gname, aggregationSetID); err != nil {
+				a.logger.Warnf("Failed to delete aggregation set: queue=%q, group=%q, setID=%q",
+					qname, gname, aggregationSetID)
+			}
+			cancel()
+		}
+	}
+}