feat: common flusher for reporting with mtu handler

enterprise/reporting/flusher/aggregator/aggregator.go

@@ -0,0 +1,12 @@
+//go:generate mockgen -destination=./aggregator_mock.go -package=aggregator -source=./aggregator.go Aggregator
+package aggregator
+
+import (
+	"context"
+	"encoding/json"
+	"time"
+)
+
+type Aggregator interface {
+	Aggregate(ctx context.Context, start, end time.Time) (jsonReports []json.RawMessage, err error)
+}

enterprise/reporting/flusher/aggregator/nop.go

@@ -0,0 +1,13 @@
+package aggregator
+
+import (
+	"context"
+	"encoding/json"
+	"time"
+)
+
+type NOP struct{}
+
+func (n *NOP) Aggregate(ctx context.Context, start, end time.Time) (jsonReports []json.RawMessage, err error) {
+	return nil, nil
+}

enterprise/reporting/flusher/aggregator/tracked_users_inapp.go

@@ -0,0 +1,120 @@
+package aggregator
+
+import (
+	"context"
+	"database/sql"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"time"
+
+	"github.com/segmentio/go-hll"
+
+	"github.com/rudderlabs/rudder-go-kit/config"
+	"github.com/rudderlabs/rudder-go-kit/stats"
+)
+
+const tableName = `tracked_users_reports`
+
+type TrackedUsersInAppAggregator struct {
+	db    *sql.DB
+	stats stats.Stats
+
+	reportsCounter    stats.Measurement
+	aggReportsCounter stats.Measurement
+}
+
+func NewTrackedUsersInAppAggregator(db *sql.DB, s stats.Stats, conf *config.Config, module string) *TrackedUsersInAppAggregator {
+	t := TrackedUsersInAppAggregator{db: db, stats: s}
+	tags := stats.Tags{
+		"instance": conf.GetString("INSTANCE_ID", "1"),
+		"table":    tableName,
+		"module":   module,
+	}
+	t.reportsCounter = t.stats.NewTaggedStat("reporting_flusher_get_reports_count", stats.HistogramType, tags)
+	t.aggReportsCounter = t.stats.NewTaggedStat("reporting_flusher_get_aggregated_reports_count", stats.HistogramType, tags)
+	return &t
+}
+
+func (a *TrackedUsersInAppAggregator) Aggregate(ctx context.Context, start, end time.Time) (jsonReports []json.RawMessage, err error) {
+	query := `SELECT reported_at, workspace_id, source_id, instance_id, userid_hll, anonymousid_hll, identified_anonymousid_hll FROM ` + tableName + `  WHERE reported_at >= $1 AND reported_at < $2 ORDER BY reported_at`
+
+	rows, err := a.db.Query(query, start, end)
+	if err != nil {
+		return nil, fmt.Errorf("cannot read reports %w", err)
+	}
+	defer rows.Close()
+
+	total := 0
+	aggReportsMap := make(map[string]*TrackedUsersReport)
+	for rows.Next() {
+		total += 1
+		r := TrackedUsersReport{}
+		err := rows.Scan(&r.ReportedAt, &r.WorkspaceID, &r.SourceID, &r.InstanceID, &r.UserIDHLLHex, &r.AnonymousIDHLLHex, &r.IdentifiedAnonymousIDHLLHex)
+		if err != nil {
+			return nil, fmt.Errorf("cannot scan row %w", err)
+		}
+		r.UserIDHLL, err = a.decodeHLL(r.UserIDHLLHex)
+		if err != nil {
+			return nil, fmt.Errorf("cannot decode hll %w", err)
+		}
+		r.AnonymousIDHLL, err = a.decodeHLL(r.AnonymousIDHLLHex)
+		if err != nil {
+			return nil, fmt.Errorf("cannot decode hll %w", err)
+		}
+		r.IdentifiedAnonymousIDHLL, err = a.decodeHLL(r.IdentifiedAnonymousIDHLLHex)
+		if err != nil {
+			return nil, fmt.Errorf("cannot decode hll %w", err)
+		}
+
+		k := fmt.Sprintf("%s-%s-%s", r.WorkspaceID, r.SourceID, r.InstanceID)
+
+		if agg, exists := aggReportsMap[k]; exists {
+			agg.UserIDHLL.Union(*r.UserIDHLL)
+			agg.AnonymousIDHLL.Union(*r.AnonymousIDHLL)
+			agg.IdentifiedAnonymousIDHLL.Union(*r.IdentifiedAnonymousIDHLL)
+			aggReportsMap[k] = agg
+		} else {
+			aggReportsMap[k] = &r
+		}
+
+	}
+
+	if err := rows.Err(); err != nil {
+		return nil, fmt.Errorf("rows errors %w", err)
+	}
+
+	jsonReports, err = marshalReports(aggReportsMap)
+	if err != nil {
+		return nil, fmt.Errorf("error marshalling reports %w", err)
+	}
+
+	a.reportsCounter.Observe(float64(total))
+	a.aggReportsCounter.Observe(float64(len(jsonReports)))
+
+	return jsonReports, nil
+}
+
+func (a *TrackedUsersInAppAggregator) decodeHLL(encoded string) (*hll.Hll, error) {
+	data, err := hex.DecodeString(encoded)
+	if err != nil {
+		return nil, err
+	}
+	hll, err := hll.FromBytes(data)
+	if err != nil {
+		return nil, err
+	}
+	return &hll, nil
+}
+
+func marshalReports(aggReportsMap map[string]*TrackedUsersReport) ([]json.RawMessage, error) {
+	jsonReports := make([]json.RawMessage, 0, len(aggReportsMap))
+	for _, v := range aggReportsMap {
+		data, err := json.Marshal(v)
+		if err != nil {
+			return nil, err
+		}
+		jsonReports = append(jsonReports, data)
+	}
+	return jsonReports, nil
+}

