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

     1  package state
     2  
     3  import (
     4  	"context"
     5  	"net/netip"
     6  	"time"
     7  
     8  	"github.com/ClickHouse/clickhouse-go/v2"
     9  	castaipb "github.com/castai/kvisor/api/v1/runtime"
    10  	"github.com/castai/kvisor/pkg/logging"
    11  	"github.com/castai/kvisor/pkg/metrics"
    12  )
    13  
    14  func NewClickhouseNetflowExporter(log *logging.Logger, conn clickhouse.Conn, queueSize int) *ClickHouseNetflowExporter {
    15  	return &ClickHouseNetflowExporter{
    16  		log:   log.WithField("component", "clickhouse_netflow_exporter"),
    17  		conn:  conn,
    18  		queue: make(chan *castaipb.Netflow, queueSize),
    19  	}
    20  }
    21  
    22  type ClickHouseNetflowExporter struct {
    23  	log   *logging.Logger
    24  	conn  clickhouse.Conn
    25  	queue chan *castaipb.Netflow
    26  }
    27  
    28  func (c *ClickHouseNetflowExporter) Run(ctx context.Context) error {
    29  	c.log.Info("running export loop")
    30  	defer c.log.Info("export loop done")
    31  
    32  	sendErrorMetric := metrics.AgentExporterSendErrorsTotal.WithLabelValues("clickhouse_netflow")
    33  	sendMetric := metrics.AgentExporterSendTotal.WithLabelValues("clickhouse_netflow")
    34  
    35  	for {
    36  		select {
    37  		case <-ctx.Done():
    38  			return ctx.Err()
    39  		case e := <-c.queue:
    40  			if err := c.asyncWrite(ctx, false, e); err != nil {
    41  				sendErrorMetric.Inc()
    42  				continue
    43  			}
    44  			sendMetric.Inc()
    45  		}
    46  	}
    47  }
    48  
    49  func (c *ClickHouseNetflowExporter) Enqueue(e *castaipb.Netflow) {
    50  	select {
    51  	case c.queue <- e:
    52  	default:
    53  		metrics.AgentExporterQueueDroppedTotal.WithLabelValues("clickhouse_netflow").Inc()
    54  	}
    55  }
    56  
    57  func (c *ClickHouseNetflowExporter) asyncWrite(ctx context.Context, wait bool, e *castaipb.Netflow) error {
    58  	q := `INSERT INTO netflows(
    59  			start,
    60  			end,
    61  			protocol,
    62  			process,
    63  			container_name,
    64  			pod_name,
    65  			namespace,
    66  			zone,
    67  			workload_name,
    68  			workload_kind,
    69  			addr,
    70  			port,
    71  			dst_addr,
    72  			dst_port,
    73  			dst_domain,
    74  			dst_pod_name,
    75  			dst_namespace,
    76  			dst_zone,
    77  			dst_workload_name,
    78  			dst_workload_kind,
    79  			tx_bytes,
    80  			tx_packets,
    81  			rx_bytes,
    82  			rx_packets
    83  			) VALUES(
    84  			?,
    85  			?,
    86  			?,
    87  			?,
    88  			?,
    89  			?,
    90  			?,
    91  			?,
    92  			?,
    93  			?,
    94  			?,
    95  			?,
    96  			?,
    97  			?,
    98  			?,
    99  			?,
   100  			?,
   101  			?,
   102  			?,
   103  			?,
   104  			?,
   105  			?,
   106  			?,
   107  			?
   108  			)`
   109  
   110  	for _, dst := range e.Destinations {
   111  		srcAddr, _ := netip.AddrFromSlice(e.Addr)
   112  		dstAddr, _ := netip.AddrFromSlice(dst.Addr)
   113  
   114  		if err := c.conn.AsyncInsert(
   115  			ctx,
   116  			q,
   117  			wait,
   118  
   119  			time.UnixMicro(int64(e.StartTs)/1000),
   120  			time.UnixMicro(int64(e.EndTs)/1000),
   121  			toDBProtocol(e.Protocol),
   122  			e.ProcessName,
   123  			e.ContainerName,
   124  			e.PodName,
   125  			e.Namespace,
   126  			e.Zone,
   127  			e.WorkloadName,
   128  			e.WorkloadKind,
   129  			srcAddr.Unmap(),
   130  			e.Port,
   131  			dstAddr.Unmap(),
   132  			dst.Port,
   133  			dst.DnsQuestion,
   134  			dst.PodName,
   135  			dst.Namespace,
   136  			dst.Zone,
   137  			dst.WorkloadName,
   138  			dst.WorkloadKind,
   139  			dst.TxBytes,
   140  			dst.TxPackets,
   141  			dst.RxBytes,
   142  			dst.RxPackets,
   143  		); err != nil {
   144  			return err
   145  		}
   146  	}
   147  	return nil
   148  }
   149  
   150  func toDBProtocol(proto castaipb.NetflowProtocol) string {
   151  	switch proto {
   152  	case castaipb.NetflowProtocol_NETFLOW_PROTOCOL_TCP:
   153  		return "tcp"
   154  	case castaipb.NetflowProtocol_NETFLOW_PROTOCOL_UDP:
   155  		return "udp"
   156  	default:
   157  		return "unknown"
   158  	}
   159  }
   160  
   161  func ClickhouseNetflowSchema() string {
   162  	return `
   163  CREATE TABLE IF NOT EXISTS netflows
   164  (
   165  	start DateTime('UTC'),
   166  	end DateTime('UTC'),
   167  	protocol LowCardinality(String),
   168  	process LowCardinality(String),
   169  	container_name LowCardinality(String),
   170  	pod_name LowCardinality(String),
   171  	namespace LowCardinality(String),
   172  	zone LowCardinality(String),
   173  	workload_name LowCardinality(String),
   174  	workload_kind LowCardinality(String),
   175  	addr IPv6,
   176  	port UInt16,
   177  	dst_addr IPv6,
   178  	dst_port UInt16,
   179  	dst_domain String,
   180  	dst_pod_name LowCardinality(String),
   181  	dst_namespace LowCardinality(String),
   182  	dst_zone LowCardinality(String),
   183  	dst_workload_name LowCardinality(String),
   184  	dst_workload_kind LowCardinality(String),
   185  	tx_bytes UInt64,
   186  	tx_packets UInt64,
   187  	rx_bytes UInt64,
   188  	rx_packets UInt64
   189  )
   190  ENGINE = MergeTree()
   191  ORDER BY (start, end, namespace, container_name)
   192  TTL toDateTime(start) + INTERVAL 3 HOUR;
   193  `
   194  }