Add asynchronous ACK handling to S3 and SQS inputs

libbeat/common/fifo/fifo.go

@@ -0,0 +1,73 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you 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 fifo
+
+// FIFO is a minimal first-in first-out queue based on a singly linked list.
+type FIFO[T any] struct {
+	first *node[T]
+	last  *node[T]
+}
+
+type node[T any] struct {
+	next  *node[T]
+	value T
+}
+
+func (f *FIFO[T]) Add(value T) {
+	newNode := &node[T]{value: value}
+	if f.first == nil {
+		f.first = newNode
+	} else {
+		f.last.next = newNode
+	}
+	f.last = newNode
+}
+
+func (f *FIFO[T]) Empty() bool {
+	return f.first == nil
+}
+
+// Return the first value (if present) without removing it from the queue.
+// If the queue is empty a default value is returned, to detect this case
+// use f.Empty().
+func (f *FIFO[T]) First() T {
+	if f.first == nil {
+		var none T
+		return none
+	}
+	return f.first.value
+}
+
+// Return the first value (if present) and remove it from the queue.
+// If the queue is empty a default value is returned, to detect this case
+// use f.Empty().
+func (f *FIFO[T]) ConsumeFirst() T {
+	result := f.First()
+	f.Remove()
+	return result
+}
+
+// Remove the first entry in the queue. Does nothing if the FIFO is empty.
+func (f *FIFO[T]) Remove() {
+	if f.first != nil {
+		f.first = f.first.next
+		if f.first == nil {
+			f.last = nil
+		}
+	}
+}

x-pack/filebeat/input/awss3/acks.go

@@ -0,0 +1,86 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License;
+// you may not use this file except in compliance with the Elastic License.
+
+package awss3
+
+import (
+	"github.com/elastic/beats/v7/libbeat/beat"
+	"github.com/elastic/beats/v7/libbeat/common/acker"
+	"github.com/elastic/beats/v7/libbeat/common/fifo"
+)
+
+type awsACKHandler struct {
+	pending    fifo.FIFO[pendingACK]
+	ackedCount int
+
+	pendingChan chan pendingACK
+	ackChan     chan int
+}
+
+type pendingACK struct {
+	eventCount  int
+	ackCallback func()
+}
+
+func newAWSACKHandler() *awsACKHandler {
+	handler := &awsACKHandler{
+		pendingChan: make(chan pendingACK, 10),
+		ackChan:     make(chan int, 10),
+	}
+	go handler.run()
+	return handler
+}
+
+func (ah *awsACKHandler) Add(eventCount int, ackCallback func()) {
+	ah.pendingChan <- pendingACK{
+		eventCount:  eventCount,
+		ackCallback: ackCallback,
+	}
+}
+
+// Called when a worker is closing, to indicate to the ack handler that it
+// should shut down as soon as the current pending list is acknowledged.
+func (ah *awsACKHandler) Close() {
+	close(ah.pendingChan)
+}
+
+func (ah *awsACKHandler) pipelineEventListener() beat.EventListener {
+	return acker.TrackingCounter(func(_ int, total int) {
+		// Notify the ack handler goroutine
+		ah.ackChan <- total
+	})
+}
+
+// Listener that handles both incoming metadata and ACK
+// confirmations.
+func (ah *awsACKHandler) run() {
+	for {
+		select {
+		case result, ok := <-ah.pendingChan:
+			if ok {
+				ah.pending.Add(result)
+			} else {
+				// Channel is closed, reset so we don't receive any more values
+				ah.pendingChan = nil
+			}
+		case count := <-ah.ackChan:
+			ah.ackedCount += count
+		}
+
+		// Finalize any objects that are now completed
+		for !ah.pending.Empty() && ah.ackedCount >= ah.pending.First().eventCount {
+			result := ah.pending.ConsumeFirst()
+			ah.ackedCount -= result.eventCount
+			// Run finalization asynchronously so we don't block the SQS worker
+			// or the queue by ignoring the ack handler's input channels. Ordering
+			// is no longer important at this point.
+			go result.ackCallback()
+		}
+
+		// If the input is closed and all acks are completed, we're done
+		if ah.pending.Empty() && ah.pendingChan == nil {
+			return
+		}
+	}
+}

x-pack/filebeat/input/awss3/input_benchmark_test.go

