github.com/castai/kvisor@v1.7.1-0.20240516114728-b3572a2607b5/cmd/agent/daemon/enrichment/service.go (about)

     1  package enrichment
     2  
     3  import (
     4  	"context"
     5  	"runtime/debug"
     6  
     7  	castpb "github.com/castai/kvisor/api/v1/runtime"
     8  	"github.com/castai/kvisor/pkg/ebpftracer/types"
     9  	"github.com/castai/kvisor/pkg/logging"
    10  	"github.com/castai/kvisor/pkg/metrics"
    11  )
    12  
    13  type EnrichRequest struct {
    14  	Event     *castpb.Event
    15  	EbpfEvent *types.Event
    16  }
    17  
    18  type EventEnricher interface {
    19  	// Enrich will add additional data to the provided Event.
    20  	Enrich(context.Context, *EnrichRequest)
    21  
    22  	// EventsTypes returns a slice of event types, this enricher reacts to.
    23  	EventTypes() []castpb.EventType
    24  }
    25  
    26  type Config struct {
    27  	WorkerCount    int
    28  	EventEnrichers []EventEnricher
    29  }
    30  
    31  func NewService(log *logging.Logger, cfg Config) *Service {
    32  	return &Service{
    33  		log:            log.WithField("component", "enrichment"),
    34  		eventsQueue:    make(chan *EnrichRequest, 1000),
    35  		outQueue:       make(chan *castpb.Event, 1000),
    36  		cfg:            cfg,
    37  		eventEnrichers: groupEventEnrichers(cfg.EventEnrichers),
    38  	}
    39  }
    40  
    41  func groupEventEnrichers(enrichers []EventEnricher) map[castpb.EventType][]EventEnricher {
    42  	result := map[castpb.EventType][]EventEnricher{}
    43  
    44  	for _, enricher := range enrichers {
    45  		for _, t := range enricher.EventTypes() {
    46  			result[t] = append(result[t], enricher)
    47  		}
    48  	}
    49  
    50  	return result
    51  }
    52  
    53  type Service struct {
    54  	log         *logging.Logger
    55  	cfg         Config
    56  	eventsQueue chan *EnrichRequest
    57  	outQueue    chan *castpb.Event
    58  
    59  	eventEnrichers map[castpb.EventType][]EventEnricher
    60  }
    61  
    62  func (s *Service) Events() <-chan *castpb.Event {
    63  	return s.outQueue
    64  }
    65  
    66  func (s *Service) Enqueue(e *EnrichRequest) bool {
    67  	if _, found := s.eventEnrichers[e.Event.EventType]; !found {
    68  		return false
    69  	}
    70  
    71  	select {
    72  	case s.eventsQueue <- e:
    73  		return true
    74  	default:
    75  		metrics.AgentAnalyzersQueueDroppedEventsTotal.Inc()
    76  		return false
    77  	}
    78  }
    79  
    80  func (s *Service) Run(ctx context.Context) error {
    81  	s.log.Infof("running, workers=%d", s.cfg.WorkerCount)
    82  	defer s.log.Infof("stopping")
    83  
    84  	// Events processing workers loops.
    85  	for i := 0; i < s.cfg.WorkerCount; i++ {
    86  		go func() {
    87  			for {
    88  				select {
    89  				case <-ctx.Done():
    90  					return
    91  				case e := <-s.eventsQueue:
    92  					s.processEvent(ctx, e)
    93  				}
    94  			}
    95  		}()
    96  	}
    97  
    98  	for {
    99  		select {
   100  		case <-ctx.Done():
   101  			return ctx.Err()
   102  		}
   103  	}
   104  }
   105  
   106  func (s *Service) processEvent(ctx context.Context, e *EnrichRequest) {
   107  	defer func() {
   108  		if perr := recover(); perr != nil {
   109  			stack := string(debug.Stack())
   110  			s.log.Errorf("panic while enriching event: %v, stack=%s", perr, stack)
   111  		}
   112  	}()
   113  
   114  	enrichers := s.eventEnrichers[e.Event.EventType]
   115  
   116  	for _, enricher := range enrichers {
   117  		enricher.Enrich(ctx, e)
   118  	}
   119  
   120  	s.outQueue <- e.Event
   121  }