From 1d643f71c43972e69e093b781e9136c46cac89a0 Mon Sep 17 00:00:00 2001
From: ahmed <ahmed.zahran@trufflesec.com>
Date: Tue, 13 Aug 2024 13:35:48 -0400
Subject: [PATCH] new metrics for scan job reports

---
 main.go                          | 29 +++++++++++++++-
 metrics.go                       | 59 ++++++++++++++++++++++++++++++++
 pkg/sources/job_progress_hook.go |  1 +
 pkg/sources/metrics.go           |  7 ++++
 4 files changed, 95 insertions(+), 1 deletion(-)
 create mode 100644 metrics.go

diff --git a/main.go b/main.go
index 0dbf1f8d1..5f1874aea 100644
--- a/main.go
+++ b/main.go
@@ -13,6 +13,7 @@ import (
 	"strings"
 	"sync"
 	"syscall"
+	"time"
 
 	"github.com/alecthomas/kingpin/v2"
 	"github.com/felixge/fgprof"
@@ -542,7 +543,23 @@ func runSingleScan(ctx context.Context, cmd string, cfg engine.Config) (metrics,
 				}
 			}()
 
+			startTime := time.Now()
+			var totalBytes int64
+
 			for metrics := range finishedMetrics {
+				reportGenerationTotal.Inc()
+
+				unitType := "unknown"
+				if metrics.Unit != nil {
+					_, unitKind := metrics.Unit.SourceUnitID()
+					unitType = string(unitKind)
+				}
+				unitMetricsTotal.WithLabelValues(unitType).Inc()
+
+				if len(metrics.Errors) > 0 {
+					unitMetricsWithErrors.WithLabelValues(unitType).Inc()
+				}
+
 				metrics.Errors = common.ExportErrors(metrics.Errors...)
 				details, err := json.Marshal(map[string]any{
 					"version": 1,
@@ -550,12 +567,22 @@ func runSingleScan(ctx context.Context, cmd string, cfg engine.Config) (metrics,
 				})
 				if err != nil {
 					ctx.Logger().Error(err, "error marshalling job details")
+					reportGenerationErrors.WithLabelValues("marshal").Inc()
 					continue
 				}
-				if _, err := jobReportWriter.Write(append(details, '\n')); err != nil {
+
+				details = append(details, '\n')
+				totalBytes += int64(len(details))
+
+				if _, err := jobReportWriter.Write(details); err != nil {
 					ctx.Logger().Error(err, "error writing to file")
+					reportGenerationErrors.WithLabelValues("write").Inc()
 				}
 			}
+
+			duration := time.Since(startTime)
+			reportGenerationDuration.Observe(duration.Seconds())
+			reportFileSize.Set(float64(totalBytes))
 		}()
 	}
 
diff --git a/metrics.go b/metrics.go
new file mode 100644
index 000000000..f8ad05f1b
--- /dev/null
+++ b/metrics.go
@@ -0,0 +1,59 @@
+package main
+
+import (
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+
+	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
+)
+
+var (
+	reportGenerationTotal = promauto.NewCounter(prometheus.CounterOpts{
+		Namespace: common.MetricsNamespace,
+		Subsystem: common.MetricsSubsystem,
+		Name:      "report_generation_total",
+		Help:      "Total number of report generation attempts",
+	})
+
+	reportGenerationErrors = promauto.NewCounterVec(prometheus.CounterOpts{
+		Namespace: common.MetricsNamespace,
+		Subsystem: common.MetricsSubsystem,
+		Name:      "report_generation_errors",
+		Help:      "Total number of errors encountered during report generation, by error type",
+	},
+		[]string{"error_type"},
+	)
+
+	unitMetricsTotal = promauto.NewCounterVec(prometheus.CounterOpts{
+		Namespace: common.MetricsNamespace,
+		Subsystem: common.MetricsSubsystem,
+		Name:      "unit_metrics_total",
+		Help:      "Total number of UnitMetrics generated, by unit type",
+	},
+		[]string{"unit_type"},
+	)
+
+	unitMetricsWithErrors = promauto.NewCounterVec(prometheus.CounterOpts{
+		Namespace: common.MetricsNamespace,
+		Subsystem: common.MetricsSubsystem,
+		Name:      "unit_metrics_with_errors_total",
+		Help:      "Total number of UnitMetrics containing errors, by unit type",
+	},
+		[]string{"unit_type"},
+	)
+
+	reportFileSize = promauto.NewGauge(prometheus.GaugeOpts{
+		Namespace: common.MetricsNamespace,
+		Subsystem: common.MetricsSubsystem,
+		Name:      "report_file_size_bytes",
+		Help:      "Size of the generated report file in bytes",
+	})
+
+	reportGenerationDuration = promauto.NewHistogram(prometheus.HistogramOpts{
+		Namespace: common.MetricsNamespace,
+		Subsystem: common.MetricsSubsystem,
+		Name:      "report_generation_duration_seconds",
+		Help:      "Time taken to generate a report",
+		Buckets:   prometheus.ExponentialBuckets(0.1, 2, 10),
+	})
+)
diff --git a/pkg/sources/job_progress_hook.go b/pkg/sources/job_progress_hook.go
index 6f94feff2..43d52b0fa 100644
--- a/pkg/sources/job_progress_hook.go
+++ b/pkg/sources/job_progress_hook.go
@@ -39,6 +39,7 @@ func NewUnitHook(ctx context.Context, opts ...UnitHookOpt) (*UnitHook, <-chan Un
 		logBackPressure: func() {
 			once.Do(func() {
 				ctx.Logger().Info("back pressure detected in unit hook")
+				backPressureOccurrences.Inc()
 			})
 		},
 	}
diff --git a/pkg/sources/metrics.go b/pkg/sources/metrics.go
index 394771dc6..ad829ecef 100644
--- a/pkg/sources/metrics.go
+++ b/pkg/sources/metrics.go
@@ -21,4 +21,11 @@ var (
 		Name:      "hooks_channel_size",
 		Help:      "Total number of metrics waiting in the finished channel.",
 	}, nil)
+
+	backPressureOccurrences = promauto.NewCounter(prometheus.CounterOpts{
+		Namespace: common.MetricsNamespace,
+		Subsystem: common.MetricsSubsystem,
+		Name:      "back_pressure_occurrences_total",
+		Help:      "Total number of times back pressure was detected",
+	})
 )