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

     1  package state
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	castpb "github.com/castai/kvisor/api/v1/runtime"
     8  	"github.com/castai/kvisor/pkg/castai"
     9  	"github.com/castai/kvisor/pkg/logging"
    10  	"github.com/castai/kvisor/pkg/metrics"
    11  	"github.com/prometheus/client_golang/prometheus"
    12  	"google.golang.org/grpc"
    13  )
    14  
    15  func NewCastaiEventsExporter(log *logging.Logger, apiClient *castai.Client, queueSize int) *CastaiEventsExporter {
    16  	return &CastaiEventsExporter{
    17  		log:                         log.WithField("component", "castai_events_exporter"),
    18  		apiClient:                   apiClient,
    19  		queue:                       make(chan *castpb.Event, queueSize),
    20  		writeStreamCreateRetryDelay: 2 * time.Second,
    21  	}
    22  }
    23  
    24  type CastaiEventsExporter struct {
    25  	log                         *logging.Logger
    26  	apiClient                   *castai.Client
    27  	queue                       chan *castpb.Event
    28  	writeStreamCreateRetryDelay time.Duration
    29  }
    30  
    31  func (c *CastaiEventsExporter) Run(ctx context.Context) error {
    32  	c.log.Info("running export loop")
    33  	defer c.log.Info("export loop done")
    34  
    35  	ws := castai.NewWriteStream[*castpb.Event, *castpb.WriteStreamResponse](ctx, func(ctx context.Context) (grpc.ClientStream, error) {
    36  		return c.apiClient.GRPC.EventsWriteStream(ctx)
    37  	})
    38  	defer ws.Close()
    39  	ws.ReopenDelay = c.writeStreamCreateRetryDelay
    40  
    41  	sendErrorMetric := metrics.AgentExporterSendErrorsTotal.WithLabelValues("castai_events")
    42  	sendMetric := metrics.AgentExporterSendTotal.WithLabelValues("castai_events")
    43  
    44  	for {
    45  		select {
    46  		case <-ctx.Done():
    47  			return ctx.Err()
    48  		case e := <-c.queue:
    49  			if err := ws.Send(e); err != nil {
    50  				sendErrorMetric.Inc()
    51  				continue
    52  			}
    53  			sendMetric.Inc()
    54  			metrics.AgentExportedEventsTotal.With(prometheus.Labels{metrics.EventTypeLabel: e.GetEventType().String()}).Inc()
    55  		}
    56  	}
    57  }
    58  
    59  func (c *CastaiEventsExporter) Enqueue(e *castpb.Event) {
    60  	select {
    61  	case c.queue <- e:
    62  	default:
    63  		metrics.AgentExporterQueueDroppedTotal.WithLabelValues("castai_events").Inc()
    64  	}
    65  }