acquisition: set S3 anon creds during tests, prevent deadlock (#3961)

Author: mmetc

.golangci.yml

@@ -197,9 +197,12 @@ linters:
       block-size: 6
 
     nolintlint:
-      require-explanation: false  # don't require an explanation for nolint directives
-      require-specific: true     # don't require nolint directives to be specific about which linter is being skipped
-      allow-unused: false         # report any unused nolint directives
+      # don't require an explanation for nolint directives
+      require-explanation: false
+      # don't require nolint directives to be specific about which linter is being skipped
+      require-specific: true
+      # report any unused nolint directives
+      allow-unused: false
 
     revive:
       severity: error
@@ -443,12 +446,6 @@ linters:
         path: pkg/types/utils.go
         text: 'argument-limit: .*'
 
-      # need some cleanup first: to create db in memory and share the client, not the config
-      - linters:
-          - usetesting
-        path: (.+)_test.go
-        text: context.Background.*
-
       - linters:
           - usetesting
         path: pkg/apiserver/(.+)_test.go
@@ -469,18 +466,27 @@ linters:
         path: pkg/acquisition/modules/s3/s3.go
         text: found a struct that contains a context.Context field
 
+      - linters:
+          - contextcheck
+        text: Function `Configure->newS3Client` should pass the context parameter
+
+      - linters:
+          - contextcheck
+        text: Function `ConfigureByDSN->newS3Client` should pass the context parameter
+
       # migrate over time
 
       - linters:
           - noctx
+        path: pkg/parser/enrich_dns.go
         text: "net.LookupAddr must not be called"
 
       - linters:
           - noctx
+        path: pkg/exprhelpers/helpers.go
         text: "net.LookupHost must not be called"
 
     paths:
-      - pkg/metabase
       - third_party$
       - builtin$
       - examples$

pkg/acquisition/modules/kinesis/creds.go

@@ -1,9 +1,11 @@
-//go:build !test
-
 package kinesisacquisition
 
 import "github.com/aws/aws-sdk-go-v2/aws"
 
+var defaultCredsFunc = func() aws.CredentialsProvider {
+	return nil
+}
+
 func defaultCreds() aws.CredentialsProvider {
-    return nil
+	return defaultCredsFunc()
 }

pkg/acquisition/modules/kinesis/creds_test.go

@@ -1,9 +1,10 @@
-//go:build test
-
 package kinesisacquisition
 
 import "github.com/aws/aws-sdk-go-v2/aws"
 
-func defaultCreds() aws.CredentialsProvider {
-    return aws.AnonymousCredentials{}
+//nolint:gochecknoinits
+func init() {
+	defaultCredsFunc = func() aws.CredentialsProvider {
+		return aws.AnonymousCredentials{}
+	}
 }

pkg/acquisition/modules/s3/creds.go

@@ -0,0 +1,11 @@
+package s3acquisition
+
+import "github.com/aws/aws-sdk-go-v2/aws"
+
+var defaultCredsFunc = func() aws.CredentialsProvider {
+	return nil
+}
+
+func defaultCreds() aws.CredentialsProvider {
+	return defaultCredsFunc()
+}

pkg/acquisition/modules/s3/creds_test.go

@@ -0,0 +1,10 @@
+package s3acquisition
+
+import "github.com/aws/aws-sdk-go-v2/aws"
+
+//nolint:gochecknoinits
+func init() {
+	defaultCredsFunc = func() aws.CredentialsProvider {
+		return aws.AnonymousCredentials{}
+	}
+}

pkg/acquisition/modules/s3/s3.go

@@ -118,11 +118,7 @@ const (
 	SQSFormatSNS            = "sns"
 )
 
-func (s *S3Source) newS3Client() error {
-	if s.s3Client != nil {
-		return nil
-	}
-
+func (s *S3Source) newS3Client() (*s3.Client, error) {
 	var loadOpts []func(*config.LoadOptions) error
 	if s.Config.AwsProfile != nil && *s.Config.AwsProfile != "" {
 		loadOpts = append(loadOpts, config.WithSharedConfigProfile(*s.Config.AwsProfile))
@@ -134,11 +130,14 @@ func (s *S3Source) newS3Client() error {
 	}
 
 	loadOpts = append(loadOpts, config.WithRegion(region))
-	loadOpts = append(loadOpts, config.WithCredentialsProvider(aws.AnonymousCredentials{}))
 
-	cfg, err := config.LoadDefaultConfig(s.ctx, loadOpts...)
+	if c := defaultCreds(); c != nil {
+		loadOpts = append(loadOpts, config.WithCredentialsProvider(c))
+	}
+
+	cfg, err := config.LoadDefaultConfig(context.TODO(), loadOpts...)
 	if err != nil {
-		return fmt.Errorf("failed to load aws config: %w", err)
+		return nil, fmt.Errorf("failed to load aws config: %w", err)
 	}
 
 	var clientOpts []func(*s3.Options)
@@ -148,16 +147,10 @@ func (s *S3Source) newS3Client() error {
 		})
 	}
 
-	s.s3Client = s3.NewFromConfig(cfg, clientOpts...)
-
-	return nil
+	return s3.NewFromConfig(cfg, clientOpts...), nil
 }
 
-func (s *S3Source) newSQSClient() error {
-	if s.sqsClient != nil {
-		return nil
-	}
-
+func (s *S3Source) newSQSClient() (*sqs.Client, error) {
 	var loadOpts []func(*config.LoadOptions) error
 	if s.Config.AwsProfile != nil && *s.Config.AwsProfile != "" {
 		loadOpts = append(loadOpts, config.WithSharedConfigProfile(*s.Config.AwsProfile))
@@ -169,21 +162,22 @@ func (s *S3Source) newSQSClient() error {
 	}
 
 	loadOpts = append(loadOpts, config.WithRegion(region))
-	loadOpts = append(loadOpts, config.WithCredentialsProvider(aws.AnonymousCredentials{}))
 
-	cfg, err := config.LoadDefaultConfig(s.ctx, loadOpts...)
+	if c := defaultCreds(); c != nil {
+		loadOpts = append(loadOpts, config.WithCredentialsProvider(c))
+	}
+
+	cfg, err := config.LoadDefaultConfig(context.TODO(), loadOpts...)
 	if err != nil {
-		return fmt.Errorf("failed to load aws config: %w", err)
+		return nil, fmt.Errorf("failed to load aws config: %w", err)
 	}
 
 	var clientOpts []func(*sqs.Options)
 	if s.Config.AwsEndpoint != "" {
 		clientOpts = append(clientOpts, func(o *sqs.Options) { o.BaseEndpoint = aws.String(s.Config.AwsEndpoint) })
 	}
 
-	s.sqsClient = sqs.NewFromConfig(cfg, clientOpts...)
-
-	return nil
+	return sqs.NewFromConfig(cfg, clientOpts...), nil
 }
 
 func (s *S3Source) readManager() {
@@ -207,7 +201,7 @@ func (s *S3Source) readManager() {
 
 func (s *S3Source) getBucketContent() ([]s3types.Object, error) {
 	logger := s.logger.WithField("method", "getBucketContent")
-	logger.Debugf("Getting bucket content for %s", s.Config.BucketName)
+	logger.Debugf("Getting bucket content")
 
 	bucketObjects := make([]s3types.Object, 0)
 
@@ -274,10 +268,17 @@ func (s *S3Source) listPoll() error {
 
 				logger.Debugf("Found new object %s", *bucketObjects[i].Key)
 
-				s.readerChan <- S3Object{
+				obj := S3Object{
 					Bucket: s.Config.BucketName,
 					Key:    *bucketObjects[i].Key,
 				}
+
+				select {
+				case s.readerChan <- obj:
+				case <-s.t.Dying():
+					logger.Debug("tomb is dying, dropping object send")
+					return nil
+				}
 			}
 
 			if newObject {
@@ -391,6 +392,9 @@ func (s *S3Source) sqsPoll() error {
 				WaitTimeSeconds:     20, // Probably no need to make it configurable ?
 			})
 			if err != nil {
+				if errors.Is(err, context.Canceled) {
+					return nil
+				}
 				logger.Errorf("Error while polling SQS: %s", err)
 				continue
 			}
@@ -421,7 +425,13 @@ func (s *S3Source) sqsPoll() error {
 
 				logger.Debugf("Received SQS message for object %s/%s", bucket, key)
 
-				s.readerChan <- S3Object{Key: key, Bucket: bucket}
+				// don't block if readManager has quit
+				select {
+				case s.readerChan <- S3Object{Key: key, Bucket: bucket}:
+				case <-s.t.Dying():
+					logger.Debug("tomb is dying, dropping object send")
+					return nil
+				}
 
 				_, err = s.sqsClient.DeleteMessage(s.ctx,
 					&sqs.DeleteMessageInput{
@@ -516,7 +526,13 @@ func (s *S3Source) readFile(bucket string, key string) error {
 			evt := types.MakeEvent(s.Config.UseTimeMachine, types.LOG, true)
 			evt.Line = l
 
-			s.out <- evt
+			// don't block in shutdown
+			select {
+			case s.out <-evt:
+			case <-s.t.Dying():
+				s.logger.Infof("tomb is dying, dropping event for %s/%s", bucket, key)
+				return nil
+			}
 		}
 	}
 
@@ -615,16 +631,20 @@ func (s *S3Source) Configure(yamlConfig []byte, logger *log.Entry, metricsLevel
 		s.logger.Warning("Polling method is set to list. This is not recommended as it will not scale well. Consider using SQS instead.")
 	}
 
-	err = s.newS3Client()
+	client, err := s.newS3Client()
 	if err != nil {
 		return err
 	}
 
+	s.s3Client = client
+
 	if s.Config.PollingMethod == PollMethodSQS {
-		err = s.newSQSClient()
+		sqsClient, err := s.newSQSClient()
 		if err != nil {
 			return err
 		}
+
+		s.sqsClient = sqsClient
 	}
 
 	return nil
@@ -707,11 +727,13 @@ func (s *S3Source) ConfigureByDSN(dsn string, labels map[string]string, logger *
 		return fmt.Errorf("invalid DSN %s for S3 source", dsn)
 	}
 
-	err := s.newS3Client()
+	client, err := s.newS3Client()
 	if err != nil {
 		return err
 	}
 
+	s.s3Client = client
+
 	return nil
 }
 
@@ -768,21 +790,11 @@ func (s *S3Source) StreamingAcquisition(ctx context.Context, out chan types.Even
 
 	if s.Config.PollingMethod == PollMethodSQS {
 		t.Go(func() error {
-			err := s.sqsPoll()
-			if err != nil {
-				return err
-			}
-
-			return nil
+			return s.sqsPoll()
 		})
 	} else {
 		t.Go(func() error {
-			err := s.listPoll()
-			if err != nil {
-				return err
-			}
-
-			return nil
+			return s.listPoll()
 		})
 	}
 

pkg/acquisition/modules/s3/s3_test.go

@@ -268,9 +268,11 @@ func TestDSNAcquis(t *testing.T) {
 			linesRead := 0
 			f := S3Source{}
 			logger := log.NewEntry(log.New())
-			f.s3Client = mockS3Client{}
 			err := f.ConfigureByDSN(test.dsn, map[string]string{"foo": "bar"}, logger, "")
 			require.NoError(t, err)
+
+			f.s3Client = mockS3Client{}
+
 			assert.Equal(t, test.expectedBucketName, f.Config.BucketName)
 			assert.Equal(t, test.expectedPrefix, f.Config.Prefix)
 
@@ -338,12 +340,10 @@ prefix: foo/
 			logger := log.NewEntry(log.New())
 			logger.Logger.SetLevel(log.TraceLevel)
 
-			f.s3Client = mockS3Client{}
-
 			err := f.Configure([]byte(test.config), logger, metrics.AcquisitionMetricsLevelNone)
-			if err != nil {
-				t.Fatalf("unexpected error: %s", err.Error())
-			}
+			require.NoError(t, err)
+
+			f.s3Client = mockS3Client{}
 
 			if f.Config.PollingMethod != PollMethodList {
 				t.Fatalf("expected list polling, got %s", f.Config.PollingMethod)
@@ -423,10 +423,11 @@ sqs_name: test
 			linesRead := 0
 			f := S3Source{}
 			logger := log.NewEntry(log.New())
-			f.s3Client = mockS3Client{}
 			err := f.Configure([]byte(test.config), logger, metrics.AcquisitionMetricsLevelNone)
 			require.NoError(t, err)
 
+			f.s3Client = mockS3Client{}
+
 			if f.Config.PollingMethod != PollMethodSQS {
 				t.Fatalf("expected sqs polling, got %s", f.Config.PollingMethod)
 			}