feat(metrics): add metrics collection mechanism

.env.example

@@ -142,5 +142,12 @@ SENTRY_DSN=
 # OFFLOADING_SAVE_POINTS=1000
 # OFFLOADING_WRITES_PER_SEC=200
 
+# Configure metrics collection.
+# By default, the metrics collection is enabled
+# DISABLED_TELEMETRY=false
+# METRICS_COLLECTION_JOB_INTERVAL=1h
+
 # Abort and exit if license cannot be validated.
 KILL_IF_READ_LICENSE_ERROR=false
+
+# BIGQUERY_PROJECT_ID=

api/handle_metrics_ingestion.go

@@ -0,0 +1,34 @@
+package api
+
+import (
+	"net/http"
+
+	"github.com/checkmarble/marble-backend/dto"
+	"github.com/checkmarble/marble-backend/usecases"
+	"github.com/checkmarble/marble-backend/utils"
+	"github.com/gin-gonic/gin"
+)
+
+// Handle metrics ingestion from the metrics collection worker.
+func handleMetricsIngestion(uc usecases.Usecases) func(c *gin.Context) {
+	return func(c *gin.Context) {
+		logger := utils.LoggerFromContext(c.Request.Context())
+		var metricsCollectionDto dto.MetricsCollectionDto
+		if err := c.ShouldBindJSON(&metricsCollectionDto); err != nil {
+			c.Status(http.StatusBadRequest)
+			logger.WarnContext(c.Request.Context(), "Failed to bind metrics collection", "error", err.Error())
+			return
+		}
+
+		metricsCollection := dto.AdaptMetricsCollection(metricsCollectionDto)
+
+		usecase := uc.NewMetricsIngestionUsecase()
+		err := usecase.IngestMetrics(c.Request.Context(), metricsCollection)
+		if presentError(c.Request.Context(), c, err) {
+			return
+		}
+
+		// Success response
+		c.Status(http.StatusOK)
+	}
+}

api/routes.go

