Add clusterUUID column to S3Uploader and ClickHouseExporter

build/yamls/flow-visibility-e2e.yml

@@ -90,6 +90,7 @@ data:
             throughputFromDestinationNode UInt64,
             reverseThroughputFromSourceNode UInt64,
             reverseThroughputFromDestinationNode UInt64,
+            clusterUUID String,
             trusted UInt8 DEFAULT 0
         ) engine=MergeTree
         ORDER BY (timeInserted, flowEndSeconds)

pkg/flowaggregator/clickhouseclient/clickhouseclient.go

@@ -81,9 +81,10 @@ const (
                    throughputFromSourceNode,
                    throughputFromDestinationNode,
                    reverseThroughputFromSourceNode,
-                   reverseThroughputFromDestinationNode) 
+                   reverseThroughputFromDestinationNode,
+                   clusterUUID)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 
-                           ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`
+                           ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`
 )
 
 // PrepareClickHouseConnection is used for unit testing
@@ -118,7 +119,8 @@ type ClickHouseExportProcess struct {
 	commitTicker         *time.Ticker
 	exportProcessRunning bool
 	// mutex protects configuration state from concurrent access
-	mutex sync.Mutex
+	mutex       sync.Mutex
+	clusterUUID string
 }
 
 type ClickHouseInput struct {
@@ -156,7 +158,7 @@ func (ci *ClickHouseInput) GetDataSourceName() (string, error) {
 	return sb.String(), nil
 }
 
-func NewClickHouseClient(input ClickHouseInput) (*ClickHouseExportProcess, error) {
+func NewClickHouseClient(input ClickHouseInput, clusterUUID string) (*ClickHouseExportProcess, error) {
 	dsn, connect, err := PrepareClickHouseConnection(input)
 	if err != nil {
 		return nil, err
@@ -168,6 +170,7 @@ func NewClickHouseClient(input ClickHouseInput) (*ClickHouseExportProcess, error
 		deque:          deque.New(),
 		queueSize:      maxQueueSize,
 		commitInterval: input.CommitInterval,
+		clusterUUID:    clusterUUID,
 	}
 	return chClient, nil
 }
@@ -343,7 +346,8 @@ func (ch *ClickHouseExportProcess) batchCommitAll(ctx context.Context) (int, err
 			record.ThroughputFromSourceNode,
 			record.ThroughputFromDestinationNode,
 			record.ReverseThroughputFromSourceNode,
-			record.ReverseThroughputFromDestinationNode)
+			record.ReverseThroughputFromDestinationNode,
+			ch.clusterUUID)
 
 		if err != nil {
 			klog.ErrorS(err, "Error when adding record")

pkg/flowaggregator/clickhouseclient/clickhouseclient_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/DATA-DOG/go-sqlmock"
 	"github.com/gammazero/deque"
 	"github.com/golang/mock/gomock"
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	ipfixentitiestesting "github.com/vmware/go-ipfix/pkg/entities/testing"
@@ -40,6 +41,8 @@ func init() {
 	registry.LoadRegistry()
 }
 
+var fakeClusterUUID = uuid.New().String()
+
 func TestGetDataSourceName(t *testing.T) {
 	chInput := ClickHouseInput{
 		Username:       "username",
@@ -103,9 +106,10 @@ func TestBatchCommitAll(t *testing.T) {
 	defer db.Close()
 
 	chExportProc := ClickHouseExportProcess{
-		db:        db,
-		deque:     deque.New(),
-		queueSize: maxQueueSize,
+		db:          db,
+		deque:       deque.New(),
+		queueSize:   maxQueueSize,
+		clusterUUID: fakeClusterUUID,
 	}
 
 	recordRow := flowrecordtesting.PrepareTestFlowRecord()
@@ -161,7 +165,8 @@ func TestBatchCommitAll(t *testing.T) {
 			15902813473,
 			15902813474,
 			12381345,
-			12381346).
+			12381346,
+			fakeClusterUUID).
 		WillReturnResult(sqlmock.NewResult(0, 1))
 	mock.ExpectCommit()
 
@@ -183,7 +188,7 @@ func TestBatchCommitAllMultiRecord(t *testing.T) {
 		queueSize: maxQueueSize,
 	}
 	recordRow := flowrecord.FlowRecord{}
-	fieldCount := reflect.TypeOf(recordRow).NumField()
+	fieldCount := reflect.TypeOf(recordRow).NumField() + 1
 	argList := make([]driver.Value, fieldCount)
 	for i := 0; i < len(argList); i++ {
 		argList[i] = sqlmock.AnyArg()
@@ -216,7 +221,7 @@ func TestBatchCommitAllError(t *testing.T) {
 	}
 	recordRow := flowrecord.FlowRecord{}
 	chExportProc.deque.PushBack(&recordRow)
-	fieldCount := reflect.TypeOf(recordRow).NumField()
+	fieldCount := reflect.TypeOf(recordRow).NumField() + 1
 	argList := make([]driver.Value, fieldCount)
 	for i := 0; i < len(argList); i++ {
 		argList[i] = sqlmock.AnyArg()

pkg/flowaggregator/exporter/clickhouse.go

@@ -18,6 +18,7 @@ import (
 	"os"
 
 	ipfixentities "github.com/vmware/go-ipfix/pkg/entities"
+	"k8s.io/client-go/kubernetes"
 	"k8s.io/klog/v2"
 
 	"antrea.io/antrea/pkg/flowaggregator/clickhouseclient"
@@ -41,10 +42,14 @@ func buildClickHouseInput(opt *options.Options) clickhouseclient.ClickHouseInput
 	}
 }
 
-func NewClickHouseExporter(opt *options.Options) (*ClickHouseExporter, error) {
+func NewClickHouseExporter(k8sClient kubernetes.Interface, opt *options.Options) (*ClickHouseExporter, error) {
 	chInput := buildClickHouseInput(opt)
 	klog.InfoS("ClickHouse configuration", "database", chInput.Database, "databaseURL", chInput.DatabaseURL, "debug", chInput.Debug, "compress", *chInput.Compress, "commitInterval", chInput.CommitInterval)
-	chExportProcess, err := clickhouseclient.NewClickHouseClient(chInput)
+	clusterUUID, err := getClusterUUID(k8sClient)
+	if err != nil {
+		return nil, err
+	}
+	chExportProcess, err := clickhouseclient.NewClickHouseClient(chInput, clusterUUID.String())
 	if err != nil {
 		return nil, err
 	}

pkg/flowaggregator/exporter/clickhouse_test.go

@@ -20,6 +20,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -54,8 +55,10 @@ func TestClickHouse_UpdateOptions(t *testing.T) {
 		},
 		ClickHouseCommitInterval: 8 * time.Second,
 	}
-	clickHouseExporter, err := NewClickHouseExporter(opt)
+	chInput := buildClickHouseInput(opt)
+	chExportProcess, err := clickhouseclient.NewClickHouseClient(chInput, uuid.New().String())
 	require.NoError(t, err)
+	clickHouseExporter := ClickHouseExporter{chInput: &chInput, chExportProcess: chExportProcess}
 	clickHouseExporter.Start()
 	assert.Equal(t, clickHouseExporter.chExportProcess.GetDsn(), "tcp://clickhouse-clickhouse.flow-visibility.svc:9000?username=default&password=default&database=default&debug=true&compress=false")
 	assert.Equal(t, clickHouseExporter.chExportProcess.GetCommitInterval().String(), "8s")

pkg/flowaggregator/exporter/ipfix.go

@@ -18,17 +18,14 @@ import (
 	"fmt"
 	"hash/fnv"
 	"sync"
-	"time"
 
 	"github.com/google/uuid"
 	ipfixentities "github.com/vmware/go-ipfix/pkg/entities"
 	"github.com/vmware/go-ipfix/pkg/exporter"
 	ipfixregistry "github.com/vmware/go-ipfix/pkg/registry"
-	"k8s.io/apimachinery/pkg/util/wait"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/klog/v2"
 
-	"antrea.io/antrea/pkg/clusteridentity"
 	"antrea.io/antrea/pkg/flowaggregator/infoelements"
 	"antrea.io/antrea/pkg/flowaggregator/options"
 	"antrea.io/antrea/pkg/ipfix"
@@ -59,31 +56,11 @@ type IPFIXExporter struct {
 // genObservationDomainID generates an IPFIX Observation Domain ID when one is not provided by the
 // user through the flow aggregator configuration. It will first try to generate one
 // deterministically based on the cluster UUID (if available, with a timeout of 10s). Otherwise, it
-// will generate a random one. The cluster UUID should be available if Antrea is deployed to the
-// cluster ahead of the flow aggregator, which is the expectation since when deploying flow
-// aggregator as a Pod, networking needs to be configured by the CNI plugin.
+// will generate a random one.
 func genObservationDomainID(k8sClient kubernetes.Interface) uint32 {
-	const retryInterval = time.Second
-	const timeout = 10 * time.Second
-	const defaultAntreaNamespace = "kube-system"
-
-	clusterIdentityProvider := clusteridentity.NewClusterIdentityProvider(
-		defaultAntreaNamespace,
-		clusteridentity.DefaultClusterIdentityConfigMapName,
-		k8sClient,
-	)
-	var clusterUUID uuid.UUID
-	if err := wait.PollImmediate(retryInterval, timeout, func() (bool, error) {
-		clusterIdentity, _, err := clusterIdentityProvider.Get()
-		if err != nil {
-			return false, nil
-		}
-		clusterUUID = clusterIdentity.UUID
-		return true, nil
-	}); err != nil {
-		klog.InfoS(
-			"Unable to retrieve cluster UUID; will generate a random observation domain ID", "timeout", timeout, "ConfigMapNameSpace", defaultAntreaNamespace, "ConfigMapName", clusteridentity.DefaultClusterIdentityConfigMapName,
-		)
+	clusterUUID, err := getClusterUUID(k8sClient)
+	if err != nil {
+		klog.ErrorS(err, "Error when retrieving cluster UUID; will generate a random observation domain ID")
 		clusterUUID = uuid.New()
 	}
 	h := fnv.New32()

pkg/flowaggregator/exporter/s3.go

@@ -16,6 +16,7 @@ package exporter
 
 import (
 	ipfixentities "github.com/vmware/go-ipfix/pkg/entities"
+	"k8s.io/client-go/kubernetes"
 	"k8s.io/klog/v2"
 
 	"antrea.io/antrea/pkg/flowaggregator/options"
@@ -34,10 +35,14 @@ func buildS3Input(opt *options.Options) s3uploader.S3Input {
 	}
 }
 
-func NewS3Exporter(opt *options.Options) (*S3Exporter, error) {
+func NewS3Exporter(k8sClient kubernetes.Interface, opt *options.Options) (*S3Exporter, error) {
 	s3Input := buildS3Input(opt)
 	klog.InfoS("S3Uploader configuration", "bucketName", s3Input.Config.BucketName, "bucketPrefix", s3Input.Config.BucketPrefix, "region", s3Input.Config.Region, "recordFormat", s3Input.Config.RecordFormat, "compress", *s3Input.Config.Compress, "maxRecordsPerFile", s3Input.Config.MaxRecordsPerFile, "uploadInterval", s3Input.UploadInterval)
-	s3UploadProcess, err := s3uploader.NewS3UploadProcess(s3Input)
+	clusterUUID, err := getClusterUUID(k8sClient)
+	if err != nil {
+		return nil, err
+	}
+	s3UploadProcess, err := s3uploader.NewS3UploadProcess(s3Input, clusterUUID.String())
 	if err != nil {
 		return nil, err
 	}

pkg/flowaggregator/exporter/s3_test.go

@@ -18,11 +18,13 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"antrea.io/antrea/pkg/config/flowaggregator"
 	"antrea.io/antrea/pkg/flowaggregator/options"
+	"antrea.io/antrea/pkg/flowaggregator/s3uploader"
 )
 
 func TestS3_UpdateOptions(t *testing.T) {
@@ -40,8 +42,10 @@ func TestS3_UpdateOptions(t *testing.T) {
 		},
 		S3UploadInterval: 8 * time.Second,
 	}
-	s3Exporter, err := NewS3Exporter(opt)
+	s3Input := buildS3Input(opt)
+	s3UploadProcess, err := s3uploader.NewS3UploadProcess(s3Input, uuid.New().String())
 	require.NoError(t, err)
+	s3Exporter := S3Exporter{s3Input: &s3Input, s3UploadProcess: s3UploadProcess}
 	s3Exporter.Start()
 	assert.Equal(t, s3Exporter.s3UploadProcess.GetBucketName(), "defaultBucketName")
 	assert.Equal(t, s3Exporter.s3UploadProcess.GetBucketPrefix(), "defaultBucketPrefix")

pkg/flowaggregator/exporter/utils.go

@@ -0,0 +1,55 @@
+// Copyright 2022 Antrea Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package exporter
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	"k8s.io/apimachinery/pkg/util/wait"
+	"k8s.io/client-go/kubernetes"
+
+	"antrea.io/antrea/pkg/clusteridentity"
+)
+
+// getClusterUUID retrieves the cluster UUID (if available, with a timeout of 10s).
+// Otherwise, it returns an empty cluster UUID and error. The cluster UUID should
+// be available if Antrea is deployed to the cluster ahead of the flow aggregator,
+// which is the expectation since when deploying flow aggregator as a Pod,
+// networking needs to be configured by the CNI plugin.
+func getClusterUUID(k8sClient kubernetes.Interface) (uuid.UUID, error) {
+	const retryInterval = time.Second
+	const timeout = 10 * time.Second
+	const defaultAntreaNamespace = "kube-system"
+
+	clusterIdentityProvider := clusteridentity.NewClusterIdentityProvider(
+		defaultAntreaNamespace,
+		clusteridentity.DefaultClusterIdentityConfigMapName,
+		k8sClient,
+	)
+	var clusterUUID uuid.UUID
+	if err := wait.PollImmediate(retryInterval, timeout, func() (bool, error) {
+		clusterIdentity, _, err := clusterIdentityProvider.Get()
+		if err != nil {
+			return false, nil
+		}
+		clusterUUID = clusterIdentity.UUID
+		return true, nil
+	}); err != nil {
+		return clusterUUID, fmt.Errorf("unable to retrieve cluster UUID from ConfigMap '%s/%s': timeout after %v", defaultAntreaNamespace, clusteridentity.DefaultClusterIdentityConfigMapName, timeout)
+	}
+	return clusterUUID, nil
+}

pkg/flowaggregator/flowaggregator.go

@@ -93,11 +93,11 @@ var (
 	newIPFIXExporter = func(k8sClient kubernetes.Interface, opt *options.Options, registry ipfix.IPFIXRegistry) exporter.Interface {
 		return exporter.NewIPFIXExporter(k8sClient, opt, registry)
 	}
-	newClickHouseExporter = func(opt *options.Options) (exporter.Interface, error) {
-		return exporter.NewClickHouseExporter(opt)
+	newClickHouseExporter = func(k8sClient kubernetes.Interface, opt *options.Options) (exporter.Interface, error) {
+		return exporter.NewClickHouseExporter(k8sClient, opt)
 	}
-	newS3Exporter = func(opt *options.Options) (exporter.Interface, error) {
-		return exporter.NewS3Exporter(opt)
+	newS3Exporter = func(k8sClient kubernetes.Interface, opt *options.Options) (exporter.Interface, error) {
+		return exporter.NewS3Exporter(k8sClient, opt)
 	}
 )
 
@@ -183,14 +183,14 @@ func NewFlowAggregator(
 	}
 	if opt.Config.ClickHouse.Enable {
 		var err error
-		fa.clickHouseExporter, err = newClickHouseExporter(opt)
+		fa.clickHouseExporter, err = newClickHouseExporter(k8sClient, opt)
 		if err != nil {
 			return nil, fmt.Errorf("error when creating ClickHouse export process: %v", err)
 		}
 	}
 	if opt.Config.S3Uploader.Enable {
 		var err error
-		fa.s3Exporter, err = newS3Exporter(opt)
+		fa.s3Exporter, err = newS3Exporter(k8sClient, opt)
 		if err != nil {
 			return nil, fmt.Errorf("error when creating S3 export process: %v", err)
 		}
@@ -586,7 +586,7 @@ func (fa *flowAggregator) updateFlowAggregator(opt *options.Options) {
 		if fa.clickHouseExporter == nil {
 			klog.InfoS("Enabling ClickHouse")
 			var err error
-			fa.clickHouseExporter, err = newClickHouseExporter(opt)
+			fa.clickHouseExporter, err = newClickHouseExporter(fa.k8sClient, opt)
 			if err != nil {
 				klog.ErrorS(err, "Error when creating ClickHouse export process")
 				return
@@ -608,7 +608,7 @@ func (fa *flowAggregator) updateFlowAggregator(opt *options.Options) {
 		if fa.s3Exporter == nil {
 			klog.InfoS("Enabling S3Uploader")
 			var err error
-			fa.s3Exporter, err = newS3Exporter(opt)
+			fa.s3Exporter, err = newS3Exporter(fa.k8sClient, opt)
 			if err != nil {
 				klog.ErrorS(err, "Error when creating S3 export process")
 				return

pkg/flowaggregator/flowaggregator_test.go

@@ -264,10 +264,10 @@ func TestFlowAggregator_updateFlowAggregator(t *testing.T) {
 	newIPFIXExporter = func(kubernetes.Interface, *options.Options, ipfix.IPFIXRegistry) exporter.Interface {
 		return mockIPFIXExporter
 	}
-	newClickHouseExporter = func(*options.Options) (exporter.Interface, error) {
+	newClickHouseExporter = func(kubernetes.Interface, *options.Options) (exporter.Interface, error) {
 		return mockClickHouseExporter, nil
 	}
-	newS3Exporter = func(*options.Options) (exporter.Interface, error) {
+	newS3Exporter = func(kubernetes.Interface, *options.Options) (exporter.Interface, error) {
 		return mockS3Exporter, nil
 	}
 
@@ -405,10 +405,10 @@ func TestFlowAggregator_Run(t *testing.T) {
 	newIPFIXExporter = func(kubernetes.Interface, *options.Options, ipfix.IPFIXRegistry) exporter.Interface {
 		return mockIPFIXExporter
 	}
-	newClickHouseExporter = func(*options.Options) (exporter.Interface, error) {
+	newClickHouseExporter = func(kubernetes.Interface, *options.Options) (exporter.Interface, error) {
 		return mockClickHouseExporter, nil
 	}
-	newS3Exporter = func(*options.Options) (exporter.Interface, error) {
+	newS3Exporter = func(kubernetes.Interface, *options.Options) (exporter.Interface, error) {
 		return mockS3Exporter, nil
 	}