From b042f8e8fad1e1b4835da1c95ee8daf6ab3c0b0e Mon Sep 17 00:00:00 2001
From: lesovsky <lesovsky@gmail.com>
Date: Wed, 8 Jul 2020 13:50:27 +0500
Subject: [PATCH] Add '--duration' flag, determines duration of the test in
 seconds.

---
 cmd/app.go  | 7 ++++++-
 cmd/main.go | 2 ++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/cmd/app.go b/cmd/app.go
index 3258e4c..e30e654 100644
--- a/cmd/app.go
+++ b/cmd/app.go
@@ -13,6 +13,7 @@ import (
 	"github.com/lesovsky/noisia/waitxacts"
 	"github.com/rs/zerolog"
 	"sync"
+	"time"
 )
 
 type config struct {
@@ -20,6 +21,7 @@ type config struct {
 	doCleanup             bool
 	postgresConninfo      string
 	jobs                  uint16 // max 65535
+	duration              int
 	idleXacts             bool
 	idleXactsNaptimeMin   int
 	idleXactsNaptimeMax   int
@@ -41,12 +43,15 @@ type config struct {
 }
 
 func runApplication(ctx context.Context, c *config, log zerolog.Logger) error {
-
 	if c.doCleanup {
 		log.Info().Msg("do cleanup")
 		return noisia.Cleanup(ctx, c.postgresConninfo)
 	}
 
+	timeout := time.Duration(c.duration) * time.Second
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
 	var wg sync.WaitGroup
 
 	if c.idleXacts {
diff --git a/cmd/main.go b/cmd/main.go
index 794280f..20e7356 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -21,6 +21,7 @@ func main() {
 		doCleanup             = kingpin.Flag("cleanup", "do cleanup of workloads related tables from database").Default("false").Envar("NOISIA_CLEANUP").Bool()
 		postgresConninfo      = kingpin.Flag("conninfo", "Postgres connection string (DSN or URL), must be specified explicitly").Default("").Envar("NOISIA_POSTGRES_CONNINFO").String()
 		jobs                  = kingpin.Flag("jobs", "Run workload with specified number of workers").Default("1").Envar("NOISIA_JOBS").Uint16()
+		duration              = kingpin.Flag("duration", "Duration of tests in seconds").Default("10").Envar("NOISIA_DURATION").Int()
 		idleXacts             = kingpin.Flag("idle-xacts", "Run idle transactions workload").Default("false").Envar("NOISIA_IDLE_XACTS").Bool()
 		idleXactsNaptimeMin   = kingpin.Flag("idle-xacts.naptime-min", "Min transactions naptime, in seconds").Default("5").Envar("NOISIA_IDLE_XACTS_NAPTIME_MIN").Int()
 		idleXactsNaptimeMax   = kingpin.Flag("idle-xacts.naptime-max", "Max transactions naptime, in seconds").Default("20").Envar("NOISIA_IDLE_XACTS_NAPTIME_MAX").Int()
@@ -54,6 +55,7 @@ func main() {
 		doCleanup:             *doCleanup,
 		postgresConninfo:      *postgresConninfo,
 		jobs:                  *jobs,
+		duration:              *duration,
 		idleXacts:             *idleXacts,
 		idleXactsNaptimeMin:   *idleXactsNaptimeMin,
 		idleXactsNaptimeMax:   *idleXactsNaptimeMax,