@@ -211,7 +211,6 @@ file_selectors:
 func benchmarkInputSQS(t *testing.T, maxMessagesInflight int) testing.BenchmarkResult {
 	return testing.Benchmark(func(b *testing.B) {
 		var err error
-		pipeline := &fakePipeline{}
 
 		config := makeBenchmarkConfig(t)
 		config.MaxNumberOfMessages = maxMessagesInflight
@@ -221,7 +220,7 @@ func benchmarkInputSQS(t *testing.T, maxMessagesInflight int) testing.BenchmarkR
 		sqsReader.sqs, err = newConstantSQS()
 		require.NoError(t, err)
 		sqsReader.s3 = newConstantS3(t)
-		sqsReader.msgHandler, err = sqsReader.createEventProcessor(pipeline)
+		sqsReader.msgHandler, err = sqsReader.createEventProcessor()
 		require.NoError(t, err, "createEventProcessor must succeed")
 
 		ctx, cancel := context.WithCancel(context.Background())
@@ -307,6 +306,7 @@ func benchmarkInputS3(t *testing.T, numberOfWorkers int) testing.BenchmarkResult
 		client := pubtest.NewChanClientWithCallback(100, func(event beat.Event) {
 			event.Private.(*awscommon.EventACKTracker).ACK()
 		})
+		pipeline := pubtest.PublisherWithClient(client)
 
 		defer func() {
 			_ = client.Close()
@@ -342,13 +342,13 @@ func benchmarkInputS3(t *testing.T, numberOfWorkers int) testing.BenchmarkResult
 				states, err := newStates(nil, store)
 				assert.NoError(t, err, "states creation should succeed")
 
-				s3EventHandlerFactory := newS3ObjectProcessorFactory(log.Named("s3"), metrics, s3API, config.FileSelectors, backupConfig{})
+				s3EventHandlerFactory := newS3ObjectProcessorFactory(metrics, s3API, config.FileSelectors, backupConfig{})
 				s3Poller := &s3PollerInput{
 					log:             logp.NewLogger(inputName),
 					config:          config,
 					metrics:         metrics,
 					s3:              s3API,
-					client:          client,
+					pipeline:        pipeline,
 					s3ObjectHandler: s3EventHandlerFactory,
 					states:          states,
 					provider:        "provider",

x-pack/filebeat/input/awss3/interfaces.go

@@ -17,7 +17,6 @@ import (
 	"github.com/aws/smithy-go/middleware"
 
 	"github.com/elastic/beats/v7/libbeat/beat"
-	awscommon "github.com/elastic/beats/v7/x-pack/libbeat/common/aws"
 
 	awssdk "github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/s3"
@@ -41,25 +40,9 @@ import (
 const s3RequestURLMetadataKey = `x-beat-s3-request-url`
 
 type sqsAPI interface {
-	sqsReceiver
-	sqsDeleter
-	sqsVisibilityChanger
-	sqsAttributeGetter
-}
-
-type sqsReceiver interface {
 	ReceiveMessage(ctx context.Context, maxMessages int) ([]types.Message, error)
-}
-
-type sqsDeleter interface {
 	DeleteMessage(ctx context.Context, msg *types.Message) error
-}
-
-type sqsVisibilityChanger interface {
 	ChangeMessageVisibility(ctx context.Context, msg *types.Message, timeout time.Duration) error
-}
-
-type sqsAttributeGetter interface {
 	GetQueueAttributes(ctx context.Context, attr []types.QueueAttributeName) (map[string]string, error)
 }
 
@@ -68,7 +51,7 @@ type sqsProcessor interface {
 	// given message and is responsible for updating the message's visibility
 	// timeout while it is being processed and for deleting it when processing
 	// completes successfully.
-	ProcessSQS(ctx context.Context, msg *types.Message) error
+	ProcessSQS(ctx context.Context, msg *types.Message, eventCallback func(e beat.Event)) sqsProcessingResult
 }
 
 // ------
@@ -103,25 +86,18 @@ type s3ObjectHandlerFactory interface {
 	// Create returns a new s3ObjectHandler that can be used to process the
 	// specified S3 object. If the handler is not configured to process the
 	// given S3 object (based on key name) then it will return nil.
-	Create(ctx context.Context, log *logp.Logger, client beat.Client, acker *awscommon.EventACKTracker, obj s3EventV2) s3ObjectHandler
+	Create(ctx context.Context, obj s3EventV2) s3ObjectHandler
 }
 
 type s3ObjectHandler interface {
 	// ProcessS3Object downloads the S3 object, parses it, creates events, and
-	// publishes them. It returns when processing finishes or when it encounters
-	// an unrecoverable error. It does not wait for the events to be ACKed by
-	// the publisher before returning (use eventACKTracker's Wait() method to
-	// determine this).
-	ProcessS3Object() error
+	// passes to the given callback. It returns when processing finishes or
+	// when it encounters an unrecoverable error.
+	ProcessS3Object(log *logp.Logger, eventCallback func(e beat.Event)) error
 
 	// FinalizeS3Object finalizes processing of an S3 object after the current
 	// batch is finished.
 	FinalizeS3Object() error
-
-	// Wait waits for every event published by ProcessS3Object() to be ACKed
-	// by the publisher before returning. Internally it uses the
-	// s3ObjectHandler eventACKTracker's Wait() method
-	Wait()
 }
 
 // ------