components-contrib/bindings/aws/kinesis/kinesis.go

/*
Copyright 2021 The Dapr 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 kinesis

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/google/uuid"
	"github.com/vmware/vmware-go-kcl/clientlibrary/config"
	"github.com/vmware/vmware-go-kcl/clientlibrary/interfaces"
	"github.com/vmware/vmware-go-kcl/clientlibrary/worker"

	aws_auth "github.com/dapr/components-contrib/authentication/aws"
	"github.com/dapr/components-contrib/bindings"
	"github.com/dapr/kit/logger"
)

// AWSKinesis allows receiving and sending data to/from AWS Kinesis stream.
type AWSKinesis struct {
	client   *kinesis.Kinesis
	metadata *kinesisMetadata

	worker       *worker.Worker
	workerConfig *config.KinesisClientLibConfiguration

	streamARN   *string
	consumerARN *string
	logger      logger.Logger
}

type kinesisMetadata struct {
	StreamName          string              `json:"streamName"`
	ConsumerName        string              `json:"consumerName"`
	Region              string              `json:"region"`
	Endpoint            string              `json:"endpoint"`
	AccessKey           string              `json:"accessKey"`
	SecretKey           string              `json:"secretKey"`
	SessionToken        string              `json:"sessionToken"`
	KinesisConsumerMode kinesisConsumerMode `json:"mode"`
}

type kinesisConsumerMode string

const (
	// ExtendedFanout - dedicated throughput through data stream api.
	ExtendedFanout kinesisConsumerMode = "extended"

	// SharedThroughput - shared throughput using checkpoint and monitoring.
	SharedThroughput kinesisConsumerMode = "shared"

	partitionKeyName = "partitionKey"
)

// recordProcessorFactory.
type recordProcessorFactory struct {
	logger  logger.Logger
	handler func(*bindings.ReadResponse) ([]byte, error)
}

type recordProcessor struct {
	logger  logger.Logger
	handler func(*bindings.ReadResponse) ([]byte, error)
}

// NewAWSKinesis returns a new AWS Kinesis instance.
func NewAWSKinesis(logger logger.Logger) *AWSKinesis {
	return &AWSKinesis{logger: logger}
}

// Init does metadata parsing and connection creation.
func (a *AWSKinesis) Init(metadata bindings.Metadata) error {
	m, err := a.parseMetadata(metadata)
	if err != nil {
		return err
	}

	if m.KinesisConsumerMode == "" {
		m.KinesisConsumerMode = SharedThroughput
	}

	if m.KinesisConsumerMode != SharedThroughput && m.KinesisConsumerMode != ExtendedFanout {
		return fmt.Errorf("%s invalid \"mode\" field %s", "aws.kinesis", m.KinesisConsumerMode)
	}

	client, err := a.getClient(m)
	if err != nil {
		return err
	}

	streamName := aws.String(m.StreamName)
	stream, err := client.DescribeStream(&kinesis.DescribeStreamInput{
		StreamName: streamName,
	})
	if err != nil {
		return err
	}

	if m.KinesisConsumerMode == SharedThroughput {
		kclConfig := config.NewKinesisClientLibConfigWithCredential(m.ConsumerName,
			m.StreamName, m.Region, m.ConsumerName,
			credentials.NewStaticCredentials(m.AccessKey, m.SecretKey, ""))
		a.workerConfig = kclConfig
	}

	a.streamARN = stream.StreamDescription.StreamARN
	a.metadata = m
	a.client = client

	return nil
}

func (a *AWSKinesis) Operations() []bindings.OperationKind {
	return []bindings.OperationKind{bindings.CreateOperation}
}

func (a *AWSKinesis) Invoke(req *bindings.InvokeRequest) (*bindings.InvokeResponse, error) {
	partitionKey := req.Metadata[partitionKeyName]
	if partitionKey == "" {
		partitionKey = uuid.New().String()
	}
	_, err := a.client.PutRecord(&kinesis.PutRecordInput{
		StreamName:   &a.metadata.StreamName,
		Data:         req.Data,
		PartitionKey: &partitionKey,
	})

	return nil, err
}

func (a *AWSKinesis) Read(handler func(*bindings.ReadResponse) ([]byte, error)) error {
	if a.metadata.KinesisConsumerMode == SharedThroughput {
		a.worker = worker.NewWorker(a.recordProcessorFactory(handler), a.workerConfig)
		err := a.worker.Start()
		if err != nil {
			return err
		}
	} else if a.metadata.KinesisConsumerMode == ExtendedFanout {
		ctx := context.Background()
		stream, err := a.client.DescribeStream(&kinesis.DescribeStreamInput{StreamName: &a.metadata.StreamName})
		if err != nil {
			return err
		}
		go a.Subscribe(ctx, *stream.StreamDescription, handler)
	}

	exitChan := make(chan os.Signal, 1)
	signal.Notify(exitChan, os.Interrupt, syscall.SIGTERM)
	<-exitChan

	if a.metadata.KinesisConsumerMode == SharedThroughput {
		go a.worker.Shutdown()
	} else if a.metadata.KinesisConsumerMode == ExtendedFanout {
		go a.deregisterConsumer(a.streamARN, a.consumerARN)
	}

	return nil
}

// Subscribe to all shards.
func (a *AWSKinesis) Subscribe(ctx context.Context, streamDesc kinesis.StreamDescription, handler func(*bindings.ReadResponse) ([]byte, error)) error {
	consumerARN, err := a.ensureConsumer(streamDesc.StreamARN)
	if err != nil {
		a.logger.Error(err)

		return err
	}

	a.consumerARN = consumerARN

	for {
		var wg sync.WaitGroup
		wg.Add(len(streamDesc.Shards))
		for i, shard := range streamDesc.Shards {
			go func(idx int, s *kinesis.Shard) error {
				defer wg.Done()
				sub, err := a.client.SubscribeToShardWithContext(ctx, &kinesis.SubscribeToShardInput{
					ConsumerARN:      consumerARN,
					ShardId:          s.ShardId,
					StartingPosition: &kinesis.StartingPosition{Type: aws.String(kinesis.ShardIteratorTypeLatest)},
				})
				if err != nil {
					a.logger.Error(err)

					return err
				}

				for event := range sub.EventStream.Events() {
					switch e := event.(type) {
					case *kinesis.SubscribeToShardEvent:
						for _, rec := range e.Records {
							handler(&bindings.ReadResponse{
								Data: rec.Data,
							})
						}
					}
				}

				return nil
			}(i, shard)
		}
		wg.Wait()
		time.Sleep(time.Minute * 5)
	}
}

func (a *AWSKinesis) ensureConsumer(streamARN *string) (*string, error) {
	consumer, err := a.client.DescribeStreamConsumer(&kinesis.DescribeStreamConsumerInput{
		ConsumerName: &a.metadata.ConsumerName,
		StreamARN:    streamARN,
	})
	if err != nil {
		arn, err := a.registerConsumer(streamARN)

		return arn, err
	}

	return consumer.ConsumerDescription.ConsumerARN, nil
}

func (a *AWSKinesis) registerConsumer(streamARN *string) (*string, error) {
	consumer, err := a.client.RegisterStreamConsumer(&kinesis.RegisterStreamConsumerInput{
		ConsumerName: &a.metadata.ConsumerName,
		StreamARN:    streamARN,
	})
	if err != nil {
		return nil, err
	}

	err = a.waitUntilConsumerExists(context.Background(), &kinesis.DescribeStreamConsumerInput{
		ConsumerName: &a.metadata.ConsumerName,
		StreamARN:    streamARN,
	})

	if err != nil {
		return nil, err
	}

	return consumer.Consumer.ConsumerARN, nil
}

func (a *AWSKinesis) deregisterConsumer(streamARN *string, consumerARN *string) error {
	if a.consumerARN != nil {
		_, err := a.client.DeregisterStreamConsumer(&kinesis.DeregisterStreamConsumerInput{
			ConsumerARN:  consumerARN,
			StreamARN:    streamARN,
			ConsumerName: &a.metadata.ConsumerName,
		})

		return err
	}

	return nil
}

func (a *AWSKinesis) waitUntilConsumerExists(ctx aws.Context, input *kinesis.DescribeStreamConsumerInput, opts ...request.WaiterOption) error {
	w := request.Waiter{
		Name:        "WaitUntilConsumerExists",
		MaxAttempts: 18,
		Delay:       request.ConstantWaiterDelay(10 * time.Second),
		Acceptors: []request.WaiterAcceptor{
			{
				State:   request.SuccessWaiterState,
				Matcher: request.PathWaiterMatch, Argument: "ConsumerDescription.ConsumerStatus",
				Expected: "ACTIVE",
			},
		},
		NewRequest: func(opts []request.Option) (*request.Request, error) {
			var inCpy *kinesis.DescribeStreamConsumerInput
			if input != nil {
				tmp := *input
				inCpy = &tmp
			}
			req, _ := a.client.DescribeStreamConsumerRequest(inCpy)
			req.SetContext(ctx)
			req.ApplyOptions(opts...)

			return req, nil
		},
	}
	w.ApplyOptions(opts...)

	return w.WaitWithContext(ctx)
}

func (a *AWSKinesis) getClient(metadata *kinesisMetadata) (*kinesis.Kinesis, error) {
	sess, err := aws_auth.GetClient(metadata.AccessKey, metadata.SecretKey, metadata.SessionToken, metadata.Region, metadata.Endpoint)
	if err != nil {
		return nil, err
	}
	k := kinesis.New(sess)

	return k, nil
}

func (a *AWSKinesis) parseMetadata(metadata bindings.Metadata) (*kinesisMetadata, error) {
	b, err := json.Marshal(metadata.Properties)
	if err != nil {
		return nil, err
	}

	var m kinesisMetadata
	err = json.Unmarshal(b, &m)
	if err != nil {
		return nil, err
	}

	return &m, nil
}

func (a *AWSKinesis) recordProcessorFactory(handler func(*bindings.ReadResponse) ([]byte, error)) interfaces.IRecordProcessorFactory {
	return &recordProcessorFactory{logger: a.logger, handler: handler}
}

func (r *recordProcessorFactory) CreateProcessor() interfaces.IRecordProcessor {
	return &recordProcessor{logger: r.logger, handler: r.handler}
}

func (p *recordProcessor) Initialize(input *interfaces.InitializationInput) {
	p.logger.Infof("Processing ShardId: %v at checkpoint: %v", input.ShardId, aws.StringValue(input.ExtendedSequenceNumber.SequenceNumber))
}

func (p *recordProcessor) ProcessRecords(input *interfaces.ProcessRecordsInput) {
	// don't process empty record
	if len(input.Records) == 0 {
		return
	}

	for _, v := range input.Records {
		p.handler(&bindings.ReadResponse{
			Data: v.Data,
		})
	}

	// checkpoint it after processing this batch
	lastRecordSequenceNumber := input.Records[len(input.Records)-1].SequenceNumber
	input.Checkpointer.Checkpoint(lastRecordSequenceNumber)
}

func (p *recordProcessor) Shutdown(input *interfaces.ShutdownInput) {
	if input.ShutdownReason == interfaces.TERMINATE {
		input.Checkpointer.Checkpoint(nil)
	}
}