Add netflows send retry

cmd/agent/daemon/app/app.go

@@ -173,7 +173,7 @@ func (a *App) Run(ctx context.Context) error {
 			exporters.Stats = append(exporters.Stats, state.NewCastaiStatsExporter(log, castaiClient, a.cfg.ExportersQueueSize))
 		}
 		if cfg.Netflow.Enabled {
-			exporters.Netflow = append(exporters.Netflow, state.NewCastaiNetflowExporter(log, castaiClient, a.cfg.ExportersQueueSize))
+			exporters.Netflow = append(exporters.Netflow, state.NewCastaiNetflowExporter(log, castaiClient.GRPC, a.cfg.ExportersQueueSize))
 		}
 		if cfg.ProcessTree.Enabled {
 			exporter := state.NewCastaiProcessTreeExporter(log, castaiClient, a.cfg.ExportersQueueSize)

cmd/agent/daemon/state/castai_events_exporter.go

@@ -78,7 +78,6 @@ func (s *eventSender) send(e *castpb.Event, retry bool) {
 
 			return
 		}
-
 		s.sendErrorMetric.Inc()
 		return
 	}

cmd/agent/daemon/state/castai_netflow_exporter.go

@@ -12,18 +12,22 @@ import (
 	"google.golang.org/grpc/encoding/gzip"
 )
 
-func NewCastaiNetflowExporter(log *logging.Logger, apiClient *castai.Client, queueSize int) *CastaiNetflowExporter {
+type castaiNetflowExporterClient interface {
+	NetflowWriteStream(ctx context.Context, opts ...grpc.CallOption) (castpb.RuntimeSecurityAgentAPI_NetflowWriteStreamClient, error)
+}
+
+func NewCastaiNetflowExporter(log *logging.Logger, apiClient castaiNetflowExporterClient, queueSize int) *CastaiNetflowExporter {
 	return &CastaiNetflowExporter{
 		log:                         log.WithField("component", "castai_netflow_exporter"),
 		apiClient:                   apiClient,
 		queue:                       make(chan *castpb.Netflow, queueSize),
-		writeStreamCreateRetryDelay: 2 * time.Second,
+		writeStreamCreateRetryDelay: 1 * time.Second,
 	}
 }
 
 type CastaiNetflowExporter struct {
 	log                         *logging.Logger
-	apiClient                   *castai.Client
+	apiClient                   castaiNetflowExporterClient
 	queue                       chan *castpb.Netflow
 	writeStreamCreateRetryDelay time.Duration
 }
@@ -33,24 +37,31 @@ func (c *CastaiNetflowExporter) Run(ctx context.Context) error {
 	defer c.log.Info("export loop done")
 
 	ws := castai.NewWriteStream[*castpb.Netflow, *castpb.WriteStreamResponse](ctx, func(ctx context.Context) (grpc.ClientStream, error) {
-		return c.apiClient.GRPC.NetflowWriteStream(ctx, grpc.UseCompressor(gzip.Name))
+		return c.apiClient.NetflowWriteStream(ctx, grpc.UseCompressor(gzip.Name))
 	})
 	defer ws.Close()
 	ws.ReopenDelay = c.writeStreamCreateRetryDelay
 
 	sendErrorMetric := metrics.AgentExporterSendErrorsTotal.WithLabelValues("castai_netflow")
 	sendMetric := metrics.AgentExporterSendTotal.WithLabelValues("castai_netflow")
 
+	sendWithRetry := func(e *castpb.Netflow) {
+		for range 2 {
+			if err := ws.Send(e); err != nil {
+				continue
+			}
+			sendMetric.Inc()
+			return
+		}
+		sendErrorMetric.Inc()
+	}
+
 	for {
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
 		case e := <-c.queue:
-			if err := ws.Send(e); err != nil {
-				sendErrorMetric.Inc()
-				continue
-			}
-			sendMetric.Inc()
+			sendWithRetry(e)
 		}
 	}
 }

cmd/agent/daemon/state/castai_netflow_exporter_test.go

