Adding upstream version 1.34.4.

Signed-off-by: Daniel Baumann <daniel@debian.org>

plugins/inputs/kinesis_consumer/consumer.go

@@ -0,0 +1,355 @@
+package kinesis_consumer
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/kinesis"
+	"github.com/aws/aws-sdk-go-v2/service/kinesis/types"
+
+	"github.com/influxdata/telegraf"
+)
+
+type recordHandler func(ctx context.Context, shard string, r *types.Record)
+
+type shardConsumer struct {
+	seqnr    string
+	interval time.Duration
+	log      telegraf.Logger
+
+	client *kinesis.Client
+	params *kinesis.GetShardIteratorInput
+
+	onMessage recordHandler
+}
+
+func (c *shardConsumer) consume(ctx context.Context, shard string) ([]types.ChildShard, error) {
+	ticker := time.NewTicker(c.interval)
+	defer ticker.Stop()
+
+	// Get the first shard iterator
+	iter, err := c.iterator(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("getting first shard iterator failed: %w", err)
+	}
+
+	for {
+		// Get new records from the shard
+		resp, err := c.client.GetRecords(ctx, &kinesis.GetRecordsInput{
+			ShardIterator: iter,
+		})
+		if err != nil {
+			// Handle recoverable errors
+			var throughputErr *types.ProvisionedThroughputExceededException
+			var expiredIterErr *types.ExpiredIteratorException
+			switch {
+			case errors.As(err, &throughputErr):
+				// Wait a second before trying again as suggested by
+				// https://docs.aws.amazon.com/streams/latest/dev/service-sizes-and-limits.html
+				c.log.Tracef("throughput exceeded when getting records for shard %s...", shard)
+				time.Sleep(time.Second)
+				continue
+			case errors.As(err, &expiredIterErr):
+				c.log.Tracef("iterator expired for shard %s...", shard)
+				if iter, err = c.iterator(ctx); err != nil {
+					return nil, fmt.Errorf("getting shard iterator failed: %w", err)
+				}
+				continue
+			case errors.Is(err, context.Canceled):
+				return nil, nil
+			default:
+				c.log.Tracef("get-records error is of type %T", err)
+				return nil, fmt.Errorf("getting records failed: %w", err)
+			}
+		}
+		c.log.Tracef("read %d records for shard %s...", len(resp.Records), shard)
+
+		// Check if we fully read the shard
+		if resp.NextShardIterator == nil {
+			return resp.ChildShards, nil
+		}
+		iter = resp.NextShardIterator
+
+		// Process the records and keep track of the last sequence number
+		// consumed for recreating the iterator.
+		for _, r := range resp.Records {
+			c.onMessage(ctx, shard, &r)
+			c.seqnr = *r.SequenceNumber
+			if errors.Is(ctx.Err(), context.Canceled) {
+				return nil, nil
+			}
+		}
+
+		// Wait for the poll interval to pass or cancel
+		select {
+		case <-ctx.Done():
+			return nil, nil
+		case <-ticker.C:
+			continue
+		}
+	}
+}
+
+func (c *shardConsumer) iterator(ctx context.Context) (*string, error) {
+	for {
+		resp, err := c.client.GetShardIterator(ctx, c.params)
+		if err != nil {
+			var throughputErr *types.ProvisionedThroughputExceededException
+			if errors.As(err, &throughputErr) {
+				// We called the function too often and should wait a bit
+				// until trying again
+				c.log.Tracef("throughput exceeded when getting iterator for shard %s...", *c.params.ShardId)
+				time.Sleep(time.Second)
+				continue
+			}
+
+			return nil, err
+		}
+		c.log.Tracef("successfully updated iterator for shard %s (%s)...", *c.params.ShardId, c.seqnr)
+		return resp.ShardIterator, nil
+	}
+}
+
+type consumer struct {
+	config              aws.Config
+	stream              string
+	iterType            types.ShardIteratorType
+	pollInterval        time.Duration
+	shardUpdateInterval time.Duration
+	log                 telegraf.Logger
+
+	onMessage recordHandler
+	position  func(shard string) string
+
+	client *kinesis.Client
+
+	shardsConsumed map[string]bool
+	shardConsumers map[string]*shardConsumer
+
+	wg sync.WaitGroup
+
+	sync.Mutex
+}
+
+func (c *consumer) init() error {
+	if c.stream == "" {
+		return errors.New("stream cannot be empty")
+	}
+	if c.pollInterval <= 0 {
+		return errors.New("invalid poll interval")
+	}
+
+	if c.onMessage == nil {
+		return errors.New("message handler is undefined")
+	}
+
+	c.shardsConsumed = make(map[string]bool)
+	c.shardConsumers = make(map[string]*shardConsumer)
+
+	return nil
+}
+
+func (c *consumer) start(ctx context.Context) {
+	// Setup the client
+	c.client = kinesis.NewFromConfig(c.config)
+
+	// Do the initial discovery of shards
+	if err := c.updateShardConsumers(ctx); err != nil {
+		c.log.Errorf("Initializing shards failed: %v", err)
+	}
+
+	// If the consumer has a shard-update interval, use a ticker to update
+	// available shards on a regular basis
+	if c.shardUpdateInterval <= 0 {
+		return
+	}
+	ticker := time.NewTicker(c.shardUpdateInterval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			if err := c.updateShardConsumers(ctx); err != nil {
+				c.log.Errorf("Updating shards failed: %v", err)
+			}
+		}
+	}
+}
+
+func (c *consumer) updateShardConsumers(ctx context.Context) error {
+	// List all shards of the given stream
+	var availableShards []types.Shard
+	req := &kinesis.ListShardsInput{StreamName: aws.String(c.stream)}
+	for {
+		resp, err := c.client.ListShards(ctx, req)
+		if err != nil {
+			return fmt.Errorf("listing shards failed: %w", err)
+		}
+		availableShards = append(availableShards, resp.Shards...)
+
+		if resp.NextToken == nil {
+			break
+		}
+
+		req = &kinesis.ListShardsInput{NextToken: resp.NextToken}
+	}
+	c.log.Tracef("got %d shards during update", len(availableShards))
+
+	// All following operations need to be locked to create a consistent
+	// state of the shards and consumers
+	c.Lock()
+	defer c.Unlock()
+
+	// Filter out all shards actively consumed already
+	inactiveShards := make([]types.Shard, 0, len(availableShards))
+	for _, shard := range availableShards {
+		id := *shard.ShardId
+		if _, found := c.shardConsumers[id]; found {
+			c.log.Tracef("shard %s is actively consumed...", id)
+			continue
+		}
+		c.log.Tracef("shard %s is not actively consumed...", id)
+		inactiveShards = append(inactiveShards, shard)
+	}
+
+	// Fill the shards already consumed and get the positions if the consumer
+	// is backed by an iterator store
+	newShards := make([]types.Shard, 0, len(inactiveShards))
+	seqnrs := make(map[string]string, len(inactiveShards))
+	for _, shard := range inactiveShards {
+		id := *shard.ShardId
+
+		if c.shardsConsumed[id] {
+			c.log.Tracef("shard %s is already fully consumed...", id)
+			continue
+		}
+		c.log.Tracef("shard %s is not fully consumed...", id)
+
+		// Retrieve the shard position from the store
+		if c.position != nil {
+			seqnr := c.position(id)
+			if seqnr == "" {
+				// A truely new shard
+				newShards = append(newShards, shard)
+				c.log.Tracef("shard %s is new...", id)
+				continue
+			}
+			seqnrs[id] = seqnr
+
+			// Check if we already fully consumed for closed shards
+			end := shard.SequenceNumberRange.EndingSequenceNumber
+			if end != nil && *end == seqnr {
+				c.log.Tracef("shard %s is closed and already fully consumed...", id)
+				c.shardsConsumed[id] = true
+				continue
+			}
+			c.log.Tracef("shard %s is not yet fully consumed...", id)
+		}
+
+		// The shard is not fully consumed yet so save the sequence number
+		// and the shard as "new".
+		newShards = append(newShards, shard)
+	}
+
+	// Filter all shards already fully consumed and create a new consumer for
+	// every remaining new shard respecting resharding artifacts
+	for _, shard := range newShards {
+		id := *shard.ShardId
+
+		// Handle resharding by making sure all parents are consumed already
+		// before starting a consumer on a child shard. If parents are not
+		// consumed fully we ignore this shard here as it will be reported
+		// by the call to `GetRecords` as a child later.
+		if shard.ParentShardId != nil && *shard.ParentShardId != "" {
+			pid := *shard.ParentShardId
+			if !c.shardsConsumed[pid] {
+				c.log.Tracef("shard %s has parent %s which is not fully consumed yet...", id, pid)
+				continue
+			}
+		}
+		if shard.AdjacentParentShardId != nil && *shard.AdjacentParentShardId != "" {
+			pid := *shard.AdjacentParentShardId
+			if !c.shardsConsumed[pid] {
+				c.log.Tracef("shard %s has adjacent parent %s which is not fully consumed yet...", id, pid)
+				continue
+			}
+		}
+
+		// Create a new consumer and start it
+		c.wg.Add(1)
+		go func(shardID string) {
+			defer c.wg.Done()
+			c.startShardConsumer(ctx, shardID, seqnrs[shardID])
+		}(id)
+	}
+
+	return nil
+}
+
+func (c *consumer) startShardConsumer(ctx context.Context, id, seqnr string) {
+	c.log.Tracef("starting consumer for shard %s at sequence number %q...", id, seqnr)
+	sc := &shardConsumer{
+		seqnr:     seqnr,
+		interval:  c.pollInterval,
+		log:       c.log,
+		onMessage: c.onMessage,
+		client:    c.client,
+		params: &kinesis.GetShardIteratorInput{
+			ShardId:           &id,
+			ShardIteratorType: c.iterType,
+			StreamName:        &c.stream,
+		},
+	}
+	if seqnr != "" {
+		sc.params.ShardIteratorType = types.ShardIteratorTypeAfterSequenceNumber
+		sc.params.StartingSequenceNumber = &seqnr
+	}
+	c.shardConsumers[id] = sc
+
+	childs, err := sc.consume(ctx, id)
+	if err != nil {
+		c.log.Errorf("Consuming shard %s failed: %v", id, err)
+		return
+	}
+	c.log.Tracef("finished consuming shard %s", id)
+
+	c.Lock()
+	defer c.Unlock()
+
+	c.shardsConsumed[id] = true
+	delete(c.shardConsumers, id)
+
+	for _, shard := range childs {
+		cid := *shard.ShardId
+
+		startable := true
+		for _, pid := range shard.ParentShards {
+			startable = startable && c.shardsConsumed[pid]
+		}
+		if !startable {
+			c.log.Tracef("child shard %s of shard %s is not startable as parents are fully consumed yet...", cid, id)
+			continue
+		}
+		c.log.Tracef("child shard %s of shard %s is startable...", cid, id)
+
+		var cseqnr string
+		if c.position != nil {
+			cseqnr = c.position(cid)
+		}
+		c.wg.Add(1)
+		go func() {
+			defer c.wg.Done()
+			c.startShardConsumer(ctx, cid, cseqnr)
+		}()
+	}
+}
+
+func (c *consumer) stop() {
+	c.wg.Wait()
+}