├── .gitreview ├── .gitignore ├── logger ├── zerolog │ ├── zerolog_test.go │ └── zerolog.go ├── zap │ ├── zap_test.go │ └── zap.go ├── logger_test.go ├── logger.go └── logrus.go ├── LICENSE ├── clientlibrary ├── utils │ ├── uuid.go │ ├── random_test.go │ └── random.go ├── config │ ├── initial-stream-pos.go │ └── config_test.go ├── interfaces │ ├── sequence-number.go │ ├── record-processor.go │ ├── inputs.go │ └── record-processor-checkpointer.go ├── worker │ ├── record-processor-checkpointer.go │ ├── worker-fan-out.go │ ├── fan-out-shard-consumer.go │ ├── common-shard-consumer.go │ ├── polling-shard-consumer_test.go │ └── polling-shard-consumer.go ├── metrics │ ├── interfaces.go │ ├── prometheus │ │ └── prometheus.go │ └── cloudwatch │ │ └── cloudwatch.go ├── partition │ └── partition.go └── checkpoint │ ├── checkpointer.go │ ├── mock-dynamodb_test.go │ └── dynamodb-api.go ├── .github └── workflows │ └── vmware-go-kcl-v2-ci.yml ├── Makefile ├── _support └── scripts │ ├── sonar-scan.sh │ └── ci.sh ├── go.mod ├── CONTRIBUTING.md ├── README.md ├── test ├── logger_test.go ├── record_processor_test.go ├── worker_lease_stealing_test.go ├── worker_custom_test.go ├── lease_stealing_util_test.go ├── record_publisher_test.go └── worker_test.go └── CODE_OF_CONDUCT.md /.gitreview: -------------------------------------------------------------------------------- 1 | [gerrit] 2 | host=review.ec.eng.vmware.com 3 | port=29418 4 | project=cascade-kinesis-client 5 | defaultbranch=develop 6 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /gen 2 | /vendor 3 | !/vendor/manifest 4 | /bin 5 | /pkg 6 | /tmp 7 | /log 8 | /vms 9 | /run 10 | /go 11 | .hmake 12 | .hmakerc 13 | .project 14 | .idea 15 | .vscode 16 | *_mock_test.go 17 | filenames 18 | 19 | .DS_Store 20 | -------------------------------------------------------------------------------- /logger/zerolog/zerolog_test.go: -------------------------------------------------------------------------------- 1 | package zerolog 2 | 3 | import ( 4 | "github.com/vmware/vmware-go-kcl-v2/logger" 5 | "testing" 6 | ) 7 | 8 | func TestZeroLogLoggerWithConfig(t *testing.T) { 9 | config := logger.Configuration{ 10 | EnableConsole: true, 11 | ConsoleLevel: logger.Debug, 12 | ConsoleJSONFormat: true, 13 | EnableFile: true, 14 | FileLevel: logger.Info, 15 | FileJSONFormat: false, 16 | Filename: "/tmp/kcl-zerolog-log.log", 17 | } 18 | 19 | log := NewZerologLoggerWithConfig(config) 20 | 21 | contextLogger := log.WithFields(logger.Fields{"key1": "value1"}) 22 | contextLogger.Debugf("Starting with rs zerolog") 23 | contextLogger.Infof("Rs zerolog is awesome") 24 | } 25 | 26 | func TestZeroLogLogger(t *testing.T) { 27 | log := NewZerologLogger() 28 | 29 | contextLogger := log.WithFields(logger.Fields{"key1": "value1"}) 30 | contextLogger.Debugf("Starting with zerolog") 31 | contextLogger.Infof("Zerolog is awesome") 32 | } 33 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 VMware, Inc. 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /logger/zap/zap_test.go: -------------------------------------------------------------------------------- 1 | package zap_test 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | "github.com/vmware/vmware-go-kcl-v2/logger" 8 | "github.com/vmware/vmware-go-kcl-v2/logger/zap" 9 | uzap "go.uber.org/zap" 10 | ) 11 | 12 | func TestZapLoggerWithConfig(t *testing.T) { 13 | config := logger.Configuration{ 14 | EnableConsole: true, 15 | ConsoleLevel: logger.Debug, 16 | ConsoleJSONFormat: true, 17 | EnableFile: false, 18 | FileLevel: logger.Info, 19 | FileJSONFormat: true, 20 | Filename: "log.log", 21 | } 22 | 23 | log := zap.NewZapLoggerWithConfig(config) 24 | 25 | contextLogger := log.WithFields(logger.Fields{"key1": "value1"}) 26 | contextLogger.Debugf("Starting with zap") 27 | contextLogger.Infof("Zap is awesome") 28 | } 29 | 30 | func TestZapLogger(t *testing.T) { 31 | zapLogger, err := uzap.NewProduction() 32 | assert.Nil(t, err) 33 | 34 | log := zap.NewZapLogger(zapLogger.Sugar()) 35 | 36 | contextLogger := log.WithFields(logger.Fields{"key1": "value1"}) 37 | contextLogger.Debugf("Starting with zap") 38 | contextLogger.Infof("Zap is awesome") 39 | } 40 | -------------------------------------------------------------------------------- /clientlibrary/utils/uuid.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package utils 21 | package utils 22 | 23 | import ( 24 | guuid "github.com/google/uuid" 25 | ) 26 | 27 | // MustNewUUID generates a new UUID and panics if failed 28 | func MustNewUUID() string { 29 | id, err := guuid.NewUUID() 30 | if err != nil { 31 | panic(err) 32 | } 33 | 34 | return id.String() 35 | } 36 | -------------------------------------------------------------------------------- /.github/workflows/vmware-go-kcl-v2-ci.yml: -------------------------------------------------------------------------------- 1 | name: vmware-go-kcl-v2 2 | 3 | on: 4 | push: 5 | branches: [ main ] 6 | paths-ignore: [ README.md ] 7 | pull_request: 8 | branches: [ main ] 9 | paths-ignore: [ README.md ] 10 | 11 | jobs: 12 | build: 13 | name: Continous Integration 14 | runs-on: ubuntu-latest 15 | steps: 16 | - name: Check out code into the Go module directory 17 | uses: actions/checkout@v2 18 | 19 | - name: Set up Go 1.17.x 20 | uses: actions/setup-go@v2 21 | with: 22 | go-version: ^1.17 23 | id: go 24 | 25 | - name: Build 26 | shell: bash 27 | run: | 28 | make build 29 | 30 | - name: Test 31 | shell: bash 32 | run: | 33 | make test 34 | 35 | scans: 36 | name: Checks, Lints and Scans 37 | runs-on: ubuntu-latest 38 | steps: 39 | - name: Check out code into the Go module directory 40 | uses: actions/checkout@v2 41 | 42 | - name: Set up Go 1.17.x 43 | uses: actions/setup-go@v2 44 | with: 45 | go-version: ^1.17 46 | id: go 47 | 48 | - name: Format Check 49 | shell: bash 50 | run: | 51 | make format-check 52 | 53 | - name: Lint 54 | shell: bash 55 | run: | 56 | make lint-docker 57 | 58 | - name: Run Gosec Security Scanner 59 | uses: securego/gosec@master 60 | with: 61 | # let the report trigger content trigger a failure using the GitHub Security features. 62 | args: '-no-fail -fmt sarif -out results.sarif -exclude-dir internal -exclude-dir vendor -severity high ./...' 63 | 64 | - name: Upload SARIF file 65 | uses: github/codeql-action/upload-sarif@v1 66 | with: 67 | # path to SARIF file relative to the root of the repository 68 | sarif_file: results.sarif 69 | -------------------------------------------------------------------------------- /clientlibrary/utils/random_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package utils 21 | package utils 22 | 23 | import ( 24 | "fmt" 25 | "testing" 26 | "time" 27 | ) 28 | 29 | func TestRandom(t *testing.T) { 30 | for i := 0; i < 10; i++ { 31 | s1 := RandStringBytesMaskImpr(10) 32 | s2 := RandStringBytesMaskImpr(10) 33 | if s1 == s2 { 34 | t.Fatalf("failed in generating random string. s1: %s, s2: %s", s1, s2) 35 | } 36 | fmt.Println(s1) 37 | fmt.Println(s2) 38 | } 39 | } 40 | 41 | func TestRandomNum(t *testing.T) { 42 | for i := 0; i < 10; i++ { 43 | seed := time.Now().UTC().Second() 44 | s1 := RandStringBytesMaskImpr(seed) 45 | s2 := RandStringBytesMaskImpr(seed) 46 | if s1 == s2 { 47 | t.Fatalf("failed in generating random string. s1: %s, s2: %s", s1, s2) 48 | } 49 | fmt.Println(s1) 50 | fmt.Println(s2) 51 | } 52 | } 53 | -------------------------------------------------------------------------------- /logger/logger_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | // Note: The implementation comes from https://www.mountedthoughts.com/golang-logger-interface/ 20 | 21 | package logger 22 | 23 | import ( 24 | "testing" 25 | 26 | "github.com/sirupsen/logrus" 27 | ) 28 | 29 | func TestLogrusLoggerWithConfig(t *testing.T) { 30 | config := Configuration{ 31 | EnableConsole: true, 32 | ConsoleLevel: Debug, 33 | ConsoleJSONFormat: false, 34 | EnableFile: false, 35 | FileLevel: Info, 36 | FileJSONFormat: true, 37 | } 38 | 39 | log := NewLogrusLoggerWithConfig(config) 40 | 41 | contextLogger := log.WithFields(Fields{"key1": "value1"}) 42 | contextLogger.Debugf("Starting with logrus") 43 | contextLogger.Infof("Logrus is awesome") 44 | } 45 | 46 | func TestLogrusLogger(t *testing.T) { 47 | // adapts to Logger interface 48 | log := NewLogrusLogger(logrus.StandardLogger()) 49 | 50 | contextLogger := log.WithFields(Fields{"key1": "value1"}) 51 | contextLogger.Debugf("Starting with logrus") 52 | contextLogger.Infof("Logrus is awesome") 53 | } 54 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: help 2 | help: ## - Show this help message 3 | @printf "\033[32m\xE2\x9c\x93 usage: make [target]\n\n\033[0m" 4 | @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-30s\033[0m %s\n", $$1, $$2}' 5 | 6 | .PHONY: up 7 | up: ## - start docker compose 8 | @ cd _support/docker && docker-compose -f docker-compose.yml up 9 | 10 | .PHONY: build-common 11 | build-common: ## - execute build common tasks clean and mod tidy 12 | @ go version 13 | @ go clean 14 | @ go mod download && go mod tidy 15 | @ go mod verify 16 | 17 | .PHONY: build 18 | build: build-common ## - build a debug binary to the current platform (windows, linux or darwin(mac)) 19 | @ echo building 20 | @ go build -v ./... 21 | @ echo "done" 22 | 23 | .PHONY: format-check 24 | format-check: ## - check files format using gofmt 25 | @ ./_support/scripts/ci.sh fmtCheck 26 | 27 | .PHONY: format-check 28 | format: ## - apply golang file format using gofmt 29 | @ ./_support/scripts/ci.sh format 30 | 31 | .PHONY: test 32 | test: build-common ## - execute go test command for unit and mocked tests 33 | @ ./_support/scripts/ci.sh unitTest 34 | 35 | .PHONY: integration-test 36 | integration-test: ## - execute go test command for integration tests (aws credentials needed) 37 | @ go test -v -cover -race ./test 38 | 39 | .PHONY: scan 40 | scan: ## - execute static code analysis 41 | @ ./_support/scripts/ci.sh scan 42 | 43 | .PHONY: local-scan 44 | local-scan: ## - execute static code analysis locally 45 | @ ./_support/scripts/ci.sh localScan 46 | 47 | .PHONY: lint 48 | lint: ## - runs golangci-lint 49 | @ ./_support/scripts/ci.sh lint 50 | 51 | .PHONY: lint-docker 52 | lint-docker: ## - runs golangci-lint with docker container 53 | @ ./_support/scripts/ci.sh lintDocker 54 | 55 | .PHONY: sonar-scan 56 | sonar-scan: ## - start sonar qube locally with docker (you will need docker installed in your machine) 57 | @ # after start, setup a new project with the name sms-local and a new token sms-token, fill the token against the -Dsonar.login= parameter. 58 | @ # login with user: admin pwd: vmware 59 | @ $(SHELL) _support/scripts/sonar-scan.sh 60 | 61 | .PHONY: sonar-stop 62 | sonar-stop: ## - stop sonar qube docker container 63 | @ docker stop sonarqube 64 | -------------------------------------------------------------------------------- /clientlibrary/utils/random.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package utils 21 | package utils 22 | 23 | import ( 24 | "crypto/rand" 25 | "math/big" 26 | "time" 27 | ) 28 | 29 | const letterBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" 30 | const ( 31 | letterIdxBits = 6 // 6 bits to represent a letter index 32 | letterIdxMask = 1<= 0; { 42 | if remain == 0 { 43 | rnd, _ = rand.Int(rand.Reader, big.NewInt(seed)) 44 | cache, remain = rnd.Int64(), letterIdxMax 45 | } 46 | if idx := int(cache & letterIdxMask); idx < len(letterBytes) { 47 | b[i] = letterBytes[idx] 48 | i-- 49 | } 50 | cache >>= letterIdxBits 51 | remain-- 52 | } 53 | 54 | return string(b) 55 | } 56 | -------------------------------------------------------------------------------- /clientlibrary/config/initial-stream-pos.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package config 21 | // The implementation is derived from https://github.com/awslabs/amazon-kinesis-client 22 | /* 23 | * Copyright 2014-2015 Amazon.com, Inc. or its affiliates. All Rights Reserved. 24 | * 25 | * Licensed under the Amazon Software License (the "License"). 26 | * You may not use this file except in compliance with the License. 27 | * A copy of the License is located at 28 | * 29 | * http://aws.amazon.com/asl/ 30 | * 31 | * or in the "license" file accompanying this file. This file is distributed 32 | * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either 33 | * express or implied. See the License for the specific language governing 34 | * permissions and limitations under the License. 35 | */ 36 | 37 | package config 38 | 39 | import ( 40 | "time" 41 | ) 42 | 43 | func newInitialPositionAtTimestamp(timestamp *time.Time) *InitialPositionInStreamExtended { 44 | return &InitialPositionInStreamExtended{Position: AT_TIMESTAMP, Timestamp: timestamp} 45 | } 46 | 47 | func newInitialPosition(position InitialPositionInStream) *InitialPositionInStreamExtended { 48 | return &InitialPositionInStreamExtended{Position: position, Timestamp: nil} 49 | } 50 | -------------------------------------------------------------------------------- /_support/scripts/sonar-scan.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | ######################## 4 | # requirements: # 5 | # 0. docker # 6 | # 1. wget # 7 | # 2. curl # 8 | # 3. jq # 9 | # 4. sonar-scanner # 10 | ######################## 11 | 12 | set -e 13 | 14 | projectKey="vmware-go-kcl-v2" 15 | user_tokenName="local_token" 16 | username="admin" 17 | user_password="admin" 18 | new_password="vmware" 19 | url="http://localhost" 20 | port="9000" 21 | 22 | if [[ "$( docker container inspect -f '{{.State.Running}}' sonarqube )" == "true" ]]; 23 | then 24 | docker ps 25 | else 26 | docker run --rm -d --name sonarqube -e SONAR_ES_BOOTSTRAP_CHECKS_DISABLE=true -p 9000:9000 sonarqube 27 | fi 28 | 29 | echo "waiting for sonarqube starts..." 30 | wget -q -O - "$@" http://localhost:9000 | awk '/STARTING/{ print $0 }' | xargs 31 | 32 | STATUS="$(wget -q -O - "$@" http://localhost:9000 | awk '/UP/{ print $0 }')" 33 | while [ -z "$STATUS" ] 34 | do 35 | sleep 2 36 | STATUS="$(wget -q -O - "$@" http://localhost:9000 | awk '/UP/{ print $0 }')" 37 | printf "." 38 | done 39 | 40 | printf '\n %s' "${STATUS}" | xargs 41 | echo "" 42 | 43 | # change the default password to avoid create a new one when login for the very first time 44 | curl -u ${username}:${user_password} -X POST "${url}:${port}/api/users/change_password?login=${username}&previousPassword=${user_password}&password=${new_password}" 45 | 46 | # search the specific user tokens for SonarQube 47 | hasToken=$(curl --silent -u ${username}:${new_password} -X GET "${url}:${port}/api/user_tokens/search") 48 | if [[ -n "${hasToken}" ]]; then 49 | # Revoke the user token for SonarQube 50 | curl -X POST -H "Content-Type: application/x-www-form-urlencoded" -d "name=${user_tokenName}" -u ${username}:${new_password} "${url}:${port}"/api/user_tokens/revoke 51 | fi 52 | 53 | # generate new token 54 | token=$(curl --silent -X POST -H "Content-Type: application/x-www-form-urlencoded" -d "name=${user_tokenName}" -u ${username}:${new_password} "${url}:${port}"/api/user_tokens/generate | jq '.token' | xargs) 55 | 56 | # scan and push the results to localhost docker container 57 | sonar-scanner -Dsonar.projectKey="${projectKey}" \ 58 | -Dsonar.projectName="${projectKey}" \ 59 | -Dsonar.sources=. \ 60 | -Dsonar.exclusions="internal/records/**, test/**" \ 61 | -Dsonar.host.url="${url}:${port}" \ 62 | -Dsonar.login="${token}" 63 | 64 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/vmware/vmware-go-kcl-v2 2 | 3 | go 1.21 4 | 5 | require ( 6 | github.com/aws/aws-sdk-go-v2 v1.25.2 7 | github.com/aws/aws-sdk-go-v2/config v1.27.5 8 | github.com/aws/aws-sdk-go-v2/credentials v1.17.5 9 | github.com/aws/aws-sdk-go-v2/service/cloudwatch v1.36.1 10 | github.com/aws/aws-sdk-go-v2/service/dynamodb v1.30.2 11 | github.com/aws/aws-sdk-go-v2/service/kinesis v1.27.1 12 | github.com/awslabs/kinesis-aggregation/go/v2 v2.0.0-20211222152315-953b66f67407 13 | github.com/golang/protobuf v1.5.2 14 | github.com/google/uuid v1.3.0 15 | github.com/prometheus/client_golang v1.11.1 16 | github.com/prometheus/common v0.32.1 17 | github.com/rs/zerolog v1.26.1 18 | github.com/sirupsen/logrus v1.8.1 19 | github.com/stretchr/testify v1.8.1 20 | go.uber.org/zap v1.20.0 21 | gopkg.in/natefinch/lumberjack.v2 v2.0.0 22 | ) 23 | 24 | require ( 25 | github.com/BurntSushi/toml v0.4.1 // indirect 26 | github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.1 // indirect 27 | github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.15.2 // indirect 28 | github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.2 // indirect 29 | github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.2 // indirect 30 | github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 // indirect 31 | github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.1 // indirect 32 | github.com/aws/aws-sdk-go-v2/service/internal/endpoint-discovery v1.9.3 // indirect 33 | github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.3 // indirect 34 | github.com/aws/aws-sdk-go-v2/service/sso v1.20.1 // indirect 35 | github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.1 // indirect 36 | github.com/aws/aws-sdk-go-v2/service/sts v1.28.2 // indirect 37 | github.com/aws/smithy-go v1.20.1 // indirect 38 | github.com/beorn7/perks v1.0.1 // indirect 39 | github.com/cespare/xxhash/v2 v2.1.2 // indirect 40 | github.com/davecgh/go-spew v1.1.1 // indirect 41 | github.com/jmespath/go-jmespath v0.4.0 // indirect 42 | github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect 43 | github.com/pmezard/go-difflib v1.0.0 // indirect 44 | github.com/prometheus/client_model v0.2.0 // indirect 45 | github.com/prometheus/procfs v0.7.3 // indirect 46 | github.com/stretchr/objx v0.5.0 // indirect 47 | go.uber.org/atomic v1.9.0 // indirect 48 | go.uber.org/multierr v1.7.0 // indirect 49 | golang.org/x/sys v0.1.0 // indirect 50 | google.golang.org/protobuf v1.27.1 // indirect 51 | gopkg.in/yaml.v3 v3.0.1 // indirect 52 | ) 53 | -------------------------------------------------------------------------------- /clientlibrary/interfaces/sequence-number.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package interfaces 21 | // The implementation is derived from https://github.com/awslabs/amazon-kinesis-client 22 | /* 23 | * Copyright 2014-2015 Amazon.com, Inc. or its affiliates. All Rights Reserved. 24 | * 25 | * Licensed under the Amazon Software License (the "License"). 26 | * You may not use this file except in compliance with the License. 27 | * A copy of the License is located at 28 | * 29 | * http://aws.amazon.com/asl/ 30 | * 31 | * or in the "license" file accompanying this file. This file is distributed 32 | * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either 33 | * express or implied. See the License for the specific language governing 34 | * permissions and limitations under the License. 35 | */ 36 | 37 | package interfaces 38 | 39 | // ExtendedSequenceNumber represents a two-part sequence number for records aggregated by the Kinesis Producer Library. 40 | // 41 | // The KPL combines multiple user records into a single Kinesis record. Each user record therefore has an integer 42 | // sub-sequence number, in addition to the regular sequence number of the Kinesis record. The sub-sequence number 43 | // is used to checkpoint within an aggregated record. 44 | type ExtendedSequenceNumber struct { 45 | SequenceNumber *string 46 | SubSequenceNumber int64 47 | } 48 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing to vmware-go-kcl-v2 2 | 3 | The vmware-go-kcl-v2 project team welcomes contributions from the community. Before you start working with vmware-go-kcl-v2, please 4 | read our [Developer Certificate of Origin](https://cla.vmware.com/dco). All contributions to this repository must be 5 | signed as described on that page. Your signature certifies that you wrote the patch or have the right to pass it on 6 | as an open-source patch. 7 | 8 | ## Community 9 | 10 | ## Contribution Flow 11 | 12 | This is a rough outline of what a contributor's workflow looks like: 13 | 14 | - Create a topic branch from where you want to base your work 15 | - Make commits of logical units 16 | - Make sure your commit messages are in the proper format (see below) 17 | - Push your changes to a topic branch in your fork of the repository 18 | - Submit a pull request 19 | 20 | Example: 21 | 22 | ``` shell 23 | git remote add upstream https://github.com/vmware/@(project).git 24 | git checkout -b my-new-feature main 25 | git commit -a 26 | git push origin my-new-feature 27 | ``` 28 | 29 | ### Staying In Sync With Upstream 30 | 31 | When your branch gets out of sync with the vmware/main branch, use the following to update: 32 | 33 | ``` shell 34 | git checkout my-new-feature 35 | git fetch -a 36 | git pull --rebase upstream master 37 | git push --force-with-lease origin my-new-feature 38 | ``` 39 | 40 | ### Updating pull requests 41 | 42 | If your PR fails to pass CI or needs changes based on code review, you'll most likely want to squash these changes into 43 | existing commits. 44 | 45 | If your pull request contains a single commit or your changes are related to the most recent commit, you can simply 46 | amend the commit. 47 | 48 | ``` shell 49 | git add . 50 | git commit --amend 51 | git push --force-with-lease origin my-new-feature 52 | ``` 53 | 54 | If you need to squash changes into an earlier commit, you can use: 55 | 56 | ``` shell 57 | git add . 58 | git commit --fixup 59 | git rebase -i --autosquash main 60 | git push --force-with-lease origin my-new-feature 61 | ``` 62 | 63 | Be sure to add a comment to the PR indicating your new changes are ready to review, as GitHub does not generate a 64 | notification when you git push. 65 | 66 | ### Formatting Commit Messages 67 | 68 | We follow the conventions on [How to Write a Git Commit Message](http://chris.beams.io/posts/git-commit/). 69 | 70 | Be sure to include any related GitHub issue references in the commit message. See 71 | [GFM syntax](https://guides.github.com/features/mastering-markdown/#GitHub-flavored-markdown) for referencing issues 72 | and commits. 73 | 74 | ## Reporting Bugs and Creating Issues 75 | 76 | When opening a new issue, try to roughly follow the commit message format conventions above. 77 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # VMWare Go KCL v2 2 | 3 | ![technology Go](https://img.shields.io/badge/technology-go-blue.svg) 4 | [![Go Report Card](https://goreportcard.com/badge/github.com/vmware/vmware-go-kcl-v2)](https://goreportcard.com/report/github.com/vmware/vmware-go-kcl-v2) 5 | [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT) 6 | [![vmware-go-kcl-v2](https://github.com/vmware/vmware-go-kcl-v2/actions/workflows/vmware-go-kcl-v2-ci.yml/badge.svg)](https://github.com/vmware/vmware-go-kcl-v2/actions/workflows/vmware-go-kcl-v2-ci.yml) 7 | 8 | ## Overview 9 | 10 | VMware-Go-KCL-V2 is a native open-source Go library for Amazon Kinesis Data Stream (KDS) consumption. It allows developers 11 | to program KDS consumers in lightweight Go language and still take advantage of the features presented by the native 12 | KDS Java API libraries. 13 | 14 | [vmware-go-kcl-v2](https://github.com/vmware/vmware-go-kcl-v2) is a VMWare originated open-source project for AWS Kinesis 15 | Client Library in Go. Within VMware, we have seen adoption in vSecureState and Carbon Black. In addition, Carbon Black 16 | has contributed to the vmware-go-kcl codebase and heavily used it in the product. Besides, 17 | [vmware-go-kcl-v2](https://github.com/vmware/vmware-go-kcl-v2) has got 18 | [recognition](https://www.linkedin.com/posts/adityakrish_vmware-go-kcl-a-native-open-source-go-programming-activity-6810626798133616640-B6W8/), 19 | and [contributions](https://github.com/vmware/vmware-go-kcl-v2/graphs/contributors) from the industry. 20 | 21 | `vmware-go-kcl-v2` is the v2 version of VMWare KCL for the Go programming language by utilizing [AWS Go SDK V2](https://github.com/aws/aws-sdk-go-v2). 22 | 23 | ## Try it out 24 | 25 | ### Prerequisites 26 | 27 | * [aws-sdk-go-v2](https://github.com/aws/aws-sdk-go-v2) 28 | * The v2 SDK requires a minimum version of `Go 1.17`. 29 | * [gosec](https://github.com/securego/gosec) 30 | 31 | ### Build & Run 32 | 33 | 1. Initialize Project 34 | 35 | 2. Build 36 | > `make build` 37 | 38 | 3. Test 39 | > `make test` 40 | 41 | ## Documentation 42 | 43 | VMware-Go-KCL matches exactly the same interface and programming model from original Amazon KCL, the best place for getting reference, tutorial is from Amazon itself: 44 | 45 | * [Developing Consumers Using the Kinesis Client Library](https://docs.aws.amazon.com/streams/latest/dev/developing-consumers-with-kcl.html) 46 | * [Troubleshooting](https://docs.aws.amazon.com/streams/latest/dev/troubleshooting-consumers.html) 47 | * [Advanced Topics](https://docs.aws.amazon.com/streams/latest/dev/advanced-consumers.html) 48 | 49 | ## Contributing 50 | 51 | The vmware-go-kcl-v2 project team welcomes contributions from the community. Before you start working with vmware-go-kcl-v2, please 52 | read our [Developer Certificate of Origin](https://cla.vmware.com/dco). All contributions to this repository must be 53 | signed as described on that page. Your signature certifies that you wrote the patch or have the right to pass it on 54 | as an open-source patch. For more detailed information, refer to [CONTRIBUTING.md](CONTRIBUTING.md). 55 | 56 | ## License 57 | 58 | MIT License 59 | -------------------------------------------------------------------------------- /clientlibrary/worker/record-processor-checkpointer.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package worker 21 | package worker 22 | 23 | import ( 24 | "github.com/aws/aws-sdk-go-v2/aws" 25 | chk "github.com/vmware/vmware-go-kcl-v2/clientlibrary/checkpoint" 26 | kcl "github.com/vmware/vmware-go-kcl-v2/clientlibrary/interfaces" 27 | par "github.com/vmware/vmware-go-kcl-v2/clientlibrary/partition" 28 | ) 29 | 30 | type ( 31 | 32 | // PreparedCheckpointer 33 | /* 34 | * Objects of this class are prepared to checkpoint at a specific sequence number. They use an 35 | * IRecordProcessorCheckpointer to do the actual checkpointing, so their checkpoint is subject to the same 'didn't go 36 | * backwards' validation as a normal checkpoint. 37 | */ 38 | PreparedCheckpointer struct { 39 | pendingCheckpointSequenceNumber *kcl.ExtendedSequenceNumber 40 | checkpointer kcl.IRecordProcessorCheckpointer 41 | } 42 | 43 | //RecordProcessorCheckpointer 44 | /* 45 | * This class is used to enable RecordProcessors to checkpoint their progress. 46 | * The Amazon Kinesis Client Library will instantiate an object and provide a reference to the application 47 | * RecordProcessor instance. Amazon Kinesis Client Library will create one instance per shard assignment. 48 | */ 49 | RecordProcessorCheckpointer struct { 50 | shard *par.ShardStatus 51 | checkpoint chk.Checkpointer 52 | } 53 | ) 54 | 55 | func NewRecordProcessorCheckpoint(shard *par.ShardStatus, checkpoint chk.Checkpointer) kcl.IRecordProcessorCheckpointer { 56 | return &RecordProcessorCheckpointer{ 57 | shard: shard, 58 | checkpoint: checkpoint, 59 | } 60 | } 61 | 62 | func (pc *PreparedCheckpointer) GetPendingCheckpoint() *kcl.ExtendedSequenceNumber { 63 | return pc.pendingCheckpointSequenceNumber 64 | } 65 | 66 | func (pc *PreparedCheckpointer) Checkpoint() error { 67 | return pc.checkpointer.Checkpoint(pc.pendingCheckpointSequenceNumber.SequenceNumber) 68 | } 69 | 70 | func (rc *RecordProcessorCheckpointer) Checkpoint(sequenceNumber *string) error { 71 | // checkpoint the last sequence of a closed shard 72 | if sequenceNumber == nil { 73 | rc.shard.SetCheckpoint(chk.ShardEnd) 74 | } else { 75 | rc.shard.SetCheckpoint(aws.ToString(sequenceNumber)) 76 | } 77 | 78 | return rc.checkpoint.CheckpointSequence(rc.shard) 79 | } 80 | 81 | func (rc *RecordProcessorCheckpointer) PrepareCheckpoint(_ *string) (kcl.IPreparedCheckpointer, error) { 82 | return &PreparedCheckpointer{}, nil 83 | } 84 | -------------------------------------------------------------------------------- /test/logger_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | // Note: The implementation comes from https://www.mountedthoughts.com/golang-logger-interface/ 20 | 21 | package test 22 | 23 | import ( 24 | "github.com/stretchr/testify/assert" 25 | 26 | "testing" 27 | 28 | "github.com/sirupsen/logrus" 29 | "go.uber.org/zap" 30 | 31 | "github.com/vmware/vmware-go-kcl-v2/logger" 32 | zaplogger "github.com/vmware/vmware-go-kcl-v2/logger/zap" 33 | ) 34 | 35 | func TestZapLoggerWithConfig(t *testing.T) { 36 | config := logger.Configuration{ 37 | EnableConsole: true, 38 | ConsoleLevel: logger.Debug, 39 | ConsoleJSONFormat: true, 40 | EnableFile: true, 41 | FileLevel: logger.Info, 42 | FileJSONFormat: true, 43 | Filename: "log.log", 44 | } 45 | 46 | log := zaplogger.NewZapLoggerWithConfig(config) 47 | 48 | contextLogger := log.WithFields(logger.Fields{"key1": "value1"}) 49 | contextLogger.Debugf("Starting with zap") 50 | contextLogger.Infof("Zap is awesome") 51 | } 52 | 53 | func TestZapLogger(t *testing.T) { 54 | zapLogger, err := zap.NewProduction() 55 | assert.Nil(t, err) 56 | 57 | log := zaplogger.NewZapLogger(zapLogger.Sugar()) 58 | 59 | contextLogger := log.WithFields(logger.Fields{"key1": "value1"}) 60 | contextLogger.Debugf("Starting with zap") 61 | contextLogger.Infof("Zap is awesome") 62 | } 63 | 64 | func TestLogrusLoggerWithConfig(t *testing.T) { 65 | config := logger.Configuration{ 66 | EnableConsole: true, 67 | ConsoleLevel: logger.Debug, 68 | ConsoleJSONFormat: false, 69 | EnableFile: true, 70 | FileLevel: logger.Info, 71 | FileJSONFormat: true, 72 | Filename: "log.log", 73 | } 74 | 75 | log := logger.NewLogrusLoggerWithConfig(config) 76 | 77 | contextLogger := log.WithFields(logger.Fields{"key1": "value1"}) 78 | contextLogger.Debugf("Starting with logrus") 79 | contextLogger.Infof("Logrus is awesome") 80 | } 81 | 82 | func TestLogrusLogger(t *testing.T) { 83 | // adapts to Logger interface from *logrus.Logger 84 | log := logger.NewLogrusLogger(logrus.StandardLogger()) 85 | 86 | contextLogger := log.WithFields(logger.Fields{"key1": "value1"}) 87 | contextLogger.Debugf("Starting with logrus") 88 | contextLogger.Infof("Logrus is awesome") 89 | } 90 | 91 | func TestLogrusLoggerWithFieldsAtInit(t *testing.T) { 92 | // adapts to Logger interface from *logrus.Entry 93 | fieldLogger := logrus.StandardLogger().WithField("key0", "value0") 94 | log := logger.NewLogrusLogger(fieldLogger) 95 | 96 | contextLogger := log.WithFields(logger.Fields{"key1": "value1"}) 97 | contextLogger.Debugf("Starting with logrus") 98 | contextLogger.Infof("Structured logging is awesome") 99 | } 100 | -------------------------------------------------------------------------------- /clientlibrary/metrics/interfaces.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package metrics 21 | // The implementation is derived from https://github.com/patrobinson/gokini 22 | // 23 | // Copyright 2018 Patrick robinson. 24 | // 25 | // Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 26 | // 27 | // The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 28 | // 29 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 30 | package metrics 31 | 32 | type MonitoringService interface { 33 | Init(appName, streamName, workerID string) error 34 | Start() error 35 | IncrRecordsProcessed(shard string, count int) 36 | IncrBytesProcessed(shard string, count int64) 37 | MillisBehindLatest(shard string, milliSeconds float64) 38 | DeleteMetricMillisBehindLatest(shard string) 39 | LeaseGained(shard string) 40 | LeaseLost(shard string) 41 | LeaseRenewed(shard string) 42 | RecordGetRecordsTime(shard string, time float64) 43 | RecordProcessRecordsTime(shard string, time float64) 44 | Shutdown() 45 | } 46 | 47 | // NoopMonitoringService implements MonitoringService by does nothing. 48 | type NoopMonitoringService struct{} 49 | 50 | func (NoopMonitoringService) Init(_, _, _ string) error { return nil } 51 | func (NoopMonitoringService) Start() error { return nil } 52 | func (NoopMonitoringService) Shutdown() {} 53 | 54 | func (NoopMonitoringService) IncrRecordsProcessed(_ string, _ int) {} 55 | func (NoopMonitoringService) IncrBytesProcessed(_ string, _ int64) {} 56 | func (NoopMonitoringService) MillisBehindLatest(_ string, _ float64) {} 57 | func (NoopMonitoringService) DeleteMetricMillisBehindLatest(_ string) {} 58 | func (NoopMonitoringService) LeaseGained(_ string) {} 59 | func (NoopMonitoringService) LeaseLost(_ string) {} 60 | func (NoopMonitoringService) LeaseRenewed(_ string) {} 61 | func (NoopMonitoringService) RecordGetRecordsTime(_ string, _ float64) {} 62 | func (NoopMonitoringService) RecordProcessRecordsTime(_ string, _ float64) {} 63 | -------------------------------------------------------------------------------- /test/record_processor_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2020 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package test 21 | package test 22 | 23 | import ( 24 | "testing" 25 | 26 | "github.com/aws/aws-sdk-go-v2/aws" 27 | "github.com/stretchr/testify/assert" 28 | 29 | kc "github.com/vmware/vmware-go-kcl-v2/clientlibrary/interfaces" 30 | ) 31 | 32 | // Record processor factory is used to create RecordProcessor 33 | func recordProcessorFactory(t *testing.T) kc.IRecordProcessorFactory { 34 | return &dumpRecordProcessorFactory{t: t} 35 | } 36 | 37 | // simple record processor and dump everything 38 | type dumpRecordProcessorFactory struct { 39 | t *testing.T 40 | } 41 | 42 | func (d *dumpRecordProcessorFactory) CreateProcessor() kc.IRecordProcessor { 43 | return &dumpRecordProcessor{ 44 | t: d.t, 45 | } 46 | } 47 | 48 | // Create a dump record processor for printing out all data from record. 49 | type dumpRecordProcessor struct { 50 | t *testing.T 51 | count int 52 | } 53 | 54 | func (dd *dumpRecordProcessor) Initialize(input *kc.InitializationInput) { 55 | dd.t.Logf("Processing SharId: %v at checkpoint: %v", input.ShardId, aws.ToString(input.ExtendedSequenceNumber.SequenceNumber)) 56 | shardID = input.ShardId 57 | dd.count = 0 58 | } 59 | 60 | func (dd *dumpRecordProcessor) ProcessRecords(input *kc.ProcessRecordsInput) { 61 | dd.t.Log("Processing Records...") 62 | 63 | // don't process empty record 64 | if len(input.Records) == 0 { 65 | return 66 | } 67 | 68 | for _, v := range input.Records { 69 | dd.t.Logf("Record = %s", v.Data) 70 | assert.Equal(dd.t, specstr, string(v.Data)) 71 | dd.count++ 72 | } 73 | 74 | // checkpoint it after processing this batch. 75 | // Especially, for processing de-aggregated KPL records, checkpointing has to happen at the end of batch 76 | // because de-aggregated records share the same sequence number. 77 | lastRecordSequenceNumber := input.Records[len(input.Records)-1].SequenceNumber 78 | // Calculate the time taken from polling records and delivering to record processor for a batch. 79 | diff := input.CacheExitTime.Sub(*input.CacheEntryTime) 80 | dd.t.Logf("Checkpoint progress at: %v, MillisBehindLatest = %v, KCLProcessTime = %v", lastRecordSequenceNumber, input.MillisBehindLatest, diff) 81 | _ = input.Checkpointer.Checkpoint(lastRecordSequenceNumber) 82 | } 83 | 84 | func (dd *dumpRecordProcessor) Shutdown(input *kc.ShutdownInput) { 85 | dd.t.Logf("Shutdown Reason: %v", aws.ToString(kc.ShutdownReasonMessage(input.ShutdownReason))) 86 | dd.t.Logf("Processed Record Count = %d", dd.count) 87 | 88 | // When the value of {@link ShutdownInput#getShutdownReason()} is 89 | // {@link com.amazonaws.services.kinesis.clientlibrary.lib.worker.ShutdownReason#TERMINATE} it is required that you 90 | // checkpoint. Failure to do so will result in an IllegalArgumentException, and the KCL no longer making progress. 91 | if input.ShutdownReason == kc.TERMINATE { 92 | _ = input.Checkpointer.Checkpoint(nil) 93 | } 94 | 95 | assert.True(dd.t, dd.count > 0) 96 | } 97 | -------------------------------------------------------------------------------- /clientlibrary/partition/partition.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package partition 21 | // The implementation is derived from https://github.com/patrobinson/gokini 22 | // 23 | // Copyright 2018 Patrick robinson. 24 | // 25 | // Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 26 | // 27 | // The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 28 | // 29 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 30 | package partition 31 | 32 | import ( 33 | "sync" 34 | "time" 35 | 36 | "github.com/vmware/vmware-go-kcl-v2/clientlibrary/config" 37 | ) 38 | 39 | type ShardStatus struct { 40 | ID string 41 | ParentShardId string 42 | Checkpoint string 43 | AssignedTo string 44 | Mux *sync.RWMutex 45 | LeaseTimeout time.Time 46 | // Shard Range 47 | StartingSequenceNumber string 48 | // child shard doesn't have end sequence number 49 | EndingSequenceNumber string 50 | ClaimRequest string 51 | } 52 | 53 | func (ss *ShardStatus) GetLeaseOwner() string { 54 | ss.Mux.RLock() 55 | defer ss.Mux.RUnlock() 56 | return ss.AssignedTo 57 | } 58 | 59 | func (ss *ShardStatus) SetLeaseOwner(owner string) { 60 | ss.Mux.Lock() 61 | defer ss.Mux.Unlock() 62 | ss.AssignedTo = owner 63 | } 64 | 65 | func (ss *ShardStatus) GetCheckpoint() string { 66 | ss.Mux.RLock() 67 | defer ss.Mux.RUnlock() 68 | return ss.Checkpoint 69 | } 70 | 71 | func (ss *ShardStatus) SetCheckpoint(c string) { 72 | ss.Mux.Lock() 73 | defer ss.Mux.Unlock() 74 | ss.Checkpoint = c 75 | } 76 | 77 | func (ss *ShardStatus) GetLeaseTimeout() time.Time { 78 | ss.Mux.Lock() 79 | defer ss.Mux.Unlock() 80 | return ss.LeaseTimeout 81 | } 82 | 83 | func (ss *ShardStatus) SetLeaseTimeout(timeout time.Time) { 84 | ss.Mux.Lock() 85 | defer ss.Mux.Unlock() 86 | ss.LeaseTimeout = timeout 87 | } 88 | 89 | func (ss *ShardStatus) IsClaimRequestExpired(kclConfig *config.KinesisClientLibConfiguration) bool { 90 | if leaseTimeout := ss.GetLeaseTimeout(); leaseTimeout.IsZero() { 91 | return false 92 | } else { 93 | return leaseTimeout. 94 | Before(time.Now().UTC().Add(time.Duration(-kclConfig.LeaseStealingClaimTimeoutMillis) * time.Millisecond)) 95 | } 96 | } 97 | -------------------------------------------------------------------------------- /logger/logger.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | // Note: The implementation comes from https://www.mountedthoughts.com/golang-logger-interface/ 20 | // https://github.com/amitrai48/logger 21 | 22 | package logger 23 | 24 | import ( 25 | "github.com/sirupsen/logrus" 26 | ) 27 | 28 | // Fields Type to pass when we want to call WithFields for structured logging 29 | type Fields map[string]interface{} 30 | 31 | const ( 32 | //Debug has verbose message 33 | Debug = "debug" 34 | //Info is default log level 35 | Info = "info" 36 | //Warn is for logging messages about possible issues 37 | Warn = "warn" 38 | //Error is for logging errors 39 | Error = "error" 40 | //Fatal is for logging fatal messages. The sytem shutsdown after logging the message. 41 | Fatal = "fatal" 42 | ) 43 | 44 | // Logger is the common interface for logging. 45 | type Logger interface { 46 | Debugf(format string, args ...interface{}) 47 | 48 | Infof(format string, args ...interface{}) 49 | 50 | Warnf(format string, args ...interface{}) 51 | 52 | Errorf(format string, args ...interface{}) 53 | 54 | Fatalf(format string, args ...interface{}) 55 | 56 | Panicf(format string, args ...interface{}) 57 | 58 | WithFields(keyValues Fields) Logger 59 | } 60 | 61 | // Configuration stores the config for the logger 62 | // For some loggers there can only be one level across writers, for such the level of Console is picked by default 63 | type Configuration struct { 64 | EnableConsole bool 65 | ConsoleJSONFormat bool 66 | ConsoleLevel string 67 | EnableFile bool 68 | FileJSONFormat bool 69 | FileLevel string 70 | 71 | // Filename is the file to write logs to. Backup log files will be retained 72 | // in the same directory. It uses -lumberjack.log in 73 | // os.TempDir() if empty. 74 | Filename string 75 | 76 | // MaxSize is the maximum size in megabytes of the log file before it gets 77 | // rotated. It defaults to 100 megabytes. 78 | MaxSizeMB int 79 | 80 | // MaxAge is the maximum number of days to retain old log files based on the 81 | // timestamp encoded in their filename. Note that a day is defined as 24 82 | // hours and may not exactly correspond to calendar days due to daylight 83 | // savings, leap seconds, etc. The default is 7 days. 84 | MaxAgeDays int 85 | 86 | // MaxBackups is the maximum number of old log files to retain. The default 87 | // is to retain all old log files (though MaxAge may still cause them to get 88 | // deleted.) 89 | MaxBackups int 90 | 91 | // LocalTime determines if the time used for formatting the timestamps in 92 | // backup files is the computer's local time. The default is to use UTC 93 | // time. 94 | LocalTime bool 95 | } 96 | 97 | // GetDefaultLogger creates a default logger. 98 | func GetDefaultLogger() Logger { 99 | return NewLogrusLogger(logrus.StandardLogger()) 100 | } 101 | 102 | // normalizeConfig to enforce default value in configuration. 103 | func normalizeConfig(config *Configuration) { 104 | if config.MaxSizeMB <= 0 { 105 | config.MaxSizeMB = 100 106 | } 107 | 108 | if config.MaxAgeDays <= 0 { 109 | config.MaxAgeDays = 7 110 | } 111 | 112 | if config.MaxBackups < 0 { 113 | config.MaxBackups = 0 114 | } 115 | } 116 | -------------------------------------------------------------------------------- /clientlibrary/interfaces/record-processor.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package interfaces 21 | // The implementation is derived from https://github.com/awslabs/amazon-kinesis-client 22 | /* 23 | * Copyright 2014-2015 Amazon.com, Inc. or its affiliates. All Rights Reserved. 24 | * 25 | * Licensed under the Amazon Software License (the "License"). 26 | * You may not use this file except in compliance with the License. 27 | * A copy of the License is located at 28 | * 29 | * http://aws.amazon.com/asl/ 30 | * 31 | * or in the "license" file accompanying this file. This file is distributed 32 | * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either 33 | * express or implied. See the License for the specific language governing 34 | * permissions and limitations under the License. 35 | */ 36 | package interfaces 37 | 38 | type ( 39 | // IRecordProcessor is the interface for some callback functions invoked by KCL will 40 | // The main task of using KCL is to provide implementation on IRecordProcessor interface. 41 | // Note: This is exactly the same interface as Amazon KCL IRecordProcessor v2 42 | IRecordProcessor interface { 43 | // Initialize 44 | /* 45 | * Invoked by the Amazon Kinesis Client Library before data records are delivered to the RecordProcessor instance 46 | * (via processRecords). 47 | * 48 | * @param initializationInput Provides information related to initialization 49 | */ 50 | Initialize(initializationInput *InitializationInput) 51 | 52 | // ProcessRecords 53 | /* 54 | * Process data records. The Amazon Kinesis Client Library will invoke this method to deliver data records to the 55 | * application. 56 | * Upon fail over, the new instance will get records with sequence number > checkpoint position 57 | * for each partition key. 58 | * 59 | * @param processRecordsInput Provides the records to be processed as well as information and capabilities related 60 | * to them (eg checkpointing). 61 | */ 62 | ProcessRecords(processRecordsInput *ProcessRecordsInput) 63 | 64 | // Shutdown 65 | /* 66 | * Invoked by the Amazon Kinesis Client Library to indicate it will no longer send data records to this 67 | * RecordProcessor instance. 68 | * 69 | *

