github.com/castai/kvisor@v1.7.1-0.20240516114728-b3572a2607b5/cmd/agent/daemon/state/castai_container_stats_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  	"google.golang.org/grpc"
    12  )
    13  
    14  func NewCastaiContainerStatsExporter(log *logging.Logger, apiClient *castai.Client, queueSize int) *CastaiContainerStatsExporter {
    15  	return &CastaiContainerStatsExporter{
    16  		log:                         log.WithField("component", "castai_container_stats_exporter"),
    17  		apiClient:                   apiClient,
    18  		queue:                       make(chan *castpb.ContainerStatsBatch, queueSize),
    19  		writeStreamCreateRetryDelay: 2 * time.Second,
    20  	}
    21  }
    22  
    23  type CastaiContainerStatsExporter struct {
    24  	log                         *logging.Logger
    25  	apiClient                   *castai.Client
    26  	queue                       chan *castpb.ContainerStatsBatch
    27  	writeStreamCreateRetryDelay time.Duration
    28  }
    29  
    30  func (c *CastaiContainerStatsExporter) Run(ctx context.Context) error {
    31  	c.log.Info("running export loop")
    32  	defer c.log.Info("export loop done")
    33  
    34  	ws := castai.NewWriteStream[*castpb.ContainerStatsBatch, *castpb.WriteStreamResponse](ctx, func(ctx context.Context) (grpc.ClientStream, error) {
    35  		return c.apiClient.GRPC.ContainerStatsWriteStream(ctx)
    36  	})
    37  	defer ws.Close()
    38  	ws.ReopenDelay = c.writeStreamCreateRetryDelay
    39  
    40  	sendErrorMetric := metrics.AgentExporterSendErrorsTotal.WithLabelValues("castai_container_stats")
    41  	sendMetric := metrics.AgentExporterSendTotal.WithLabelValues("castai_container_stats")
    42  
    43  	for {
    44  		select {
    45  		case <-ctx.Done():
    46  			return ctx.Err()
    47  		case e := <-c.queue:
    48  			if err := ws.Send(e); err != nil {
    49  				sendErrorMetric.Inc()
    50  				continue
    51  			}
    52  			sendMetric.Inc()
    53  		}
    54  	}
    55  }
    56  
    57  func (c *CastaiContainerStatsExporter) Enqueue(e *castpb.ContainerStatsBatch) {
    58  	select {
    59  	case c.queue <- e:
    60  	default:
    61  		metrics.AgentExporterQueueDroppedTotal.WithLabelValues("castai_container_stats").Inc()
    62  	}
    63  }