tetragon: Allow to specify rb-* size options with size suffix

cmd/tetragon/main.go

@@ -709,8 +709,9 @@ func execute() error {
 		Use:   "tetragon",
 		Short: "Tetragon - eBPF-based Security Observability and Runtime Enforcement",
 		Run: func(cmd *cobra.Command, args []string) {
-			option.ReadAndSetFlags()
-
+			if err := option.ReadAndSetFlags(); err != nil {
+				log.WithError(err).Fatal("Failed to start gops")
+			}
 			if err := startGopsServer(); err != nil {
 				log.WithError(err).Fatal("Failed to start gops")
 			}
@@ -792,8 +793,8 @@ func execute() error {
 	flags.Bool(option.KeyDisableKprobeMulti, false, "Allow to disable kprobe multi interface")
 
 	// Allow to specify perf ring buffer size
-	flags.Int(option.KeyRBSizeTotal, 0, "Set perf ring buffer size in total for all cpus (default 65k per cpu)")
-	flags.Int(option.KeyRBSize, 0, "Set perf ring buffer size for single cpu (default 65k)")
+	flags.String(option.KeyRBSizeTotal, "0", "Set perf ring buffer size in total for all cpus (default 65k per cpu, allows K/M/G suffix)")
+	flags.String(option.KeyRBSize, "0", "Set perf ring buffer size for single cpu (default 65k, allows K/M/G suffix)")
 
 	// Provide option to remove existing pinned BPF programs and maps in Tetragon's
 	// observer dir on startup. Useful for doing upgrades/downgrades. Set to false to
@@ -812,7 +813,7 @@ func execute() error {
 
 	flags.StringSlice(option.KeyKmods, []string{}, "List of kernel modules to load symbols from")
 
-	flags.Int(option.KeyRBQueueSize, 65535, "Set size of channel between ring buffer and sensor go routines (default 65k)")
+	flags.String(option.KeyRBQueueSize, "65535", "Set size of channel between ring buffer and sensor go routines (default 65k, allows K/M/G suffix)")
 
 	flags.Bool(option.KeyEnablePodInfo, false, "Enable PodInfo custom resource")
 

pkg/observer/observer.go

@@ -165,6 +165,18 @@ func perfBufferSize(perCPUBuffer int) int {
 	return nPages * pageSize
 }
 
+func sizeWithSuffix(size int) string {
+	suffix := [4]string{"", "K", "M", "G"}
+
+	i := 0
+	for size > 1024 && i < 3 {
+		size = size / 1024
+		i++
+	}
+
+	return fmt.Sprintf("%d%s", size, suffix[i])
+}
+
 func (k *Observer) getRBSize(cpus int) int {
 	var size int
 
@@ -179,7 +191,9 @@ func (k *Observer) getRBSize(cpus int) int {
 	cpuSize := perfBufferSize(size)
 	totalSize := cpuSize * cpus
 
-	k.log.WithField("percpu", cpuSize).WithField("total", totalSize).Info("Perf ring buffer size (bytes)")
+	k.log.WithField("percpu", sizeWithSuffix(cpuSize)).
+		WithField("total", sizeWithSuffix(totalSize)).
+		Info("Perf ring buffer size (bytes)")
 	return size
 }
 
@@ -188,7 +202,8 @@ func (k *Observer) getRBQueueSize() int {
 	if size == 0 {
 		size = 65535
 	}
-	k.log.WithField("size", size).Info("Perf ring buffer events queue size (events)")
+	k.log.WithField("size", sizeWithSuffix(size)).
+		Info("Perf ring buffer events queue size (events)")
 	return size
 }
 

pkg/option/flags.go

@@ -4,9 +4,11 @@
 package option
 
 import (
+	"fmt"
 	"strings"
 
 	"github.com/cilium/tetragon/pkg/logger"
+	"github.com/cilium/tetragon/pkg/strutils"
 	"github.com/spf13/viper"
 )
 
@@ -86,7 +88,7 @@ const (
 	KeyExposeKernelAddresses = "expose-kernel-addresses"
 )
 
-func ReadAndSetFlags() {
+func ReadAndSetFlags() error {
 	Config.HubbleLib = viper.GetString(KeyHubbleLib)
 	Config.BTF = viper.GetString(KeyBTF)
 	Config.ProcFS = viper.GetString(KeyProcFS)
@@ -103,9 +105,17 @@ func ReadAndSetFlags() {
 
 	Config.DisableKprobeMulti = viper.GetBool(KeyDisableKprobeMulti)
 
-	Config.RBSize = viper.GetInt(KeyRBSize)
-	Config.RBSizeTotal = viper.GetInt(KeyRBSizeTotal)
-	Config.RBQueueSize = viper.GetInt(KeyRBQueueSize)
+	var err error
+
+	if Config.RBSize, err = strutils.ParseSize(viper.GetString(KeyRBSize)); err != nil {
+		return fmt.Errorf("failed to parse rb-size value: %s", err)
+	}
+	if Config.RBSizeTotal, err = strutils.ParseSize(viper.GetString(KeyRBSizeTotal)); err != nil {
+		return fmt.Errorf("failed to parse rb-size-total value: %s", err)
+	}
+	if Config.RBQueueSize, err = strutils.ParseSize(viper.GetString(KeyRBQueueSize)); err != nil {
+		return fmt.Errorf("failed to parse rb-queue-size value: %s", err)
+	}
 
 	Config.GopsAddr = viper.GetString(KeyGopsAddr)
 
@@ -154,6 +164,8 @@ func ReadAndSetFlags() {
 	Config.TracingPolicy = viper.GetString(KeyTracingPolicy)
 
 	Config.ExposeKernelAddresses = viper.GetBool(KeyExposeKernelAddresses)
+
+	return nil
 }
 
 func ParseMetricsLabelFilter(labels string) map[string]interface{} {

pkg/strutils/strutls.go

@@ -3,7 +3,10 @@
 
 package strutils
 
-import "strings"
+import (
+	"strconv"
+	"strings"
+)
 
 // UTF8FromBPFBytes transforms bpf (C) strings to valid utf-8 strings
 //
@@ -22,3 +25,28 @@ import "strings"
 func UTF8FromBPFBytes(b []byte) string {
 	return strings.ToValidUTF8(string(b), "�")
 }
+
+func ParseSize(str string) (int, error) {
+	suffix := str[len(str)-1:]
+
+	if !strings.Contains("KMG", suffix) {
+		return strconv.Atoi(str)
+	}
+
+	val, err := strconv.Atoi(str[0 : len(str)-1])
+	if err != nil {
+		return 0, err
+	}
+
+	switch suffix {
+	case "K":
+		return val * 1024, nil
+	case "M":
+		return val * 1024 * 1024, nil
+	case "G":
+		return val * 1024 * 1024 * 1024, nil
+	}
+
+	// never reached
+	return 0, nil
+}

pkg/strutils/strutls_test.go

@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright Authors of Tetragon
+
+package strutils
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type parseSize struct {
+	str string
+	err bool
+	val int
+}
+
+func TestParseSize(t *testing.T) {
+	var tests = []parseSize{
+		parseSize{"1K", false, 1024},
+		parseSize{"256M", false, 256 * 1024 * 1024},
+		parseSize{"10G", false, 10 * 1024 * 1024 * 1024},
+		parseSize{"10k", true, 0},
+		parseSize{"abc", true, 0},
+		parseSize{"abcM", true, 0},
+	}
+
+	for idx := range tests {
+		test := tests[idx]
+		val, err := ParseSize(test.str)
+		assert.Equal(t, val, test.val)
+		assert.Equal(t, err != nil, test.err)
+	}
+}