Warning

70 | * 71 | * When the value of {@link ShutdownInput#getShutdownReason()} is 72 | * {@link com.amazonaws.services.kinesis.clientlibrary.lib.worker.ShutdownReason#TERMINATE} it is required that you 73 | * checkpoint. Failure to do so will result in an IllegalArgumentException, and the KCL no longer making progress. 74 | * 75 | * @param shutdownInput 76 | * Provides information and capabilities (eg checkpointing) related to shutdown of this record processor. 77 | */ 78 | Shutdown(shutdownInput *ShutdownInput) 79 | } 80 | 81 | // IRecordProcessorFactory is interface for creating IRecordProcessor. Each Worker can have multiple threads 82 | // for processing shard. Client can choose either creating one processor per shard or sharing them. 83 | IRecordProcessorFactory interface { 84 | 85 | // CreateProcessor 86 | /* 87 | * Returns a record processor to be used for processing data records for a (assigned) shard. 88 | * 89 | * @return Returns a processor object. 90 | */ 91 | CreateProcessor() IRecordProcessor 92 | } 93 | ) 94 | -------------------------------------------------------------------------------- /clientlibrary/worker/worker-fan-out.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2021 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package worker 21 | package worker 22 | 23 | import ( 24 | "context" 25 | "errors" 26 | "fmt" 27 | "math" 28 | "time" 29 | 30 | "github.com/aws/aws-sdk-go-v2/service/kinesis" 31 | "github.com/aws/aws-sdk-go-v2/service/kinesis/types" 32 | ) 33 | 34 | // fetchConsumerARNWithRetry tries to fetch consumer ARN. Retries 10 times with exponential backoff in case of an error 35 | func (w *Worker) fetchConsumerARNWithRetry() (string, error) { 36 | for retry := 0; ; retry++ { 37 | consumerARN, err := w.fetchConsumerARN() 38 | if err == nil { 39 | return consumerARN, nil 40 | } 41 | if retry < 10 { 42 | sleepDuration := time.Duration(math.Exp2(float64(retry))*100) * time.Millisecond 43 | w.kclConfig.Logger.Errorf("Could not get consumer ARN: %v, retrying after: %s", err, sleepDuration) 44 | time.Sleep(sleepDuration) 45 | continue 46 | } 47 | return consumerARN, err 48 | } 49 | } 50 | 51 | // fetchConsumerARN gets enhanced fan-out consumerARN. 52 | // Registers enhanced fan-out consumer if the consumer is not found 53 | func (w *Worker) fetchConsumerARN() (string, error) { 54 | log := w.kclConfig.Logger 55 | log.Debugf("Fetching stream consumer ARN") 56 | 57 | streamDescription, err := w.kc.DescribeStream(context.TODO(), &kinesis.DescribeStreamInput{ 58 | StreamName: &w.kclConfig.StreamName, 59 | }) 60 | 61 | if err != nil { 62 | log.Errorf("Could not describe stream: %v", err) 63 | return "", err 64 | } 65 | 66 | streamConsumerDescription, err := w.kc.DescribeStreamConsumer(context.TODO(), &kinesis.DescribeStreamConsumerInput{ 67 | ConsumerName: &w.kclConfig.EnhancedFanOutConsumerName, 68 | StreamARN: streamDescription.StreamDescription.StreamARN, 69 | }) 70 | 71 | if err == nil { 72 | log.Infof("Enhanced fan-out consumer found, consumer status: %s", streamConsumerDescription.ConsumerDescription.ConsumerStatus) 73 | if streamConsumerDescription.ConsumerDescription.ConsumerStatus != types.ConsumerStatusActive { 74 | return "", fmt.Errorf("consumer is not in active status yet, current status: %s", streamConsumerDescription.ConsumerDescription.ConsumerStatus) 75 | } 76 | return *streamConsumerDescription.ConsumerDescription.ConsumerARN, nil 77 | } 78 | 79 | //aws-sdk-go-v2 https://github.com/aws/aws-sdk-go-v2/blob/main/CHANGELOG.md#error-handling 80 | var notFoundErr *types.ResourceNotFoundException 81 | if errors.As(err, ¬FoundErr) { 82 | log.Infof("Enhanced fan-out consumer not found, registering new consumer with name: %s", w.kclConfig.EnhancedFanOutConsumerName) 83 | out, err := w.kc.RegisterStreamConsumer(context.TODO(), &kinesis.RegisterStreamConsumerInput{ 84 | ConsumerName: &w.kclConfig.EnhancedFanOutConsumerName, 85 | StreamARN: streamDescription.StreamDescription.StreamARN, 86 | }) 87 | if err != nil { 88 | log.Errorf("Could not register enhanced fan-out consumer: %v", err) 89 | return "", err 90 | } 91 | if out.Consumer.ConsumerStatus != types.ConsumerStatusActive { 92 | return "", fmt.Errorf("consumer is not in active status yet, current status: %s", out.Consumer.ConsumerStatus) 93 | } 94 | return *out.Consumer.ConsumerARN, nil 95 | } 96 | 97 | log.Errorf("Could not describe stream consumer: %v", err) //%w should we unwrap the underlying error? 98 | 99 | return "", err 100 | } 101 | -------------------------------------------------------------------------------- /_support/scripts/ci.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | function local_go_pkgs() { 4 | find './clientlibrary' -name '*.go' | \ 5 | grep -Fv '/vendor/' | \ 6 | grep -Fv '/go/' | \ 7 | grep -Fv '/gen/' | \ 8 | grep -Fv '/tmp/' | \ 9 | grep -Fv '/run/' | \ 10 | grep -Fv '/tests/' | \ 11 | sed -r 's|(.+)/[^/]+\.go$|\1|g' | \ 12 | sort -u 13 | } 14 | 15 | function checkfmt() { 16 | local files="" 17 | files="$(find . -type f -iname "*.go" -exec gofmt -l {} \;)" 18 | 19 | if [ -n "$files" ]; then 20 | echo "You need to run \"gofmt -w ./\" to fix your formatting." 21 | echo "$files" >&2 22 | return 1 23 | fi 24 | } 25 | 26 | function goFormat() { 27 | echo "go formatting..." 28 | gofmt -w ./ 29 | echo "done" 30 | } 31 | 32 | function lint() { 33 | # golangci-lint run --enable-all -D forbidigo -D gochecknoglobals -D gofumpt -D gofmt -D nlreturn 34 | 35 | golangci-lint run \ 36 | --skip-files=_mock.go \ 37 | --skip-dirs=test \ 38 | --skip-dirs=internal \ 39 | --timeout=600s \ 40 | --verbose 41 | } 42 | 43 | function lintDocker() { 44 | lintVersion="1.41.1" 45 | lintImage="golangci/golangci-lint:v$lintVersion-alpine" 46 | 47 | docker run --rm -v "${PWD}":/app -w /app "$lintImage" golangci-lint run \ 48 | --skip-files=_mock.go \ 49 | --skip-dirs=test \ 50 | --skip-dirs=internal \ 51 | --timeout=600s \ 52 | --verbose 53 | } 54 | 55 | function unitTest() { 56 | go list ./... | grep -v /test | \ 57 | xargs -L 1 -I% bash -c 'echo -e "\n**************** Package: % ****************" && go test % -v -cover -race ./...' 58 | } 59 | 60 | function scanast() { 61 | gosec version 62 | gosec ./... > security.log 2>&1 63 | 64 | local issues="" 65 | issues=$(grep -c 'Severity: MEDIUM' security.log | grep -v deaggregator | grep -c _) 66 | if [ -n "$issues" ] && [ "$issues" -gt 0 ]; then 67 | echo "" 68 | echo "Medium Severity Issues:" 69 | grep -e "Severity: MEDIUM" -A 1 security.log 70 | echo "$issues" "medium severity issues found." 71 | fi 72 | 73 | local issues="" 74 | local issues_count="" 75 | issues="$(grep -E 'Severity: HIGH' security.log | grep -v vendor)" 76 | issues_count="$(grep -E 'Severity: HIGH' security.log | grep -v vendor | grep -c _)" 77 | if [ -n "$issues_count" ] && [ "$issues_count" -gt 0 ]; then 78 | echo "" 79 | echo "High Severity Issues:" 80 | grep -E "Severity: HIGH" -A 1 security.log 81 | echo "$issues_count" "high severity issues found." 82 | echo "$issues" 83 | echo "You need to resolve the high severity issues at the least." 84 | exit 1 85 | fi 86 | 87 | local issues="" 88 | local issues_count="" 89 | issues="$(grep -E 'Errors unhandled' security.log | grep -v vendor | grep -v /src/go/src)" 90 | issues_count="$(grep -E 'Errors unhandled' security.log | grep -v vendor | grep -v /src/go/src | grep -c _)" 91 | if [ -n "$issues_count" ] && [ "$issues_count" -gt 0 ]; then 92 | echo "" 93 | echo "Unhandled errors:" 94 | grep -E "Errors unhandled" security.log 95 | echo "$issues_count" "unhandled errors, please indicate with the right comment that this case is ok, or handle the error." 96 | echo "$issues" 97 | echo "You need to resolve the all unhandled errors." 98 | exit 1 99 | fi 100 | 101 | rm -f security.log 102 | } 103 | 104 | function scan() { 105 | gosec -fmt=sarif -out=results.sarif -exclude-dir=internal -exclude-dir=vendor -severity=high ./... 106 | } 107 | 108 | function localScan() { 109 | # you can use the vs code plugin https://marketplace.visualstudio.com/items?itemName=MS-SarifVSCode.sarif-viewer 110 | # to navigate against the issues 111 | gosec -fmt=sarif -out=results.sarif -exclude-dir=internal -exclude-dir=vendor ./... 112 | } 113 | 114 | function usage() { 115 | echo "check.sh fmt|lint" >&2 116 | exit 2 117 | } 118 | 119 | case "$1" in 120 | fmtCheck) checkfmt ;; 121 | format) goFormat ;; 122 | lint) lint ;; 123 | lintDocker) lintDocker ;; 124 | unitTest) unitTest ;; 125 | scan) scan ;; 126 | localScan) localScan ;; 127 | *) usage ;; 128 | esac 129 | -------------------------------------------------------------------------------- /test/worker_lease_stealing_test.go: -------------------------------------------------------------------------------- 1 | package test 2 | 3 | import ( 4 | chk "github.com/vmware/vmware-go-kcl-v2/clientlibrary/checkpoint" 5 | "testing" 6 | 7 | cfg "github.com/vmware/vmware-go-kcl-v2/clientlibrary/config" 8 | wk "github.com/vmware/vmware-go-kcl-v2/clientlibrary/worker" 9 | "github.com/vmware/vmware-go-kcl-v2/logger" 10 | ) 11 | 12 | func TestLeaseStealing(t *testing.T) { 13 | config := &TestClusterConfig{ 14 | numShards: 4, 15 | numWorkers: 2, 16 | appName: appName, 17 | streamName: streamName, 18 | regionName: regionName, 19 | workerIDTemplate: workerID + "-%v", 20 | } 21 | test := NewLeaseStealingTest(t, config, newLeaseStealingWorkerFactory(t)) 22 | test.Run(LeaseStealingAssertions{ 23 | expectedLeasesForInitialWorker: config.numShards, 24 | expectedLeasesPerWorker: config.numShards / config.numWorkers, 25 | }) 26 | } 27 | 28 | type leaseStealingWorkerFactory struct { 29 | t *testing.T 30 | } 31 | 32 | func newLeaseStealingWorkerFactory(t *testing.T) *leaseStealingWorkerFactory { 33 | return &leaseStealingWorkerFactory{t} 34 | } 35 | 36 | func (wf *leaseStealingWorkerFactory) CreateKCLConfig(workerID string, config *TestClusterConfig) *cfg.KinesisClientLibConfiguration { 37 | log := logger.NewLogrusLoggerWithConfig(logger.Configuration{ 38 | EnableConsole: true, 39 | ConsoleLevel: logger.Error, 40 | ConsoleJSONFormat: false, 41 | EnableFile: true, 42 | FileLevel: logger.Info, 43 | FileJSONFormat: true, 44 | Filename: "log.log", 45 | }) 46 | 47 | log.WithFields(logger.Fields{"worker": workerID}) 48 | 49 | return cfg.NewKinesisClientLibConfig(config.appName, config.streamName, config.regionName, workerID). 50 | WithInitialPositionInStream(cfg.LATEST). 51 | WithMaxRecords(10). 52 | WithShardSyncIntervalMillis(5000). 53 | WithFailoverTimeMillis(10000). 54 | WithLeaseStealing(true). 55 | WithLogger(log) 56 | } 57 | 58 | func (wf *leaseStealingWorkerFactory) CreateWorker(_ string, kclConfig *cfg.KinesisClientLibConfiguration) *wk.Worker { 59 | worker := wk.NewWorker(recordProcessorFactory(wf.t), kclConfig) 60 | return worker 61 | } 62 | 63 | func TestLeaseStealingInjectCheckpointer(t *testing.T) { 64 | config := &TestClusterConfig{ 65 | numShards: 4, 66 | numWorkers: 2, 67 | appName: appName, 68 | streamName: streamName, 69 | regionName: regionName, 70 | workerIDTemplate: workerID + "-%v", 71 | } 72 | test := NewLeaseStealingTest(t, config, newleaseStealingWorkerFactoryCustomChk(t)) 73 | test.Run(LeaseStealingAssertions{ 74 | expectedLeasesForInitialWorker: config.numShards, 75 | expectedLeasesPerWorker: config.numShards / config.numWorkers, 76 | }) 77 | } 78 | 79 | type leaseStealingWorkerFactoryCustom struct { 80 | *leaseStealingWorkerFactory 81 | } 82 | 83 | func newleaseStealingWorkerFactoryCustomChk(t *testing.T) *leaseStealingWorkerFactoryCustom { 84 | return &leaseStealingWorkerFactoryCustom{ 85 | newLeaseStealingWorkerFactory(t), 86 | } 87 | } 88 | 89 | func (wfc *leaseStealingWorkerFactoryCustom) CreateWorker(workerID string, kclConfig *cfg.KinesisClientLibConfiguration) *wk.Worker { 90 | worker := wfc.leaseStealingWorkerFactory.CreateWorker(workerID, kclConfig) 91 | checkpointer := chk.NewDynamoCheckpoint(kclConfig) 92 | return worker.WithCheckpointer(checkpointer) 93 | } 94 | 95 | func TestLeaseStealingWithMaxLeasesForWorker(t *testing.T) { 96 | config := &TestClusterConfig{ 97 | numShards: 4, 98 | numWorkers: 2, 99 | appName: appName, 100 | streamName: streamName, 101 | regionName: regionName, 102 | workerIDTemplate: workerID + "-%v", 103 | } 104 | test := NewLeaseStealingTest(t, config, newLeaseStealingWorkerFactoryMaxLeases(t, config.numShards-1)) 105 | test.Run(LeaseStealingAssertions{ 106 | expectedLeasesForInitialWorker: config.numShards - 1, 107 | expectedLeasesPerWorker: 2, 108 | }) 109 | } 110 | 111 | type leaseStealingWorkerFactoryMaxLeases struct { 112 | maxLeases int 113 | *leaseStealingWorkerFactory 114 | } 115 | 116 | func newLeaseStealingWorkerFactoryMaxLeases(t *testing.T, maxLeases int) *leaseStealingWorkerFactoryMaxLeases { 117 | return &leaseStealingWorkerFactoryMaxLeases{ 118 | maxLeases, 119 | newLeaseStealingWorkerFactory(t), 120 | } 121 | } 122 | 123 | func (wfm *leaseStealingWorkerFactoryMaxLeases) CreateKCLConfig(workerID string, config *TestClusterConfig) *cfg.KinesisClientLibConfiguration { 124 | kclConfig := wfm.leaseStealingWorkerFactory.CreateKCLConfig(workerID, config) 125 | kclConfig.WithMaxLeasesForWorker(wfm.maxLeases) 126 | return kclConfig 127 | } 128 | -------------------------------------------------------------------------------- /clientlibrary/checkpoint/checkpointer.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package checkpoint 21 | // The implementation is derived from https://github.com/patrobinson/gokini 22 | // 23 | // Copyright 2018 Patrick robinson. 24 | // 25 | // Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 26 | // 27 | // The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 28 | // 29 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 30 | package checkpoint 31 | 32 | import ( 33 | "errors" 34 | "fmt" 35 | 36 | par "github.com/vmware/vmware-go-kcl-v2/clientlibrary/partition" 37 | ) 38 | 39 | const ( 40 | LeaseKeyKey = "ShardID" 41 | LeaseOwnerKey = "AssignedTo" 42 | LeaseTimeoutKey = "LeaseTimeout" 43 | SequenceNumberKey = "Checkpoint" 44 | ParentShardIdKey = "ParentShardId" 45 | ClaimRequestKey = "ClaimRequest" 46 | 47 | // ShardEnd We've completely processed all records in this shard. 48 | ShardEnd = "SHARD_END" 49 | 50 | // ErrShardClaimed is returned when shard is claimed 51 | ErrShardClaimed = "shard is already claimed by another node" 52 | ) 53 | 54 | type ErrLeaseNotAcquired struct { 55 | cause string 56 | } 57 | 58 | func (e ErrLeaseNotAcquired) Error() string { 59 | return fmt.Sprintf("lease not acquired: %s", e.cause) 60 | } 61 | 62 | // Checkpointer handles checkpointing when a record has been processed 63 | type Checkpointer interface { 64 | // Init initialises the Checkpoint 65 | Init() error 66 | 67 | // GetLease attempts to gain a lock on the given shard 68 | GetLease(*par.ShardStatus, string) error 69 | 70 | // CheckpointSequence writes a checkpoint at the designated sequence ID 71 | CheckpointSequence(*par.ShardStatus) error 72 | 73 | // FetchCheckpoint retrieves the checkpoint for the given shard 74 | FetchCheckpoint(*par.ShardStatus) error 75 | 76 | // RemoveLeaseInfo to remove lease info for shard entry because the shard no longer exists 77 | RemoveLeaseInfo(string) error 78 | 79 | // RemoveLeaseOwner to remove lease owner for the shard entry to make the shard available for reassignment 80 | RemoveLeaseOwner(string) error 81 | 82 | // GetLeaseOwner to get current owner of lease for shard 83 | GetLeaseOwner(string) (string, error) 84 | 85 | // ListActiveWorkers returns active workers and their shards (New Lease Stealing Methods) 86 | ListActiveWorkers(map[string]*par.ShardStatus) (map[string][]*par.ShardStatus, error) 87 | 88 | // ClaimShard claims a shard for stealing 89 | ClaimShard(*par.ShardStatus, string) error 90 | } 91 | 92 | // ErrSequenceIDNotFound is returned by FetchCheckpoint when no SequenceID is found 93 | var ErrSequenceIDNotFound = errors.New("SequenceIDNotFoundForShard") 94 | 95 | // ErrShardNotAssigned is returned by ListActiveWorkers when no AssignedTo is found 96 | var ErrShardNotAssigned = errors.New("AssignedToNotFoundForShard") 97 | -------------------------------------------------------------------------------- /clientlibrary/config/config_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | package config 20 | 21 | import ( 22 | "testing" 23 | 24 | "github.com/stretchr/testify/assert" 25 | 26 | "github.com/vmware/vmware-go-kcl-v2/logger" 27 | ) 28 | 29 | func TestConfig(t *testing.T) { 30 | kclConfig := NewKinesisClientLibConfig("appName", "StreamName", "us-west-2", "workerId"). 31 | WithFailoverTimeMillis(500). 32 | WithMaxRecords(100). 33 | WithInitialPositionInStream(TRIM_HORIZON). 34 | WithIdleTimeBetweenReadsInMillis(20). 35 | WithCallProcessRecordsEvenForEmptyRecordList(true). 36 | WithTaskBackoffTimeMillis(10). 37 | WithEnhancedFanOutConsumerName("fan-out-consumer") 38 | 39 | assert.Equal(t, "appName", kclConfig.ApplicationName) 40 | assert.Equal(t, 500, kclConfig.FailoverTimeMillis) 41 | assert.Equal(t, 10, kclConfig.TaskBackoffTimeMillis) 42 | 43 | assert.True(t, kclConfig.EnableEnhancedFanOutConsumer) 44 | assert.Equal(t, "fan-out-consumer", kclConfig.EnhancedFanOutConsumerName) 45 | 46 | assert.Equal(t, false, kclConfig.EnableLeaseStealing) 47 | assert.Equal(t, 5000, kclConfig.LeaseStealingIntervalMillis) 48 | 49 | contextLogger := kclConfig.Logger.WithFields(logger.Fields{"key1": "value1"}) 50 | contextLogger.Debugf("Starting with default logger") 51 | contextLogger.Infof("Default logger is awesome") 52 | } 53 | 54 | func TestConfigLeaseStealing(t *testing.T) { 55 | kclConfig := NewKinesisClientLibConfig("appName", "StreamName", "us-west-2", "workerId"). 56 | WithFailoverTimeMillis(500). 57 | WithMaxRecords(100). 58 | WithInitialPositionInStream(TRIM_HORIZON). 59 | WithIdleTimeBetweenReadsInMillis(20). 60 | WithCallProcessRecordsEvenForEmptyRecordList(true). 61 | WithTaskBackoffTimeMillis(10). 62 | WithLeaseStealing(true). 63 | WithLeaseStealingIntervalMillis(10000) 64 | 65 | assert.Equal(t, "appName", kclConfig.ApplicationName) 66 | assert.Equal(t, 500, kclConfig.FailoverTimeMillis) 67 | assert.Equal(t, 10, kclConfig.TaskBackoffTimeMillis) 68 | assert.Equal(t, true, kclConfig.EnableLeaseStealing) 69 | assert.Equal(t, 10000, kclConfig.LeaseStealingIntervalMillis) 70 | 71 | contextLogger := kclConfig.Logger.WithFields(logger.Fields{"key1": "value1"}) 72 | contextLogger.Debugf("Starting with default logger") 73 | contextLogger.Infof("Default logger is awesome") 74 | } 75 | 76 | func TestConfigDefaultEnhancedFanOutConsumerName(t *testing.T) { 77 | kclConfig := NewKinesisClientLibConfig("appName", "StreamName", "us-west-2", "workerId") 78 | 79 | assert.Equal(t, "appName", kclConfig.ApplicationName) 80 | assert.False(t, kclConfig.EnableEnhancedFanOutConsumer) 81 | assert.Equal(t, "appName", kclConfig.EnhancedFanOutConsumerName) 82 | } 83 | 84 | func TestEmptyEnhancedFanOutConsumerName(t *testing.T) { 85 | assert.PanicsWithValue(t, "Non-empty value expected for EnhancedFanOutConsumerName, actual: ", func() { 86 | NewKinesisClientLibConfig("app", "stream", "us-west-2", "worker").WithEnhancedFanOutConsumerName("") 87 | }) 88 | } 89 | 90 | func TestConfigWithEnhancedFanOutConsumerARN(t *testing.T) { 91 | kclConfig := NewKinesisClientLibConfig("app", "stream", "us-west-2", "worker"). 92 | WithEnhancedFanOutConsumerARN("consumer:arn") 93 | 94 | assert.True(t, kclConfig.EnableEnhancedFanOutConsumer) 95 | assert.Equal(t, "consumer:arn", kclConfig.EnhancedFanOutConsumerARN) 96 | } 97 | 98 | func TestEmptyEnhancedFanOutConsumerARN(t *testing.T) { 99 | assert.PanicsWithValue(t, "Non-empty value expected for EnhancedFanOutConsumerARN, actual: ", func() { 100 | NewKinesisClientLibConfig("app", "stream", "us-west-2", "worker").WithEnhancedFanOutConsumerARN("") 101 | }) 102 | } 103 | -------------------------------------------------------------------------------- /logger/zerolog/zerolog.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | // Note: The implementation comes from https://www.mountedthoughts.com/golang-logger-interface/ 20 | // https://github.com/amitrai48/logger 21 | 22 | // Package zerolog implements the KCL logger using RS Zerolog logger 23 | package zerolog 24 | 25 | import ( 26 | "github.com/rs/zerolog" 27 | "github.com/vmware/vmware-go-kcl-v2/logger" 28 | "gopkg.in/natefinch/lumberjack.v2" 29 | "os" 30 | ) 31 | 32 | type zeroLogger struct { 33 | log zerolog.Logger 34 | } 35 | 36 | // NewZerologLogger creates a new logger.Logger backed by RS Zerolog using a default config 37 | func NewZerologLogger() logger.Logger { 38 | return NewZerologLoggerWithConfig(logger.Configuration{ 39 | EnableConsole: true, 40 | ConsoleJSONFormat: true, 41 | ConsoleLevel: logger.Info, 42 | EnableFile: false, 43 | FileJSONFormat: false, 44 | FileLevel: logger.Info, 45 | Filename: "", 46 | MaxSizeMB: 0, 47 | MaxAgeDays: 0, 48 | MaxBackups: 0, 49 | LocalTime: true, 50 | }) 51 | } 52 | 53 | // NewZerologLoggerWithConfig creates a new logger.Logger backed by RS Zerolog using the provided config 54 | func NewZerologLoggerWithConfig(config logger.Configuration) logger.Logger { 55 | var consoleHandler *zerolog.ConsoleWriter 56 | var fileHandler *lumberjack.Logger 57 | var finalLogger zerolog.Logger 58 | 59 | normalizeConfig(&config) 60 | 61 | if config.EnableConsole { 62 | consoleHandler = &zerolog.ConsoleWriter{Out: os.Stdout} 63 | } 64 | 65 | if config.EnableFile { 66 | fileHandler = &lumberjack.Logger{ 67 | Filename: config.Filename, 68 | MaxSize: config.MaxSizeMB, 69 | Compress: true, 70 | MaxAge: config.MaxAgeDays, 71 | MaxBackups: config.MaxBackups, 72 | LocalTime: config.LocalTime, 73 | } 74 | } 75 | 76 | if config.EnableConsole && config.EnableFile { 77 | multi := zerolog.MultiLevelWriter(consoleHandler, fileHandler) 78 | finalLogger = zerolog.New(multi).Level(getZeroLogLevel(config.ConsoleLevel)).With().Timestamp().Logger() 79 | } else if config.EnableFile { 80 | finalLogger = zerolog.New(fileHandler).Level(getZeroLogLevel(config.FileLevel)).With().Timestamp().Logger() 81 | } else { 82 | finalLogger = zerolog.New(consoleHandler).Level(getZeroLogLevel(config.ConsoleLevel)).With().Timestamp().Logger() 83 | } 84 | 85 | return &zeroLogger{log: finalLogger} 86 | } 87 | 88 | func (z *zeroLogger) Debugf(format string, args ...interface{}) { 89 | z.log.Debug().Msgf(format, args...) 90 | } 91 | 92 | func (z *zeroLogger) Infof(format string, args ...interface{}) { 93 | z.log.Info().Msgf(format, args...) 94 | } 95 | 96 | func (z *zeroLogger) Warnf(format string, args ...interface{}) { 97 | z.log.Warn().Msgf(format, args...) 98 | } 99 | 100 | func (z *zeroLogger) Errorf(format string, args ...interface{}) { 101 | z.log.Error().Msgf(format, args...) 102 | } 103 | 104 | func (z *zeroLogger) Fatalf(format string, args ...interface{}) { 105 | z.log.Fatal().Msgf(format, args...) 106 | } 107 | 108 | func (z *zeroLogger) Panicf(format string, args ...interface{}) { 109 | z.log.Panic().Msgf(format, args...) 110 | } 111 | 112 | func (z *zeroLogger) WithFields(keyValues logger.Fields) logger.Logger { 113 | newLogger := z.log.With() 114 | for k, v := range keyValues { 115 | newLogger.Interface(k, v) 116 | } 117 | 118 | return &zeroLogger{ 119 | log: newLogger.Logger(), 120 | } 121 | } 122 | 123 | func getZeroLogLevel(level string) zerolog.Level { 124 | switch level { 125 | case logger.Info: 126 | return zerolog.InfoLevel 127 | case logger.Warn: 128 | return zerolog.WarnLevel 129 | case logger.Debug: 130 | return zerolog.DebugLevel 131 | case logger.Error: 132 | return zerolog.ErrorLevel 133 | case logger.Fatal: 134 | return zerolog.FatalLevel 135 | default: 136 | return zerolog.InfoLevel 137 | } 138 | } 139 | 140 | func normalizeConfig(config *logger.Configuration) { 141 | if config.MaxSizeMB <= 0 { 142 | config.MaxSizeMB = 100 143 | } 144 | 145 | if config.MaxAgeDays <= 0 { 146 | config.MaxAgeDays = 7 147 | } 148 | 149 | if config.MaxBackups < 0 { 150 | config.MaxBackups = 0 151 | } 152 | } 153 | -------------------------------------------------------------------------------- /logger/zap/zap.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | // Note: The implementation comes from https://www.mountedthoughts.com/golang-logger-interface/ 20 | // https://github.com/amitrai48/logger 21 | 22 | package zap 23 | 24 | import ( 25 | "os" 26 | 27 | "github.com/vmware/vmware-go-kcl-v2/logger" 28 | uzap "go.uber.org/zap" 29 | "go.uber.org/zap/zapcore" 30 | lumberjack "gopkg.in/natefinch/lumberjack.v2" 31 | ) 32 | 33 | type ZapLogger struct { 34 | sugaredLogger *uzap.SugaredLogger 35 | } 36 | 37 | // NewZapLogger adapts existing sugared zap logger to Logger interface. 38 | // The call is responsible for configuring sugard zap logger appropriately. 39 | // 40 | // Note: Sugar wraps the Logger to provide a more ergonomic, but slightly slower, 41 | // API. Sugaring a Logger is quite inexpensive, so it's reasonable for a 42 | // single application to use both Loggers and SugaredLoggers, converting 43 | // between them on the boundaries of performance-sensitive code. 44 | // 45 | // Base zap logger can be convert to SugaredLogger by calling to add a wrapper: 46 | // sugaredLogger := log.Sugar() 47 | func NewZapLogger(logger *uzap.SugaredLogger) logger.Logger { 48 | return &ZapLogger{ 49 | sugaredLogger: logger, 50 | } 51 | } 52 | 53 | // NewZapLoggerWithConfig creates and configs Logger instance backed by 54 | // zap Sugared logger. 55 | func NewZapLoggerWithConfig(config logger.Configuration) logger.Logger { 56 | cores := []zapcore.Core{} 57 | 58 | if config.EnableConsole { 59 | level := getZapLevel(config.ConsoleLevel) 60 | writer := zapcore.Lock(os.Stdout) 61 | core := zapcore.NewCore(getEncoder(config.ConsoleJSONFormat), writer, level) 62 | cores = append(cores, core) 63 | } 64 | 65 | if config.EnableFile { 66 | level := getZapLevel(config.FileLevel) 67 | writer := zapcore.AddSync(&lumberjack.Logger{ 68 | Filename: config.Filename, 69 | MaxSize: config.MaxSizeMB, 70 | Compress: true, 71 | MaxAge: config.MaxAgeDays, 72 | MaxBackups: config.MaxBackups, 73 | LocalTime: config.LocalTime, 74 | }) 75 | core := zapcore.NewCore(getEncoder(config.FileJSONFormat), writer, level) 76 | cores = append(cores, core) 77 | } 78 | 79 | combinedCore := zapcore.NewTee(cores...) 80 | 81 | // AddCallerSkip skips 2 number of callers, this is important else the file that gets 82 | // logged will always be the wrapped file. In our case zap.go 83 | logger := uzap.New(combinedCore, 84 | uzap.AddCallerSkip(2), 85 | uzap.AddCaller(), 86 | ).Sugar() 87 | 88 | return &ZapLogger{ 89 | sugaredLogger: logger, 90 | } 91 | } 92 | 93 | func (l *ZapLogger) Debugf(format string, args ...interface{}) { 94 | l.sugaredLogger.Debugf(format, args...) 95 | } 96 | 97 | func (l *ZapLogger) Infof(format string, args ...interface{}) { 98 | l.sugaredLogger.Infof(format, args...) 99 | } 100 | 101 | func (l *ZapLogger) Warnf(format string, args ...interface{}) { 102 | l.sugaredLogger.Warnf(format, args...) 103 | } 104 | 105 | func (l *ZapLogger) Errorf(format string, args ...interface{}) { 106 | l.sugaredLogger.Errorf(format, args...) 107 | } 108 | 109 | func (l *ZapLogger) Fatalf(format string, args ...interface{}) { 110 | l.sugaredLogger.Fatalf(format, args...) 111 | } 112 | 113 | func (l *ZapLogger) Panicf(format string, args ...interface{}) { 114 | l.sugaredLogger.Fatalf(format, args...) 115 | } 116 | 117 | func (l *ZapLogger) WithFields(fields logger.Fields) logger.Logger { 118 | var f = make([]interface{}, 0) 119 | for k, v := range fields { 120 | f = append(f, k) 121 | f = append(f, v) 122 | } 123 | newLogger := l.sugaredLogger.With(f...) 124 | return &ZapLogger{newLogger} 125 | } 126 | 127 | func getEncoder(isJSON bool) zapcore.Encoder { 128 | encoderConfig := uzap.NewProductionEncoderConfig() 129 | encoderConfig.EncodeTime = zapcore.ISO8601TimeEncoder 130 | if isJSON { 131 | return zapcore.NewJSONEncoder(encoderConfig) 132 | } 133 | return zapcore.NewConsoleEncoder(encoderConfig) 134 | } 135 | 136 | func getZapLevel(level string) zapcore.Level { 137 | switch level { 138 | case logger.Info: 139 | return zapcore.InfoLevel 140 | case logger.Warn: 141 | return zapcore.WarnLevel 142 | case logger.Debug: 143 | return zapcore.DebugLevel 144 | case logger.Error: 145 | return zapcore.ErrorLevel 146 | case logger.Fatal: 147 | return zapcore.FatalLevel 148 | default: 149 | return zapcore.InfoLevel 150 | } 151 | } 152 | -------------------------------------------------------------------------------- /logger/logrus.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2019 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | // Note: The implementation comes from https://www.mountedthoughts.com/golang-logger-interface/ 20 | // https://github.com/amitrai48/logger 21 | 22 | package logger 23 | 24 | import ( 25 | "io" 26 | "os" 27 | 28 | "github.com/sirupsen/logrus" 29 | lumberjack "gopkg.in/natefinch/lumberjack.v2" 30 | ) 31 | 32 | type LogrusLogEntry struct { 33 | entry *logrus.Entry 34 | } 35 | 36 | type LogrusLogger struct { 37 | logger logrus.FieldLogger 38 | } 39 | 40 | // NewLogrusLogger adapts existing logrus logger to Logger interface. 41 | // The call is responsible for configuring logrus logger appropriately. 42 | func NewLogrusLogger(lLogger logrus.FieldLogger) Logger { 43 | return &LogrusLogger{ 44 | logger: lLogger, 45 | } 46 | } 47 | 48 | // NewLogrusLoggerWithConfig creates and configs Logger instance backed by 49 | // logrus logger. 50 | func NewLogrusLoggerWithConfig(config Configuration) Logger { 51 | logLevel := config.ConsoleLevel 52 | if logLevel == "" { 53 | logLevel = config.FileLevel 54 | } 55 | 56 | level, err := logrus.ParseLevel(logLevel) 57 | if err != nil { 58 | // fallback to InfoLevel 59 | level = logrus.InfoLevel 60 | } 61 | 62 | normalizeConfig(&config) 63 | 64 | stdOutHandler := os.Stdout 65 | fileHandler := &lumberjack.Logger{ 66 | Filename: config.Filename, 67 | MaxSize: config.MaxSizeMB, 68 | Compress: true, 69 | MaxAge: config.MaxAgeDays, 70 | MaxBackups: config.MaxBackups, 71 | LocalTime: config.LocalTime, 72 | } 73 | lLogger := &logrus.Logger{ 74 | Out: stdOutHandler, 75 | Formatter: getFormatter(config.ConsoleJSONFormat), 76 | Hooks: make(logrus.LevelHooks), 77 | Level: level, 78 | } 79 | 80 | if config.EnableConsole && config.EnableFile { 81 | lLogger.SetOutput(io.MultiWriter(stdOutHandler, fileHandler)) 82 | } else { 83 | if config.EnableFile { 84 | lLogger.SetOutput(fileHandler) 85 | lLogger.SetFormatter(getFormatter(config.FileJSONFormat)) 86 | } 87 | } 88 | 89 | return &LogrusLogger{ 90 | logger: lLogger, 91 | } 92 | } 93 | 94 | func (l *LogrusLogger) Debugf(format string, args ...interface{}) { 95 | l.logger.Debugf(format, args...) 96 | } 97 | 98 | func (l *LogrusLogger) Infof(format string, args ...interface{}) { 99 | l.logger.Infof(format, args...) 100 | } 101 | 102 | func (l *LogrusLogger) Warnf(format string, args ...interface{}) { 103 | l.logger.Warnf(format, args...) 104 | } 105 | 106 | func (l *LogrusLogger) Errorf(format string, args ...interface{}) { 107 | l.logger.Errorf(format, args...) 108 | } 109 | 110 | func (l *LogrusLogger) Fatalf(format string, args ...interface{}) { 111 | l.logger.Fatalf(format, args...) 112 | } 113 | 114 | func (l *LogrusLogger) Panicf(format string, args ...interface{}) { 115 | l.logger.Fatalf(format, args...) 116 | } 117 | 118 | func (l *LogrusLogger) WithFields(fields Fields) Logger { 119 | return &LogrusLogEntry{ 120 | entry: l.logger.WithFields(convertToLogrusFields(fields)), 121 | } 122 | } 123 | 124 | func (l *LogrusLogEntry) Debugf(format string, args ...interface{}) { 125 | l.entry.Debugf(format, args...) 126 | } 127 | 128 | func (l *LogrusLogEntry) Infof(format string, args ...interface{}) { 129 | l.entry.Infof(format, args...) 130 | } 131 | 132 | func (l *LogrusLogEntry) Warnf(format string, args ...interface{}) { 133 | l.entry.Warnf(format, args...) 134 | } 135 | 136 | func (l *LogrusLogEntry) Errorf(format string, args ...interface{}) { 137 | l.entry.Errorf(format, args...) 138 | } 139 | 140 | func (l *LogrusLogEntry) Fatalf(format string, args ...interface{}) { 141 | l.entry.Fatalf(format, args...) 142 | } 143 | 144 | func (l *LogrusLogEntry) Panicf(format string, args ...interface{}) { 145 | l.entry.Fatalf(format, args...) 146 | } 147 | 148 | func (l *LogrusLogEntry) WithFields(fields Fields) Logger { 149 | return &LogrusLogEntry{ 150 | entry: l.entry.WithFields(convertToLogrusFields(fields)), 151 | } 152 | } 153 | 154 | func getFormatter(isJSON bool) logrus.Formatter { 155 | if isJSON { 156 | return &logrus.JSONFormatter{} 157 | } 158 | return &logrus.TextFormatter{ 159 | FullTimestamp: true, 160 | DisableLevelTruncation: true, 161 | } 162 | } 163 | 164 | func convertToLogrusFields(fields Fields) logrus.Fields { 165 | logrusFields := logrus.Fields{} 166 | for index, val := range fields { 167 | logrusFields[index] = val 168 | } 169 | return logrusFields 170 | } 171 | -------------------------------------------------------------------------------- /clientlibrary/checkpoint/mock-dynamodb_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2021 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // The implementation is derived from https://github.com/patrobinson/gokini 21 | // 22 | // Copyright 2018 Patrick robinson. 23 | // 24 | // Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 25 | // 26 | // The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 27 | // 28 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 29 | package checkpoint 30 | 31 | import ( 32 | "context" 33 | "github.com/aws/aws-sdk-go-v2/aws" 34 | "github.com/aws/aws-sdk-go-v2/service/dynamodb" 35 | "github.com/aws/aws-sdk-go-v2/service/dynamodb/types" 36 | ) 37 | 38 | type mockDynamoDB struct { 39 | client *dynamodb.Client 40 | tableExist bool 41 | item map[string]types.AttributeValue 42 | conditionalExpression string 43 | expressionAttributeValues map[string]types.AttributeValue 44 | } 45 | 46 | func (m *mockDynamoDB) Scan(ctx context.Context, params *dynamodb.ScanInput, optFns ...func(*dynamodb.Options)) (*dynamodb.ScanOutput, error) { 47 | return &dynamodb.ScanOutput{}, nil 48 | } 49 | 50 | func (m *mockDynamoDB) DescribeTable(ctx context.Context, params *dynamodb.DescribeTableInput, optFns ...func(*dynamodb.Options)) (*dynamodb.DescribeTableOutput, error) { 51 | if !m.tableExist { 52 | return &dynamodb.DescribeTableOutput{}, &types.ResourceNotFoundException{Message: aws.String("doesNotExist")} 53 | } 54 | 55 | return &dynamodb.DescribeTableOutput{}, nil 56 | } 57 | 58 | func (m *mockDynamoDB) CreateTable(ctx context.Context, params *dynamodb.CreateTableInput, optFns ...func(*dynamodb.Options)) (*dynamodb.CreateTableOutput, error) { 59 | return &dynamodb.CreateTableOutput{}, nil 60 | } 61 | 62 | func (m *mockDynamoDB) PutItem(ctx context.Context, params *dynamodb.PutItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.PutItemOutput, error) { 63 | item := params.Item 64 | 65 | if shardID, ok := item[LeaseKeyKey]; ok { 66 | m.item[LeaseKeyKey] = shardID 67 | } 68 | 69 | if owner, ok := item[LeaseOwnerKey]; ok { 70 | m.item[LeaseOwnerKey] = owner 71 | } 72 | 73 | if timeout, ok := item[LeaseTimeoutKey]; ok { 74 | m.item[LeaseTimeoutKey] = timeout 75 | } 76 | 77 | if checkpoint, ok := item[SequenceNumberKey]; ok { 78 | m.item[SequenceNumberKey] = checkpoint 79 | } 80 | 81 | if parent, ok := item[ParentShardIdKey]; ok { 82 | m.item[ParentShardIdKey] = parent 83 | } 84 | 85 | if claimRequest, ok := item[ClaimRequestKey]; ok { 86 | m.item[ClaimRequestKey] = claimRequest 87 | } 88 | 89 | if params.ConditionExpression != nil { 90 | m.conditionalExpression = *params.ConditionExpression 91 | } 92 | 93 | m.expressionAttributeValues = params.ExpressionAttributeValues 94 | 95 | return nil, nil 96 | } 97 | 98 | func (m *mockDynamoDB) GetItem(ctx context.Context, params *dynamodb.GetItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.GetItemOutput, error) { 99 | return &dynamodb.GetItemOutput{ 100 | Item: m.item, 101 | }, nil 102 | } 103 | 104 | func (m *mockDynamoDB) UpdateItem(ctx context.Context, params *dynamodb.UpdateItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.UpdateItemOutput, error) { 105 | exp := params.UpdateExpression 106 | 107 | if aws.ToString(exp) == "remove "+LeaseOwnerKey { 108 | delete(m.item, LeaseOwnerKey) 109 | } 110 | 111 | return nil, nil 112 | } 113 | 114 | func (m *mockDynamoDB) DeleteItem(ctx context.Context, params *dynamodb.DeleteItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.DeleteItemOutput, error) { 115 | return &dynamodb.DeleteItemOutput{}, nil 116 | } 117 | -------------------------------------------------------------------------------- /clientlibrary/interfaces/inputs.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2020 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package interfaces 21 | // The implementation is derived from https://github.com/awslabs/amazon-kinesis-client 22 | /* 23 | * Copyright 2014-2015 Amazon.com, Inc. or its affiliates. All Rights Reserved. 24 | * 25 | * Licensed under the Amazon Software License (the "License"). 26 | * You may not use this file except in compliance with the License. 27 | * A copy of the License is located at 28 | * 29 | * http://aws.amazon.com/asl/ 30 | * 31 | * or in the "license" file accompanying this file. This file is distributed 32 | * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either 33 | * express or implied. See the License for the specific language governing 34 | * permissions and limitations under the License. 35 | */ 36 | 37 | package interfaces 38 | 39 | import ( 40 | "time" 41 | 42 | "github.com/aws/aws-sdk-go-v2/aws" 43 | "github.com/aws/aws-sdk-go-v2/service/kinesis/types" 44 | ) 45 | 46 | const ( 47 | /* 48 | * REQUESTED Indicates that the entire application is being shutdown, and if desired the record processor will be given a 49 | * final chance to checkpoint. This state will not trigger a direct call to 50 | * {@link com.amazonaws.services.kinesis.clientlibrary.interfaces.v2.IRecordProcessor#shutdown(ShutdownInput)}, but 51 | * instead depend on a different interface for backward compatibility. 52 | */ 53 | REQUESTED ShutdownReason = iota + 1 54 | 55 | /* 56 | * Terminate processing for this RecordProcessor (resharding use case). 57 | * Indicates that the shard is closed and all records from the shard have been delivered to the application. 58 | * Applications SHOULD checkpoint their progress to indicate that they have successfully processed all records 59 | * from this shard and processing of child shards can be started. 60 | */ 61 | TERMINATE 62 | 63 | /* 64 | * Processing will be moved to a different record processor (fail over, load balancing use cases). 65 | * Applications SHOULD NOT checkpoint their progress (as another record processor may have already started 66 | * processing data). 67 | */ 68 | ZOMBIE 69 | ) 70 | 71 | // Containers for the parameters to the IRecordProcessor 72 | type ( 73 | /* 74 | * Reason the RecordProcessor is being shutdown. 75 | * Used to distinguish between a fail-over vs. a termination (shard is closed and all records have been delivered). 76 | * In case of a fail-over, applications should NOT checkpoint as part of shutdown, 77 | * since another record processor may have already started processing records for that shard. 78 | * In case of termination (resharding use case), applications SHOULD keep checkpointing their progress to indicate 79 | * that they have successfully processed all the records (processing of child shards can then begin). 80 | */ 81 | ShutdownReason int 82 | 83 | InitializationInput struct { 84 | // The shardId that the record processor is being initialized for. 85 | ShardId string 86 | 87 | // The last extended sequence number that was successfully checkpointed by the previous record processor. 88 | ExtendedSequenceNumber *ExtendedSequenceNumber 89 | } 90 | 91 | ProcessRecordsInput struct { 92 | // The time that this batch of records was received by the KCL. 93 | CacheEntryTime *time.Time 94 | 95 | // The time that this batch of records was prepared to be provided to the RecordProcessor. 96 | CacheExitTime *time.Time 97 | 98 | // The records received from Kinesis. These records may have been de-aggregated if they were published by the KPL. 99 | Records []types.Record 100 | 101 | // A checkpointer that the RecordProcessor can use to checkpoint its progress. 102 | Checkpointer IRecordProcessorCheckpointer 103 | 104 | // How far behind this batch of records was when received from Kinesis. 105 | MillisBehindLatest int64 106 | } 107 | 108 | ShutdownInput struct { 109 | // ShutdownReason shows why RecordProcessor is going to be shutdown. 110 | ShutdownReason ShutdownReason 111 | 112 | // Checkpointer is used to record the current progress. 113 | Checkpointer IRecordProcessorCheckpointer 114 | } 115 | ) 116 | 117 | var shutdownReasonMap = map[ShutdownReason]*string{ 118 | REQUESTED: aws.String("REQUESTED"), 119 | TERMINATE: aws.String("TERMINATE"), 120 | ZOMBIE: aws.String("ZOMBIE"), 121 | } 122 | 123 | func ShutdownReasonMessage(reason ShutdownReason) *string { 124 | return shutdownReasonMap[reason] 125 | } 126 | -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # Contributor Covenant Code of Conduct 2 | 3 | ## Our Pledge 4 | 5 | We as members, contributors, and leaders pledge to make participation in vmware-go-kcl-v2 project and our 6 | community a harassment-free experience for everyone, regardless of age, body 7 | size, visible or invisible disability, ethnicity, sex characteristics, gender 8 | identity and expression, level of experience, education, socio-economic status, 9 | nationality, personal appearance, race, religion, or sexual identity 10 | and orientation. 11 | 12 | We pledge to act and interact in ways that contribute to an open, welcoming, 13 | diverse, inclusive, and healthy community. 14 | 15 | ## Our Standards 16 | 17 | Examples of behavior that contributes to a positive environment for our 18 | community include: 19 | 20 | * Demonstrating empathy and kindness toward other people 21 | * Being respectful of differing opinions, viewpoints, and experiences 22 | * Giving and gracefully accepting constructive feedback 23 | * Accepting responsibility and apologizing to those affected by our mistakes, 24 | and learning from the experience 25 | * Focusing on what is best not just for us as individuals, but for the 26 | overall community 27 | 28 | Examples of unacceptable behavior include: 29 | 30 | * The use of sexualized language or imagery, and sexual attention or 31 | advances of any kind 32 | * Trolling, insulting or derogatory comments, and personal or political attacks 33 | * Public or private harassment 34 | * Publishing others' private information, such as a physical or email 35 | address, without their explicit permission 36 | * Other conduct which could reasonably be considered inappropriate in a 37 | professional setting 38 | 39 | ## Enforcement Responsibilities 40 | 41 | Community leaders are responsible for clarifying and enforcing our standards of 42 | acceptable behavior and will take appropriate and fair corrective action in 43 | response to any behavior that they deem inappropriate, threatening, offensive, 44 | or harmful. 45 | 46 | Community leaders have the right and responsibility to remove, edit, or reject 47 | comments, commits, code, wiki edits, issues, and other contributions that are 48 | not aligned to this Code of Conduct, and will communicate reasons for moderation 49 | decisions when appropriate. 50 | 51 | ## Scope 52 | 53 | This Code of Conduct applies within all community spaces, and also applies when 54 | an individual is officially representing the community in public spaces. 55 | Examples of representing our community include using an official e-mail address, 56 | posting via an official social media account, or acting as an appointed 57 | representative at an online or offline event. 58 | 59 | ## Enforcement 60 | 61 | Instances of abusive, harassing, or otherwise unacceptable behavior may be 62 | reported to the community leaders responsible for enforcement at oss-coc@@vmware.com. 63 | All complaints will be reviewed and investigated promptly and fairly. 64 | 65 | All community leaders are obligated to respect the privacy and security of the 66 | reporter of any incident. 67 | 68 | ## Enforcement Guidelines 69 | 70 | Community leaders will follow these Community Impact Guidelines in determining 71 | the consequences for any action they deem in violation of this Code of Conduct: 72 | 73 | ### 1. Correction 74 | 75 | **Community Impact**: Use of inappropriate language or other behavior deemed 76 | unprofessional or unwelcome in the community. 77 | 78 | **Consequence**: A private, written warning from community leaders, providing 79 | clarity around the nature of the violation and an explanation of why the 80 | behavior was inappropriate. A public apology may be requested. 81 | 82 | ### 2. Warning 83 | 84 | **Community Impact**: A violation through a single incident or series 85 | of actions. 86 | 87 | **Consequence**: A warning with consequences for continued behavior. No 88 | interaction with the people involved, including unsolicited interaction with 89 | those enforcing the Code of Conduct, for a specified period of time. This 90 | includes avoiding interactions in community spaces as well as external channels 91 | like social media. Violating these terms may lead to a temporary or 92 | permanent ban. 93 | 94 | ### 3. Temporary Ban 95 | 96 | **Community Impact**: A serious violation of community standards, including 97 | sustained inappropriate behavior. 98 | 99 | **Consequence**: A temporary ban from any sort of interaction or public 100 | communication with the community for a specified period of time. No public or 101 | private interaction with the people involved, including unsolicited interaction 102 | with those enforcing the Code of Conduct, is allowed during this period. 103 | Violating these terms may lead to a permanent ban. 104 | 105 | ### 4. Permanent Ban 106 | 107 | **Community Impact**: Demonstrating a pattern of violation of community 108 | standards, including sustained inappropriate behavior, harassment of an 109 | individual, or aggression toward or disparagement of classes of individuals. 110 | 111 | **Consequence**: A permanent ban from any sort of public interaction within 112 | the community. 113 | 114 | ## Attribution 115 | 116 | This Code of Conduct is adapted from the [Contributor Covenant][homepage], 117 | version 2.0, available at 118 | https://www.contributor-covenant.org/version/2/0/code_of_conduct.html. 119 | 120 | Community Impact Guidelines were inspired by [Mozilla's code of conduct 121 | enforcement ladder](https://github.com/mozilla/diversity). 122 | 123 | [homepage]: https://www.contributor-covenant.org 124 | 125 | For answers to common questions about this code of conduct, see the FAQ at 126 | https://www.contributor-covenant.org/faq. Translations are available at 127 | https://www.contributor-covenant.org/translations. 128 | -------------------------------------------------------------------------------- /test/worker_custom_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | package test 20 | 21 | import ( 22 | "context" 23 | chk "github.com/vmware/vmware-go-kcl-v2/clientlibrary/checkpoint" 24 | "os" 25 | "sync" 26 | "testing" 27 | "time" 28 | 29 | "github.com/aws/aws-sdk-go-v2/config" 30 | "github.com/aws/aws-sdk-go-v2/service/kinesis" 31 | log "github.com/sirupsen/logrus" 32 | "github.com/stretchr/testify/assert" 33 | 34 | cfg "github.com/vmware/vmware-go-kcl-v2/clientlibrary/config" 35 | par "github.com/vmware/vmware-go-kcl-v2/clientlibrary/partition" 36 | wk "github.com/vmware/vmware-go-kcl-v2/clientlibrary/worker" 37 | ) 38 | 39 | func TestWorkerInjectCheckpointer(t *testing.T) { 40 | kclConfig := cfg.NewKinesisClientLibConfig(appName, streamName, regionName, workerID). 41 | WithInitialPositionInStream(cfg.LATEST). 42 | WithMaxRecords(10). 43 | WithMaxLeasesForWorker(1). 44 | WithShardSyncIntervalMillis(5000). 45 | WithFailoverTimeMillis(300000) 46 | log.SetOutput(os.Stdout) 47 | log.SetLevel(log.DebugLevel) 48 | 49 | assert.Equal(t, regionName, kclConfig.RegionName) 50 | assert.Equal(t, streamName, kclConfig.StreamName) 51 | 52 | // configure cloudwatch as metrics system 53 | kclConfig.WithMonitoringService(getMetricsConfig(kclConfig, metricsSystem)) 54 | 55 | // Put some data into stream. 56 | kc := NewKinesisClient(t, regionName, kclConfig.KinesisEndpoint, kclConfig.KinesisCredentials) 57 | // publishSomeData(t, kc) 58 | stop := continuouslyPublishSomeData(t, kc) 59 | defer stop() 60 | 61 | // custom checkpointer or a mock checkpointer. 62 | checkpointer := chk.NewDynamoCheckpoint(kclConfig) 63 | 64 | // Inject a custom checkpointer into the worker. 65 | worker := wk.NewWorker(recordProcessorFactory(t), kclConfig). 66 | WithCheckpointer(checkpointer) 67 | 68 | err := worker.Start() 69 | assert.Nil(t, err) 70 | 71 | // wait a few seconds before shutdown processing 72 | time.Sleep(30 * time.Second) 73 | worker.Shutdown() 74 | 75 | // verify the checkpointer after graceful shutdown 76 | status := &par.ShardStatus{ 77 | ID: shardID, 78 | Mux: &sync.RWMutex{}, 79 | } 80 | 81 | _ = checkpointer.FetchCheckpoint(status) 82 | 83 | // checkpointer should be the same 84 | assert.NotEmpty(t, status.Checkpoint) 85 | 86 | // Only the lease owner has been wiped out 87 | assert.Equal(t, "", status.GetLeaseOwner()) 88 | 89 | } 90 | 91 | func TestWorkerInjectKinesis(t *testing.T) { 92 | kclConfig := cfg.NewKinesisClientLibConfig(appName, streamName, regionName, workerID). 93 | WithInitialPositionInStream(cfg.LATEST). 94 | WithMaxRecords(10). 95 | WithMaxLeasesForWorker(1). 96 | WithShardSyncIntervalMillis(5000). 97 | WithFailoverTimeMillis(300000) 98 | 99 | log.SetOutput(os.Stdout) 100 | log.SetLevel(log.DebugLevel) 101 | 102 | assert.Equal(t, regionName, kclConfig.RegionName) 103 | assert.Equal(t, streamName, kclConfig.StreamName) 104 | 105 | // configure cloudwatch as metrics system 106 | kclConfig.WithMonitoringService(getMetricsConfig(kclConfig, metricsSystem)) 107 | 108 | defaultConfig, err := config.LoadDefaultConfig( 109 | context.TODO(), 110 | config.WithRegion(regionName), 111 | ) 112 | 113 | assert.Nil(t, err) 114 | kc := kinesis.NewFromConfig(defaultConfig) 115 | 116 | // Put some data into stream. 117 | // publishSomeData(t, kc) 118 | stop := continuouslyPublishSomeData(t, kc) 119 | defer stop() 120 | 121 | // Inject a custom checkpointer into the worker. 122 | worker := wk.NewWorker(recordProcessorFactory(t), kclConfig). 123 | WithKinesis(kc) 124 | 125 | err = worker.Start() 126 | assert.Nil(t, err) 127 | 128 | // wait a few seconds before shutdown processing 129 | time.Sleep(30 * time.Second) 130 | worker.Shutdown() 131 | } 132 | 133 | func TestWorkerInjectKinesisAndCheckpointer(t *testing.T) { 134 | kclConfig := cfg.NewKinesisClientLibConfig(appName, streamName, regionName, workerID). 135 | WithInitialPositionInStream(cfg.LATEST). 136 | WithMaxRecords(10). 137 | WithMaxLeasesForWorker(1). 138 | WithShardSyncIntervalMillis(5000). 139 | WithFailoverTimeMillis(300000) 140 | 141 | log.SetOutput(os.Stdout) 142 | log.SetLevel(log.DebugLevel) 143 | 144 | assert.Equal(t, regionName, kclConfig.RegionName) 145 | assert.Equal(t, streamName, kclConfig.StreamName) 146 | 147 | // configure cloudwatch as metrics system 148 | kclConfig.WithMonitoringService(getMetricsConfig(kclConfig, metricsSystem)) 149 | 150 | // create custom Kinesis 151 | defaultConfig, err := config.LoadDefaultConfig( 152 | context.TODO(), 153 | config.WithRegion(regionName), 154 | ) 155 | 156 | assert.Nil(t, err) 157 | kc := kinesis.NewFromConfig(defaultConfig) 158 | 159 | // Put some data into stream. 160 | // publishSomeData(t, kc) 161 | stop := continuouslyPublishSomeData(t, kc) 162 | defer stop() 163 | 164 | // custom checkpointer or a mock checkpointer. 165 | checkpointer := chk.NewDynamoCheckpoint(kclConfig) 166 | 167 | // Inject both custom checkpointer and kinesis into the worker. 168 | worker := wk.NewWorker(recordProcessorFactory(t), kclConfig). 169 | WithKinesis(kc). 170 | WithCheckpointer(checkpointer) 171 | 172 | err = worker.Start() 173 | assert.Nil(t, err) 174 | 175 | // wait a few seconds before shutdown processing 176 | time.Sleep(30 * time.Second) 177 | worker.Shutdown() 178 | } 179 | -------------------------------------------------------------------------------- /clientlibrary/interfaces/record-processor-checkpointer.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package interfaces 21 | // The implementation is derived from https://github.com/awslabs/amazon-kinesis-client 22 | /* 23 | * Copyright 2014-2015 Amazon.com, Inc. or its affiliates. All Rights Reserved. 24 | * 25 | * Licensed under the Amazon Software License (the "License"). 26 | * You may not use this file except in compliance with the License. 27 | * A copy of the License is located at 28 | * 29 | * http://aws.amazon.com/asl/ 30 | * 31 | * or in the "license" file accompanying this file. This file is distributed 32 | * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either 33 | * express or implied. See the License for the specific language governing 34 | * permissions and limitations under the License. 35 | */ 36 | 37 | package interfaces 38 | 39 | type ( 40 | IPreparedCheckpointer interface { 41 | GetPendingCheckpoint() *ExtendedSequenceNumber 42 | 43 | // Checkpoint 44 | /* 45 | * This method will record a pending checkpoint. 46 | * 47 | * @error ThrottlingError Can't store checkpoint. Can be caused by checkpointing too frequently. 48 | * Consider increasing the throughput/capacity of the checkpoint store or reducing checkpoint frequency. 49 | * @error ShutdownError The record processor instance has been shutdown. Another instance may have 50 | * started processing some of these records already. 51 | * The application should abort processing via this RecordProcessor instance. 52 | * @error InvalidStateError Can't store checkpoint. 53 | * Unable to store the checkpoint in the DynamoDB table (e.g. table doesn't exist). 54 | * @error KinesisClientLibDependencyError Encountered an issue when storing the checkpoint. The application can 55 | * backoff and retry. 56 | * @error IllegalArgumentError The sequence number being checkpointed is invalid because it is out of range, 57 | * i.e. it is smaller than the last check point value (prepared or committed), or larger than the greatest 58 | * sequence number seen by the associated record processor. 59 | */ 60 | Checkpoint() error 61 | } 62 | 63 | // IRecordProcessorCheckpointer 64 | /* 65 | * Used by RecordProcessors when they want to checkpoint their progress. 66 | * The Kinesis Client Library will pass an object implementing this interface to RecordProcessors, so they can 67 | * checkpoint their progress. 68 | */ 69 | IRecordProcessorCheckpointer interface { 70 | // Checkpoint 71 | /* 72 | * This method will checkpoint the progress at the provided sequenceNumber. This method is analogous to 73 | * {@link #checkpoint()} but provides the ability to specify the sequence number at which to 74 | * checkpoint. 75 | * 76 | * @param sequenceNumber A sequence number at which to checkpoint in this shard. Upon failover, 77 | * the Kinesis Client Library will start fetching records after this sequence number. 78 | * @error ThrottlingError Can't store checkpoint. Can be caused by checkpointing too frequently. 79 | * Consider increasing the throughput/capacity of the checkpoint store or reducing checkpoint frequency. 80 | * @error ShutdownError The record processor instance has been shutdown. Another instance may have 81 | * started processing some of these records already. 82 | * The application should abort processing via this RecordProcessor instance. 83 | * @error InvalidStateError Can't store checkpoint. 84 | * Unable to store the checkpoint in the DynamoDB table (e.g. table doesn't exist). 85 | * @error KinesisClientLibDependencyError Encountered an issue when storing the checkpoint. The application can 86 | * backoff and retry. 87 | * @error IllegalArgumentError The sequence number is invalid for one of the following reasons: 88 | * 1.) It appears to be out of range, i.e. it is smaller than the last check point value, or larger than the 89 | * greatest sequence number seen by the associated record processor. 90 | * 2.) It is not a valid sequence number for a record in this shard. 91 | */ 92 | Checkpoint(sequenceNumber *string) error 93 | 94 | // PrepareCheckpoint 95 | /** 96 | * This method will record a pending checkpoint at the provided sequenceNumber. 97 | * 98 | * @param sequenceNumber A sequence number at which to prepare checkpoint in this shard. 99 | 100 | * @return an IPreparedCheckpointer object that can be called later to persist the checkpoint. 101 | * 102 | * @error ThrottlingError Can't store pending checkpoint. Can be caused by checkpointing too frequently. 103 | * Consider increasing the throughput/capacity of the checkpoint store or reducing checkpoint frequency. 104 | * @error ShutdownError The record processor instance has been shutdown. Another instance may have 105 | * started processing some of these records already. 106 | * The application should abort processing via this RecordProcessor instance. 107 | * @error InvalidStateError Can't store pending checkpoint. 108 | * Unable to store the checkpoint in the DynamoDB table (e.g. table doesn't exist). 109 | * @error KinesisClientLibDependencyError Encountered an issue when storing the pending checkpoint. The 110 | * application can backoff and retry. 111 | * @error IllegalArgumentError The sequence number is invalid for one of the following reasons: 112 | * 1.) It appears to be out of range, i.e. it is smaller than the last check point value, or larger than the 113 | * greatest sequence number seen by the associated record processor. 114 | * 2.) It is not a valid sequence number for a record in this shard. 115 | */ 116 | PrepareCheckpoint(sequenceNumber *string) (IPreparedCheckpointer, error) 117 | } 118 | ) 119 | -------------------------------------------------------------------------------- /test/lease_stealing_util_test.go: -------------------------------------------------------------------------------- 1 | package test 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "sync" 7 | "testing" 8 | "time" 9 | 10 | "github.com/aws/aws-sdk-go-v2/aws" 11 | "github.com/aws/aws-sdk-go-v2/service/dynamodb" 12 | "github.com/aws/aws-sdk-go-v2/service/dynamodb/types" 13 | "github.com/aws/aws-sdk-go-v2/service/kinesis" 14 | "github.com/stretchr/testify/assert" 15 | 16 | chk "github.com/vmware/vmware-go-kcl-v2/clientlibrary/checkpoint" 17 | cfg "github.com/vmware/vmware-go-kcl-v2/clientlibrary/config" 18 | wk "github.com/vmware/vmware-go-kcl-v2/clientlibrary/worker" 19 | ) 20 | 21 | type LeaseStealingTest struct { 22 | t *testing.T 23 | config *TestClusterConfig 24 | cluster *TestCluster 25 | kc *kinesis.Client 26 | dc *dynamodb.Client 27 | 28 | backOffSeconds int 29 | maxRetries int 30 | } 31 | 32 | func NewLeaseStealingTest(t *testing.T, config *TestClusterConfig, workerFactory TestWorkerFactory) *LeaseStealingTest { 33 | cluster := NewTestCluster(t, config, workerFactory) 34 | clientConfig := cluster.workerFactory.CreateKCLConfig("test-client", config) 35 | return &LeaseStealingTest{ 36 | t: t, 37 | config: config, 38 | cluster: cluster, 39 | kc: NewKinesisClient(t, config.regionName, clientConfig.KinesisEndpoint, clientConfig.KinesisCredentials), 40 | dc: NewDynamoDBClient(t, config.regionName, clientConfig.DynamoDBEndpoint, clientConfig.KinesisCredentials), 41 | backOffSeconds: 5, 42 | maxRetries: 60, 43 | } 44 | } 45 | 46 | func (lst *LeaseStealingTest) WithBackoffSeconds(backoff int) *LeaseStealingTest { 47 | lst.backOffSeconds = backoff 48 | return lst 49 | } 50 | 51 | func (lst *LeaseStealingTest) WithMaxRetries(retries int) *LeaseStealingTest { 52 | lst.maxRetries = retries 53 | return lst 54 | } 55 | 56 | func (lst *LeaseStealingTest) publishSomeData() (stop func()) { 57 | done := make(chan int) 58 | wg := &sync.WaitGroup{} 59 | 60 | wg.Add(1) 61 | go func() { 62 | ticker := time.NewTicker(500 * time.Millisecond) 63 | defer wg.Done() 64 | defer ticker.Stop() 65 | for { 66 | select { 67 | case <-done: 68 | return 69 | case <-ticker.C: 70 | lst.t.Log("Coninuously publishing records") 71 | publishSomeData(lst.t, lst.kc) 72 | } 73 | } 74 | }() 75 | 76 | return func() { 77 | close(done) 78 | wg.Wait() 79 | } 80 | } 81 | 82 | func (lst *LeaseStealingTest) getShardCountByWorker() map[string]int { 83 | input := &dynamodb.ScanInput{ 84 | TableName: aws.String(lst.config.appName), 85 | } 86 | 87 | shardsByWorker := map[string]map[string]bool{} 88 | scan, err := lst.dc.Scan(context.TODO(), input) 89 | for _, result := range scan.Items { 90 | if shardID, ok := result[chk.LeaseKeyKey]; !ok { 91 | continue 92 | } else if assignedTo, ok := result[chk.LeaseOwnerKey]; !ok { 93 | continue 94 | } else { 95 | if _, ok := shardsByWorker[assignedTo.(*types.AttributeValueMemberS).Value]; !ok { 96 | shardsByWorker[assignedTo.(*types.AttributeValueMemberS).Value] = map[string]bool{} 97 | } 98 | shardsByWorker[assignedTo.(*types.AttributeValueMemberS).Value][shardID.(*types.AttributeValueMemberS).Value] = true 99 | } 100 | } 101 | assert.Nil(lst.t, err) 102 | 103 | shardCountByWorker := map[string]int{} 104 | for worker, shards := range shardsByWorker { 105 | shardCountByWorker[worker] = len(shards) 106 | } 107 | return shardCountByWorker 108 | } 109 | 110 | type LeaseStealingAssertions struct { 111 | expectedLeasesForInitialWorker int 112 | expectedLeasesPerWorker int 113 | } 114 | 115 | func (lst *LeaseStealingTest) Run(assertions LeaseStealingAssertions) { 116 | // Publish records onto stream throughout the entire duration of the test 117 | stop := lst.publishSomeData() 118 | defer stop() 119 | 120 | // Start worker 1 121 | worker1, _ := lst.cluster.SpawnWorker() 122 | 123 | // Wait until the above worker has all leases 124 | var worker1ShardCount int 125 | for i := 0; i < lst.maxRetries; i++ { 126 | time.Sleep(time.Duration(lst.backOffSeconds) * time.Second) 127 | 128 | shardCountByWorker := lst.getShardCountByWorker() 129 | if shardCount, ok := shardCountByWorker[worker1]; ok && shardCount == assertions.expectedLeasesForInitialWorker { 130 | worker1ShardCount = shardCount 131 | break 132 | } 133 | } 134 | 135 | // Assert correct number of leases 136 | assert.Equal(lst.t, assertions.expectedLeasesForInitialWorker, worker1ShardCount) 137 | 138 | // Spawn Remaining Workers 139 | for i := 0; i < lst.config.numWorkers-1; i++ { 140 | lst.cluster.SpawnWorker() 141 | } 142 | 143 | // Wait For Rebalance 144 | var shardCountByWorker map[string]int 145 | for i := 0; i < lst.maxRetries; i++ { 146 | time.Sleep(time.Duration(lst.backOffSeconds) * time.Second) 147 | 148 | shardCountByWorker = lst.getShardCountByWorker() 149 | 150 | correctCount := true 151 | for _, count := range shardCountByWorker { 152 | if count != assertions.expectedLeasesPerWorker { 153 | correctCount = false 154 | } 155 | } 156 | 157 | if correctCount { 158 | break 159 | } 160 | } 161 | 162 | // Assert Rebalanced 163 | assert.Greater(lst.t, len(shardCountByWorker), 0) 164 | for _, count := range shardCountByWorker { 165 | assert.Equal(lst.t, assertions.expectedLeasesPerWorker, count) 166 | } 167 | 168 | // Shutdown Workers 169 | time.Sleep(10 * time.Second) 170 | lst.cluster.Shutdown() 171 | } 172 | 173 | type TestWorkerFactory interface { 174 | CreateWorker(workerID string, kclConfig *cfg.KinesisClientLibConfiguration) *wk.Worker 175 | CreateKCLConfig(workerID string, config *TestClusterConfig) *cfg.KinesisClientLibConfiguration 176 | } 177 | 178 | type TestClusterConfig struct { 179 | numShards int 180 | numWorkers int 181 | 182 | appName string 183 | streamName string 184 | regionName string 185 | workerIDTemplate string 186 | } 187 | 188 | type TestCluster struct { 189 | t *testing.T 190 | config *TestClusterConfig 191 | workerFactory TestWorkerFactory 192 | workerIDs []string 193 | workers map[string]*wk.Worker 194 | } 195 | 196 | func NewTestCluster(t *testing.T, config *TestClusterConfig, workerFactory TestWorkerFactory) *TestCluster { 197 | return &TestCluster{ 198 | t: t, 199 | config: config, 200 | workerFactory: workerFactory, 201 | workerIDs: make([]string, 0), 202 | workers: make(map[string]*wk.Worker), 203 | } 204 | } 205 | 206 | func (tc *TestCluster) addWorker(workerID string, config *cfg.KinesisClientLibConfiguration) *wk.Worker { 207 | worker := tc.workerFactory.CreateWorker(workerID, config) 208 | tc.workerIDs = append(tc.workerIDs, workerID) 209 | tc.workers[workerID] = worker 210 | return worker 211 | } 212 | 213 | func (tc *TestCluster) SpawnWorker() (string, *wk.Worker) { 214 | id := len(tc.workers) 215 | workerID := fmt.Sprintf(tc.config.workerIDTemplate, id) 216 | 217 | config := tc.workerFactory.CreateKCLConfig(workerID, tc.config) 218 | worker := tc.addWorker(workerID, config) 219 | 220 | err := worker.Start() 221 | assert.Nil(tc.t, err) 222 | return workerID, worker 223 | } 224 | 225 | func (tc *TestCluster) Shutdown() { 226 | for workerID, worker := range tc.workers { 227 | tc.t.Logf("Shutting down worker: %v", workerID) 228 | worker.Shutdown() 229 | } 230 | } 231 | -------------------------------------------------------------------------------- /clientlibrary/checkpoint/dynamodb-api.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2021 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | package checkpoint 21 | 22 | import ( 23 | "context" 24 | 25 | "github.com/aws/aws-sdk-go-v2/service/dynamodb" 26 | ) 27 | 28 | // DynamoDBAPI provides an interface to enable mocking the 29 | // dynamodb.DynamoDB service client's API operation, 30 | // paginators, and waiters. This make unit testing your code that calls out 31 | // to the SDK's service client's calls easier. 32 | type DynamoDBAPI interface { 33 | // The Scan operation returns one or more items and item attributes by accessing 34 | // every item in a table or a secondary index. To have DynamoDB return fewer items, 35 | // you can provide a FilterExpression operation. If the total number of scanned 36 | // items exceeds the maximum dataset size limit of 1 MB, the scan stops and results 37 | // are returned to the user as a LastEvaluatedKey value to continue the scan in a 38 | // subsequent operation. The results also include the number of items exceeding the 39 | // limit. A scan can result in no table data meeting the filter criteria. A single 40 | // Scan operation reads up to the maximum number of items set (if using the Limit 41 | // parameter) or a maximum of 1 MB of data and then apply any filtering to the 42 | // results using FilterExpression. If LastEvaluatedKey is present in the response, 43 | // you need to paginate the result set. 44 | Scan(ctx context.Context, params *dynamodb.ScanInput, optFns ...func(*dynamodb.Options)) (*dynamodb.ScanOutput, error) 45 | 46 | // DescribeTable returns information about the table, including the current status of the table, 47 | // when it was created, the primary key schema, and any indexes on the table. If 48 | // you issue a DescribeTable request immediately after a CreateTable request, 49 | // DynamoDB might return a ResourceNotFoundException. This is because DescribeTable 50 | // uses an eventually consistent query, and the metadata for your table might not 51 | // be available at that moment. Wait for a few seconds, and then try the 52 | // DescribeTable request again. 53 | DescribeTable(ctx context.Context, params *dynamodb.DescribeTableInput, optFns ...func(*dynamodb.Options)) (*dynamodb.DescribeTableOutput, error) 54 | 55 | // The CreateTable operation adds a new table to your account. In an AWS account, 56 | // table names must be unique within each Region. That is, you can have two tables 57 | // with same name if you create the tables in different Regions. CreateTable is an 58 | // asynchronous operation. Upon receiving a CreateTable request, DynamoDB 59 | // immediately returns a response with a TableStatus of CREATING. After the table 60 | // is created, DynamoDB sets the TableStatus to ACTIVE. You can perform read and 61 | // write operations only on an ACTIVE table. You can optionally define secondary 62 | // indexes on the new table, as part of the CreateTable operation. If you want to 63 | // create multiple tables with secondary indexes on them, you must create the 64 | // tables sequentially. Only one table with secondary indexes can be in the 65 | // CREATING state at any given time. You can use the DescribeTable action to check 66 | // the table status. 67 | CreateTable(ctx context.Context, params *dynamodb.CreateTableInput, optFns ...func(*dynamodb.Options)) (*dynamodb.CreateTableOutput, error) 68 | 69 | // PutItem creates a new item, or replaces an old item with a new item. If an item that has 70 | // the same primary key as the new item already exists in the specified table, the 71 | // new item completely replaces the existing item. You can perform a conditional 72 | // put operation (add a new item if one with the specified primary key doesn't 73 | // exist), or replace an existing item if it has certain attribute values. You can 74 | // return the item's attribute values in the same operation, using the ReturnValues 75 | // parameter. 76 | PutItem(ctx context.Context, params *dynamodb.PutItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.PutItemOutput, error) 77 | 78 | // The GetItem operation returns a set of attributes for the item with the given 79 | // primary key. If there is no matching item, GetItem does not return any data and 80 | // there will be no Item element in the response. GetItem provides an eventually 81 | // consistent read by default. If your application requires a strongly consistent 82 | // read, set ConsistentRead to true. Although a strongly consistent read might take 83 | // more time than an eventually consistent read, it always returns the last updated 84 | // value. 85 | GetItem(ctx context.Context, params *dynamodb.GetItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.GetItemOutput, error) 86 | 87 | // UpdateItem edits an existing item's attributes, or adds a new item to the table if it does 88 | // not already exist. You can put, delete, or add attribute values. You can also 89 | // perform a conditional update on an existing item (insert a new attribute 90 | // name-value pair if it doesn't exist, or replace an existing name-value pair if 91 | // it has certain expected attribute values). You can also return the item's 92 | // attribute values in the same UpdateItem operation using the ReturnValues 93 | // parameter. 94 | UpdateItem(ctx context.Context, params *dynamodb.UpdateItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.UpdateItemOutput, error) 95 | 96 | // DeleteItem deletes a single item in a table by primary key. You can perform a conditional 97 | // delete operation that deletes the item if it exists, or if it has an expected 98 | // attribute value. In addition to deleting an item, you can also return the item's 99 | // attribute values in the same operation, using the ReturnValues parameter. Unless 100 | // you specify conditions, the DeleteItem is an idempotent operation; running it 101 | // multiple times on the same item or attribute does not result in an error 102 | // response. Conditional deletes are useful for deleting items only if specific 103 | // conditions are met. If those conditions are met, DynamoDB performs the delete. 104 | // Otherwise, the item is not deleted. 105 | DeleteItem(ctx context.Context, params *dynamodb.DeleteItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.DeleteItemOutput, error) 106 | } 107 | -------------------------------------------------------------------------------- /clientlibrary/worker/fan-out-shard-consumer.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2021 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package worker 21 | package worker 22 | 23 | import ( 24 | "context" 25 | "errors" 26 | "time" 27 | 28 | "github.com/aws/aws-sdk-go-v2/aws" 29 | "github.com/aws/aws-sdk-go-v2/service/kinesis" 30 | "github.com/aws/aws-sdk-go-v2/service/kinesis/types" 31 | 32 | chk "github.com/vmware/vmware-go-kcl-v2/clientlibrary/checkpoint" 33 | kcl "github.com/vmware/vmware-go-kcl-v2/clientlibrary/interfaces" 34 | ) 35 | 36 | // FanOutShardConsumer is responsible for consuming data records of a (specified) shard. 37 | // Note: FanOutShardConsumer only deal with one shard. 38 | // For more info see: https://docs.aws.amazon.com/streams/latest/dev/enhanced-consumers.html 39 | type FanOutShardConsumer struct { 40 | commonShardConsumer 41 | consumerARN string 42 | consumerID string 43 | stop *chan struct{} 44 | } 45 | 46 | // getRecords subscribes to a shard and reads events from it. 47 | // Precondition: it currently has the lease on the shard. 48 | func (sc *FanOutShardConsumer) getRecords() error { 49 | defer sc.releaseLease(sc.shard.ID) 50 | 51 | log := sc.kclConfig.Logger 52 | 53 | // If the shard is child shard, need to wait until the parent finished. 54 | if err := sc.waitOnParentShard(); err != nil { 55 | // If parent shard has been deleted by Kinesis system already, just ignore the error. 56 | if err != chk.ErrSequenceIDNotFound { 57 | log.Errorf("Error in waiting for parent shard: %v to finish. Error: %+v", sc.shard.ParentShardId, err) 58 | return err 59 | } 60 | } 61 | 62 | shardSub, err := sc.subscribeToShard() 63 | if err != nil { 64 | log.Errorf("Unable to subscribe to shard %s: %v", sc.shard.ID, err) 65 | return err 66 | } 67 | defer func() { 68 | if shardSub == nil || shardSub.GetStream() == nil { 69 | log.Debugf("Nothing to close, EventStream is nil") 70 | return 71 | } 72 | err = shardSub.GetStream().Close() 73 | if err != nil { 74 | log.Errorf("Unable to close event stream for %s: %v", sc.shard.ID, err) 75 | } 76 | }() 77 | 78 | input := &kcl.InitializationInput{ 79 | ShardId: sc.shard.ID, 80 | ExtendedSequenceNumber: &kcl.ExtendedSequenceNumber{SequenceNumber: aws.String(sc.shard.GetCheckpoint())}, 81 | } 82 | sc.recordProcessor.Initialize(input) 83 | recordCheckpointer := NewRecordProcessorCheckpoint(sc.shard, sc.checkpointer) 84 | 85 | var continuationSequenceNumber *string 86 | refreshLeaseTimer := time.After(time.Until(sc.shard.LeaseTimeout.Add(-time.Duration(sc.kclConfig.LeaseRefreshPeriodMillis) * time.Millisecond))) 87 | for { 88 | getRecordsStartTime := time.Now() 89 | select { 90 | case <-*sc.stop: 91 | shutdownInput := &kcl.ShutdownInput{ShutdownReason: kcl.REQUESTED, Checkpointer: recordCheckpointer} 92 | sc.recordProcessor.Shutdown(shutdownInput) 93 | return nil 94 | case <-refreshLeaseTimer: 95 | log.Debugf("Refreshing lease on shard: %s for worker: %s", sc.shard.ID, sc.consumerID) 96 | err = sc.checkpointer.GetLease(sc.shard, sc.consumerID) 97 | if err != nil { 98 | if errors.As(err, &chk.ErrLeaseNotAcquired{}) { 99 | log.Warnf("Failed in acquiring lease on shard: %s for worker: %s", sc.shard.ID, sc.consumerID) 100 | return nil 101 | } 102 | log.Errorf("Error in refreshing lease on shard: %s for worker: %s. Error: %+v", sc.shard.ID, sc.consumerID, err) 103 | return err 104 | } 105 | refreshLeaseTimer = time.After(time.Until(sc.shard.LeaseTimeout.Add(-time.Duration(sc.kclConfig.LeaseRefreshPeriodMillis) * time.Millisecond))) 106 | // log metric for renewed lease for worker 107 | sc.mService.LeaseRenewed(sc.shard.ID) 108 | case event, ok := <-shardSub.GetStream().Events(): 109 | if !ok { 110 | // need to resubscribe to shard 111 | log.Debugf("Event stream ended, refreshing subscription on shard: %s for worker: %s", sc.shard.ID, sc.consumerID) 112 | if continuationSequenceNumber == nil || *continuationSequenceNumber == "" { 113 | log.Debugf("No continuation sequence number") 114 | return nil 115 | } 116 | shardSub, err = sc.resubscribe(shardSub, continuationSequenceNumber) 117 | if err != nil { 118 | return err 119 | } 120 | continue 121 | } 122 | subEvent, ok := event.(*types.SubscribeToShardEventStreamMemberSubscribeToShardEvent) 123 | if !ok { 124 | log.Errorf("Received unexpected event type: %T", event) 125 | continue 126 | } 127 | continuationSequenceNumber = subEvent.Value.ContinuationSequenceNumber 128 | sc.processRecords(getRecordsStartTime, subEvent.Value.Records, subEvent.Value.MillisBehindLatest, recordCheckpointer) 129 | 130 | // The shard has been closed, so no new records can be read from it 131 | if continuationSequenceNumber == nil { 132 | log.Infof("Shard %s closed", sc.shard.ID) 133 | shutdownInput := &kcl.ShutdownInput{ShutdownReason: kcl.TERMINATE, Checkpointer: recordCheckpointer} 134 | sc.recordProcessor.Shutdown(shutdownInput) 135 | return nil 136 | } 137 | } 138 | } 139 | } 140 | 141 | func (sc *FanOutShardConsumer) subscribeToShard() (*kinesis.SubscribeToShardOutput, error) { 142 | startPosition, err := sc.getStartingPosition() 143 | if err != nil { 144 | return nil, err 145 | } 146 | 147 | return sc.kc.SubscribeToShard(context.TODO(), &kinesis.SubscribeToShardInput{ 148 | ConsumerARN: &sc.consumerARN, 149 | ShardId: &sc.shard.ID, 150 | StartingPosition: startPosition, 151 | }) 152 | } 153 | 154 | func (sc *FanOutShardConsumer) resubscribe(shardSub *kinesis.SubscribeToShardOutput, continuationSequence *string) (*kinesis.SubscribeToShardOutput, error) { 155 | err := shardSub.GetStream().Close() 156 | if err != nil { 157 | sc.kclConfig.Logger.Errorf("Unable to close event stream for %s: %v", sc.shard.ID, err) 158 | return nil, err 159 | } 160 | startPosition := &types.StartingPosition{ 161 | Type: types.ShardIteratorTypeAfterSequenceNumber, 162 | SequenceNumber: continuationSequence, 163 | } 164 | shardSub, err = sc.kc.SubscribeToShard(context.TODO(), &kinesis.SubscribeToShardInput{ 165 | ConsumerARN: &sc.consumerARN, 166 | ShardId: &sc.shard.ID, 167 | StartingPosition: startPosition, 168 | }) 169 | if err != nil { 170 | sc.kclConfig.Logger.Errorf("Unable to resubscribe to shard %s: %v", sc.shard.ID, err) 171 | return nil, err 172 | } 173 | return shardSub, nil 174 | } 175 | -------------------------------------------------------------------------------- /clientlibrary/worker/common-shard-consumer.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2021 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package worker 21 | package worker 22 | 23 | import ( 24 | "context" 25 | "sync" 26 | "time" 27 | 28 | "github.com/aws/aws-sdk-go-v2/aws" 29 | "github.com/aws/aws-sdk-go-v2/service/kinesis" 30 | "github.com/aws/aws-sdk-go-v2/service/kinesis/types" 31 | deagg "github.com/awslabs/kinesis-aggregation/go/v2/deaggregator" 32 | 33 | chk "github.com/vmware/vmware-go-kcl-v2/clientlibrary/checkpoint" 34 | "github.com/vmware/vmware-go-kcl-v2/clientlibrary/config" 35 | kcl "github.com/vmware/vmware-go-kcl-v2/clientlibrary/interfaces" 36 | "github.com/vmware/vmware-go-kcl-v2/clientlibrary/metrics" 37 | par "github.com/vmware/vmware-go-kcl-v2/clientlibrary/partition" 38 | ) 39 | 40 | type shardConsumer interface { 41 | getRecords() error 42 | } 43 | 44 | type KinesisSubscriberGetter interface { 45 | SubscribeToShard(ctx context.Context, params *kinesis.SubscribeToShardInput, optFns ...func(*kinesis.Options)) (*kinesis.SubscribeToShardOutput, error) 46 | GetShardIterator(ctx context.Context, params *kinesis.GetShardIteratorInput, optFns ...func(*kinesis.Options)) (*kinesis.GetShardIteratorOutput, error) 47 | GetRecords(ctx context.Context, params *kinesis.GetRecordsInput, optFns ...func(*kinesis.Options)) (*kinesis.GetRecordsOutput, error) 48 | } 49 | 50 | // commonShardConsumer implements common functionality for regular and enhanced fan-out consumers 51 | type commonShardConsumer struct { 52 | shard *par.ShardStatus 53 | kc KinesisSubscriberGetter 54 | checkpointer chk.Checkpointer 55 | recordProcessor kcl.IRecordProcessor 56 | kclConfig *config.KinesisClientLibConfiguration 57 | mService metrics.MonitoringService 58 | } 59 | 60 | // Cleanup the internal lease cache 61 | func (sc *commonShardConsumer) releaseLease(shard string) { 62 | log := sc.kclConfig.Logger 63 | log.Infof("Release lease for shard %s", sc.shard.ID) 64 | sc.shard.SetLeaseOwner("") 65 | 66 | // Release the lease by wiping out the lease owner for the shard 67 | // Note: we don't need to do anything in case of error here and shard lease will eventually be expired. 68 | if err := sc.checkpointer.RemoveLeaseOwner(sc.shard.ID); err != nil { 69 | log.Debugf("Failed to release shard lease or shard: %s Error: %+v", sc.shard.ID, err) 70 | } 71 | 72 | // reporting lease lose metrics 73 | sc.mService.DeleteMetricMillisBehindLatest(shard) 74 | sc.mService.LeaseLost(sc.shard.ID) 75 | } 76 | 77 | // getStartingPosition gets kinesis stating position. 78 | // First try to fetch checkpoint. If checkpoint is not found use InitialPositionInStream 79 | func (sc *commonShardConsumer) getStartingPosition() (*types.StartingPosition, error) { 80 | err := sc.checkpointer.FetchCheckpoint(sc.shard) 81 | if err != nil && err != chk.ErrSequenceIDNotFound { 82 | return nil, err 83 | } 84 | 85 | checkpoint := sc.shard.GetCheckpoint() 86 | if checkpoint != "" { 87 | sc.kclConfig.Logger.Debugf("Start shard: %v at checkpoint: %v", sc.shard.ID, checkpoint) 88 | return &types.StartingPosition{ 89 | Type: types.ShardIteratorTypeAfterSequenceNumber, 90 | SequenceNumber: &checkpoint, 91 | }, nil 92 | } 93 | 94 | shardIteratorType := config.InitalPositionInStreamToShardIteratorType(sc.kclConfig.InitialPositionInStream) 95 | sc.kclConfig.Logger.Debugf("No checkpoint recorded for shard: %v, starting with: %v", sc.shard.ID, aws.ToString(shardIteratorType)) 96 | if sc.kclConfig.InitialPositionInStream == config.AT_TIMESTAMP { 97 | return &types.StartingPosition{ 98 | Type: types.ShardIteratorTypeAtTimestamp, 99 | Timestamp: sc.kclConfig.InitialPositionInStreamExtended.Timestamp, 100 | }, nil 101 | } 102 | 103 | if *shardIteratorType == "TRIM_HORIZON" { 104 | return &types.StartingPosition{ 105 | Type: types.ShardIteratorTypeTrimHorizon, 106 | }, nil 107 | } 108 | 109 | return &types.StartingPosition{ 110 | Type: types.ShardIteratorTypeLatest, 111 | }, nil 112 | } 113 | 114 | // Need to wait until the parent shard finished 115 | func (sc *commonShardConsumer) waitOnParentShard() error { 116 | if len(sc.shard.ParentShardId) == 0 { 117 | return nil 118 | } 119 | 120 | pshard := &par.ShardStatus{ 121 | ID: sc.shard.ParentShardId, 122 | Mux: &sync.RWMutex{}, 123 | } 124 | 125 | for { 126 | if err := sc.checkpointer.FetchCheckpoint(pshard); err != nil { 127 | return err 128 | } 129 | 130 | // Parent shard is finished. 131 | if pshard.GetCheckpoint() == chk.ShardEnd { 132 | return nil 133 | } 134 | 135 | time.Sleep(time.Duration(sc.kclConfig.ParentShardPollIntervalMillis) * time.Millisecond) 136 | } 137 | } 138 | 139 | func (sc *commonShardConsumer) processRecords(getRecordsStartTime time.Time, records []types.Record, millisBehindLatest *int64, recordCheckpointer kcl.IRecordProcessorCheckpointer) { 140 | log := sc.kclConfig.Logger 141 | 142 | getRecordsTime := time.Since(getRecordsStartTime).Milliseconds() 143 | sc.mService.RecordGetRecordsTime(sc.shard.ID, float64(getRecordsTime)) 144 | 145 | log.Debugf("Received %d original records.", len(records)) 146 | 147 | // De-aggregate the records if they were published by the KPL. 148 | dars, err := deagg.DeaggregateRecords(records) 149 | if err != nil { 150 | // The error is caused by bad KPL publisher and just skip the bad records 151 | // instead of being stuck here. 152 | log.Errorf("Error in de-aggregating KPL records: %+v", err) 153 | } 154 | 155 | input := &kcl.ProcessRecordsInput{ 156 | Records: dars, 157 | MillisBehindLatest: *millisBehindLatest, 158 | Checkpointer: recordCheckpointer, 159 | } 160 | 161 | recordLength := len(input.Records) 162 | recordBytes := int64(0) 163 | log.Debugf("Received %d de-aggregated records, MillisBehindLatest: %v", recordLength, input.MillisBehindLatest) 164 | 165 | for _, r := range input.Records { 166 | recordBytes += int64(len(r.Data)) 167 | } 168 | 169 | if recordLength > 0 || sc.kclConfig.CallProcessRecordsEvenForEmptyRecordList { 170 | processRecordsStartTime := time.Now() 171 | 172 | // Delivery the events to the record processor 173 | input.CacheEntryTime = &getRecordsStartTime 174 | input.CacheExitTime = &processRecordsStartTime 175 | sc.recordProcessor.ProcessRecords(input) 176 | processedRecordsTiming := time.Since(processRecordsStartTime).Milliseconds() 177 | sc.mService.RecordProcessRecordsTime(sc.shard.ID, float64(processedRecordsTiming)) 178 | } 179 | 180 | sc.mService.IncrRecordsProcessed(sc.shard.ID, recordLength) 181 | sc.mService.IncrBytesProcessed(sc.shard.ID, recordBytes) 182 | sc.mService.MillisBehindLatest(sc.shard.ID, float64(*millisBehindLatest)) 183 | } 184 | -------------------------------------------------------------------------------- /clientlibrary/metrics/prometheus/prometheus.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package prometheus 21 | // The implementation is derived from https://github.com/patrobinson/gokini 22 | // 23 | // Copyright 2018 Patrick robinson. 24 | // 25 | // Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 26 | // 27 | // The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 28 | // 29 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 30 | package prometheus 31 | 32 | import ( 33 | "net/http" 34 | 35 | prom "github.com/prometheus/client_golang/prometheus" 36 | "github.com/prometheus/client_golang/prometheus/promhttp" 37 | 38 | "github.com/vmware/vmware-go-kcl-v2/logger" 39 | ) 40 | 41 | // MonitoringService publishes kcl metrics to Prometheus. 42 | // It might be trick if the service onboarding with KCL already uses Prometheus. 43 | type MonitoringService struct { 44 | listenAddress string 45 | namespace string 46 | streamName string 47 | workerID string 48 | region string 49 | logger logger.Logger 50 | 51 | processedRecords *prom.CounterVec 52 | processedBytes *prom.CounterVec 53 | behindLatestMillis *prom.GaugeVec 54 | leasesHeld *prom.GaugeVec 55 | leaseRenewals *prom.CounterVec 56 | getRecordsTime *prom.HistogramVec 57 | processRecordsTime *prom.HistogramVec 58 | } 59 | 60 | // NewMonitoringService returns a Monitoring service publishing metrics to Prometheus. 61 | func NewMonitoringService(listenAddress, region string, logger logger.Logger) *MonitoringService { 62 | return &MonitoringService{ 63 | listenAddress: listenAddress, 64 | region: region, 65 | logger: logger, 66 | } 67 | } 68 | 69 | func (p *MonitoringService) Init(appName, streamName, workerID string) error { 70 | p.namespace = appName 71 | p.streamName = streamName 72 | p.workerID = workerID 73 | 74 | p.processedBytes = prom.NewCounterVec(prom.CounterOpts{ 75 | Name: p.namespace + `_processed_bytes`, 76 | Help: "Number of bytes processed", 77 | }, []string{"kinesisStream", "shard"}) 78 | p.processedRecords = prom.NewCounterVec(prom.CounterOpts{ 79 | Name: p.namespace + `_processed_records`, 80 | Help: "Number of records processed", 81 | }, []string{"kinesisStream", "shard"}) 82 | p.behindLatestMillis = prom.NewGaugeVec(prom.GaugeOpts{ 83 | Name: p.namespace + `_behind_latest_millis`, 84 | Help: "The amount of milliseconds processing is behind", 85 | }, []string{"kinesisStream", "shard"}) 86 | p.leasesHeld = prom.NewGaugeVec(prom.GaugeOpts{ 87 | Name: p.namespace + `_leases_held`, 88 | Help: "The number of leases held by the worker", 89 | }, []string{"kinesisStream", "shard", "workerID"}) 90 | p.leaseRenewals = prom.NewCounterVec(prom.CounterOpts{ 91 | Name: p.namespace + `_lease_renewals`, 92 | Help: "The number of successful lease renewals", 93 | }, []string{"kinesisStream", "shard", "workerID"}) 94 | p.getRecordsTime = prom.NewHistogramVec(prom.HistogramOpts{ 95 | Name: p.namespace + `_get_records_duration_milliseconds`, 96 | Help: "The time taken to fetch records and process them", 97 | }, []string{"kinesisStream", "shard"}) 98 | p.processRecordsTime = prom.NewHistogramVec(prom.HistogramOpts{ 99 | Name: p.namespace + `_process_records_duration_milliseconds`, 100 | Help: "The time taken to process records", 101 | }, []string{"kinesisStream", "shard"}) 102 | 103 | metrics := []prom.Collector{ 104 | p.processedBytes, 105 | p.processedRecords, 106 | p.behindLatestMillis, 107 | p.leasesHeld, 108 | p.leaseRenewals, 109 | p.getRecordsTime, 110 | p.processRecordsTime, 111 | } 112 | for _, metric := range metrics { 113 | err := prom.Register(metric) 114 | if err != nil { 115 | return err 116 | } 117 | } 118 | 119 | return nil 120 | } 121 | 122 | func (p *MonitoringService) Start() error { 123 | http.Handle("/metrics", promhttp.Handler()) 124 | go func() { 125 | p.logger.Infof("Starting Prometheus listener on %s", p.listenAddress) 126 | err := http.ListenAndServe(p.listenAddress, nil) 127 | if err != nil { 128 | p.logger.Errorf("Error starting Prometheus metrics endpoint. %+v", err) 129 | } 130 | p.logger.Infof("Stopped metrics server") 131 | }() 132 | 133 | return nil 134 | } 135 | 136 | func (p *MonitoringService) Shutdown() {} 137 | 138 | func (p *MonitoringService) IncrRecordsProcessed(shard string, count int) { 139 | p.processedRecords.With(prom.Labels{"shard": shard, "kinesisStream": p.streamName}).Add(float64(count)) 140 | } 141 | 142 | func (p *MonitoringService) IncrBytesProcessed(shard string, count int64) { 143 | p.processedBytes.With(prom.Labels{"shard": shard, "kinesisStream": p.streamName}).Add(float64(count)) 144 | } 145 | 146 | func (p *MonitoringService) MillisBehindLatest(shard string, millSeconds float64) { 147 | p.behindLatestMillis.With(prom.Labels{"shard": shard, "kinesisStream": p.streamName}).Set(millSeconds) 148 | } 149 | 150 | func (p *MonitoringService) DeleteMetricMillisBehindLatest(shard string) { 151 | p.behindLatestMillis.Delete(prom.Labels{"shard": shard, "kinesisStream": p.streamName}) 152 | } 153 | 154 | func (p *MonitoringService) LeaseGained(shard string) { 155 | p.leasesHeld.With(prom.Labels{"shard": shard, "kinesisStream": p.streamName, "workerID": p.workerID}).Inc() 156 | } 157 | 158 | func (p *MonitoringService) LeaseLost(shard string) { 159 | p.leasesHeld.With(prom.Labels{"shard": shard, "kinesisStream": p.streamName, "workerID": p.workerID}).Dec() 160 | } 161 | 162 | func (p *MonitoringService) LeaseRenewed(shard string) { 163 | p.leaseRenewals.With(prom.Labels{"shard": shard, "kinesisStream": p.streamName, "workerID": p.workerID}).Inc() 164 | } 165 | 166 | func (p *MonitoringService) RecordGetRecordsTime(shard string, time float64) { 167 | p.getRecordsTime.With(prom.Labels{"shard": shard, "kinesisStream": p.streamName}).Observe(time) 168 | } 169 | 170 | func (p *MonitoringService) RecordProcessRecordsTime(shard string, time float64) { 171 | p.processRecordsTime.With(prom.Labels{"shard": shard, "kinesisStream": p.streamName}).Observe(time) 172 | } 173 | -------------------------------------------------------------------------------- /test/record_publisher_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2020 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | package test 20 | 21 | import ( 22 | "context" 23 | "crypto/md5" 24 | "fmt" 25 | "sync" 26 | "testing" 27 | "time" 28 | 29 | "github.com/aws/aws-sdk-go-v2/aws" 30 | "github.com/aws/aws-sdk-go-v2/aws/retry" 31 | awsConfig "github.com/aws/aws-sdk-go-v2/config" 32 | "github.com/aws/aws-sdk-go-v2/service/dynamodb" 33 | "github.com/aws/aws-sdk-go-v2/service/kinesis" 34 | "github.com/aws/aws-sdk-go-v2/service/kinesis/types" 35 | rec "github.com/awslabs/kinesis-aggregation/go/v2/records" 36 | "github.com/golang/protobuf/proto" 37 | 38 | "github.com/vmware/vmware-go-kcl-v2/clientlibrary/utils" 39 | ) 40 | 41 | const specstr = `{"name":"kube-qQyhk","networking":{"containerNetworkCidr":"10.2.0.0/16"},"orgName":"BVT-Org-cLQch","projectName":"project-tDSJd","serviceLevel":"DEVELOPER","size":{"count":1},"version":"1.8.1-4"}` 42 | 43 | // NewKinesisClient to create a Kinesis Client. 44 | func NewKinesisClient(t *testing.T, regionName, endpoint string, creds aws.CredentialsProvider) *kinesis.Client { 45 | // create session for Kinesis 46 | t.Logf("Creating Kinesis client") 47 | 48 | resolver := aws.EndpointResolverWithOptionsFunc(func(service, region string, options ...interface{}) (aws.Endpoint, error) { 49 | return aws.Endpoint{ 50 | PartitionID: "aws", 51 | URL: endpoint, 52 | SigningRegion: regionName, 53 | }, nil 54 | }) 55 | 56 | cfg, err := awsConfig.LoadDefaultConfig( 57 | context.TODO(), 58 | awsConfig.WithRegion(regionName), 59 | awsConfig.WithCredentialsProvider(creds), 60 | awsConfig.WithEndpointResolverWithOptions(resolver), 61 | awsConfig.WithRetryer(func() aws.Retryer { 62 | return retry.AddWithMaxBackoffDelay(retry.NewStandard(), retry.DefaultMaxBackoff) 63 | }), 64 | ) 65 | 66 | if err != nil { 67 | // no need to move forward 68 | t.Fatalf("Failed in loading Kinesis default config for creating Worker: %+v", err) 69 | } 70 | 71 | return kinesis.NewFromConfig(cfg) 72 | } 73 | 74 | // NewDynamoDBClient to create a Kinesis Client. 75 | func NewDynamoDBClient(t *testing.T, regionName, endpoint string, creds aws.CredentialsProvider) *dynamodb.Client { 76 | resolver := aws.EndpointResolverWithOptionsFunc(func(service, region string, options ...interface{}) (aws.Endpoint, error) { 77 | return aws.Endpoint{ 78 | PartitionID: "aws", 79 | URL: endpoint, 80 | SigningRegion: regionName, 81 | }, nil 82 | }) 83 | 84 | cfg, err := awsConfig.LoadDefaultConfig( 85 | context.TODO(), 86 | awsConfig.WithRegion(regionName), 87 | awsConfig.WithCredentialsProvider(creds), 88 | awsConfig.WithEndpointResolverWithOptions(resolver), 89 | awsConfig.WithRetryer(func() aws.Retryer { 90 | return retry.AddWithMaxBackoffDelay(retry.NewStandard(), retry.DefaultMaxBackoff) 91 | }), 92 | ) 93 | 94 | if err != nil { 95 | t.Fatalf("unable to load SDK config, %v", err) 96 | } 97 | 98 | return dynamodb.NewFromConfig(cfg) 99 | } 100 | 101 | func continuouslyPublishSomeData(t *testing.T, kc *kinesis.Client) func() { 102 | var shards []types.Shard 103 | var nextToken *string 104 | for { 105 | out, err := kc.ListShards(context.TODO(), &kinesis.ListShardsInput{ 106 | StreamName: aws.String(streamName), 107 | NextToken: nextToken, 108 | }) 109 | if err != nil { 110 | t.Errorf("Error in ListShards. %+v", err) 111 | } 112 | 113 | shards = append(shards, out.Shards...) 114 | if out.NextToken == nil { 115 | break 116 | } 117 | nextToken = out.NextToken 118 | } 119 | 120 | done := make(chan int) 121 | wg := &sync.WaitGroup{} 122 | 123 | wg.Add(1) 124 | go func() { 125 | defer wg.Done() 126 | ticker := time.NewTicker(500 * time.Millisecond) 127 | for { 128 | select { 129 | case <-done: 130 | return 131 | case <-ticker.C: 132 | publishToAllShards(t, kc, shards) 133 | publishSomeData(t, kc) 134 | } 135 | } 136 | }() 137 | 138 | return func() { 139 | close(done) 140 | wg.Wait() 141 | } 142 | } 143 | 144 | func publishToAllShards(t *testing.T, kc *kinesis.Client, shards []types.Shard) { 145 | // Put records to all shards 146 | for i := 0; i < 10; i++ { 147 | for _, shard := range shards { 148 | publishRecord(t, kc, shard.HashKeyRange.StartingHashKey) 149 | } 150 | } 151 | } 152 | 153 | // publishSomeData to put some records into Kinesis stream 154 | func publishSomeData(t *testing.T, kc *kinesis.Client) { 155 | // Put some data into stream. 156 | t.Log("Putting data into stream using PutRecord API...") 157 | for i := 0; i < 50; i++ { 158 | publishRecord(t, kc, nil) 159 | } 160 | t.Log("Done putting data into stream using PutRecord API.") 161 | 162 | // Put some data into stream using PutRecords API 163 | t.Log("Putting data into stream using PutRecords API...") 164 | for i := 0; i < 10; i++ { 165 | publishRecords(t, kc) 166 | } 167 | t.Log("Done putting data into stream using PutRecords API.") 168 | 169 | // Put some data into stream using KPL Aggregate Record format 170 | t.Log("Putting data into stream using KPL Aggregate Record ...") 171 | for i := 0; i < 10; i++ { 172 | publishAggregateRecord(t, kc) 173 | } 174 | t.Log("Done putting data into stream using KPL Aggregate Record.") 175 | } 176 | 177 | // publishRecord to put a record into Kinesis stream using PutRecord API. 178 | func publishRecord(t *testing.T, kc *kinesis.Client, hashKey *string) { 179 | input := &kinesis.PutRecordInput{ 180 | Data: []byte(specstr), 181 | StreamName: aws.String(streamName), 182 | PartitionKey: aws.String(utils.RandStringBytesMaskImpr(10)), 183 | } 184 | if hashKey != nil { 185 | input.ExplicitHashKey = hashKey 186 | } 187 | // Use random string as partition key to ensure even distribution across shards 188 | _, err := kc.PutRecord(context.TODO(), input) 189 | 190 | if err != nil { 191 | t.Errorf("Error in PutRecord. %+v", err) 192 | } 193 | } 194 | 195 | // publishRecord to put a record into Kinesis stream using PutRecords API. 196 | func publishRecords(t *testing.T, kc *kinesis.Client) { 197 | // Use random string as partition key to ensure even distribution across shards 198 | records := make([]types.PutRecordsRequestEntry, 5) 199 | 200 | for i := 0; i < 5; i++ { 201 | record := types.PutRecordsRequestEntry{ 202 | Data: []byte(specstr), 203 | PartitionKey: aws.String(utils.RandStringBytesMaskImpr(10)), 204 | } 205 | records[i] = record 206 | } 207 | 208 | _, err := kc.PutRecords(context.TODO(), &kinesis.PutRecordsInput{ 209 | Records: records, 210 | StreamName: aws.String(streamName), 211 | }) 212 | 213 | if err != nil { 214 | t.Errorf("Error in PutRecords. %+v", err) 215 | } 216 | } 217 | 218 | // publishRecord to put a record into Kinesis stream using PutRecord API. 219 | func publishAggregateRecord(t *testing.T, kc *kinesis.Client) { 220 | data := generateAggregateRecord(5, specstr) 221 | // Use random string as partition key to ensure even distribution across shards 222 | _, err := kc.PutRecord(context.TODO(), &kinesis.PutRecordInput{ 223 | Data: data, 224 | StreamName: aws.String(streamName), 225 | PartitionKey: aws.String(utils.RandStringBytesMaskImpr(10)), 226 | }) 227 | 228 | if err != nil { 229 | t.Errorf("Error in PutRecord. %+v", err) 230 | } 231 | } 232 | 233 | // generateAggregateRecord generates an aggregate record in the correct AWS-specified format used by KPL. 234 | // https://github.com/awslabs/amazon-kinesis-producer/blob/master/aggregation-format.md 235 | // copy from: https://github.com/awslabs/kinesis-aggregation/blob/master/go/deaggregator/deaggregator_test.go 236 | func generateAggregateRecord(numRecords int, content string) []byte { 237 | aggr := &rec.AggregatedRecord{} 238 | // Start with the magic header 239 | aggRecord := []byte("\xf3\x89\x9a\xc2") 240 | partKeyTable := make([]string, 0) 241 | 242 | // Create proto record with numRecords length 243 | for i := 0; i < numRecords; i++ { 244 | var partKey uint64 245 | var hashKey uint64 246 | partKey = uint64(i) 247 | hashKey = uint64(i) * uint64(10) 248 | r := &rec.Record{ 249 | PartitionKeyIndex: &partKey, 250 | ExplicitHashKeyIndex: &hashKey, 251 | Data: []byte(content), 252 | Tags: make([]*rec.Tag, 0), 253 | } 254 | 255 | aggr.Records = append(aggr.Records, r) 256 | partKeyVal := fmt.Sprint(i) 257 | partKeyTable = append(partKeyTable, partKeyVal) 258 | } 259 | 260 | aggr.PartitionKeyTable = partKeyTable 261 | // Marshal to protobuf record, create md5 sum from proto record 262 | // and append both to aggRecord with magic header 263 | data, _ := proto.Marshal(aggr) 264 | md5Hash := md5.Sum(data) 265 | aggRecord = append(aggRecord, data...) 266 | aggRecord = append(aggRecord, md5Hash[:]...) 267 | return aggRecord 268 | } 269 | -------------------------------------------------------------------------------- /clientlibrary/worker/polling-shard-consumer_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2023 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | package worker 21 | 22 | import ( 23 | "context" 24 | "errors" 25 | "testing" 26 | "time" 27 | 28 | "github.com/aws/aws-sdk-go-v2/aws" 29 | "github.com/aws/aws-sdk-go-v2/service/kinesis" 30 | "github.com/stretchr/testify/assert" 31 | "github.com/stretchr/testify/mock" 32 | ) 33 | 34 | var ( 35 | testGetRecordsError = errors.New("GetRecords Error") 36 | ) 37 | 38 | func TestCallGetRecordsAPI(t *testing.T) { 39 | // basic happy path 40 | m1 := MockKinesisSubscriberGetter{} 41 | ret := kinesis.GetRecordsOutput{} 42 | m1.On("GetRecords", mock.Anything, mock.Anything, mock.Anything).Return(&ret, nil) 43 | psc := PollingShardConsumer{ 44 | commonShardConsumer: commonShardConsumer{kc: &m1}, 45 | } 46 | gri := kinesis.GetRecordsInput{ 47 | ShardIterator: aws.String("shard-iterator-01"), 48 | } 49 | out, _, err := psc.callGetRecordsAPI(&gri) 50 | assert.Nil(t, err) 51 | assert.Equal(t, &ret, out) 52 | m1.AssertExpectations(t) 53 | 54 | // check that localTPSExceededError is thrown when trying more than 5 TPS 55 | m2 := MockKinesisSubscriberGetter{} 56 | psc2 := PollingShardConsumer{ 57 | commonShardConsumer: commonShardConsumer{kc: &m2}, 58 | callsLeft: 0, 59 | } 60 | rateLimitTimeSince = func(t time.Time) time.Duration { 61 | return 500 * time.Millisecond 62 | } 63 | out2, _, err2 := psc2.callGetRecordsAPI(&gri) 64 | assert.Nil(t, out2) 65 | assert.ErrorIs(t, err2, localTPSExceededError) 66 | m2.AssertExpectations(t) 67 | 68 | // check that getRecords is called normally in bytesRead = 0 case 69 | m3 := MockKinesisSubscriberGetter{} 70 | ret3 := kinesis.GetRecordsOutput{} 71 | m3.On("GetRecords", mock.Anything, mock.Anything, mock.Anything).Return(&ret3, nil) 72 | psc3 := PollingShardConsumer{ 73 | commonShardConsumer: commonShardConsumer{kc: &m3}, 74 | callsLeft: 2, 75 | bytesRead: 0, 76 | } 77 | rateLimitTimeSince = func(t time.Time) time.Duration { 78 | return 2 * time.Second 79 | } 80 | out3, checkSleepVal, err3 := psc3.callGetRecordsAPI(&gri) 81 | assert.Nil(t, err3) 82 | assert.Equal(t, checkSleepVal, 0) 83 | assert.Equal(t, &ret3, out3) 84 | m3.AssertExpectations(t) 85 | 86 | // check that correct cool off period is taken for 10mb in 1 second 87 | testTime := time.Now() 88 | m4 := MockKinesisSubscriberGetter{} 89 | ret4 := kinesis.GetRecordsOutput{Records: nil} 90 | m4.On("GetRecords", mock.Anything, mock.Anything, mock.Anything).Return(&ret4, nil) 91 | psc4 := PollingShardConsumer{ 92 | commonShardConsumer: commonShardConsumer{kc: &m4}, 93 | callsLeft: 2, 94 | bytesRead: MaxBytes, 95 | lastCheckTime: testTime, 96 | remBytes: MaxBytes, 97 | } 98 | rateLimitTimeSince = func(t time.Time) time.Duration { 99 | return 2 * time.Second 100 | } 101 | rateLimitTimeNow = func() time.Time { 102 | return testTime.Add(time.Second) 103 | } 104 | out4, checkSleepVal2, err4 := psc4.callGetRecordsAPI(&gri) 105 | assert.Nil(t, err4) 106 | assert.Equal(t, &ret4, out4) 107 | m4.AssertExpectations(t) 108 | if checkSleepVal2 != 0 { 109 | t.Errorf("Incorrect Cool Off Period: %v", checkSleepVal2) 110 | } 111 | 112 | // check that no cool off period is taken for 6mb in 3 seconds 113 | testTime2 := time.Now() 114 | m5 := MockKinesisSubscriberGetter{} 115 | ret5 := kinesis.GetRecordsOutput{} 116 | m5.On("GetRecords", mock.Anything, mock.Anything, mock.Anything).Return(&ret5, nil) 117 | psc5 := PollingShardConsumer{ 118 | commonShardConsumer: commonShardConsumer{kc: &m5}, 119 | callsLeft: 2, 120 | bytesRead: MaxBytesPerSecond * 3, 121 | lastCheckTime: testTime2, 122 | remBytes: MaxBytes, 123 | } 124 | rateLimitTimeSince = func(t time.Time) time.Duration { 125 | return 3 * time.Second 126 | } 127 | rateLimitTimeNow = func() time.Time { 128 | return testTime2.Add(time.Second * 3) 129 | } 130 | out5, checkSleepVal3, err5 := psc5.callGetRecordsAPI(&gri) 131 | assert.Nil(t, err5) 132 | assert.Equal(t, checkSleepVal3, 0) 133 | assert.Equal(t, &ret5, out5) 134 | m5.AssertExpectations(t) 135 | 136 | // check for correct cool off period with 8mb in .2 seconds with 6mb remaining 137 | testTime3 := time.Now() 138 | m6 := MockKinesisSubscriberGetter{} 139 | ret6 := kinesis.GetRecordsOutput{Records: nil} 140 | m6.On("GetRecords", mock.Anything, mock.Anything, mock.Anything).Return(&ret6, nil) 141 | psc6 := PollingShardConsumer{ 142 | commonShardConsumer: commonShardConsumer{kc: &m6}, 143 | callsLeft: 2, 144 | bytesRead: MaxBytesPerSecond * 4, 145 | lastCheckTime: testTime3, 146 | remBytes: MaxBytesPerSecond * 3, 147 | } 148 | rateLimitTimeSince = func(t time.Time) time.Duration { 149 | return 3 * time.Second 150 | } 151 | rateLimitTimeNow = func() time.Time { 152 | return testTime3.Add(time.Second / 5) 153 | } 154 | out6, checkSleepVal4, err6 := psc6.callGetRecordsAPI(&gri) 155 | assert.Nil(t, err6) 156 | assert.Equal(t, &ret6, out6) 157 | m5.AssertExpectations(t) 158 | if checkSleepVal4 != 0 { 159 | t.Errorf("Incorrect Cool Off Period: %v", checkSleepVal4) 160 | } 161 | 162 | // case where getRecords throws error 163 | m7 := MockKinesisSubscriberGetter{} 164 | ret7 := kinesis.GetRecordsOutput{Records: nil} 165 | m7.On("GetRecords", mock.Anything, mock.Anything, mock.Anything).Return(&ret7, testGetRecordsError) 166 | psc7 := PollingShardConsumer{ 167 | commonShardConsumer: commonShardConsumer{kc: &m7}, 168 | callsLeft: 2, 169 | bytesRead: 0, 170 | } 171 | rateLimitTimeSince = func(t time.Time) time.Duration { 172 | return 2 * time.Second 173 | } 174 | out7, checkSleepVal7, err7 := psc7.callGetRecordsAPI(&gri) 175 | assert.Equal(t, err7, testGetRecordsError) 176 | assert.Equal(t, checkSleepVal7, 0) 177 | assert.Equal(t, out7, &ret7) 178 | m7.AssertExpectations(t) 179 | 180 | // restore original func 181 | rateLimitTimeNow = time.Now 182 | rateLimitTimeSince = time.Since 183 | 184 | } 185 | 186 | type MockKinesisSubscriberGetter struct { 187 | mock.Mock 188 | } 189 | 190 | func (m *MockKinesisSubscriberGetter) GetRecords(ctx context.Context, params *kinesis.GetRecordsInput, optFns ...func(*kinesis.Options)) (*kinesis.GetRecordsOutput, error) { 191 | ret := m.Called(ctx, params, optFns) 192 | 193 | return ret.Get(0).(*kinesis.GetRecordsOutput), ret.Error(1) 194 | } 195 | 196 | func (m *MockKinesisSubscriberGetter) GetShardIterator(ctx context.Context, params *kinesis.GetShardIteratorInput, optFns ...func(*kinesis.Options)) (*kinesis.GetShardIteratorOutput, error) { 197 | return nil, nil 198 | } 199 | 200 | func (m *MockKinesisSubscriberGetter) SubscribeToShard(ctx context.Context, params *kinesis.SubscribeToShardInput, optFns ...func(*kinesis.Options)) (*kinesis.SubscribeToShardOutput, error) { 201 | return nil, nil 202 | } 203 | 204 | func TestPollingShardConsumer_checkCoolOffPeriod(t *testing.T) { 205 | refTime := time.Now() 206 | type fields struct { 207 | lastCheckTime time.Time 208 | remBytes int 209 | bytesRead int 210 | } 211 | tests := []struct { 212 | name string 213 | fields fields 214 | timeNow time.Time 215 | want int 216 | wantErr bool 217 | }{ 218 | { 219 | "zero time max bytes to spend", 220 | fields{ 221 | time.Time{}, 222 | 0, 223 | 0, 224 | }, 225 | refTime, 226 | 0, 227 | false, 228 | }, 229 | { 230 | "same second, bytes still left to spend", 231 | fields{ 232 | refTime, 233 | MaxBytesPerSecond, 234 | MaxBytesPerSecond - 1, 235 | }, 236 | refTime, 237 | 0, 238 | false, 239 | }, 240 | { 241 | "same second, not many but some bytes still left to spend", 242 | fields{ 243 | refTime, 244 | 8, 245 | MaxBytesPerSecond, 246 | }, 247 | refTime, 248 | 0, 249 | false, 250 | }, 251 | { 252 | "same second, 1 byte still left to spend", 253 | fields{ 254 | refTime, 255 | 1, 256 | MaxBytesPerSecond, 257 | }, 258 | refTime, 259 | 0, 260 | false, 261 | }, 262 | { 263 | "next second, bytes still left to spend", 264 | fields{ 265 | refTime, 266 | 42, 267 | 1024, 268 | }, 269 | refTime.Add(1 * time.Second), 270 | 0, 271 | false, 272 | }, 273 | { 274 | "same second, max bytes per second already spent", 275 | fields{ 276 | refTime, 277 | 0, 278 | MaxBytesPerSecond, 279 | }, 280 | refTime, 281 | 1, 282 | true, 283 | }, 284 | { 285 | "same second, more than max bytes per second already spent", 286 | fields{ 287 | refTime, 288 | 0, 289 | MaxBytesPerSecond + 1, 290 | }, 291 | refTime, 292 | 2, 293 | true, 294 | }, 295 | 296 | // Kinesis prevents reading more than 10 MiB at once 297 | { 298 | "same second, 10 MiB read all at once", 299 | fields{ 300 | refTime, 301 | 0, 302 | 10 * 1024 * 1024, 303 | }, 304 | refTime, 305 | 6, 306 | true, 307 | }, 308 | 309 | { 310 | "same second, 10 MB read all at once", 311 | fields{ 312 | refTime, 313 | 0, 314 | 10 * 1000 * 1000, 315 | }, 316 | refTime, 317 | 5, 318 | true, 319 | }, 320 | { 321 | "5 seconds ago, 10 MB read all at once", 322 | fields{ 323 | refTime, 324 | 0, 325 | 10 * 1000 * 1000, 326 | }, 327 | refTime.Add(5 * time.Second), 328 | 0, 329 | false, 330 | }, 331 | } 332 | for _, tt := range tests { 333 | t.Run(tt.name, func(t *testing.T) { 334 | sc := &PollingShardConsumer{ 335 | lastCheckTime: tt.fields.lastCheckTime, 336 | remBytes: tt.fields.remBytes, 337 | bytesRead: tt.fields.bytesRead, 338 | } 339 | rateLimitTimeNow = func() time.Time { 340 | return tt.timeNow 341 | } 342 | got, err := sc.checkCoolOffPeriod() 343 | if (err != nil) != tt.wantErr { 344 | t.Errorf("PollingShardConsumer.checkCoolOffPeriod() error = %v, wantErr %v", err, tt.wantErr) 345 | return 346 | } 347 | if got != tt.want { 348 | t.Errorf("PollingShardConsumer.checkCoolOffPeriod() = %v, want %v", got, tt.want) 349 | } 350 | }) 351 | } 352 | 353 | // restore original time.Now 354 | rateLimitTimeNow = time.Now 355 | } 356 | -------------------------------------------------------------------------------- /test/worker_test.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | package test 20 | 21 | import ( 22 | "net/http" 23 | "os" 24 | "os/signal" 25 | "syscall" 26 | "testing" 27 | "time" 28 | 29 | "github.com/aws/aws-sdk-go-v2/credentials" 30 | "github.com/prometheus/common/expfmt" 31 | "github.com/stretchr/testify/assert" 32 | 33 | cfg "github.com/vmware/vmware-go-kcl-v2/clientlibrary/config" 34 | "github.com/vmware/vmware-go-kcl-v2/clientlibrary/metrics" 35 | "github.com/vmware/vmware-go-kcl-v2/clientlibrary/metrics/cloudwatch" 36 | "github.com/vmware/vmware-go-kcl-v2/clientlibrary/metrics/prometheus" 37 | wk "github.com/vmware/vmware-go-kcl-v2/clientlibrary/worker" 38 | "github.com/vmware/vmware-go-kcl-v2/logger" 39 | zaplogger "github.com/vmware/vmware-go-kcl-v2/logger/zap" 40 | ) 41 | 42 | const ( 43 | appName = "appName" 44 | streamName = "kcl-test" 45 | regionName = "us-west-2" 46 | workerID = "test-worker" 47 | consumerName = "enhanced-fan-out-consumer" 48 | kinesisEndpoint = "https://kinesis.eu-west-1.amazonaws.com" 49 | dynamoEndpoint = "https://dynamodb.eu-west-1.amazonaws.com" 50 | ) 51 | 52 | const metricsSystem = "cloudwatch" 53 | 54 | var shardID string 55 | 56 | func TestWorker(t *testing.T) { 57 | // At minimal. use standard logrus logger 58 | // log := logger.NewLogrusLogger(logrus.StandardLogger()) 59 | // 60 | // In order to have precise control over logging. Use logger with config 61 | config := logger.Configuration{ 62 | EnableConsole: true, 63 | ConsoleLevel: logger.Error, 64 | ConsoleJSONFormat: false, 65 | EnableFile: true, 66 | FileLevel: logger.Info, 67 | FileJSONFormat: true, 68 | Filename: "log.log", 69 | } 70 | // Use logrus logger 71 | log := logger.NewLogrusLoggerWithConfig(config) 72 | 73 | kclConfig := cfg.NewKinesisClientLibConfig(appName, streamName, regionName, workerID). 74 | WithInitialPositionInStream(cfg.LATEST). 75 | WithMaxRecords(8). 76 | WithMaxLeasesForWorker(1). 77 | WithShardSyncIntervalMillis(5000). 78 | WithFailoverTimeMillis(300000). 79 | WithLogger(log). 80 | WithKinesisEndpoint(kinesisEndpoint) 81 | 82 | runTest(kclConfig, false, t) 83 | } 84 | 85 | func TestWorkerWithTimestamp(t *testing.T) { 86 | // In order to have precise control over logging. Use logger with config 87 | config := logger.Configuration{ 88 | EnableConsole: true, 89 | ConsoleLevel: logger.Debug, 90 | ConsoleJSONFormat: false, 91 | } 92 | // Use logrus logger 93 | log := logger.NewLogrusLoggerWithConfig(config) 94 | 95 | ts := time.Now().Add(time.Second * 5) 96 | kclConfig := cfg.NewKinesisClientLibConfig(appName, streamName, regionName, workerID). 97 | WithTimestampAtInitialPositionInStream(&ts). 98 | WithMaxRecords(10). 99 | WithMaxLeasesForWorker(1). 100 | WithShardSyncIntervalMillis(5000). 101 | WithFailoverTimeMillis(300000). 102 | WithLogger(log). 103 | WithKinesisEndpoint(kinesisEndpoint) 104 | 105 | runTest(kclConfig, false, t) 106 | } 107 | 108 | func TestWorkerWithSigInt(t *testing.T) { 109 | // At miminal. use standard zap logger 110 | //zapLogger, err := zap.NewProduction() 111 | //assert.Nil(t, err) 112 | //log := zaplogger.NewZapLogger(zapLogger.Sugar()) 113 | // 114 | // In order to have precise control over logging. Use logger with config. 115 | config := logger.Configuration{ 116 | EnableConsole: true, 117 | ConsoleLevel: logger.Debug, 118 | ConsoleJSONFormat: true, 119 | EnableFile: true, 120 | FileLevel: logger.Info, 121 | FileJSONFormat: true, 122 | Filename: "log.log", 123 | } 124 | // use zap logger 125 | log := zaplogger.NewZapLoggerWithConfig(config) 126 | 127 | kclConfig := cfg.NewKinesisClientLibConfig(appName, streamName, regionName, workerID). 128 | WithInitialPositionInStream(cfg.LATEST). 129 | WithMaxRecords(10). 130 | WithMaxLeasesForWorker(1). 131 | WithShardSyncIntervalMillis(5000). 132 | WithFailoverTimeMillis(300000). 133 | WithLogger(log). 134 | WithKinesisEndpoint(kinesisEndpoint) 135 | 136 | runTest(kclConfig, true, t) 137 | } 138 | 139 | func TestWorkerStatic(t *testing.T) { 140 | //t.Skip("Need to provide actual credentials") 141 | 142 | // Fill in the credentials for accessing Kinesis and DynamoDB. 143 | // Note: use empty string as SessionToken for long-term credentials. 144 | kinesisCreds := credentials.NewStaticCredentialsProvider("", "", "") 145 | dynamoCreds := credentials.NewStaticCredentialsProvider("", "", "") 146 | 147 | kclConfig := cfg.NewKinesisClientLibConfigWithCredentials(appName, streamName, regionName, workerID, &kinesisCreds, &dynamoCreds). 148 | WithInitialPositionInStream(cfg.LATEST). 149 | WithMaxRecords(10). 150 | WithMaxLeasesForWorker(1). 151 | WithShardSyncIntervalMillis(5000). 152 | WithFailoverTimeMillis(300000). 153 | WithKinesisEndpoint(kinesisEndpoint). 154 | WithDynamoDBEndpoint(dynamoEndpoint) 155 | 156 | runTest(kclConfig, false, t) 157 | } 158 | 159 | func TestWorkerAssumeRole(t *testing.T) { 160 | t.Skip("Need to provide actual roleARN") 161 | 162 | // Initial credentials loaded from SDK's default credential chain. Such as 163 | // the environment, shared credentials (~/.aws/credentials), or EC2 Instance 164 | // Role. These credentials will be used to make the STS Assume Role API. 165 | //sess := session.Must(session.NewSession()) 166 | 167 | // Create the credentials from AssumeRoleProvider to assume the role 168 | // referenced by the "myRoleARN" ARN. 169 | //kinesisCreds := stscreds.NewAssumeRoleProvider(sess, "arn:aws:iam::*:role/kcl-test-publisher") 170 | kinesisCreds := credentials.NewStaticCredentialsProvider("", "", "") 171 | dynamoCreds := credentials.NewStaticCredentialsProvider("", "", "") 172 | 173 | kclConfig := cfg.NewKinesisClientLibConfigWithCredentials(appName, streamName, regionName, workerID, &kinesisCreds, &dynamoCreds). 174 | WithInitialPositionInStream(cfg.LATEST). 175 | WithMaxRecords(10). 176 | WithMaxLeasesForWorker(1). 177 | WithShardSyncIntervalMillis(5000). 178 | WithFailoverTimeMillis(300000). 179 | WithKinesisEndpoint(kinesisEndpoint). 180 | WithDynamoDBEndpoint(dynamoEndpoint) 181 | 182 | runTest(kclConfig, false, t) 183 | } 184 | 185 | func TestEnhancedFanOutConsumer(t *testing.T) { 186 | // At minimal, use standard logrus logger 187 | // log := logger.NewLogrusLogger(logrus.StandardLogger()) 188 | // 189 | // In order to have precise control over logging. Use logger with config 190 | config := logger.Configuration{ 191 | EnableConsole: true, 192 | ConsoleLevel: logger.Debug, 193 | ConsoleJSONFormat: false, 194 | EnableFile: true, 195 | FileLevel: logger.Info, 196 | FileJSONFormat: true, 197 | Filename: "log.log", 198 | } 199 | // Use logrus logger 200 | log := logger.NewLogrusLoggerWithConfig(config) 201 | 202 | kclConfig := cfg.NewKinesisClientLibConfig(appName, streamName, regionName, workerID). 203 | WithInitialPositionInStream(cfg.LATEST). 204 | WithEnhancedFanOutConsumerName(consumerName). 205 | WithMaxRecords(10). 206 | WithMaxLeasesForWorker(1). 207 | WithShardSyncIntervalMillis(5000). 208 | WithFailoverTimeMillis(300000). 209 | WithLogger(log) 210 | 211 | runTest(kclConfig, false, t) 212 | } 213 | 214 | func TestEnhancedFanOutConsumerDefaultConsumerName(t *testing.T) { 215 | // At minimal, use standard logrus logger 216 | // log := logger.NewLogrusLogger(logrus.StandardLogger()) 217 | // 218 | // In order to have precise control over logging. Use logger with config 219 | config := logger.Configuration{ 220 | EnableConsole: true, 221 | ConsoleLevel: logger.Debug, 222 | ConsoleJSONFormat: false, 223 | EnableFile: true, 224 | FileLevel: logger.Info, 225 | FileJSONFormat: true, 226 | Filename: "log.log", 227 | } 228 | // Use logrus logger 229 | log := logger.NewLogrusLoggerWithConfig(config) 230 | 231 | kclConfig := cfg.NewKinesisClientLibConfig(appName, streamName, regionName, workerID). 232 | WithInitialPositionInStream(cfg.LATEST). 233 | WithEnhancedFanOutConsumer(true). 234 | WithMaxRecords(10). 235 | WithMaxLeasesForWorker(1). 236 | WithShardSyncIntervalMillis(5000). 237 | WithFailoverTimeMillis(300000). 238 | WithLogger(log) 239 | 240 | runTest(kclConfig, false, t) 241 | } 242 | 243 | func TestEnhancedFanOutConsumerARN(t *testing.T) { 244 | t.Skip("Need to provide actual consumerARN") 245 | 246 | consumerARN := "arn:aws:kinesis:*:stream/kcl-test/consumer/fanout-poc-consumer-test:*" 247 | // At minimal, use standard logrus logger 248 | // log := logger.NewLogrusLogger(logrus.StandardLogger()) 249 | // 250 | // In order to have precise control over logging. Use logger with config 251 | config := logger.Configuration{ 252 | EnableConsole: true, 253 | ConsoleLevel: logger.Debug, 254 | ConsoleJSONFormat: false, 255 | EnableFile: true, 256 | FileLevel: logger.Info, 257 | FileJSONFormat: true, 258 | Filename: "log.log", 259 | } 260 | // Use logrus logger 261 | log := logger.NewLogrusLoggerWithConfig(config) 262 | 263 | kclConfig := cfg.NewKinesisClientLibConfig(appName, streamName, regionName, workerID). 264 | WithInitialPositionInStream(cfg.LATEST). 265 | WithEnhancedFanOutConsumerARN(consumerARN). 266 | WithMaxRecords(10). 267 | WithMaxLeasesForWorker(1). 268 | WithShardSyncIntervalMillis(5000). 269 | WithFailoverTimeMillis(300000). 270 | WithLogger(log) 271 | 272 | runTest(kclConfig, false, t) 273 | } 274 | 275 | func runTest(kclConfig *cfg.KinesisClientLibConfiguration, triggersig bool, t *testing.T) { 276 | assert.Equal(t, regionName, kclConfig.RegionName) 277 | assert.Equal(t, streamName, kclConfig.StreamName) 278 | 279 | // configure cloudwatch as metrics system 280 | kclConfig.WithMonitoringService(getMetricsConfig(kclConfig, metricsSystem)) 281 | 282 | // Put some data into stream. 283 | kc := NewKinesisClient(t, regionName, kclConfig.KinesisEndpoint, kclConfig.KinesisCredentials) 284 | // publishSomeData(t, kc) 285 | stop := continuouslyPublishSomeData(t, kc) 286 | defer stop() 287 | 288 | worker := wk.NewWorker(recordProcessorFactory(t), kclConfig) 289 | err := worker.Start() 290 | assert.Nil(t, err) 291 | 292 | sigs := make(chan os.Signal, 1) 293 | signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM) 294 | 295 | // Signal processing. 296 | go func() { 297 | sig := <-sigs 298 | t.Logf("Received signal %s. Exiting", sig) 299 | worker.Shutdown() 300 | // some other processing before exit. 301 | //os.Exit(0) 302 | }() 303 | 304 | if triggersig { 305 | t.Log("Trigger signal SIGINT") 306 | p, _ := os.FindProcess(os.Getpid()) 307 | _ = p.Signal(os.Interrupt) 308 | } 309 | 310 | // wait a few seconds before shutdown processing 311 | time.Sleep(30 * time.Second) 312 | 313 | switch metricsSystem { 314 | case "prometheus": 315 | res, err := http.Get("http://localhost:8080/metrics") 316 | if err != nil { 317 | t.Fatalf("Error scraping Prometheus endpoint %s", err) 318 | } 319 | 320 | var parser expfmt.TextParser 321 | parsed, err := parser.TextToMetricFamilies(res.Body) 322 | _ = res.Body.Close() 323 | if err != nil { 324 | t.Errorf("Error reading monitoring response %s", err) 325 | } 326 | 327 | t.Logf("Prometheus: %+v", parsed) 328 | } 329 | 330 | t.Log("Calling normal shutdown at the end of application.") 331 | worker.Shutdown() 332 | } 333 | 334 | // configure different metrics system 335 | func getMetricsConfig(kclConfig *cfg.KinesisClientLibConfiguration, service string) metrics.MonitoringService { 336 | 337 | if service == "cloudwatch" { 338 | return cloudwatch.NewMonitoringServiceWithOptions(kclConfig.RegionName, 339 | kclConfig.KinesisCredentials, 340 | kclConfig.Logger, 341 | cloudwatch.DefaultCloudwatchMetricsBufferDuration) 342 | } 343 | 344 | if service == "prometheus" { 345 | return prometheus.NewMonitoringService(":8080", regionName, kclConfig.Logger) 346 | } 347 | 348 | return nil 349 | } 350 | -------------------------------------------------------------------------------- /clientlibrary/metrics/cloudwatch/cloudwatch.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package cloudwatch 21 | // The implementation is derived from https://github.com/patrobinson/gokini 22 | // 23 | // Copyright 2018 Patrick robinson. 24 | // 25 | // Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 26 | // 27 | // The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 28 | // 29 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 30 | package cloudwatch 31 | 32 | import ( 33 | "context" 34 | "sync" 35 | "time" 36 | 37 | "github.com/aws/aws-sdk-go-v2/aws" 38 | cwatch "github.com/aws/aws-sdk-go-v2/service/cloudwatch" 39 | "github.com/aws/aws-sdk-go-v2/service/cloudwatch/types" 40 | 41 | "github.com/vmware/vmware-go-kcl-v2/logger" 42 | ) 43 | 44 | // DefaultCloudwatchMetricsBufferDuration Buffer metrics for at most this long before publishing to CloudWatch. 45 | const DefaultCloudwatchMetricsBufferDuration = 10 * time.Second 46 | 47 | type MonitoringService struct { 48 | appName string 49 | streamName string 50 | workerID string 51 | region string 52 | credentials aws.CredentialsProvider 53 | logger logger.Logger 54 | 55 | // control how often to publish to CloudWatch 56 | bufferDuration time.Duration 57 | 58 | stop *chan struct{} 59 | waitGroup *sync.WaitGroup 60 | svc *cwatch.Client 61 | shardMetrics *sync.Map 62 | } 63 | 64 | type cloudWatchMetrics struct { 65 | sync.Mutex 66 | 67 | processedRecords int64 68 | processedBytes int64 69 | behindLatestMillis []float64 70 | leasesHeld int64 71 | leaseRenewals int64 72 | getRecordsTime []float64 73 | processRecordsTime []float64 74 | } 75 | 76 | // NewMonitoringService returns a Monitoring service publishing metrics to CloudWatch. 77 | func NewMonitoringService(region string, creds aws.CredentialsProvider) *MonitoringService { 78 | return NewMonitoringServiceWithOptions(region, creds, logger.GetDefaultLogger(), DefaultCloudwatchMetricsBufferDuration) 79 | } 80 | 81 | // NewMonitoringServiceWithOptions returns a Monitoring service publishing metrics to 82 | // CloudWatch with the provided credentials, buffering duration and logger. 83 | func NewMonitoringServiceWithOptions(region string, creds aws.CredentialsProvider, logger logger.Logger, bufferDur time.Duration) *MonitoringService { 84 | return &MonitoringService{ 85 | region: region, 86 | credentials: creds, 87 | logger: logger, 88 | bufferDuration: bufferDur, 89 | } 90 | } 91 | 92 | func (cw *MonitoringService) Init(appName, streamName, workerID string) error { 93 | cw.appName = appName 94 | cw.streamName = streamName 95 | cw.workerID = workerID 96 | 97 | cfg := &aws.Config{Region: cw.region} 98 | cfg.Credentials = cw.credentials 99 | 100 | cw.svc = cwatch.NewFromConfig(*cfg) 101 | cw.shardMetrics = &sync.Map{} 102 | 103 | stopChan := make(chan struct{}) 104 | cw.stop = &stopChan 105 | wg := sync.WaitGroup{} 106 | cw.waitGroup = &wg 107 | 108 | return nil 109 | } 110 | 111 | func (cw *MonitoringService) Start() error { 112 | cw.waitGroup.Add(1) 113 | // entering eventloop for sending metrics to CloudWatch 114 | go cw.eventloop() 115 | return nil 116 | } 117 | 118 | func (cw *MonitoringService) Shutdown() { 119 | cw.logger.Infof("Shutting down cloudwatch metrics system...") 120 | close(*cw.stop) 121 | cw.waitGroup.Wait() 122 | cw.logger.Infof("Cloudwatch metrics system has been shutdown.") 123 | } 124 | 125 | // eventloop start daemon to flush metrics periodically 126 | func (cw *MonitoringService) eventloop() { 127 | defer cw.waitGroup.Done() 128 | 129 | for { 130 | if err := cw.flush(); err != nil { 131 | cw.logger.Errorf("Error sending metrics to CloudWatch. %+v", err) 132 | } 133 | 134 | select { 135 | case <-*cw.stop: 136 | cw.logger.Infof("Shutting down monitoring system") 137 | if err := cw.flush(); err != nil { 138 | cw.logger.Errorf("Error sending metrics to CloudWatch. %+v", err) 139 | } 140 | return 141 | case <-time.After(cw.bufferDuration): 142 | } 143 | } 144 | } 145 | 146 | func (cw *MonitoringService) flushShard(shard string, metric *cloudWatchMetrics) bool { 147 | metric.Lock() 148 | defaultDimensions := []types.Dimension{ 149 | { 150 | Name: aws.String("Shard"), 151 | Value: &shard, 152 | }, 153 | { 154 | Name: aws.String("KinesisStreamName"), 155 | Value: &cw.streamName, 156 | }, 157 | } 158 | 159 | leaseDimensions := []types.Dimension{ 160 | { 161 | Name: aws.String("Shard"), 162 | Value: &shard, 163 | }, 164 | { 165 | Name: aws.String("KinesisStreamName"), 166 | Value: &cw.streamName, 167 | }, 168 | { 169 | Name: aws.String("WorkerID"), 170 | Value: &cw.workerID, 171 | }, 172 | } 173 | metricTimestamp := time.Now() 174 | 175 | data := []types.MetricDatum{ 176 | { 177 | Dimensions: defaultDimensions, 178 | MetricName: aws.String("RecordsProcessed"), 179 | Unit: types.StandardUnitCount, 180 | Timestamp: &metricTimestamp, 181 | Value: aws.Float64(float64(metric.processedRecords)), 182 | }, 183 | { 184 | Dimensions: defaultDimensions, 185 | MetricName: aws.String("DataBytesProcessed"), 186 | Unit: types.StandardUnitBytes, 187 | Timestamp: &metricTimestamp, 188 | Value: aws.Float64(float64(metric.processedBytes)), 189 | }, 190 | { 191 | Dimensions: leaseDimensions, 192 | MetricName: aws.String("RenewLease.Success"), 193 | Unit: types.StandardUnitCount, 194 | Timestamp: &metricTimestamp, 195 | Value: aws.Float64(float64(metric.leaseRenewals)), 196 | }, 197 | { 198 | Dimensions: leaseDimensions, 199 | MetricName: aws.String("CurrentLeases"), 200 | Unit: types.StandardUnitCount, 201 | Timestamp: &metricTimestamp, 202 | Value: aws.Float64(float64(metric.leasesHeld)), 203 | }, 204 | } 205 | 206 | if len(metric.behindLatestMillis) > 0 { 207 | data = append(data, types.MetricDatum{ 208 | Dimensions: defaultDimensions, 209 | MetricName: aws.String("MillisBehindLatest"), 210 | Unit: types.StandardUnitMilliseconds, 211 | Timestamp: &metricTimestamp, 212 | StatisticValues: &types.StatisticSet{ 213 | SampleCount: aws.Float64(float64(len(metric.behindLatestMillis))), 214 | Sum: sumFloat64(metric.behindLatestMillis), 215 | Maximum: maxFloat64(metric.behindLatestMillis), 216 | Minimum: minFloat64(metric.behindLatestMillis), 217 | }}) 218 | } 219 | 220 | if len(metric.getRecordsTime) > 0 { 221 | data = append(data, types.MetricDatum{ 222 | Dimensions: defaultDimensions, 223 | MetricName: aws.String("KinesisDataFetcher.getRecords.Time"), 224 | Unit: types.StandardUnitMilliseconds, 225 | Timestamp: &metricTimestamp, 226 | StatisticValues: &types.StatisticSet{ 227 | SampleCount: aws.Float64(float64(len(metric.getRecordsTime))), 228 | Sum: sumFloat64(metric.getRecordsTime), 229 | Maximum: maxFloat64(metric.getRecordsTime), 230 | Minimum: minFloat64(metric.getRecordsTime), 231 | }}) 232 | } 233 | 234 | if len(metric.processRecordsTime) > 0 { 235 | data = append(data, types.MetricDatum{ 236 | Dimensions: defaultDimensions, 237 | MetricName: aws.String("RecordProcessor.processRecords.Time"), 238 | Unit: types.StandardUnitMilliseconds, 239 | Timestamp: &metricTimestamp, 240 | StatisticValues: &types.StatisticSet{ 241 | SampleCount: aws.Float64(float64(len(metric.processRecordsTime))), 242 | Sum: sumFloat64(metric.processRecordsTime), 243 | Maximum: maxFloat64(metric.processRecordsTime), 244 | Minimum: minFloat64(metric.processRecordsTime), 245 | }}) 246 | } 247 | 248 | // Publish metrics data to cloud watch 249 | _, err := cw.svc.PutMetricData(context.TODO(), &cwatch.PutMetricDataInput{ 250 | Namespace: aws.String(cw.appName), 251 | MetricData: data, 252 | }) 253 | 254 | if err == nil { 255 | metric.processedRecords = 0 256 | metric.processedBytes = 0 257 | metric.behindLatestMillis = []float64{} 258 | metric.leaseRenewals = 0 259 | metric.getRecordsTime = []float64{} 260 | metric.processRecordsTime = []float64{} 261 | } else { 262 | cw.logger.Errorf("Error in publishing cloudwatch metrics. Error: %+v", err) 263 | } 264 | 265 | metric.Unlock() 266 | return true 267 | } 268 | 269 | func (cw *MonitoringService) flush() error { 270 | cw.logger.Debugf("Flushing metrics data. Stream: %s, Worker: %s", cw.streamName, cw.workerID) 271 | // publish per shard metrics 272 | cw.shardMetrics.Range(func(k, v interface{}) bool { 273 | shard, metric := k.(string), v.(*cloudWatchMetrics) 274 | return cw.flushShard(shard, metric) 275 | }) 276 | 277 | return nil 278 | } 279 | 280 | func (cw *MonitoringService) IncrRecordsProcessed(shard string, count int) { 281 | m := cw.getOrCreatePerShardMetrics(shard) 282 | m.Lock() 283 | defer m.Unlock() 284 | m.processedRecords += int64(count) 285 | } 286 | 287 | func (cw *MonitoringService) IncrBytesProcessed(shard string, count int64) { 288 | m := cw.getOrCreatePerShardMetrics(shard) 289 | m.Lock() 290 | defer m.Unlock() 291 | m.processedBytes += count 292 | } 293 | 294 | func (cw *MonitoringService) MillisBehindLatest(shard string, millSeconds float64) { 295 | m := cw.getOrCreatePerShardMetrics(shard) 296 | m.Lock() 297 | defer m.Unlock() 298 | m.behindLatestMillis = append(m.behindLatestMillis, millSeconds) 299 | } 300 | 301 | func (cw *MonitoringService) LeaseGained(shard string) { 302 | m := cw.getOrCreatePerShardMetrics(shard) 303 | m.Lock() 304 | defer m.Unlock() 305 | m.leasesHeld++ 306 | } 307 | 308 | func (cw *MonitoringService) LeaseLost(shard string) { 309 | m := cw.getOrCreatePerShardMetrics(shard) 310 | m.Lock() 311 | defer m.Unlock() 312 | m.leasesHeld-- 313 | } 314 | 315 | func (cw *MonitoringService) LeaseRenewed(shard string) { 316 | m := cw.getOrCreatePerShardMetrics(shard) 317 | m.Lock() 318 | defer m.Unlock() 319 | m.leaseRenewals++ 320 | } 321 | 322 | func (cw *MonitoringService) RecordGetRecordsTime(shard string, time float64) { 323 | m := cw.getOrCreatePerShardMetrics(shard) 324 | m.Lock() 325 | defer m.Unlock() 326 | m.getRecordsTime = append(m.getRecordsTime, time) 327 | } 328 | func (cw *MonitoringService) RecordProcessRecordsTime(shard string, time float64) { 329 | m := cw.getOrCreatePerShardMetrics(shard) 330 | m.Lock() 331 | defer m.Unlock() 332 | m.processRecordsTime = append(m.processRecordsTime, time) 333 | } 334 | 335 | func (cw *MonitoringService) DeleteMetricMillisBehindLatest(shard string) { 336 | // not implemented 337 | } 338 | 339 | func (cw *MonitoringService) getOrCreatePerShardMetrics(shard string) *cloudWatchMetrics { 340 | var i interface{} 341 | var ok bool 342 | if i, ok = cw.shardMetrics.Load(shard); !ok { 343 | m := &cloudWatchMetrics{} 344 | cw.shardMetrics.Store(shard, m) 345 | return m 346 | } 347 | 348 | return i.(*cloudWatchMetrics) 349 | } 350 | 351 | func sumFloat64(slice []float64) *float64 { 352 | sum := float64(0) 353 | for _, num := range slice { 354 | sum += num 355 | } 356 | return &sum 357 | } 358 | 359 | func maxFloat64(slice []float64) *float64 { 360 | if len(slice) < 1 { 361 | return aws.Float64(0) 362 | } 363 | max := slice[0] 364 | for _, num := range slice { 365 | if num > max { 366 | max = num 367 | } 368 | } 369 | return &max 370 | } 371 | 372 | func minFloat64(slice []float64) *float64 { 373 | if len(slice) < 1 { 374 | return aws.Float64(0) 375 | } 376 | min := slice[0] 377 | for _, num := range slice { 378 | if num < min { 379 | min = num 380 | } 381 | } 382 | return &min 383 | } 384 | -------------------------------------------------------------------------------- /clientlibrary/worker/polling-shard-consumer.go: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2018 VMware, Inc. 3 | * 4 | * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and 5 | * associated documentation files (the "Software"), to deal in the Software without restriction, including 6 | * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | * copies of the Software, and to permit persons to whom the Software is furnished to do 8 | * so, subject to the following conditions: 9 | * 10 | * The above copyright notice and this permission notice shall be included in all copies or substantial 11 | * portions of the Software. 12 | * 13 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT 14 | * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. 15 | * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 16 | * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE 17 | * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 18 | */ 19 | 20 | // Package worker 21 | // The implementation is derived from https://github.com/patrobinson/gokini 22 | // 23 | // Copyright 2018 Patrick robinson. 24 | // 25 | // Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 26 | // 27 | // The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 28 | // 29 | // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 30 | package worker 31 | 32 | import ( 33 | "context" 34 | "errors" 35 | log "github.com/sirupsen/logrus" 36 | "math" 37 | "time" 38 | 39 | "github.com/aws/aws-sdk-go-v2/aws" 40 | "github.com/aws/aws-sdk-go-v2/service/kinesis" 41 | "github.com/aws/aws-sdk-go-v2/service/kinesis/types" 42 | 43 | chk "github.com/vmware/vmware-go-kcl-v2/clientlibrary/checkpoint" 44 | kcl "github.com/vmware/vmware-go-kcl-v2/clientlibrary/interfaces" 45 | "github.com/vmware/vmware-go-kcl-v2/clientlibrary/metrics" 46 | ) 47 | 48 | const ( 49 | kinesisReadTPSLimit = 5 50 | MaxBytes = 10000000 51 | MaxBytesPerSecond = 2000000 52 | BytesToMbConversion = 1000000 53 | ) 54 | 55 | var ( 56 | rateLimitTimeNow = time.Now 57 | rateLimitTimeSince = time.Since 58 | localTPSExceededError = errors.New("Error GetRecords TPS Exceeded") 59 | maxBytesExceededError = errors.New("Error GetRecords Max Bytes For Call Period Exceeded") 60 | ) 61 | 62 | // PollingShardConsumer is responsible for polling data records from a (specified) shard. 63 | // Note: PollingShardConsumer only deal with one shard. 64 | type PollingShardConsumer struct { 65 | commonShardConsumer 66 | streamName string 67 | stop *chan struct{} 68 | consumerID string 69 | mService metrics.MonitoringService 70 | currTime time.Time 71 | callsLeft int 72 | remBytes int 73 | lastCheckTime time.Time 74 | bytesRead int 75 | } 76 | 77 | func (sc *PollingShardConsumer) getShardIterator() (*string, error) { 78 | startPosition, err := sc.getStartingPosition() 79 | if err != nil { 80 | return nil, err 81 | } 82 | 83 | shardIterArgs := &kinesis.GetShardIteratorInput{ 84 | ShardId: &sc.shard.ID, 85 | ShardIteratorType: startPosition.Type, 86 | StartingSequenceNumber: startPosition.SequenceNumber, 87 | Timestamp: startPosition.Timestamp, 88 | StreamName: &sc.streamName, 89 | } 90 | 91 | iterResp, err := sc.kc.GetShardIterator(context.TODO(), shardIterArgs) 92 | if err != nil { 93 | return nil, err 94 | } 95 | 96 | return iterResp.ShardIterator, nil 97 | } 98 | 99 | // getRecords continuously poll one shard for data record 100 | // Precondition: it currently has the lease on the shard. 101 | func (sc *PollingShardConsumer) getRecords() error { 102 | ctx, cancelFunc := context.WithCancel(context.Background()) 103 | defer func() { 104 | // cancel renewLease() 105 | cancelFunc() 106 | sc.releaseLease(sc.shard.ID) 107 | }() 108 | 109 | log := sc.kclConfig.Logger 110 | 111 | // If the shard is child shard, need to wait until the parent finished. 112 | if err := sc.waitOnParentShard(); err != nil { 113 | // If parent shard has been deleted by Kinesis system already, just ignore the error. 114 | if err != chk.ErrSequenceIDNotFound { 115 | log.Errorf("Error in waiting for parent shard: %v to finish. Error: %+v", sc.shard.ParentShardId, err) 116 | return err 117 | } 118 | } 119 | 120 | shardIterator, err := sc.getShardIterator() 121 | if err != nil { 122 | log.Errorf("Unable to get shard iterator for %s: %v", sc.shard.ID, err) 123 | return err 124 | } 125 | 126 | // Start processing events and notify record processor on shard and starting checkpoint 127 | input := &kcl.InitializationInput{ 128 | ShardId: sc.shard.ID, 129 | ExtendedSequenceNumber: &kcl.ExtendedSequenceNumber{SequenceNumber: aws.String(sc.shard.GetCheckpoint())}, 130 | } 131 | sc.recordProcessor.Initialize(input) 132 | 133 | recordCheckpointer := NewRecordProcessorCheckpoint(sc.shard, sc.checkpointer) 134 | retriedErrors := 0 135 | 136 | // define API call rate limit starting window 137 | sc.currTime = rateLimitTimeNow() 138 | sc.callsLeft = kinesisReadTPSLimit 139 | sc.bytesRead = 0 140 | sc.remBytes = MaxBytes 141 | 142 | // starting async lease renewal thread 143 | leaseRenewalErrChan := make(chan error, 1) 144 | go func() { 145 | leaseRenewalErrChan <- sc.renewLease(ctx) 146 | }() 147 | for { 148 | getRecordsStartTime := time.Now() 149 | 150 | log.Debugf("Trying to read %d record from iterator: %v", sc.kclConfig.MaxRecords, aws.ToString(shardIterator)) 151 | 152 | // Get records from stream and retry as needed 153 | getRecordsArgs := &kinesis.GetRecordsInput{ 154 | Limit: aws.Int32(int32(sc.kclConfig.MaxRecords)), 155 | ShardIterator: shardIterator, 156 | } 157 | getResp, coolDownPeriod, err := sc.callGetRecordsAPI(getRecordsArgs) 158 | if err != nil { 159 | //aws-sdk-go-v2 https://github.com/aws/aws-sdk-go-v2/blob/main/CHANGELOG.md#error-handling 160 | var throughputExceededErr *types.ProvisionedThroughputExceededException 161 | var kmsThrottlingErr *types.KMSThrottlingException 162 | if errors.As(err, &throughputExceededErr) { 163 | retriedErrors++ 164 | if retriedErrors > sc.kclConfig.MaxRetryCount { 165 | log.Errorf("message", "Throughput Exceeded Error: "+ 166 | "reached max retry count getting records from shard", 167 | "shardId", sc.shard.ID, 168 | "retryCount", retriedErrors, 169 | "error", err) 170 | return err 171 | } 172 | // If there is insufficient provisioned throughput on the stream, 173 | // subsequent calls made within the next 1 second throw ProvisionedThroughputExceededException. 174 | // ref: https://docs.aws.amazon.com/streams/latest/dev/service-sizes-and-limits.html 175 | sc.waitASecond(sc.currTime) 176 | continue 177 | } 178 | if err == localTPSExceededError { 179 | log.Infof("localTPSExceededError so sleep for a second") 180 | sc.waitASecond(sc.currTime) 181 | continue 182 | } 183 | if err == maxBytesExceededError { 184 | log.Infof("maxBytesExceededError so sleep for %+v seconds", coolDownPeriod) 185 | time.Sleep(time.Duration(coolDownPeriod) * time.Second) 186 | continue 187 | } 188 | if errors.As(err, &kmsThrottlingErr) { 189 | log.Errorf("Error getting records from shard %v: %+v", sc.shard.ID, err) 190 | retriedErrors++ 191 | // Greater than MaxRetryCount so we get the last retry 192 | if retriedErrors > sc.kclConfig.MaxRetryCount { 193 | log.Errorf("message", "KMS Throttling Error: "+ 194 | "reached max retry count getting records from shard", 195 | "shardId", sc.shard.ID, 196 | "retryCount", retriedErrors, 197 | "error", err) 198 | return err 199 | } 200 | // exponential backoff 201 | // https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Programming.Errors.html#Programming.Errors.RetryAndBackoff 202 | time.Sleep(time.Duration(math.Exp2(float64(retriedErrors))*100) * time.Millisecond) 203 | continue 204 | } 205 | log.Errorf("Error getting records from Kinesis that cannot be retried: %+v Request: %s", err, getRecordsArgs) 206 | return err 207 | } 208 | // reset the retry count after success 209 | retriedErrors = 0 210 | 211 | sc.processRecords(getRecordsStartTime, getResp.Records, getResp.MillisBehindLatest, recordCheckpointer) 212 | 213 | // The shard has been closed, so no new records can be read from it 214 | if getResp.NextShardIterator == nil { 215 | log.Infof("Shard %s closed", sc.shard.ID) 216 | shutdownInput := &kcl.ShutdownInput{ShutdownReason: kcl.TERMINATE, Checkpointer: recordCheckpointer} 217 | sc.recordProcessor.Shutdown(shutdownInput) 218 | return nil 219 | } 220 | shardIterator = getResp.NextShardIterator 221 | 222 | // Idle between each read, the user is responsible for checkpoint the progress 223 | // This value is only used when no records are returned; if records are returned, it should immediately 224 | // retrieve the next set of records. 225 | if len(getResp.Records) == 0 && aws.ToInt64(getResp.MillisBehindLatest) < int64(sc.kclConfig.IdleTimeBetweenReadsInMillis) { 226 | time.Sleep(time.Duration(sc.kclConfig.IdleTimeBetweenReadsInMillis) * time.Millisecond) 227 | } 228 | 229 | select { 230 | case <-*sc.stop: 231 | shutdownInput := &kcl.ShutdownInput{ShutdownReason: kcl.REQUESTED, Checkpointer: recordCheckpointer} 232 | sc.recordProcessor.Shutdown(shutdownInput) 233 | return nil 234 | case leaseRenewalErr := <-leaseRenewalErrChan: 235 | return leaseRenewalErr 236 | default: 237 | } 238 | } 239 | } 240 | 241 | func (sc *PollingShardConsumer) waitASecond(timePassed time.Time) { 242 | waitTime := time.Since(timePassed) 243 | if waitTime < time.Second { 244 | time.Sleep(time.Second - waitTime) 245 | } 246 | } 247 | 248 | func (sc *PollingShardConsumer) checkCoolOffPeriod() (int, error) { 249 | // Each shard can support up to a maximum total data read rate of 2 MB per second via GetRecords. 250 | // If a call to GetRecords returns 10 MB, subsequent calls made within the next 5 seconds throw an exception. 251 | // ref: https://docs.aws.amazon.com/streams/latest/dev/service-sizes-and-limits.html 252 | // check for overspending of byte budget from getRecords call 253 | currentTime := rateLimitTimeNow() 254 | secondsPassed := currentTime.Sub(sc.lastCheckTime).Seconds() 255 | sc.lastCheckTime = currentTime 256 | sc.remBytes += int(secondsPassed * MaxBytesPerSecond) 257 | 258 | if sc.remBytes > MaxBytes { 259 | sc.remBytes = MaxBytes 260 | } 261 | if sc.remBytes < 1 { 262 | // Wait until cool down period has passed to prevent ProvisionedThroughputExceededException 263 | coolDown := sc.bytesRead / MaxBytesPerSecond 264 | if sc.bytesRead%MaxBytesPerSecond > 0 { 265 | coolDown++ 266 | } 267 | return coolDown, maxBytesExceededError 268 | } else { 269 | sc.remBytes -= sc.bytesRead 270 | } 271 | return 0, nil 272 | } 273 | 274 | func (sc *PollingShardConsumer) callGetRecordsAPI(gri *kinesis.GetRecordsInput) (*kinesis.GetRecordsOutput, int, error) { 275 | if sc.bytesRead != 0 { 276 | coolDownPeriod, err := sc.checkCoolOffPeriod() 277 | if err != nil { 278 | return nil, coolDownPeriod, err 279 | } 280 | } 281 | // every new second, we get a fresh set of calls 282 | if rateLimitTimeSince(sc.currTime) > time.Second { 283 | sc.callsLeft = kinesisReadTPSLimit 284 | sc.currTime = rateLimitTimeNow() 285 | } 286 | 287 | if sc.callsLeft < 1 { 288 | return nil, 0, localTPSExceededError 289 | } 290 | getResp, err := sc.kc.GetRecords(context.TODO(), gri) 291 | sc.callsLeft-- 292 | 293 | if err != nil { 294 | return getResp, 0, err 295 | } 296 | 297 | // Calculate size of records from read transaction 298 | sc.bytesRead = 0 299 | for _, record := range getResp.Records { 300 | sc.bytesRead += len(record.Data) 301 | } 302 | if sc.lastCheckTime.IsZero() { 303 | sc.lastCheckTime = rateLimitTimeNow() 304 | } 305 | 306 | return getResp, 0, err 307 | } 308 | 309 | func (sc *PollingShardConsumer) renewLease(ctx context.Context) error { 310 | renewDuration := time.Duration(sc.kclConfig.LeaseRefreshWaitTime) * time.Millisecond 311 | for { 312 | timer := time.NewTimer(renewDuration) 313 | select { 314 | case <-timer.C: 315 | log.Debugf("Refreshing lease on shard: %s for worker: %s", sc.shard.ID, sc.consumerID) 316 | err := sc.checkpointer.GetLease(sc.shard, sc.consumerID) 317 | if err != nil { 318 | // log and return error 319 | log.Errorf("Error in refreshing lease on shard: %s for worker: %s. Error: %+v", 320 | sc.shard.ID, sc.consumerID, err) 321 | return err 322 | } 323 | // log metric for renewed lease for worker 324 | sc.mService.LeaseRenewed(sc.shard.ID) 325 | case <-ctx.Done(): 326 | // clean up timer resources 327 | if !timer.Stop() { 328 | <-timer.C 329 | } 330 | log.Debugf("renewLease was canceled") 331 | return nil 332 | } 333 | } 334 | } 335 | --------------------------------------------------------------------------------