@@ -0,0 +1,155 @@
+package state
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"testing"
+	"time"
+
+	castaipb "github.com/castai/kvisor/api/v1/runtime"
+	"github.com/castai/kvisor/pkg/logging"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestNewCastaiNetflowExporter(t *testing.T) {
+	t.Run("send netflows", func(t *testing.T) {
+		r := require.New(t)
+
+		netflowStream := &netflowMockStream{}
+		exporter := NewCastaiNetflowExporter(logging.NewTestLog(), netflowStream, 100)
+
+		go func() {
+			if err := exporter.Run(context.Background()); err != nil {
+				panic(err)
+			}
+		}()
+		for range 100 {
+			exporter.Enqueue(&castaipb.Netflow{PodName: "p1"})
+		}
+
+		r.Eventually(func() bool {
+			netflowStream.mu.Lock()
+			streams := netflowStream.streams
+			netflowStream.mu.Unlock()
+			if len(streams) != 1 {
+				return false
+			}
+			stream := streams[0]
+			stream.mu.Lock()
+			flows := stream.flows
+			stream.mu.Unlock()
+			if len(flows) != 100 {
+				return false
+			}
+			r.Equal("p1", flows[0].PodName)
+			return true
+		}, 1*time.Second, 10*time.Millisecond)
+	})
+
+	t.Run("retry netflows", func(t *testing.T) {
+		r := require.New(t)
+
+		netflowStream := &netflowMockStream{firstSendError: true}
+		exporter := NewCastaiNetflowExporter(logging.NewTestLog(), netflowStream, 100)
+		exporter.writeStreamCreateRetryDelay = 1 * time.Millisecond
+
+		go func() {
+			if err := exporter.Run(context.Background()); err != nil {
+				panic(err)
+			}
+		}()
+		exporter.Enqueue(&castaipb.Netflow{PodName: "p1"})
+
+		r.Eventually(func() bool {
+			netflowStream.mu.Lock()
+			streams := netflowStream.streams
+			netflowStream.mu.Unlock()
+			if len(streams) != 1 {
+				return false
+			}
+			stream := streams[0]
+			stream.mu.Lock()
+			flows := stream.flows
+			stream.mu.Unlock()
+			if len(flows) != 1 {
+				return false
+			}
+			r.Equal("p1", flows[0].PodName)
+			return true
+		}, 1*time.Second, 10*time.Millisecond)
+	})
+}
+
+type netflowMockStream struct {
+	mu             sync.Mutex
+	streams        []*mockNetflowStreamClient
+	firstSendError bool
+}
+
+func (n *netflowMockStream) NetflowWriteStream(ctx context.Context, opts ...grpc.CallOption) (castaipb.RuntimeSecurityAgentAPI_NetflowWriteStreamClient, error) {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+	if n.firstSendError {
+		n.firstSendError = false
+	}
+	s := &mockNetflowStreamClient{firstSendError: n.firstSendError}
+	n.streams = append(n.streams, s)
+	return s, nil
+}
+
+type mockNetflowStreamClient struct {
+	mu sync.Mutex
+
+	flows          []*castaipb.Netflow
+	firstSendError bool
+}
+
+func (m *mockNetflowStreamClient) SendMsg(msg any) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if m.firstSendError {
+		m.firstSendError = false
+		return errors.New("can't send")
+	}
+
+	m.flows = append(m.flows, msg.(*castaipb.Netflow))
+	return nil
+}
+
+func (m *mockNetflowStreamClient) CloseSend() error {
+	return nil
+}
+
+func (m *mockNetflowStreamClient) Send(netflow *castaipb.Netflow) error {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (m *mockNetflowStreamClient) CloseAndRecv() (*castaipb.WriteStreamResponse, error) {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (m *mockNetflowStreamClient) Header() (metadata.MD, error) {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (m *mockNetflowStreamClient) Trailer() metadata.MD {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (m *mockNetflowStreamClient) Context() context.Context {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (m *mockNetflowStreamClient) RecvMsg(msg any) error {
+	//TODO implement me
+	panic("implement me")
+}

cmd/agent/daemon/state/castai_stats_exporter.go

@@ -17,7 +17,7 @@ func NewCastaiStatsExporter(log *logging.Logger, apiClient *castai.Client, queue
 		log:                         log.WithField("component", "castai_stats_exporter"),
 		apiClient:                   apiClient,
 		queue:                       make(chan *castpb.StatsBatch, queueSize),
-		writeStreamCreateRetryDelay: 2 * time.Second,
+		writeStreamCreateRetryDelay: 1 * time.Second,
 	}
 }
 
@@ -41,16 +41,23 @@ func (c *CastaiStatsExporter) Run(ctx context.Context) error {
 	sendErrorMetric := metrics.AgentExporterSendErrorsTotal.WithLabelValues("castai_stats")
 	sendMetric := metrics.AgentExporterSendTotal.WithLabelValues("castai_stats")
 
+	sendWithRetry := func(e *castpb.StatsBatch) {
+		for range 2 {
+			if err := ws.Send(e); err != nil {
+				continue
+			}
+			sendMetric.Inc()
+			return
+		}
+		sendErrorMetric.Inc()
+	}
+
 	for {
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
 		case e := <-c.queue:
-			if err := ws.Send(e); err != nil {
-				sendErrorMetric.Inc()
-				continue
-			}
-			sendMetric.Inc()
+			sendWithRetry(e)
 		}
 	}
 }

pkg/castai/write_stream.go

@@ -64,7 +64,9 @@ func (w *WriteStream[T, U]) Close() error {
 
 func (w *WriteStream[T, U]) open() error {
 	if w.wasOpened && w.ReopenDelay != 0 {
-		time.Sleep(w.ReopenDelay)
+		if err := sleep(w.rootCtx, w.ReopenDelay); err != nil {
+			return err
+		}
 	}
 	var err error
 	w.activeStreamCtx, w.activeStreamCtxCancel = context.WithCancel(w.rootCtx)
@@ -86,3 +88,15 @@ func (w *WriteStream[T, U]) close() {
 	w.activeStreamCtx = nil
 	w.activeStream = nil
 }
+
+func sleep(ctx context.Context, d time.Duration) error {
+	timer := time.NewTimer(d)
+	defer timer.Stop()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.C:
+		return nil
+	}
+}