Add initial implementation of aggregator

internal/base/base.go

@@ -215,14 +215,20 @@ func GroupKey(qname, gkey string) string {
 	return fmt.Sprintf("%s%s", GroupKeyPrefix(qname), gkey)
 }
 
+// AggregationSetKey returns a redis key used for an aggregation set.
+func AggregationSetKey(qname, gname, setID string) string {
+	return fmt.Sprintf("%s:%s", GroupKey(qname, gname), setID)
+}
+
 // AllGroups return a redis key used to store all group keys used in a given queue.
 func AllGroups(qname string) string {
 	return fmt.Sprintf("%sgroups", QueueKeyPrefix(qname))
 }
 
-// AllStagedGroups returns a redis key used to store all groups staged to be aggregated in a given queue.
-func AllStagedGroups(qname string) string {
-	return fmt.Sprintf("%sstaged_groups", QueueKeyPrefix(qname))
+// AllAggregationSets returns a redis key used to store all aggregation sets (set of tasks staged to be aggregated)
+// in a given queue.
+func AllAggregationSets(qname string) string {
+	return fmt.Sprintf("%saggregation_sets", QueueKeyPrefix(qname))
 }
 
 // TaskMessage is the internal representation of a task with additional metadata fields.
@@ -708,6 +714,7 @@ func (l *Lease) IsValid() bool {
 // See rdb.RDB as a reference implementation.
 type Broker interface {
 	Ping() error
+	Close() error
 	Enqueue(ctx context.Context, msg *TaskMessage) error
 	EnqueueUnique(ctx context.Context, msg *TaskMessage, ttl time.Duration) error
 	Dequeue(qnames ...string) (*TaskMessage, time.Time, error)
@@ -719,13 +726,29 @@ type Broker interface {
 	Retry(ctx context.Context, msg *TaskMessage, processAt time.Time, errMsg string, isFailure bool) error
 	Archive(ctx context.Context, msg *TaskMessage, errMsg string) error
 	ForwardIfReady(qnames ...string) error
+
+	// Group aggregation related methods
+	AddToGroup(ctx context.Context, msg *TaskMessage, gname string) error
+	AddToGroupUnique(ctx context.Context, msg *TaskMessage, groupKey string, ttl time.Duration) error
+	ListGroups(qname string) ([]string, error)
+	AggregationCheck(qname, gname string) (aggregationSetID string, err error)
+	ReadAggregationSet(qname, gname, aggregationSetID string) ([]*TaskMessage, time.Time, error)
+	DeleteAggregationSet(ctx context.Context, qname, gname, aggregationSetID string) error
+
+	// Task retention related method
 	DeleteExpiredCompletedTasks(qname string) error
+
+	// Lease related methods
 	ListLeaseExpired(cutoff time.Time, qnames ...string) ([]*TaskMessage, error)
 	ExtendLease(qname string, ids ...string) (time.Time, error)
+
+	// State snapshot related methods
 	WriteServerState(info *ServerInfo, workers []*WorkerInfo, ttl time.Duration) error
 	ClearServerState(host string, pid int, serverID string) error
+
+	// Cancelation related methods
 	CancelationPubSub() (*redis.PubSub, error) // TODO: Need to decouple from redis to support other brokers
 	PublishCancelation(id string) error
+
 	WriteResult(qname, id string, data []byte) (n int, err error)
-	Close() error
 }

internal/base/base_test.go

@@ -421,6 +421,35 @@ func TestGroupKey(t *testing.T) {
 	}
 }
 
+func TestAggregationSetKey(t *testing.T) {
+	tests := []struct {
+		qname string
+		gname string
+		setID string
+		want  string
+	}{
+		{
+			qname: "default",
+			gname: "mygroup",
+			setID: "12345",
+			want:  "asynq:{default}:g:mygroup:12345",
+		},
+		{
+			qname: "custom",
+			gname: "foo",
+			setID: "98765",
+			want:  "asynq:{custom}:g:foo:98765",
+		},
+	}
+
+	for _, tc := range tests {
+		got := AggregationSetKey(tc.qname, tc.gname, tc.setID)
+		if got != tc.want {
+			t.Errorf("AggregationSetKey(%q, %q, %q) = %q, want %q", tc.qname, tc.gname, tc.setID, got, tc.want)
+		}
+	}
+}
+
 func TestAllGroups(t *testing.T) {
 	tests := []struct {
 		qname string
@@ -444,25 +473,25 @@ func TestAllGroups(t *testing.T) {
 	}
 }
 
-func TestAllStagedGroups(t *testing.T) {
+func TestAllAggregationSets(t *testing.T) {
 	tests := []struct {
 		qname string
 		want  string
 	}{
 		{
 			qname: "default",
-			want:  "asynq:{default}:staged_groups",
+			want:  "asynq:{default}:aggregation_sets",
 		},
 		{
 			qname: "custom",
-			want:  "asynq:{custom}:staged_groups",
+			want:  "asynq:{custom}:aggregation_sets",
 		},
 	}
 
 	for _, tc := range tests {
-		got := AllStagedGroups(tc.qname)
+		got := AllAggregationSets(tc.qname)
 		if got != tc.want {
-			t.Errorf("AllStagedGroups(%q) = %q, want %q", tc.qname, got, tc.want)
+			t.Errorf("AllAggregationSets(%q) = %q, want %q", tc.qname, got, tc.want)
 		}
 	}
 }

internal/rdb/rdb.go

@@ -982,6 +982,32 @@ func (r *RDB) forwardAll(qname string) (err error) {
 	return nil
 }
 
+// ListGroups returns a list of all known groups in the given queue.
+func (r *RDB) ListGroups(qname string) ([]string, error) {
+	// TODO: Implement this with TDD
+	return nil, nil
+}
+
+// AggregationCheck checks the group identified by the given queue and group name to see if the tasks in the
+// group are ready to be aggregated. If so, it moves the tasks to be aggregated to a aggregation set and returns
+// set ID. If not, it returns an empty string for the set ID.
+func (r *RDB) AggregationCheck(qname, gname string) (string, error) {
+	// TODO: Implement this with TDD
+	return "", nil
+}
+
+// ReadAggregationSet retrieves memebers of an aggregation set and returns list of tasks and
+// the deadline for aggregating those tasks.
+func (r *RDB) ReadAggregationSet(qname, gname, setID string) ([]*base.TaskMessage, time.Time, error) {
+	// TODO: Implement this with TDD
+	return nil, time.Time{}, nil
+}
+
+// DeleteAggregationSet deletes the aggregation set identified by the parameters.
+func (r *RDB) DeleteAggregationSet(ctx context.Context, qname, gname, setID string) error {
+	return nil
+}
+
 // KEYS[1] -> asynq:{<qname>}:completed
 // ARGV[1] -> current time in unix time
 // ARGV[2] -> task key prefix