func countTempBytes(conninfo string) (int, error) {
	bytes := -1  
	conn, err := db.Connect(context.Background(), conninfo)
	if err != nil {
		return bytes, err
	}
	defer func() { _ = conn.Close() }()
	rows, err := conn.Query(context.Background(), "SELECT pg_stat_get_db_temp_bytes(oid) from pg_database where datname = current_database()")
	if err != nil {
		return bytes, err
	}
	for rows.Next() {
		err = rows.Scan(&bytes)
		if err != nil {
			return bytes, err
		}
	}
	return bytes, nil
}
countTempBytes queries current database statistics about temp bytes written.
Private context is used here, because this is auxiliary routine and is not related to
main workload.
package tempfiles
import (
	"context"
	"fmt"
	"github.com/lesovsky/noisia"
	"github.com/lesovsky/noisia/db"
	"github.com/lesovsky/noisia/log"
	"golang.org/x/time/rate"
	"sync"
)
type Config struct {
	Conninfo string
	Jobs uint16
	Rate float64
}
func (c Config) validate() error {
	if c.Jobs < 1 {
		return fmt.Errorf("jobs must be greater than zero")
	}
	if c.Rate <= 0 {
		return fmt.Errorf("temp files queries rate must be positive")
	}
	return nil
}
type workload struct {
	config Config
	logger log.Logger
	pool   db.DB
}
func NewWorkload(config Config, logger log.Logger) (noisia.Workload, error) {
	err := config.validate()
	if err != nil {
		return nil, err
	}
	return &workload{config, logger, nil}, nil
}
func (w *workload) Run(ctx context.Context) error {
	workers := int(w.config.Jobs)
	var wg sync.WaitGroup
	bytesBefore, err := countTempBytes(w.config.Conninfo)
	if err != nil {
		return err
	}
	wg.Add(workers)
	for i := 0; i < workers; i++ {
		go func() {
			err := runWorker(ctx, w.logger, w.config)
			if err != nil {
				w.logger.Warnf("start tempfiles worker failed: %s, continue", err)
			}
			wg.Done()
		}()
	}
	wg.Wait()
	bytesAfter, err := countTempBytes(w.config.Conninfo)
	if err != nil {
		return err
	}
	w.logger.Infof("generated %d temp bytes (might include temp bytes produced by concurrent workload)", bytesAfter-bytesBefore)
	return nil
}
func runWorker(ctx context.Context, log log.Logger, config Config) error {
	log.Info("start tempfiles worker")
	pool, err := db.NewPostgresDB(ctx, config.Conninfo)
	if err != nil {
		return err
	}
	defer pool.Close()
	err = startLoop(ctx, pool, log, config.Rate)
	if err != nil {
		return err
	}
	log.Infof("tempfiles worker finished")
	return nil
}
func startLoop(ctx context.Context, pool db.DB, log log.Logger, r float64) error {
	var wg sync.WaitGroup
	limiter := rate.NewLimiter(rate.Limit(r), 1)
	for {
		if limiter.Allow() {
			wg.Add(1)
			go func() {
				err := execQuery(ctx, pool)
				if err != nil && ctx.Err() == nil {
					log.Warnf("executing tempfiles query failed: %v, continue", err)
				}
				wg.Done()
			}()
		}
		select {
		case <-ctx.Done():
			wg.Wait()
			return nil
		default:
		}
	}
}
func execQuery(ctx context.Context, pool db.DB) error {
	_, _, err := pool.Exec(ctx, "SET work_mem TO '64kB'")
	if err != nil {
		return err
	}
	_, _, err = pool.Exec(ctx, "SELECT * FROM pg_class a, pg_class b ORDER BY random()")
	if err != nil {
		return err
	}
	return nil
}