@@ -6,6 +6,7 @@ import (
 	"net/url"
 	"time"
 
+	"github.com/checkmarble/marble-backend/infra"
 	"github.com/checkmarble/marble-backend/pubapi"
 	pubapiv1 "github.com/checkmarble/marble-backend/pubapi/v1"
 	"github.com/checkmarble/marble-backend/utils"
@@ -38,11 +39,15 @@ func addRoutes(r *gin.Engine, conf Configuration, uc usecases.Usecases, auth uti
 	r.GET("/health", tom, handleHealth(uc))
 	r.GET("/version", tom, handleVersion(uc))
 	r.POST("/token", tom, tokenHandler.GenerateToken)
-	r.GET("/validate-license/*license_key", tom, handleValidateLicense(uc))
 	r.GET("/config", tom, handleGetConfig(uc, conf))
 	r.GET("/is-sso-available", tom, handleIsSSOEnabled(uc))
 	r.GET("/signup-status", tom, handleSignupStatus(uc))
 
+	if infra.IsMarbleSaasProject() {
+		r.GET("/validate-license/*license_key", tom, handleValidateLicense(uc))
+		r.POST("/metrics", tom, handleMetricsIngestion(uc))
+	}
+
 	// Public API initialization
 	{
 		cfg := pubapi.Config{

cmd/server.go

@@ -110,6 +110,12 @@ func RunServer(config CompiledConfig) error {
 		LicenseKey:             utils.GetEnv("LICENSE_KEY", ""),
 		KillIfReadLicenseError: utils.GetEnv("KILL_IF_READ_LICENSE_ERROR", false),
 	}
+	bigQueryConfig := infra.BigQueryConfig{
+		ProjectID:      utils.GetEnv("BIGQUERY_PROJECT_ID", gcpConfig.ProjectId),
+		MetricsDataset: utils.GetEnv("BIGQUERY_METRICS_DATASET", infra.MetricsDataset),
+		MetricsTable:   utils.GetEnv("BIGQUERY_METRICS_TABLE", infra.MetricsTable),
+	}
+
 	serverConfig := struct {
 		batchIngestionMaxSize            int
 		caseManagerBucket                string
@@ -193,6 +199,16 @@ func RunServer(config CompiledConfig) error {
 		return err
 	}
 
+	var bigQueryInfra *infra.BigQueryInfra
+	if infra.IsMarbleSaasProject() {
+		bigQueryInfra, err = infra.InitializeBigQueryInfra(ctx, bigQueryConfig)
+		if err != nil {
+			utils.LogAndReportSentryError(ctx, err)
+			return err
+		}
+		defer bigQueryInfra.Close()
+	}
+
 	repositories := repositories.NewRepositories(
 		pool,
 		gcpConfig,
@@ -206,6 +222,7 @@ func RunServer(config CompiledConfig) error {
 		repositories.WithClientDbConfig(clientDbConfig),
 		repositories.WithTracerProvider(telemetryRessources.TracerProvider),
 		repositories.WithRiverClient(riverClient),
+		repositories.WithBigQueryInfra(bigQueryInfra),
 	)
 
 	uc := usecases.NewUsecases(repositories,
@@ -245,7 +262,8 @@ func RunServer(config CompiledConfig) error {
 
 	deps := api.InitDependencies(ctx, apiConfig, pool, marbleJwtSigningKey)
 
-	router := api.InitRouterMiddlewares(ctx, apiConfig, apiConfig.DisableSegment, deps.SegmentClient, telemetryRessources)
+	router := api.InitRouterMiddlewares(ctx, apiConfig, apiConfig.DisableSegment,
+		deps.SegmentClient, telemetryRessources)
 	server := api.NewServer(router, apiConfig, uc, deps.Authentication, deps.TokenHandler, logger)
 
 	notify, stop := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM)

cmd/worker.go

@@ -3,9 +3,11 @@ package cmd
 import (
 	"context"
 	"log/slog"
+	"maps"
 	"net/http"
 	"os"
 	"os/signal"
+	"slices"
 	"syscall"
 	"time"
 
@@ -14,6 +16,7 @@ import (
 	"github.com/checkmarble/marble-backend/models"
 	"github.com/checkmarble/marble-backend/repositories"
 	"github.com/checkmarble/marble-backend/usecases"
+	"github.com/checkmarble/marble-backend/usecases/scheduled_execution"
 	"github.com/checkmarble/marble-backend/utils"
 
 	"github.com/cockroachdb/errors"
@@ -97,6 +100,14 @@ func RunTaskQueue(apiVersion string) error {
 
 	offloadingConfig.ValidateAndFix(ctx)
 
+	metricCollectionConfig := infra.MetricCollectionConfig{
+		Disabled:            utils.GetEnv("DISABLED_TELEMETRY", false),
+		JobInterval:         utils.GetEnvDuration("METRICS_COLLECTION_JOB_INTERVAL", 1*time.Hour),
+		MetricsIngestionURL: utils.GetEnv("METRICS_INGESTION_URL", ""),
+		FallbackDuration:    utils.GetEnvDuration("METRICS_FALLBACK_DURATION", 30*24*time.Hour),
+	}
+	metricCollectionConfig.Configure()
+
 	infra.SetupSentry(workerConfig.sentryDsn, workerConfig.env, apiVersion)
 	defer sentry.Flush(3 * time.Second)
 
@@ -154,12 +165,23 @@ func RunTaskQueue(apiVersion string) error {
 		return err
 	}
 
+	// For non-org
+	nonOrgQueues := make(map[string]river.QueueConfig)
+	globalPeriodics := []*river.PeriodicJob{}
+
+	if !metricCollectionConfig.Disabled {
+		metricQueue := usecases.QueueMetrics()
+		maps.Copy(nonOrgQueues, metricQueue)
+		globalPeriodics = append(globalPeriodics,
+			scheduled_execution.NewMetricsCollectionPeriodicJob(metricCollectionConfig))
+	}
+
+	maps.Copy(queues, nonOrgQueues)
+
 	// Periodics always contain the per-org tasks retrieved above. Add other, non-organization-scoped periodics below
 	periodics := append(
 		orgPeriodics,
-		[]*river.PeriodicJob{
-			// Add periodic jobs here
-		}...,
+		globalPeriodics...,
 	)
 
 	riverClient, err = river.NewClient(riverpgxv5.New(pool), &river.Config{
@@ -183,13 +205,21 @@ func RunTaskQueue(apiVersion string) error {
 		return err
 	}
 
+	// Ensure that all global queues are active.
+	if err := ensureGlobalQueuesAreActive(ctx, riverClient, nonOrgQueues); err != nil {
+		utils.LogAndReportSentryError(ctx, err)
+		return err
+	}
+
 	uc := usecases.NewUsecases(repositories,
 		usecases.WithIngestionBucketUrl(workerConfig.ingestionBucketUrl),
 		usecases.WithOffloading(offloadingConfig),
 		usecases.WithFailedWebhooksRetryPageSize(workerConfig.failedWebhooksRetryPageSize),
 		usecases.WithLicense(license),
 		usecases.WithConvoyServer(convoyConfiguration.APIUrl),
 		usecases.WithOpensanctions(openSanctionsConfig.IsSet()),
+		usecases.WithApiVersion(apiVersion),
+		usecases.WithMetricsCollectionConfig(metricCollectionConfig),
 	)
 	adminUc := jobs.GenerateUsecaseWithCredForMarbleAdmin(ctx, uc)
 	river.AddWorker(workers, adminUc.NewAsyncDecisionWorker())
@@ -204,6 +234,9 @@ func RunTaskQueue(apiVersion string) error {
 	if offloadingConfig.Enabled {
 		river.AddWorker(workers, adminUc.NewOffloadingWorker())
 	}
+	if !metricCollectionConfig.Disabled {
+		river.AddWorker(workers, uc.NewMetricsCollectionWorker(licenseConfig))
+	}
 
 	if err := riverClient.Start(ctx); err != nil {
 		utils.LogAndReportSentryError(ctx, err)
@@ -226,7 +259,9 @@ func RunTaskQueue(apiVersion string) error {
 	logger.InfoContext(ctx, "starting worker", slog.String("version", apiVersion))
 
 	// Asynchronously keep the task queue workers up to date with the orgs in the database
-	taskQueueWorker := uc.NewTaskQueueWorker(riverClient)
+	taskQueueWorker := uc.NewTaskQueueWorker(riverClient,
+		slices.Collect(maps.Keys(nonOrgQueues)),
+	)
 	go taskQueueWorker.RefreshQueuesFromOrgIds(ctx)
 
 	// Start the cron jobs using the old entrypoint.
@@ -295,3 +330,34 @@ func cleanStop(ctx context.Context, sigintOrTerm chan os.Signal, riverClient *ri
 	}
 	// hard stop succeeded
 }
+
+// Ensure that all global queues are active in river DB, i.e. not paused. This function is more a safety net.
+// River stores the state of the queues in the DB. If a queue exists but is paused, River client will not resume it
+// We have some function which can pause queues, so we need to ensure that all global queues are active
+// cf: riverClient.QueueResume and riverClient.QueuePause
+func ensureGlobalQueuesAreActive(ctx context.Context, riverClient *river.Client[pgx.Tx],
+	nonOrgQueues map[string]river.QueueConfig,
+) error {
+	logger := utils.LoggerFromContext(ctx)
+
+	for queueName := range nonOrgQueues {
+		queueState, err := riverClient.QueueGet(ctx, queueName)
+		if err != nil {
+			if errors.Is(err, river.ErrNotFound) {
+				// Queue will be created when River starts, skip
+				continue
+			}
+			return err
+		}
+
+		// If the queue exists and is paused, resume it
+		if queueState.PausedAt != nil {
+			logger.InfoContext(ctx, "Resuming global queue at startup", "queue", queueName)
+			if err := riverClient.QueueResume(ctx, queueName, &river.QueuePauseOpts{}); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}

dto/metrics_dto.go

@@ -0,0 +1,76 @@
+package dto
+
+import (
+	"time"
+
+	"github.com/checkmarble/marble-backend/models"
+	"github.com/checkmarble/marble-backend/pure_utils"
+	"github.com/google/uuid"
+)
+
+// Be careful when changing this struct, it is used as input and output in the API.
+type MetricDataDto struct {
+	Name      string    `json:"name" binding:"required"`
+	Numeric   *float64  `json:"numeric,omitempty"`
+	Text      *string   `json:"text,omitempty"`
+	Timestamp time.Time `json:"timestamp" binding:"required"`
+	OrgID     *string   `json:"org_id,omitempty"`
+	From      time.Time `json:"from" binding:"required"`
+	To        time.Time `json:"to" binding:"required"`
+}
+
+// Be careful when changing this struct, it is used as input and output in the API.
+type MetricsCollectionDto struct {
+	CollectionID uuid.UUID       `json:"collection_id" binding:"required"`
+	Timestamp    time.Time       `json:"timestamp" binding:"required"`
+	Metrics      []MetricDataDto `json:"metrics" binding:"required"`
+	Version      string          `json:"version" binding:"required"`
+	DeploymentID uuid.UUID       `json:"deployment_id" binding:"required"`
+	LicenseKey   *string         `json:"license_key,omitempty"`
+}
+
+func AdaptMetricDataDto(metricData models.MetricData) MetricDataDto {
+	return MetricDataDto{
+		Name:      metricData.Name,
+		Numeric:   metricData.Numeric,
+		Text:      metricData.Text,
+		Timestamp: metricData.Timestamp,
+		OrgID:     metricData.OrgID,
+		From:      metricData.From,
+		To:        metricData.To,
+	}
+}
+
+func AdaptMetricsCollectionDto(metricsCollection models.MetricsCollection) MetricsCollectionDto {
+	return MetricsCollectionDto{
+		CollectionID: metricsCollection.CollectionID,
+		Timestamp:    metricsCollection.Timestamp,
+		Metrics:      pure_utils.Map(metricsCollection.Metrics, AdaptMetricDataDto),
+		Version:      metricsCollection.Version,
+		DeploymentID: metricsCollection.DeploymentID,
+		LicenseKey:   metricsCollection.LicenseKey,
+	}
+}
+
+func AdaptMetricData(metricDataDto MetricDataDto) models.MetricData {
+	return models.MetricData{
+		Name:      metricDataDto.Name,
+		Numeric:   metricDataDto.Numeric,
+		Text:      metricDataDto.Text,
+		Timestamp: metricDataDto.Timestamp,
+		OrgID:     metricDataDto.OrgID,
+		From:      metricDataDto.From,
+		To:        metricDataDto.To,
+	}
+}
+
+func AdaptMetricsCollection(metricsCollectionDto MetricsCollectionDto) models.MetricsCollection {
+	return models.MetricsCollection{
+		CollectionID: metricsCollectionDto.CollectionID,
+		Timestamp:    metricsCollectionDto.Timestamp,
+		Metrics:      pure_utils.Map(metricsCollectionDto.Metrics, AdaptMetricData),
+		Version:      metricsCollectionDto.Version,
+		DeploymentID: metricsCollectionDto.DeploymentID,
+		LicenseKey:   metricsCollectionDto.LicenseKey,
+	}
+}