Improve error handling (#12)

* Add recovery for panics in goroutines
* Add debug logging support
* Add retry on provisioning and throttling errors while writing to dynamodb

Author: Hugo Correia

copier.go

@@ -1,6 +1,9 @@
 package dynamodbcopy
 
-import "sync"
+import (
+	"fmt"
+	"sync"
+)
 
 type Copier interface {
 	Copy(readers, writers int) error
@@ -10,17 +13,20 @@ type copyService struct {
 	srcTable DynamoDBService
 	trgTable DynamoDBService
 	chans    CopierChans
+	logger   Logger
 }
 
-func NewCopier(srcTableService, trgTableService DynamoDBService, chans CopierChans) Copier {
+func NewCopier(srcTableService, trgTableService DynamoDBService, chans CopierChans, logger Logger) Copier {
 	return copyService{
 		srcTable: srcTableService,
 		trgTable: trgTableService,
 		chans:    chans,
+		logger:   logger,
 	}
 }
 
 func (service copyService) Copy(readers, writers int) error {
+	service.logger.Printf("copying table with %d readers and %d writers", readers, writers)
 	itemsChan, errChan := service.chans.Items, service.chans.Errors
 
 	wgReaders := &sync.WaitGroup{}
@@ -54,7 +60,12 @@ func (service copyService) read(
 	itemsChan chan<- []DynamoDBItem,
 	errChan chan<- error,
 ) {
-	defer wg.Done()
+	defer func() {
+		if err := recover(); err != nil {
+			errChan <- fmt.Errorf("read recovery: %s", err)
+		}
+		wg.Done()
+	}()
 
 	err := service.srcTable.Scan(totalReaders, readerID, itemsChan)
 	if err != nil {
@@ -63,13 +74,23 @@ func (service copyService) read(
 }
 
 func (service copyService) write(wg *sync.WaitGroup, itemsChan <-chan []DynamoDBItem, errChan chan<- error) {
-	defer wg.Done()
+	defer func() {
+		if err := recover(); err != nil {
+			errChan <- fmt.Errorf("write recovery: %s", err)
+		}
+		wg.Done()
+	}()
 
+	totalWritten := 0
 	for items := range itemsChan {
 		if err := service.trgTable.BatchWrite(items); err != nil {
 			errChan <- err
 		}
+
+		totalWritten += len(items)
 	}
+
+	service.logger.Printf("writer wrote a total of %d items", totalWritten)
 }
 
 // CopierChans encapsulates the chan that are used by the copier

copier_test.go

@@ -2,12 +2,15 @@ package dynamodbcopy_test
 
 import (
 	"errors"
+	"io/ioutil"
+	"log"
 	"strconv"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/dynamodb"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"github.com/uniplaces/dynamodbcopy"
 	"github.com/uniplaces/dynamodbcopy/mocks"
 )
@@ -87,6 +90,34 @@ func TestCopy(t *testing.T) {
 			3,
 			nil,
 		},
+		{
+			"ReadPanic",
+			func(src, trg *mocks.DynamoDBService, chans *dynamodbcopy.CopierChans) {
+				var readChan chan<- []dynamodbcopy.DynamoDBItem = chans.Items
+				src.On("Scan", 1, 0, readChan).Run(func(args mock.Arguments) {
+					panic("read panic")
+				}).Once()
+			},
+			1,
+			1,
+			errors.New("read recovery: read panic"),
+		},
+		{
+			"WritePanic",
+			func(src, trg *mocks.DynamoDBService, chans *dynamodbcopy.CopierChans) {
+				var readChan chan<- []dynamodbcopy.DynamoDBItem = chans.Items
+				src.On("Scan", 1, 0, readChan).Return(nil).Once()
+
+				items := buildItems(1)
+				chans.Items <- items
+				trg.On("BatchWrite", items).Run(func(args mock.Arguments) {
+					panic("write panic")
+				}).Once()
+			},
+			1,
+			1,
+			errors.New("write recovery: write panic"),
+		},
 	}
 
 	for _, testCase := range testCases {
@@ -100,7 +131,7 @@ func TestCopy(t *testing.T) {
 
 				testCase.mocker(src, trg, &copierChans)
 
-				service := dynamodbcopy.NewCopier(src, trg, copierChans)
+				service := dynamodbcopy.NewCopier(src, trg, copierChans, log.New(ioutil.Discard, "", log.Ltime))
 
 				err := service.Copy(testCase.totalReaders, testCase.totalWriters)
 

dynamodb.go

@@ -3,15 +3,22 @@ package dynamodbcopy
 import (
 	"errors"
 	"fmt"
+	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/dynamodb"
 	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
 )
 
-const maxBatchWriteSize = 25
+const (
+	maxBatchWriteSize = 25
+	maxRetryTime      = int(time.Minute) * 3
+
+	errCodeThrottlingException = "ThrottlingException"
+)
 
 // DynamoDBAPI just a wrapper over aws-sdk dynamodbiface.DynamoDBAPI interface for mocking purposes
 type DynamoDBAPI interface {
@@ -32,6 +39,7 @@ type dynamoDBSerivce struct {
 	tableName string
 	api       DynamoDBAPI
 	sleep     Sleeper
+	logger    Logger
 }
 
 func NewDynamoDBAPI(roleArn string) DynamoDBAPI {
@@ -49,8 +57,8 @@ func NewDynamoDBAPI(roleArn string) DynamoDBAPI {
 	return dynamodb.New(currentSession)
 }
 
-func NewDynamoDBService(tableName string, api DynamoDBAPI, sleepFn Sleeper) DynamoDBService {
-	return dynamoDBSerivce{tableName, api, sleepFn}
+func NewDynamoDBService(tableName string, api DynamoDBAPI, sleepFn Sleeper, logger Logger) DynamoDBService {
+	return dynamoDBSerivce{tableName, api, sleepFn, logger}
 }
 
 func (db dynamoDBSerivce) DescribeTable() (*dynamodb.TableDescription, error) {
@@ -86,6 +94,7 @@ func (db dynamoDBSerivce) UpdateCapacity(capacity Capacity) error {
 		},
 	}
 
+	db.logger.Printf("updating %s with read: %d, write: %d", db.tableName, read, write)
 	_, err := db.api.UpdateTable(input)
 	if err != nil {
 		return fmt.Errorf("unable to update table %s: %s", db.tableName, err)
@@ -95,6 +104,7 @@ func (db dynamoDBSerivce) UpdateCapacity(capacity Capacity) error {
 }
 
 func (db dynamoDBSerivce) BatchWrite(items []DynamoDBItem) error {
+	db.logger.Printf("writing batch of %d to %s", len(items), db.tableName)
 	if len(items) == 0 {
 		return nil
 	}
@@ -125,40 +135,76 @@ func (db dynamoDBSerivce) batchWriteItem(requests []*dynamodb.WriteRequest) erro
 
 	writeRequests := requests
 	for len(writeRequests) != 0 {
-		batchInput := &dynamodb.BatchWriteItemInput{
-			RequestItems: map[string][]*dynamodb.WriteRequest{
-				tableName: writeRequests,
-			},
-		}
+		retryHandler := func(attempt, elapsed int) (bool, error) {
+			batchInput := &dynamodb.BatchWriteItemInput{
+				RequestItems: map[string][]*dynamodb.WriteRequest{
+					tableName: writeRequests,
+				},
+			}
 
-		output, err := db.api.BatchWriteItem(batchInput)
-		if err != nil {
-			return fmt.Errorf("unable to batch write to table %s: %s", db.tableName, err)
+			output, err := db.api.BatchWriteItem(batchInput)
+			if err == nil {
+				writeRequests = output.UnprocessedItems[tableName]
+
+				return true, nil
+			}
+
+			if awsErr, ok := err.(awserr.Error); ok {
+				switch awsErr.Code() {
+				case dynamodb.ErrCodeProvisionedThroughputExceededException:
+					db.logger.Printf("batch write provisioning error: waited %d ms (attempt %d)", elapsed, attempt)
+					return false, nil
+				case errCodeThrottlingException:
+					db.logger.Printf("batch write throttling error: waited %d ms (attempt %d)", elapsed, attempt)
+					return false, nil
+				default:
+					return false, fmt.Errorf(
+						"aws %s error in batch write to table %s: %s",
+						awsErr.Code(),
+						db.tableName,
+						awsErr.Error(),
+					)
+				}
+			}
+
+			return false, fmt.Errorf("unable to batch write to table %s: %s", db.tableName, err)
 		}
 
-		writeRequests = output.UnprocessedItems[tableName]
+		if err := db.retry(retryHandler); err != nil {
+			return err
+		}
 	}
 
 	return nil
 }
 
 func (db dynamoDBSerivce) WaitForReadyTable() error {
-	elapsed := 0
-
-	for attempt := 0; ; attempt++ {
+	return db.retry(func(attempt, elapsed int) (bool, error) {
 		description, err := db.DescribeTable()
+		if err != nil {
+			return false, err
+		}
+
+		return *description.TableStatus == dynamodb.TableStatusActive, nil
+	})
+}
+
+func (db dynamoDBSerivce) retry(handler func(attempt, elapsed int) (bool, error)) error {
+	elapsed := 0
+	for attempt := 0; elapsed < maxRetryTime; attempt++ {
+		handled, err := handler(attempt, elapsed)
 		if err != nil {
 			return err
 		}
 
-		if *description.TableStatus == dynamodb.TableStatusActive {
-			break
+		if handled {
+			return nil
 		}
 
 		elapsed += db.sleep(elapsed * attempt)
 	}
 
-	return nil
+	return fmt.Errorf("waited for too long (%d ms) to perform operation on %s table", elapsed, db.tableName)
 }
 
 func (db dynamoDBSerivce) Scan(totalSegments, segment int, itemsChan chan<- []DynamoDBItem) error {
@@ -175,11 +221,14 @@ func (db dynamoDBSerivce) Scan(totalSegments, segment int, itemsChan chan<- []Dy
 		input.SetTotalSegments(int64(totalSegments))
 	}
 
+	totalScanned := 0
 	pagerFn := func(output *dynamodb.ScanOutput, b bool) bool {
 		var items []DynamoDBItem
 		for _, item := range output.Items {
 			items = append(items, item)
+			totalScanned++
 		}
+		db.logger.Printf("%s table scanned page with %d items (reader %d)", db.tableName, len(items), segment)
 
 		itemsChan <- items
 
@@ -190,5 +239,7 @@ func (db dynamoDBSerivce) Scan(totalSegments, segment int, itemsChan chan<- []Dy
 		return fmt.Errorf("unable to scan table %s: %s", db.tableName, err)
 	}
 
+	db.logger.Printf("%s table scanned a total of %d items (reader %d)", db.tableName, totalScanned, segment)
+
 	return nil
 }