telemetry: refactor annotator interface (#2526)

## Summary of Changes
This refactors the Annotator interface of the flow enricher to make it
less clumsy when adding annotators that require other external
dependencies. Prior, external dependencies were passed as part of the
Init method into all annotators whether they needed them or not, which
then requires touching all Init method signatures when a single
dependency is added. These dependencies are now moved into the annotator
constructors themselves.

## Testing Verification
Existing unit/integration tests should pass. There's no functional
change that should be introduced.

Author: packethog

telemetry/flow-enricher/cmd/flow-enricher/main.go

@@ -123,15 +123,16 @@ func main() {
 	rpcClient := rpc.NewWithRetries(networkConfig.LedgerPublicRPCURL, nil)
 	serviceabilityClient := serviceability.New(rpcClient, networkConfig.ServiceabilityProgramID)
 
-	enricherOpts := []enricher.EnricherOption{
+	e := enricher.NewEnricher(
 		enricher.WithClickhouseWriter(chWriter),
 		enricher.WithFlowConsumer(flowConsumer),
 		enricher.WithLogger(logger),
 		enricher.WithEnricherMetrics(enricher.NewEnricherMetrics(reg)),
 		enricher.WithServiceabilityFetcher(serviceabilityClient),
-		enricher.WithServiceabilityFetchInterval(10 * time.Second),
-	}
-	enricher := enricher.NewEnricher(enricherOpts...)
+		enricher.WithServiceabilityFetchInterval(10*time.Second),
+	)
+
+	e.AddAnnotator(enricher.NewServiceabilityAnnotator(e.ServiceabilityData))
 
 	// start prometheus
 	go func() {
@@ -141,7 +142,7 @@ func main() {
 	}()
 
 	logger.Info("starting enricher...")
-	if err := enricher.Run(ctx); err != nil {
+	if err := e.Run(ctx); err != nil {
 		logger.Error("error while running enricher", "error", err)
 		os.Exit(1)
 	}

telemetry/flow-enricher/internal/flow-enricher/enricher.go

@@ -3,8 +3,8 @@
 // records in protobuf format, enriches the flow with additional information from
 // each annotator, and writes the flows as a batch to clickhouse.
 //
-// Annotators must be registered in the RegisterAnnotators method of the enricher
-// and must implement the Annotator interface.
+// Annotators are added via the AddAnnotator method and
+// must implement the Annotator interface.
 package enricher
 
 import (
@@ -106,6 +106,11 @@ func NewEnricher(opts ...EnricherOption) *Enricher {
 	return e
 }
 
+// AddAnnotator adds an annotator to the enricher after construction.
+func (e *Enricher) AddAnnotator(a Annotator) {
+	e.annotators = append(e.annotators, a)
+}
+
 // Run starts the enricher instance and begins processing flow records.
 func (e *Enricher) Run(ctx context.Context) error {
 	if e.flowConsumer == nil {
@@ -128,8 +133,8 @@ func (e *Enricher) Run(ctx context.Context) error {
 
 	go e.fetchServiceabilityData(ctx)
 
-	// some annotators depends on serviceability data, so we need to register them after
-	if err := e.RegisterAnnotators(ctx); err != nil {
+	// initialize annotators before starting enrichment
+	if err := e.initAnnotators(ctx); err != nil {
 		return fmt.Errorf("error while initializing annotators: %v", err)
 	}
 
@@ -202,7 +207,9 @@ func (e *Enricher) fetchServiceabilityData(ctx context.Context) {
 	}
 }
 
-func (e *Enricher) serviceabilityData() serviceability.ProgramData {
+// ServiceabilityData returns a copy of the current serviceability program data.
+// This method is safe for concurrent use.
+func (e *Enricher) ServiceabilityData() serviceability.ProgramData {
 	e.programDataMutex.Lock()
 	defer e.programDataMutex.Unlock()
 	if e.programData == nil {
@@ -211,21 +218,19 @@ func (e *Enricher) serviceabilityData() serviceability.ProgramData {
 	return *e.programData
 }
 
+// Annotator defines the interface for flow enrichment annotators.
+// Dependencies should be passed via the annotator's constructor.
 type Annotator interface {
-	Init(context.Context, func() serviceability.ProgramData) error
+	// Init initializes the annotator. Called once before any Annotate calls.
+	Init(context.Context) error
 	Annotate(*FlowSample) error
 	String() string
 }
 
-// RegisterAnnotators initializes a set of annotators for use during enrichment.
-// Annotators must implement the Annotator interface.
-func (e *Enricher) RegisterAnnotators(ctx context.Context) error {
-	e.annotators = []Annotator{
-		NewServiceabilityAnnotator(),
-	}
-
+// initAnnotators initializes all registered annotators.
+func (e *Enricher) initAnnotators(ctx context.Context) error {
 	for _, a := range e.annotators {
-		if err := a.Init(ctx, e.serviceabilityData); err != nil {
+		if err := a.Init(ctx); err != nil {
 			return fmt.Errorf("error initializing annotator %s: %v", a.String(), err)
 		}
 	}

telemetry/flow-enricher/internal/flow-enricher/enricher_integration_test.go

@@ -246,6 +246,8 @@ func TestFlowEnrichment(t *testing.T) {
 		WithEnricherMetrics(NewEnricherMetrics(reg)),
 		WithServiceabilityFetcher(mockServiceability),
 	)
+	enricher.AddAnnotator(NewServiceabilityAnnotator(enricher.ServiceabilityData))
+
 	go func() {
 		if err := enricher.Run(ctx); err != nil {
 			logger.Error("error during enrichment", "error", err)

telemetry/flow-enricher/internal/flow-enricher/serviceability.go

@@ -21,14 +21,14 @@ type ServiceabilityAnnotator struct {
 	mu             sync.RWMutex
 }
 
-func NewServiceabilityAnnotator() *ServiceabilityAnnotator {
+func NewServiceabilityAnnotator(getProgramData func() serviceability.ProgramData) *ServiceabilityAnnotator {
 	return &ServiceabilityAnnotator{
-		name: "serviceability annotator",
+		name:           "serviceability annotator",
+		getProgramData: getProgramData,
 	}
 }
 
-func (s *ServiceabilityAnnotator) Init(ctx context.Context, getProgramData func() serviceability.ProgramData) error {
-	s.getProgramData = getProgramData
+func (s *ServiceabilityAnnotator) Init(ctx context.Context) error {
 	s.updateServiceabilityCache()
 
 	go func() {