Merge pull request #144 from yandex/arcadia

Http URI style ammo

cli/cli.go

@@ -25,7 +25,7 @@ import (
 	"github.com/yandex/pandora/lib/zaputil"
 )
 
-const Version = "0.3.4"
+const Version = "0.3.5"
 const defaultConfigFile = "load"
 const stdinConfigSelector = "-"
 

components/phttp/ammo/simple/provider.go

@@ -23,6 +23,7 @@ func NewProvider(fs afero.Fs, fileName string, start func(ctx context.Context, f
 		start:    start,
 		Sink:     make(chan *Ammo, 128),
 		Pool:     sync.Pool{New: func() interface{} { return &Ammo{} }},
+		Close:    func() {},
 	}
 }
 
@@ -33,6 +34,7 @@ type Provider struct {
 	Sink      chan *Ammo
 	Pool      sync.Pool
 	idCounter atomic.Int64
+	Close     func()
 	core.ProviderDeps
 }
 
@@ -49,6 +51,7 @@ func (p *Provider) Release(a core.Ammo) {
 }
 
 func (p *Provider) Run(ctx context.Context, deps core.ProviderDeps) error {
+	defer p.Close()
 	p.ProviderDeps = deps
 	defer close(p.Sink)
 	file, err := p.fs.Open(p.fileName)

components/phttp/ammo/simple/uri/provider.go

@@ -8,6 +8,7 @@ package uri
 import (
 	"bufio"
 	"context"
+	"fmt"
 
 	"github.com/pkg/errors"
 	"github.com/spf13/afero"
@@ -17,22 +18,50 @@ import (
 )
 
 type Config struct {
-	File string `validate:"required"`
+	File string
 	// Limit limits total num of ammo. Unlimited if zero.
 	Limit int `validate:"min=0"`
 	// Redefine HTTP headers
 	Headers []string
 	// Passes limits ammo file passes. Unlimited if zero.
 	Passes int `validate:"min=0"`
+	Uris   []string
 }
 
 // TODO: pass logger and metricsRegistry
 func NewProvider(fs afero.Fs, conf Config) *Provider {
+	if len(conf.Uris) > 0 {
+		if conf.File != "" {
+			panic(`One should specify either 'file' or 'uris', but not both of them.`)
+		}
+		file, err := afero.TempFile(fs, "", "generated_ammo_")
+		if err != nil {
+			panic(fmt.Sprintf(`failed to create tmp ammo file: %v`, err))
+		}
+		for _, uri := range conf.Uris {
+			_, err := file.WriteString(fmt.Sprintf("%s\n", uri))
+			if err != nil {
+				panic(fmt.Sprintf(`failed to write ammo in tmp file: %v`, err))
+			}
+		}
+		conf.File = file.Name()
+	}
+	if conf.File == "" {
+		panic(`One should specify either 'file' or 'uris'.`)
+	}
 	var p Provider
 	p = Provider{
 		Provider: simple.NewProvider(fs, conf.File, p.start),
 		Config:   conf,
 	}
+	p.Close = func() {
+		if len(conf.Uris) > 0 {
+			err := fs.Remove(conf.File)
+			if err != nil {
+				zap.L().Error("failed to delete temp file", zap.String("file name", conf.File))
+			}
+		}
+	}
 	return &p
 }