Adding retry logic for large JSONs (#32)

Co-authored-by: Yousif Akbar <[email protected]>
pbs · Sep 21, 2020 · b43c958 · b43c958
1 parent f0b4782
commit b43c958
Show file tree

Hide file tree

Showing 3 changed files with 129 additions and 12 deletions.
diff --git a/cmd/put.go b/cmd/put.go
@@ -1,8 +1,6 @@
 package cmd
 
 import (
-	"fmt"
-
 	"github.com/pbs/gorson/internal/gorson/io"
 	"github.com/pbs/gorson/internal/gorson/util"
 	"github.com/spf13/cobra"
@@ -13,9 +11,6 @@ var filename string
 func put(path string, parameters map[string]string) {
 	p := util.NewParameterStorePath(path)
 	io.WriteToParameterStore(parameters, *p)
-	for key := range parameters {
-		fmt.Println("wrote " + path + key)
-	}
 }
 
 func init() {

diff --git a/fixtures/long-file.json b/fixtures/long-file.json
@@ -0,0 +1,102 @@
+{
+    "Zero": "0",
+    "One": "1",
+    "Two": "2",
+    "Three": "3",
+    "Four": "4",
+    "Five": "5",
+    "Six": "6",
+    "Seven": "7",
+    "Eight": "8",
+    "Nine": "9",
+    "Ten": "10",
+    "Eleven": "11",
+    "Twelve": "12",
+    "Thirteen": "13",
+    "Fourteen": "14",
+    "Fifteen": "15",
+    "Sixteen": "16",
+    "Seventeen": "17",
+    "Eighteen": "18",
+    "Nineteen": "19",
+    "Twenty": "20",
+    "Twentyone": "21",
+    "Twentytwo": "22",
+    "Twentythree": "23",
+    "Twentyfour": "24",
+    "Twentyfive": "25",
+    "Twentysix": "26",
+    "Twentyseven": "27",
+    "Twentyeight": "28",
+    "Twentynine": "29",
+    "Thirty": "30",
+    "Thirtyone": "31",
+    "Thirtytwo": "32",
+    "Thirtythree": "33",
+    "Thirtyfour": "34",
+    "Thirtyfive": "35",
+    "Thirtysix": "36",
+    "Thirtyseven": "37",
+    "Thirtyeight": "38",
+    "Thirtynine": "39",
+    "Forty": "40",
+    "Fortyone": "41",
+    "Fortytwo": "42",
+    "Fortythree": "43",
+    "Fortyfour": "44",
+    "Fortyfive": "45",
+    "Fortysix": "46",
+    "Fortyseven": "47",
+    "Fortyeight": "48",
+    "Fortynine": "49",
+    "Fifty": "50",
+    "Fiftyone": "51",
+    "Fiftytwo": "52",
+    "Fiftythree": "53",
+    "Fiftyfour": "54",
+    "Fiftyfive": "55",
+    "Fiftysix": "56",
+    "Fiftyseven": "57",
+    "Fiftyeight": "58",
+    "Fiftynine": "59",
+    "Sixty": "60",
+    "Sixtyone": "61",
+    "Sixtytwo": "62",
+    "Sixtythree": "63",
+    "Sixtyfour": "64",
+    "Sixtyfive": "65",
+    "Sixtysix": "66",
+    "Sixtyseven": "67",
+    "Sixtyeight": "68",
+    "Sixtynine": "69",
+    "Seventy": "70",
+    "Seventyone": "71",
+    "Seventytwo": "72",
+    "Seventythree": "73",
+    "Seventyfour": "74",
+    "Seventyfive": "75",
+    "Seventysix": "76",
+    "Seventyseven": "77",
+    "Seventyeight": "78",
+    "Seventynine": "79",
+    "Eighty": "80",
+    "Eightyone": "81",
+    "Eightytwo": "82",
+    "Eightythree": "83",
+    "Eightyfour": "84",
+    "Eightyfive": "85",
+    "Eightysix": "86",
+    "Eightyseven": "87",
+    "Eightyeight": "88",
+    "Eightynine": "89",
+    "Ninety": "90",
+    "Ninetyone": "91",
+    "Ninetytwo": "92",
+    "Ninetythree": "93",
+    "Ninetyfour": "94",
+    "Ninetyfive": "95",
+    "Ninetysix": "96",
+    "Ninetyseven": "97",
+    "Ninetyeight": "98",
+    "Ninetynine": "99"
+}
diff --git a/internal/gorson/io/io.go b/internal/gorson/io/io.go
@@ -5,12 +5,15 @@ import (
 	"fmt"
 	"io/ioutil"
 	"log"
+	"math"
+	"math/rand"
 	"reflect"
 	"strings"
 	"time"
 
 	"github.com/pbs/gorson/internal/gorson/util"
 
+	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/ssm"
 )
@@ -65,7 +68,7 @@ func ReadFromParameterStore(path util.ParameterStorePath) map[string]string {
 	return values
 }
 
-func writeSingleParameter(c chan string, client *ssm.SSM, name string, value string) {
+func writeSingleParameter(c chan string, client *ssm.SSM, name string, value string, retryCount int) {
 	overwrite := true
 	valueType := "SecureString"
 	keyID := "alias/aws/ssm"
@@ -78,9 +81,26 @@ func writeSingleParameter(c chan string, client *ssm.SSM, name string, value str
 	}
 	_, err := client.PutParameter(&input)
 	if err != nil {
-		log.Fatal(err)
+		if awsErr, ok := err.(awserr.Error); ok {
+			if awsErr.Code() == "ThrottlingException" {
+				if retryCount < 100 {
+					// Introduce exponential backoff with jitter
+					r := math.Pow(2, float64(retryCount)) * (1 + rand.Float64())
+					time.Sleep(time.Duration(r) * time.Millisecond)
+					writeSingleParameter(c, client, name, value, retryCount+1)
+				} else {
+					fmt.Println("throttle retry limit reached for " + name)
+				}
+			} else {
+				log.Fatal(err)
+			}
+		} else {
+			log.Fatal(err)
+		}
+	} else {
+		fmt.Println("wrote " + name)
+		c <- name
 	}
-	c <- name
 }
 
 // WriteToParameterStore writes given parameters to a given parameter store path
@@ -92,8 +112,8 @@ func WriteToParameterStore(parameters map[string]string, path util.ParameterStor
 	for key, value := range parameters {
 		name := path.String() + key
 		// we pass the jobs channel into the asynchronous write function to receive
-		// success messages
-		go writeSingleParameter(jobs, client, name, value)
+		// success messages. When throttled, parameter writes wait, then retry.
+		go writeSingleParameter(jobs, client, name, value, 0)
 	}
 
 	// we keep track of the successful parameter store writes with results
@@ -111,12 +131,12 @@ func WriteToParameterStore(parameters map[string]string, path util.ParameterStor
 		}
 	}()
 
-	// we let two channels race: after 5 seconds, the channel from time.After wins,
+	// we let two channels race: after 1 minute, the channel from time.After wins,
 	// and we error out
 	select {
 	case <-done:
 		return
-	case <-time.After(5 * time.Second):
+	case <-time.After(1 * time.Minute):
 		log.Fatal("timeout")
 	}
 }