feat(prometheus): add histogram (#1141)

Signed-off-by: Miguel <[email protected]>

Author: migmartri

app/controlplane/cmd/wire_gen.go

@@ -212,7 +212,7 @@ func craftMiddleware(opts *Opts) []middleware.Middleware {
 			// 2.a - Set its workflow and organization in the context
 			usercontext.WithAttestationContextFromRobotAccount(opts.RobotAccountUseCase, opts.OrganizationUseCase, logHelper),
 			// 2.b - Set its API token and Robot Account as alternative to the user
-			usercontext.WithAttestationContextFromAPIToken(opts.APITokenUseCase, logHelper),
+			usercontext.WithAttestationContextFromAPIToken(opts.APITokenUseCase, opts.OrganizationUseCase, logHelper),
 		).Match(requireRobotAccountMatcher()).Build(),
 	)
 

app/controlplane/internal/server/grpc.go

@@ -53,6 +53,7 @@ type AttestationService struct {
 	casMappingUseCase       *biz.CASMappingUseCase
 	referrerUseCase         *biz.ReferrerUseCase
 	orgUseCase              *biz.OrganizationUseCase
+	prometheusUseCase       *biz.PrometheusUseCase
 }
 
 type NewAttestationServiceOpts struct {
@@ -68,6 +69,7 @@ type NewAttestationServiceOpts struct {
 	CASMappingUseCase  *biz.CASMappingUseCase
 	ReferrerUC         *biz.ReferrerUseCase
 	OrgUC              *biz.OrganizationUseCase
+	PromUC             *biz.PrometheusUseCase
 	Opts               []NewOpt
 }
 
@@ -86,6 +88,7 @@ func NewAttestationService(opts *NewAttestationServiceOpts) *AttestationService
 		casMappingUseCase:       opts.CASMappingUseCase,
 		referrerUseCase:         opts.ReferrerUC,
 		orgUseCase:              opts.OrgUC,
+		prometheusUseCase:       opts.PromUC,
 	}
 }
 
@@ -276,6 +279,9 @@ func (s *AttestationService) Store(ctx context.Context, req *cpAPI.AttestationSe
 		return nil, handleUseCaseErr(err, s.log)
 	}
 
+	// Record the attestation in the prometheus registry
+	_ = s.prometheusUseCase.ObserveAttestationIfNeeded(ctx, wRun, biz.WorkflowRunSuccess)
+
 	return &cpAPI.AttestationServiceStoreResponse{
 		Result: &cpAPI.AttestationServiceStoreResponse_Result{Digest: digest},
 	}, nil
@@ -306,6 +312,16 @@ func (s *AttestationService) Cancel(ctx context.Context, req *cpAPI.AttestationS
 		return nil, err
 	}
 
+	wRun, err := s.wrUseCase.GetByIDInOrgOrPublic(ctx, robotAccount.OrgID, req.WorkflowRunId)
+	if err != nil {
+		return nil, handleUseCaseErr(err, s.log)
+	} else if wRun == nil {
+		return nil, errors.NotFound("not found", "workflow run not found")
+	}
+
+	// Record the attestation in the prometheus registry
+	_ = s.prometheusUseCase.ObserveAttestationIfNeeded(ctx, wRun, status)
+
 	return &cpAPI.AttestationServiceCancelResponse{}, nil
 }
 

app/controlplane/internal/service/attestation.go

@@ -103,7 +103,7 @@ func WithCurrentAPITokenAndOrgMiddleware(apiTokenUC *biz.APITokenUseCase, orgUC
 }
 
 // WithAttestationContextFromAPIToken injects the API-Token, organization + robot account to the context
