Invoking AWS SDK using WaitGroup of configurable size

main.go

// Run N workers in a WaitGroup until they complete

package main

import (
	"flag"
	"fmt"
	"os"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/dynamodb"
)

func getSession() *session.Session {
	sess := session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))

	return sess
}

func describeTable(tableName string) error {
	dynamoDBClient := dynamodb.New(getSession())
	response, err := dynamoDBClient.DescribeTable(&dynamodb.DescribeTableInput{
		TableName: aws.String(tableName),
	})

	if err != nil {
		return err
	}

	fmt.Println(response)
	return nil
}

func worker(id int, wg *sync.WaitGroup, tableName string) {

	defer wg.Done()

	fmt.Printf("Worker %d starting\n", id)

	describeTable(tableName)
	fmt.Printf("Worker %d done\n", id)
}

func main() {

	var numWorkers int
	var tableName string
	flag.IntVar(&numWorkers, "w", 4, "Number of workers")
	flag.StringVar(&tableName, "t", "", "Table name")
	flag.Parse()

	if len(tableName) == 0 {
		fmt.Printf("You must specify a table name")
		os.Exit(1)
	}

	var wg sync.WaitGroup

	for i := 1; i <= numWorkers; i++ {
		wg.Add(1)
		go worker(i, &wg, tableName)
	}

	wg.Wait()
}

main2.go

// Run N WaitGroup workers in a loop until duration elapses

package main

import (
	"flag"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/dynamodb"
)

func getSession() *session.Session {
	sess := session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))

	return sess
}

func describeTable(tableName string) error {
	dynamoDBClient := dynamodb.New(getSession())
	response, err := dynamoDBClient.DescribeTable(&dynamodb.DescribeTableInput{
		TableName: aws.String(tableName),
	})

	if err != nil {
		return err
	}

	fmt.Println(response)
	return nil
}

func worker(id int, wg *sync.WaitGroup, tableName string) {

	defer wg.Done()

	fmt.Printf("Worker %d starting\n", id)

	for {
		describeTable(tableName)
	}
	// fmt.Printf("Worker %d done\n", id)
}

// waitTimeout waits for the waitgroup for the specified max timeout.
// Returns true if waiting timed out.
func waitTimeout(wg *sync.WaitGroup, timeout time.Duration) bool {
	c := make(chan struct{})
	go func() {
		defer close(c)
		wg.Wait()
	}()
	select {
	case <-c:
		return false // completed normally
	case <-time.After(timeout):
		return true // timed out
	}
}

func main() {

	var numWorkers int
	var tableName string
	var duration int
	flag.IntVar(&numWorkers, "w", 4, "Number of workers")
	flag.StringVar(&tableName, "t", "", "Table name")
	flag.IntVar(&duration, "d", 60, "Duration in seconds")
	flag.Parse()

	if len(tableName) == 0 {
		fmt.Printf("You must specify a table name")
		os.Exit(1)
	}

	var wg sync.WaitGroup

	for i := 1; i <= numWorkers; i++ {
		wg.Add(1)
		go worker(i, &wg, tableName)
	}

	timeout := time.Duration(duration) * time.Second
	fmt.Printf("Wait for waitgroup (up to %s)\n", timeout)

	if waitTimeout(&wg, timeout) {
		fmt.Printf("Timed out waiting for wait group after %s\n", timeout)
	} else {
		fmt.Println("Wait group finished")
	}
	fmt.Println("Done")
}