Fix rate limiting

Author: pshastry

main.go

@@ -25,6 +25,9 @@ package main
 import (
 	"encoding/json"
 	"errors"
+	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
+	"strings"
+
 	//"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
 	"github.com/thumbtack/go/lib/metrics"
 	//"github.com/thumbtack/go/lib/monitoring"
@@ -33,7 +36,6 @@ import (
 	_ "net/http/pprof"
 	"os"
 	"strconv"
-	"strings"
 	"sync"
 	"time"
 
@@ -49,18 +51,16 @@ const (
 	paramCheckpointTable    = "CHECKPOINT_DDB_TABLE"
 	paramCheckpointRegion   = "CHECKPOINT_DDB_REGION"
 	paramCheckpointEndpoint = "CHECKPOINT_DDB_ENDPOINT"
-	paramMaxRetries         = "MAX_RETRIES"
 	paramVerbose            = "VERBOSE"
 	paramPort               = "PORT"
 	paramConfigDir          = "CONFIG_DIR"
-	defaultConfigMaxRetries = 3
+	maxRetries              = 3
 )
 
 var logger = logging.New()
 var ddbTable = os.Getenv(paramCheckpointTable)
 var ddbRegion = os.Getenv(paramCheckpointRegion)
 var ddbEndpoint = os.Getenv(paramCheckpointEndpoint)
-var maxRetries = defaultConfigMaxRetries
 var ddbClient = ddbConfigConnect(ddbRegion, ddbEndpoint, maxRetries, *logger)
 var metricsClient = newMetricsClient()
 
@@ -73,16 +73,12 @@ type config struct {
 	DstEndpoint               string `json:"dst_endpoint"`
 	SrcEnv                    string `json:"src_env"`
 	DstEnv                    string `json:"dst_env"`
-	SrcAccount                string `json:"src_account"`
-	DstAccount                string `json:"dst_account"`
-	MaxConnectRetries         int    `json:"max_connect_retries"`
 	ReadWorkers               int    `json:"read_workers"`
 	WriteWorkers              int    `json:"write_workers"`
 	ReadQps                   int64  `json:"read_qps"`
 	WriteQps                  int64  `json:"write_qps"`
 	UpdateCheckpointThreshold int    `json:"update_checkpoint_threshold"`
 	EnableStreaming           *bool   `json:"enable_streaming"`
-	TruncateTable             bool   `json:"truncate_table"`
 }
 
 // Config file is read and dumped into this struct
@@ -101,13 +97,8 @@ type syncState struct {
 	timestamp             time.Time
 }
 
-func getRoleArn(env string, account string) (string) {
-	var roleType = ""
-	if strings.ToLower(account) != "admin" {
-		roleType = "NEW_" + strings.ToUpper(env) + "_ROLE"
-	} else {
-		roleType = "OLD_" + strings.ToUpper(env) + "_ROLE"
-	}
+func getRoleArn(env string) string {
+	roleType := strings.ToUpper(env) + "_ROLE"
 	logger.WithFields(logging.Fields{"Roletype": roleType}).Debug()
 	return os.Getenv(roleType)
 }
@@ -130,7 +121,7 @@ func NewSyncState(tableConfig config) *syncState {
 			aws.NewConfig().
 				WithRegion(tableConfig.SrcRegion).
 				WithEndpoint(tableConfig.SrcEndpoint).
-				WithMaxRetries(tableConfig.MaxConnectRetries).
+				WithMaxRetries(maxRetries).
 				WithHTTPClient(httpClient),
 		))
 
@@ -139,10 +130,10 @@ func NewSyncState(tableConfig config) *syncState {
 			aws.NewConfig().
 				WithRegion(tableConfig.DstRegion).
 				WithEndpoint(tableConfig.DstEndpoint).
-				WithMaxRetries(tableConfig.MaxConnectRetries),
+				WithMaxRetries(maxRetries),
 		))
-	srcRoleArn := getRoleArn(tableConfig.SrcEnv, tableConfig.SrcAccount)
-	dstRoleArn := getRoleArn(tableConfig.DstEnv, tableConfig.DstAccount)
+	srcRoleArn := getRoleArn(tableConfig.SrcEnv)
+	dstRoleArn := getRoleArn(tableConfig.DstEnv)
 
 	if srcRoleArn == "" || dstRoleArn == "" {
 		logger.WithFields(logging.Fields{}).
@@ -154,16 +145,12 @@ func NewSyncState(tableConfig config) *syncState {
 		"Src Role Arn": srcRoleArn,
 		"Dst Role Arn": dstRoleArn}).Debug("Role ARN")
 
-	/*srcCreds := stscreds.NewCredentials(srcSess, srcRoleArn)
+	srcCreds := stscreds.NewCredentials(srcSess, srcRoleArn)
 	dstCreds := stscreds.NewCredentials(dstSess, dstRoleArn)
 
 	srcDynamo = dynamodb.New(srcSess, &aws.Config{Credentials: srcCreds})
 	dstDynamo = dynamodb.New(dstSess, &aws.Config{Credentials: dstCreds})
-	stream = dynamodbstreams.New(srcSess, &aws.Config{Credentials: srcCreds})*/
-
-	srcDynamo = dynamodb.New(srcSess, &aws.Config{})
-	dstDynamo = dynamodb.New(dstSess, &aws.Config{})
-	stream = dynamodbstreams.New(srcSess, &aws.Config{})
+	stream = dynamodbstreams.New(srcSess, &aws.Config{Credentials: srcCreds})
 
 	return &syncState{
 		tableConfig:           tableConfig,
@@ -238,14 +225,6 @@ func NewApp() *appConfig {
 			logger.SetLevel(logging.DebugLevel)
 		}
 	}
-	if os.Getenv(paramMaxRetries) != "" {
-		maxRetries, err = strconv.Atoi(os.Getenv(paramMaxRetries))
-		if err != nil {
-			logger.WithFields(logging.Fields{
-				"error": err,
-			}).Fatal("Failed to parse " + paramMaxRetries)
-		}
-	}
 
 	configFile = os.Getenv(paramConfigDir) + "/config.json"
 	tableConfig, err := readConfigFile(configFile, *logger)
@@ -298,17 +277,6 @@ func setDefaults(tableConfig []config) ([]config, error) {
 			continue
 		}
 
-		if tableConfig[i].SrcAccount == "" {
-			tableConfig[i].SrcAccount = "admin"
-		}
-
-		if tableConfig[i].DstAccount == "" {
-			tableConfig[i].DstAccount = "admin"
-		}
-
-		if tableConfig[i].MaxConnectRetries == 0 {
-			tableConfig[i].MaxConnectRetries = 3
-		}
 
 		if tableConfig[i].ReadQps == 0 {
 			tableConfig[i].ReadQps = 500
@@ -354,18 +322,22 @@ func (sync *syncState) isFreshStart(key primaryKey) bool {
 	return false
 }
 
-func getPrimaryKey(sync config) (primaryKey) {
+func getPrimaryKey(sync config) primaryKey {
 	key := primaryKey{}
-	if sync.SrcAccount == "admin" {
+	delim := "_"
+
+	if !strings.Contains(sync.SrcEnv, "	new") {
 		key.sourceTable = sync.SrcTable
 	} else {
-		key.sourceTable = sync.SrcTable + ".account." + sync.SrcAccount
+		key.sourceTable = sync.SrcTable + ".account." + strings.Split(sync.SrcEnv, delim)[0]
 	}
-	if sync.DstAccount == "admin" {
+
+	if !strings.Contains(sync.DstEnv, "new") {
 		key.dstTable = sync.DstTable
 	} else {
-		key.dstTable = sync.DstTable + ".account." + sync.DstAccount
+		key.dstTable = sync.DstTable + ".account." + strings.Split(sync.DstEnv, delim)[0]
 	}
+
 	return key
 }
 
@@ -389,8 +361,9 @@ func main() {
 				"Source Table":      key.sourceTable,
 				"Destination Table": key.dstTable,
 			}).Error("Error in connecting to tables. Check config file")
-
+			return
 		}
+
 		syncWorker.readCheckpoint()
 
 		// Call a go routine to replicate for each key

replicate.go

@@ -73,7 +73,10 @@ func (sync *syncState) replicate(quit <-chan bool, key primaryKey) {
 // Once all the workers in the readWorker group are done,
 // we close the channel, and wait for the writeWorker group
 // to finish
-func (state *syncState) copyTable(key primaryKey) (error){
+func (state *syncState) copyTable(key primaryKey) error {
+	var isSourceThroughputChanged = false
+	var isDstThroughputChanged = false
+
 	logger.WithFields(logging.Fields{
 		"Source Table":      key.sourceTable,
 		"Destination Table": key.dstTable,
@@ -84,8 +87,6 @@ func (state *syncState) copyTable(key primaryKey) (error){
 	// we are done copying
 	sourceCapacity := state.getCapacity(state.tableConfig.SrcTable, state.srcDynamo)
 	dstCapacity := state.getCapacity(state.tableConfig.DstTable, state.dstDynamo)
-	var isSourceThroughputChanged = false
-	var isDstThroughputChanged = false
 	srcDynamo := state.srcDynamo
 	dstDynamo := state.dstDynamo
 
@@ -124,14 +125,14 @@ func (state *syncState) copyTable(key primaryKey) (error){
 
 	writerWG.Add(writeWorkers)
 
-	rlDst := rate.NewLimiter(rate.Limit(state.tableConfig.WriteQps), maxBatchSize)
+	rl := rate.NewLimiter(rate.Limit(state.tableConfig.WriteQps), int(state.tableConfig.WriteQps))
 	for i := 0; i < writeWorkers; i++ {
 		logger.WithFields(logging.Fields{
 			"Write Worker":      i,
 			"Source Table":      key.sourceTable,
 			"Destination Table": key.dstTable,
 		}).Debug("Starting copy table write worker..")
-		go state.writeTable(key, items, &writerWG, i, *rlDst)
+		go state.writeTable(key, items, &writerWG, i, rl)
 	}
 	readerWG.Add(readWorkers)
 	for i := 0; i < readWorkers; i++ {
@@ -213,18 +214,16 @@ func (sync *syncState) streamSync(key primaryKey, streamArn string) error {
 		if lastEvaluatedShardId != "" {
 			input.ExclusiveStartShardId = aws.String(lastEvaluatedShardId)
 		}
-		maxConnectRetries := sync.tableConfig.MaxConnectRetries
 
-		for i := 0; i < maxConnectRetries; i++ {
+		for i := 1; i <= maxRetries; i++ {
 			result, err = sync.stream.DescribeStream(input)
-			if err != nil {
-				if i == maxConnectRetries-1 {
-					return err
-				}
-				backoff(i, "Describe Stream")
-			} else {
+			if err == nil {
 				break
 			}
+			if i == maxRetries {
+				return err
+			}
+			backoff(i, "Describe Stream")
 		}
 
 		numShards += len(result.StreamDescription.Shards)

stream.go

@@ -62,12 +62,11 @@ func (sync *syncState) shardSyncStart(key primaryKey,
 	}
 
 	shardIteratorInput := sync.getShardIteratorInput(key, *shardId, streamArn)
-	maxConnectRetries := sync.tableConfig.MaxConnectRetries
 
-	for i := 0; i < maxConnectRetries; i++ {
+	for i := 1; i <= maxRetries; i++ {
 		iterator, err = sync.stream.GetShardIterator(shardIteratorInput)
 		if err != nil {
-			if i == maxConnectRetries-1 {
+			if i == maxRetries {
 				logger.WithFields(logging.Fields{
 					"Shard Id":          *shardId,
 					"Error":             err,
@@ -87,19 +86,13 @@ func (sync *syncState) shardSyncStart(key primaryKey,
 	// when nil, the shard has been closed, and the requested iterator
 	// will not return any more data
 	for shardIterator != nil {
-		for i := 0; i < maxConnectRetries; i++ {
-			logger.WithFields(logging.Fields{
-				"Source Table": key.sourceTable,
-				"Destination Table": key.dstTable,
-				"Shard Id": *shardId}).Debug("Calling GetRecords")
+		for i := 1; i <= maxRetries; i++ {
+
 			records, err = sync.stream.GetRecords(
 				&dynamodbstreams.GetRecordsInput{ShardIterator: shardIterator})
-			logger.WithFields(logging.Fields{
-				"Source Table": key.sourceTable,
-				"Destination Table": key.dstTable,
-				"Shard Id": *shardId}).Debug("Returned from GetRecords")
+
 			if err != nil {
-				if i == maxConnectRetries-1 {
+				if i == maxRetries {
 					logger.WithFields(logging.Fields{
 						"Shard Id":          *shardId,
 						"Error":             err,
@@ -209,13 +202,12 @@ func (sync *syncState) writeRecords(
 // Insert this record in the dst table
 func (sync *syncState) insertRecord(item map[string]*dynamodb.AttributeValue, key primaryKey) error {
 	var err error
-	maxConnectRetries := sync.tableConfig.MaxConnectRetries
 
 	input := &dynamodb.PutItemInput{
 		Item:      item,
 		TableName: aws.String(sync.tableConfig.DstTable),
 	}
-	for i := 0; i < maxConnectRetries; i++ {
+	for i := 1; i <= maxRetries; i++ {
 		_, err = sync.dstDynamo.PutItem(input)
 		if err == nil {
 			return nil
@@ -229,13 +221,12 @@ func (sync *syncState) insertRecord(item map[string]*dynamodb.AttributeValue, ke
 // Remove this record from the dst table
 func (sync *syncState) removeRecord(item map[string]*dynamodb.AttributeValue, key primaryKey) error {
 	var err error
-	maxConnectRetries := sync.tableConfig.MaxConnectRetries
 
 	input := &dynamodb.DeleteItemInput{
 		Key:       item,
 		TableName: aws.String(sync.tableConfig.DstTable),
 	}
-	for i := 0; i < maxConnectRetries; i++ {
+	for i := 0; i < maxRetries; i++ {
 		_, err = sync.dstDynamo.DeleteItem(input)
 		if err == nil {
 			return nil

sync_test.go

@@ -135,18 +135,17 @@ func setupTest(sync *syncState, key primaryKey) {
 
 func (sync *syncState) scanTable(key primaryKey, name string) {
 	lastEvaluatedKey := make(map[string]*dynamodb.AttributeValue, 0)
-	maxConnectRetries := sync.tableConfig.MaxConnectRetries
 	items := make([]map[string]*dynamodb.AttributeValue, 0)
 	input := &dynamodb.ScanInput{TableName: aws.String(name)}
 
 	for {
 		if len(lastEvaluatedKey) > 0 {
 			input.ExclusiveStartKey = lastEvaluatedKey
 		}
-		for i := 0; i < maxConnectRetries; i++ {
+		for i := 1; i <= maxRetries; i++ {
 			result, err := sync.srcDynamo.Scan(input)
 			if err != nil {
-				if i == maxConnectRetries-1 {
+				if i == maxRetries {
 					return
 				}
 				backoff(i, "Scan")
@@ -289,7 +288,6 @@ func TestAll(t *testing.T) {
 	os.Setenv(paramCheckpointRegion, "us-west-2")
 	os.Setenv(paramCheckpointTable, "local-dynamodb-sync.checkpoint")
 	os.Setenv(paramCheckpointEndpoint, "http://localhost:8000")
-	os.Setenv(paramMaxRetries, "3")
 
 	app := NewApp()
 	checkpointDynamo := dynamodb.New(session.Must(
@@ -322,7 +320,6 @@ func TestAll(t *testing.T) {
 	//main()
 
 	for i := 0; i < len(syncWorkers); i++ {
-		syncWorkers[i].loadCheckpointTable()
 		syncWorkers[i].testStreamSyncWait()
 		syncWorkers[i].testExpireShards()
 	}