main.go

package main

import (
  "fmt"
  "log"
  "math/rand"
  "os"
  "strconv"
  "strings"
  "sync"
  "time"

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

var endpoint = "http://localhost:9000"
var bucket = aws.String("mybucket")
var accessKey = "admin"
var secretKey = "changeme"
var prefix = "loadtest/"
var objectCount = 1000
var oneshot = true

func init() {
  rand.Seed(time.Now().UnixNano())
}

func main() {
  var err error

  if os.Getenv("LOADTEST_ENDPOINT") != "" {
    endpoint = os.Getenv("LOADTEST_ENDPOINT")
    if err != nil {
      log.Fatalf("Couldn't get LOADTEST_ENDPOINT: %v\n", err)
    }
  }

  if os.Getenv("LOADTEST_BUCKET") != "" {
    bucket = aws.String(os.Getenv("LOADTEST_BUCKET"))
    if err != nil {
      log.Fatalf("Couldn't get LOADTEST_BUCKET: %v\n", err)
    }
  }

  if os.Getenv("LOADTEST_ACCESS_KEY") != "" {
    accessKey = os.Getenv("LOADTEST_ACCESS_KEY")
    if err != nil {
      log.Fatalf("Couldn't get LOADTEST_ACCESS_KEY: %v\n", err)
    }
  }

  if os.Getenv("LOADTEST_SECRET_KEY") != "" {
	    secretKey = os.Getenv("LOADTEST_SECRET_KEY")
    if err != nil {
      log.Fatalf("Couldn't get LOADTEST_SECRET_KEY: %v\n", err)
    }
  }

  if os.Getenv("LOADTEST_OBJ_COUNT") != "" {
    objectCount, err = strconv.Atoi(os.Getenv("LOADTEST_OBJ_COUNT"))
    if err != nil {
      log.Fatalf("Couldn't parse LOADTEST_OBJ_COUNT: %v\n", err)
    }
  }

  if os.Getenv("LOADTEST_ONESHOT") != "" {
    oneshot, err = strconv.ParseBool(os.Getenv("LOADTEST_ONESHOT"))
    if err != nil {
      log.Fatalf("Couldn't parse LOADTEST_ONESHOT: %v\n", err)
    }
  }

  s3Config := &aws.Config{
    Credentials:      credentials.NewStaticCredentials(accessKey, secretKey, ""),
    Endpoint:         aws.String(endpoint),
    Region:           aws.String("us-east-1"),
    DisableSSL:       aws.Bool(true),
    S3ForcePathStyle: aws.Bool(true),
  }
  newSession := session.New(s3Config)

  s3Client := s3.New(newSession)
  for {
    var wg sync.WaitGroup
    log.Printf("Starting object uploads\n")
    for i := 0; i < objectCount; i++ {
      wg.Add(1)
      go uploadObject(&wg, s3Client)
    }
    wg.Wait()
    log.Printf("Uploaded %d objects\n", objectCount)

    if oneshot {
      break
    }

    nextTime := time.Now().Truncate(time.Minute).Add(time.Minute)
    time.Sleep(time.Until(nextTime))
  }
}

func uploadObject(wg *sync.WaitGroup, s3Client *s3.S3) {
  defer wg.Done()

  var err error
  key := aws.String(prefix + RandStringRunes(30))

  // Upload a new object "testobject" with the string "Hello World!" to our "newbucket".
  randStr := RandStringRunes(rand.Intn(1024))
  start := time.Now()
  _, err = s3Client.PutObject(&s3.PutObjectInput{
    Body:   strings.NewReader(randStr),
    Bucket: bucket,
    Key:    key,
  })
  end := time.Now()
  uploadDuration := end.Sub(start)
  if err != nil {
    fmt.Printf("Failed to upload data to %s/%s, %s\n", *bucket, *key, err.Error())
    return
  }
  log.Printf("Successfully uploaded %d bytes with key %s in %v\n", len(randStr), *key, uploadDuration)
}

// https://stackoverflow.com/questions/22892120/how-to-generate-a-random-string-of-a-fixed-length-in-go/35399339
var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")

func RandStringRunes(n int) string {
  b := make([]rune, n)
  for i := range b {
    b[i] = letterRunes[rand.Intn(len(letterRunes))]
  }
  return string(b)
}