-func WithAttestationContextFromAPIToken(apiTokenUC *biz.APITokenUseCase, logger *log.Helper) middleware.Middleware {
+func WithAttestationContextFromAPIToken(apiTokenUC *biz.APITokenUseCase, orgUC *biz.OrganizationUseCase, logger *log.Helper) middleware.Middleware {
 	return func(handler middleware.Handler) middleware.Handler {
 		return func(ctx context.Context, req interface{}) (interface{}, error) {
 			authInfo, ok := attjwtmiddleware.FromJWTAuthContext(ctx)
@@ -138,6 +138,11 @@ func WithAttestationContextFromAPIToken(apiTokenUC *biz.APITokenUseCase, logger
 				return nil, fmt.Errorf("error extracting organization from APIToken: %w", err)
 			}
 
+			ctx, err = setCurrentOrgAndAPIToken(ctx, apiTokenUC, orgUC, tokenID)
+			if err != nil {
+				return nil, fmt.Errorf("error setting current org and user: %w", err)
+			}
+
 			logger.Infow("msg", "[authN] processed credentials", "id", tokenID, "type", "API-token")
 
 			return handler(ctx, req)

app/controlplane/internal/usercontext/apitoken_middleware.go

@@ -50,6 +50,7 @@ var ProviderSet = wire.NewSet(
 	NewAttestationStateUseCase,
 	NewChainloopSigningUseCase,
 	NewPrometheusUseCase,
+	wire.Bind(new(PromObservable), new(*PrometheusUseCase)),
 	wire.Struct(new(NewIntegrationUseCaseOpts), "*"),
 	wire.Struct(new(NewUserUseCaseParams), "*"),
 )

app/controlplane/pkg/biz/biz.go

@@ -16,9 +16,13 @@
 package biz
 
 import (
+	"context"
+	"time"
+
 	conf "github.com/chainloop-dev/chainloop/app/controlplane/internal/conf/controlplane/config/v1"
-	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/metrics/prometheus"
+	chainloopprometheus "github.com/chainloop-dev/chainloop/app/controlplane/pkg/metrics/prometheus"
 	"github.com/chainloop-dev/chainloop/app/controlplane/pkg/metrics/prometheus/registry"
+	"github.com/prometheus/client_golang/prometheus"
 
 	"github.com/go-kratos/kratos/v2/log"
 )
@@ -30,7 +34,7 @@ type PrometheusUseCase struct {
 	orgUseCase        *OrganizationUseCase
 	orgMetricsUseCase *OrgMetricsUseCase
 	// Other
-	registryManager *prometheus.ChainloopRegistryManager
+	registryManager *chainloopprometheus.ChainloopRegistryManager
 }
 
 // NewPrometheusUseCase creates a new PrometheusUseCase
@@ -48,15 +52,53 @@ func NewPrometheusUseCase(conf []*conf.PrometheusIntegrationSpec, orgUseCase *Or
 }
 
 // loadPrometheusRegistries loads the prometheus registries from the configuration
-func loadPrometheusRegistries(conf []*conf.PrometheusIntegrationSpec, useCase *OrgMetricsUseCase, logger log.Logger) *prometheus.ChainloopRegistryManager {
-	rm := prometheus.NewChainloopRegistryManager()
+func loadPrometheusRegistries(conf []*conf.PrometheusIntegrationSpec, useCase *OrgMetricsUseCase, logger log.Logger) *chainloopprometheus.ChainloopRegistryManager {
+	manager := chainloopprometheus.NewChainloopRegistryManager()
 
 	for _, spec := range conf {
 		reg := registry.NewPrometheusRegistry(spec.GetOrgName(), useCase, logger)
-		rm.AddRegistry(reg)
+		manager.AddRegistry(reg)
+	}
+
+	return manager
+}
+
+// Record an attestation if the run exists and there is a registry for the organization
+func (uc *PrometheusUseCase) ObserveAttestationIfNeeded(ctx context.Context, run *WorkflowRun, status WorkflowRunStatus) bool {
+	if run == nil || run.Workflow == nil {
+		return false
+	}
+
+	workflow := run.Workflow
+	orgID := workflow.OrgID
+
+	org, err := uc.orgUseCase.FindByID(ctx, orgID.String())
+	if err != nil {
+		return false
+	}
+
+	if !uc.OrganizationHasRegistry(org.Name) {
+		return false
+	}
+
+	err = uc.observeAttestation(org.Name, workflow.Name, status, run.RunnerType, run.CreatedAt)
+	return err == nil
+}
+
+// OrganizationHasRegistry checks if an organization has a registry
+func (uc *PrometheusUseCase) observeAttestation(orgName, wfName string, status WorkflowRunStatus, runnerType string, startTime *time.Time) error {
+	if orgName == "" || wfName == "" || status == "" || startTime == nil {
+		return NewErrValidationStr("orgName, wfName, and state must be non-empty")
+	}
+
+	reg := uc.GetRegistryByOrganizationName(orgName)
+	if reg == nil {
+		return NewErrNotFound("registry not found for organization")
 	}
 
-	return rm
+	duration := time.Since(*startTime).Seconds()
+	reg.WorkflowRunDurationSeconds.With(prometheus.Labels{"org": orgName, "workflow": wfName, "status": string(status), "runner": runnerType}).Observe(duration)
+	return nil
 }
 
 // OrganizationHasRegistry checks if an organization has a registry

app/controlplane/pkg/biz/mocks/PromObservable.go

@@ -62,6 +62,7 @@ type TestingUseCases struct {
 	WorkflowContract       *biz.WorkflowContractUseCase
 	Workflow               *biz.WorkflowUseCase
 	WorkflowRun            *biz.WorkflowRunUseCase
+	Prometheus             *biz.PrometheusUseCase
 	User                   *biz.UserUseCase
 	RobotAccount           *biz.RobotAccountUseCase
 	RegisteredIntegrations sdk.AvailablePlugins
@@ -196,6 +197,10 @@ func NewConfData(db *TestDatabase, t *testing.T) *conf.Data {
 	}
 }
 
+func NewPromSpec() []*conf.PrometheusIntegrationSpec {
+	return []*conf.PrometheusIntegrationSpec{}
+}
+
 func NewDataConfig(in *conf.Data) *data.NewConfig {
 	return &data.NewConfig{
 		Driver: in.Database.Driver,

app/controlplane/pkg/biz/prometheus.go

@@ -47,6 +47,7 @@ func WireTestData(*TestDatabase, *testing.T, log.Logger, credentials.ReaderWrite
 			wire.Struct(new(TestingRepos), "*"),
 			NewConfData,
 			NewDataConfig,
+			NewPromSpec,
 			authz.NewDatabaseEnforcer,
 			wire.FieldsOf(new(*conf.Data), "Database"),
 		),