enterprise/reporting/flusher/aggregator/types.go

@@ -0,0 +1,35 @@
+package aggregator
+
+import (
+	"encoding/hex"
+	"encoding/json"
+	"time"
+
+	"github.com/segmentio/go-hll"
+)
+
+type TrackedUsersReport struct {
+	ReportedAt                  time.Time `json:"reportedAt"`
+	WorkspaceID                 string    `json:"workspaceId"`
+	SourceID                    string    `json:"sourceId"`
+	InstanceID                  string    `json:"instanceId"`
+	UserIDHLL                   *hll.Hll  `json:"-"`
+	AnonymousIDHLL              *hll.Hll  `json:"-"`
+	IdentifiedAnonymousIDHLL    *hll.Hll  `json:"-"`
+	UserIDHLLHex                string    `json:"userIdHLL"`
+	AnonymousIDHLLHex           string    `json:"anonymousIdHLL"`
+	IdentifiedAnonymousIDHLLHex string    `json:"identifiedAnonymousIdHLL"`
+}
+
+func (t *TrackedUsersReport) MarshalJSON() ([]byte, error) {
+	t.UserIDHLLHex = hex.EncodeToString(t.UserIDHLL.ToBytes())
+	t.AnonymousIDHLLHex = hex.EncodeToString(t.AnonymousIDHLL.ToBytes())
+	t.IdentifiedAnonymousIDHLLHex = hex.EncodeToString(t.IdentifiedAnonymousIDHLL.ToBytes())
+
+	type Alias TrackedUsersReport
+	return json.Marshal(&struct {
+		*Alias
+	}{
+		Alias: (*Alias)(t),
+	})
+}

enterprise/reporting/flusher/cron_runner.go

@@ -0,0 +1,130 @@
+package flusher
+
+import (
+	"context"
+	"time"
+
+	"go.uber.org/atomic"
+	"golang.org/x/sync/errgroup"
+
+	"github.com/rudderlabs/rudder-go-kit/config"
+	"github.com/rudderlabs/rudder-go-kit/logger"
+	"github.com/rudderlabs/rudder-go-kit/stats"
+	obskit "github.com/rudderlabs/rudder-observability-kit/go/labels"
+)
+
+type Runner interface {
+	Run()
+	Stop()
+}
+
+type NOPCronRunner struct{}
+
+func (c *NOPCronRunner) Run()  {}
+func (c *NOPCronRunner) Stop() {}
+
+type CronRunner struct {
+	ctx    context.Context
+	cancel context.CancelFunc
+	g      *errgroup.Group
+
+	stats stats.Stats
+	log   logger.Logger
+
+	instanceId    string
+	table         string
+	module        string
+	flusher       *Flusher
+	sleepInterval config.ValueLoader[time.Duration]
+
+	flushTimer   stats.Measurement
+	reportingLag stats.Measurement
+
+	started atomic.Bool
+}
+
+func NewCronRunner(ctx context.Context, log logger.Logger, stats stats.Stats, conf *config.Config, flusher *Flusher, table, module string) *CronRunner {
+	sleepInterval := conf.GetReloadableDurationVar(5, time.Second, "Reporting.flusher.sleepInterval")
+	instanceId := conf.GetString("INSTANCE_ID", "1")
+
+	ctx, cancel := context.WithCancel(ctx)
+	g, ctx := errgroup.WithContext(ctx)
+
+	c := &CronRunner{
+		ctx:           ctx,
+		cancel:        cancel,
+		g:             g,
+		stats:         stats,
+		log:           log,
+		instanceId:    instanceId,
+		flusher:       flusher,
+		sleepInterval: sleepInterval,
+		table:         table,
+		module:        module,
+	}
+	c.initStats()
+
+	return c
+}
+
+func (c *CronRunner) initStats() {
+	commonTags := stats.Tags{
+		"instance": c.instanceId,
+		"table":    c.table,
+		"module":   c.module,
+	}
+	c.flushTimer = c.stats.NewTaggedStat("reporting_flusher_flush_duration_seconds", stats.TimerType, commonTags)
+	c.reportingLag = c.stats.NewTaggedStat("reporting_flusher_lag_seconds", stats.GaugeType, commonTags)
+}
+
+func (c *CronRunner) Run() {
+	c.g.Go(func() error {
+		return c.startFlushing(c.ctx)
+	})
+
+	c.started.Store(true)
+
+	if err := c.g.Wait(); err != nil {
+		c.log.Errorn("error in cron-runner", obskit.Error(err))
+	}
+}
+
+func (c *CronRunner) startFlushing(ctx context.Context) error {
+	ticker := time.NewTicker(c.sleepInterval.Load())
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			s := time.Now()
+			if err := c.flusher.Flush(ctx); err != nil {
+				c.log.Errorn("error in Flush", obskit.Error(err))
+			}
+			c.flushTimer.Since(s)
+
+			if !c.flusher.ShouldFlushAggressively(ctx) {
+				select {
+				case <-ctx.Done():
+					return ctx.Err()
+				case <-ticker.C:
+				}
+			}
+		}
+	}
+}
+
+func (c *CronRunner) Stop() {
+	c.cancel()
+	err := c.g.Wait()
+	if err != nil {
+		c.log.Errorn("error in stopping cron-runner", obskit.Error(err))
+	}
+
+	err = c.flusher.CleanUp()
+	if err != nil {
+		c.log.Errorn("error in flusher cleanup", obskit.Error(err))
+	}
+	c.started.Store(false)
+}