making changes discussed in PR tsenart#534

Dockerfile

@@ -9,20 +9,12 @@ ADD go.mod /vegeta
 ADD go.sum /vegeta
 RUN go mod download
 
-#now build source code
+# now build source code
 ADD / /vegeta
 
 RUN go test -v ./...
 RUN make vegeta
-# RUN go build -v -o /bin/vegeta
 
 FROM alpine:3.12.0
 
-ENV TARGET_URL ''
-ENV DURATION '5'
-ENV REQUESTS_PER_SECOND '5'
-
 COPY --from=BUILD /vegeta/vegeta /bin/vegeta
-ADD startup.sh /
-CMD [ "/startup.sh" ]
-

README.md

@@ -96,12 +96,10 @@ attack command:
  Output file (default "stdout")
  -prom-bind
  Host to bind Prometheus service to. Defaults to 0.0.0.0
- -prom-enable
- Enable Prometheus metrics endpoint so that requests can be monitored by Prometheus. Defaults to false.
- -prom-path
- Prometheus metrics path. Defaults to /metrics
- -prom-port
- HTTP port for exposing Prometheus metrics at /metrics. Defaults to http://[host]:8880/metrics
+ -prometheus-enable
+ Enable Prometheus metrics endpoint so that requests can be monitored by Prometheus with default bind (0.0.0.0:8880). Defaults to false.
+ -prometheus-url
+ Prometheus metrics bind. Defaults to 0.0.0.0:8880
  -proxy-header value
  Proxy CONNECT header
  -rate value
@@ -735,7 +733,7 @@ It'll read and sort them by timestamp before generating reports.
 vegeta report *.bin
 ```
 
-Another way to gather results in distributed tests is to use the built-in Prometheus Exporter and configure a Prometheus Server to get test results from all Vegeta instances. See `attack` option "prom-enable" for more details and a complete example at section "Prometheus Exporter Support"
+Another way to gather results in distributed tests is to use the built-in Prometheus Exporter and configure a Prometheus Server to get test results from all Vegeta instances. See `attack` option "prom-enable" for more details and a complete example in the section "Prometheus Exporter Support"
 
 ## Usage: Real-time Analysis
 
@@ -838,13 +836,10 @@ The following metrics are exposed:
 version: '3.5'
 services:
  vegeta:
- build: .
+ image: tsenart/vegeta
  ports:
  - 8880:8880
- environment:
- - TARGET_URL=http://www.google.com
- - DURATION=300s
- - REQUESTS_PER_SECOND=1
+ command: echo "GET https://www.yahoo.com" | vegeta attack -duration=30s -rate=5 -prom-enable=true > /dev/null 2>&1
 
  prometheus:
  image: flaviostutz/prometheus:2.19.2.0
@@ -858,7 +853,7 @@ services:
 
 * Run `docker-compose up -d`
 
-* Run `curl localhost:8880/metrics` to see plain Prometheus Exporter endpoint contents
+* Run `curl localhost:8880` to see plain Prometheus Exporter endpoint contents
 
 * Open Prometheus server instance with your browser at http://localhost:9090
 

attack.go

@@ -29,9 +29,7 @@ func attackCmd() command {
  rate: vegeta.Rate{Freq: 50, Per: time.Second},
  maxBody: vegeta.DefaultMaxBody,
  promEnable: false,
- promBind: "0.0.0.0",
- promPort: 8880,
- promPath: "/metrics",
+ promURL: "0.0.0.0:8880",
  }
  fs.StringVar(&opts.name, "name", "", "Attack name")
  fs.StringVar(&opts.targetsf, "targets", "stdin", "Targets file")
@@ -61,10 +59,8 @@ func attackCmd() command {
  fs.Var(&opts.laddr, "laddr", "Local IP address")
  fs.BoolVar(&opts.keepalive, "keepalive", true, "Use persistent connections")
  fs.StringVar(&opts.unixSocket, "unix-socket", "", "Connect over a unix socket. This overrides the host address in target URLs")
- fs.BoolVar(&opts.promEnable, "prom-enable", false, "Enable Prometheus metrics endpoint so that requests can be monitored by Prometheus. Defaults to false.")
- fs.StringVar(&opts.promBind, "prom-bind", "0.0.0.0", "Host to bind Prometheus service to. Defaults to 0.0.0.0")
- fs.IntVar(&opts.promPort, "prom-port", 8880, "HTTP port for exposing Prometheus metrics at /metrics. Defaults to http://[host]:8880/metrics")
- fs.StringVar(&opts.promPath, "prom-path", "/metrics", "Prometheus metrics path. Defaults to /metrics")
+ fs.BoolVar(&opts.promEnable, "prometheus-enable", false, "Enable Prometheus metrics endpoint so that requests can be monitored by Prometheus using default parameters(bind to 0.0.0.0:8880). Defaults to false.")
+ fs.StringVar(&opts.promURL, "prometheus-url", "", "Enable Prometheus metrics with specific bind parameters in format [bind ip]:[bind port]. Example: 0.0.0.0:8880")
  systemSpecificFlags(fs, opts)
 
  return command{fs, func(args []string) error {
@@ -108,10 +104,8 @@ type attackOpts struct {
  keepalive bool
  resolvers csl
  unixSocket string
- promBind string
- promPort int
  promEnable bool
- promPath string
+ promURL  string
 }
 
 // attack validates the attack arguments, sets up the
@@ -185,9 +179,13 @@ func attack(opts *attackOpts) (err error) {
  return err
  }
 
- var promMetrics prom.PrometheusMetrics
- if opts.promEnable {
- promMetrics, err = prom.NewPrometheusMetricsWithParams(opts.promBind, opts.promPort, opts.promPath)
+ var promMetrics *prom.PrometheusMetrics
+ if opts.promEnable || opts.promURL != "" {
+ purl := opts.promURL
+ if purl == "" {
+ purl = "0.0.0.0:8880"
+ }
+ promMetrics, err = prom.NewPrometheusMetricsWithParams(purl)
  if err != nil {
  return err
  }

docker-compose.yml

@@ -7,10 +7,7 @@ services:
  image: flaviostutz/vegeta
  ports:
  - 8880:8880
- environment: 
- - TARGET_URL=https://www.yahoo.com
- - DURATION=60s
- - REQUESTS_PER_SECOND=5
+ command: echo "GET https://www.yahoo.com" | vegeta attack -duration=30s -rate=5 -prom-enable=true > /dev/null 2>&1
 
  prometheus:
  image: flaviostutz/prometheus

lib/prom/prom.go

@@ -1,12 +1,13 @@
 package prom
 
 import (
+ "context"
  "fmt"
- "net"
  "net/http"
+ "regexp"
+ "strconv"
  "time"
 
- "github.com/gorilla/mux"
  "github.com/prometheus/client_golang/prometheus"
  "github.com/prometheus/client_golang/prometheus/promauto"
  "github.com/prometheus/client_golang/prometheus/promhttp"
@@ -19,13 +20,13 @@ type PrometheusMetrics struct {
  requestBytesInCounter *prometheus.CounterVec
  requestBytesOutCounter *prometheus.CounterVec
  requestFailCounter *prometheus.CounterVec
- listenPort net.Listener
+ srv http.Server
+ registry *prometheus.Registry
 }
 
 //NewPrometheusMetrics same as NewPrometheusMetricsWithParams with default params:
-//bindHost=0.0.0.0, bindPort=8880 and metricsPath=/metrics
-func NewPrometheusMetrics() (PrometheusMetrics, error) {
- return NewPrometheusMetricsWithParams("0.0.0.0", 8880, "/metrics")
+func NewPrometheusMetrics() (*PrometheusMetrics, error) {
+ return NewPrometheusMetricsWithParams("0.0.0.0:8880")
 }
 
 // NewPrometheusMetricsWithParams start a new Prometheus observer instance for exposing
@@ -34,17 +35,36 @@ func NewPrometheusMetrics() (PrometheusMetrics, error) {
 // mechanisms of promauto.
 // Some metrics are requests/s, bytes in/out/s and failures/s
 // Options are:
-// - bindHost: host to bind the listening socket to
-// - bindPort: port to bind the listening socket to
-// - metricsPath: http path that will be used to get metrics
-// For example, after using NewPrometheusMetricsWithParams("0.0.0.0", 8880, "/metrics"),
-// during an "attack" you can call "curl http://127.0.0.0:8880/metrics" to see current metrics.
+// - bindURL: "[host]:[port]/[path]" to bind the listening socket to
+// For example, after using NewPrometheusMetricsWithParams("0.0.0.0:8880"),
+// during an "attack" you can call "curl http://127.0.0.0:8880" to see current metrics.
 // This endpoint can be configured in scrapper section of your Prometheus server.
-func NewPrometheusMetricsWithParams(bindHost string, bindPort int, metricsPath string) (PrometheusMetrics, error) {
+func NewPrometheusMetricsWithParams(bindURL string) (*PrometheusMetrics, error) {
+
+ //parse bind url elements
+ re := regexp.MustCompile("(.+):([0-9]+)")
+ rr := re.FindAllStringSubmatch(bindURL, 3)
+ bindHost := ""
+ bindPort := 0
+ var err error
+ if len(rr) == 1 {
+ if len(rr[0]) == 3 {
+ bindHost = rr[0][1]
+ bindPort, err = strconv.Atoi(rr[0][2])
+ if err != nil {
+ return nil, err
+ }
+ }
+ }
+ if bindHost == "" {
+ return nil, fmt.Errorf("Invalid bindURL %s. Must be in format '0.0.0.0:8880'", bindURL)
+ }
 
- pm := PrometheusMetrics{}
+ pm := &PrometheusMetrics{
+ registry: prometheus.NewRegistry(),
+ }
 
- pm.requestSecondsHistogram = promauto.NewHistogramVec(prometheus.HistogramOpts{
+ pm.requestSecondsHistogram = prometheus.NewHistogramVec(prometheus.HistogramOpts{
  Name: "request_seconds",
  Help: "Request latency",
  Buckets: []float64{0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0, 20, 50},
@@ -53,6 +73,7 @@ func NewPrometheusMetricsWithParams(bindHost string, bindPort int, metricsPath s
  "url",
  "status",
  })
+ pm.registry.MustRegister(pm.requestSecondsHistogram)
 
  pm.requestBytesInCounter = promauto.NewCounterVec(prometheus.CounterOpts{
  Name: "request_bytes_in",
@@ -62,6 +83,7 @@ func NewPrometheusMetricsWithParams(bindHost string, bindPort int, metricsPath s
  "url",
  "status",
  })
+ pm.registry.MustRegister(pm.requestBytesInCounter)
 
  pm.requestBytesOutCounter = promauto.NewCounterVec(prometheus.CounterOpts{
  Name: "request_bytes_out",
@@ -71,6 +93,7 @@ func NewPrometheusMetricsWithParams(bindHost string, bindPort int, metricsPath s
  "url",
  "status",
  })
+ pm.registry.MustRegister(pm.requestBytesOutCounter)
 
  pm.requestFailCounter = promauto.NewCounterVec(prometheus.CounterOpts{
  Name: "request_fail_count",
@@ -80,21 +103,16 @@ func NewPrometheusMetricsWithParams(bindHost string, bindPort int, metricsPath s
  "url",
  "message",
  })
+ pm.registry.MustRegister(pm.requestFailCounter)
 
  //setup prometheus metrics http server
- router := mux.NewRouter()
- router.Handle(metricsPath, promhttp.Handler())
-
- listen := fmt.Sprintf("%s:%d", bindHost, bindPort)
- lp, err := net.Listen("tcp", listen)
- if err != nil {
- return PrometheusMetrics{}, err
+ pm.srv = http.Server{
+ Addr: fmt.Sprintf("%s:%d", bindHost, bindPort),
+ Handler: promhttp.HandlerFor(pm.registry, promhttp.HandlerOpts{}),
  }
- pm.listenPort = lp
 
  go func() {
- http.Serve(pm.listenPort, router)
- defer pm.listenPort.Close()
+ pm.srv.ListenAndServe()
  }()
 
  return pm, nil
@@ -107,14 +125,15 @@ func (pm *PrometheusMetrics) Close() error {
  prometheus.Unregister(pm.requestBytesInCounter)
  prometheus.Unregister(pm.requestBytesOutCounter)
  prometheus.Unregister(pm.requestFailCounter)
- return pm.listenPort.Close()
+ return pm.srv.Shutdown(context.Background())
 }
 
 //Observe register metrics about hit results
 func (pm *PrometheusMetrics) Observe(res *vegeta.Result) {
- pm.requestBytesInCounter.WithLabelValues(res.Method, res.URL, fmt.Sprintf("%d", res.Code)).Add(float64(res.BytesIn))
- pm.requestBytesOutCounter.WithLabelValues(res.Method, res.URL, fmt.Sprintf("%d", res.Code)).Add(float64(res.BytesOut))
- pm.requestSecondsHistogram.WithLabelValues(res.Method, res.URL, fmt.Sprintf("%d", res.Code)).Observe(float64(res.Latency) / float64(time.Second))
+ code := strconv.FormatUint(uint64(res.Code), 10)
+ pm.requestBytesInCounter.WithLabelValues(res.Method, res.URL, code).Add(float64(res.BytesIn))
+ pm.requestBytesOutCounter.WithLabelValues(res.Method, res.URL, code).Add(float64(res.BytesOut))
+ pm.requestSecondsHistogram.WithLabelValues(res.Method, res.URL, code).Observe(float64(res.Latency) / float64(time.Second))
  if res.Error != "" {
  pm.requestFailCounter.WithLabelValues(res.Method, res.URL, res.Error)
  }

lib/prom/prom_test.go

@@ -52,7 +52,8 @@ func TestPromServerObserve(t *testing.T) {
  pm.Observe(r)
  pm.Observe(r)
 
- resp, err := http.Get("http://localhost:8880/metrics")
+ time.Sleep(1 * time.Second)
+ resp, err := http.Get("http://localhost:8880")
  assert.Nil(t, err, "Error calling prometheus metrics. err=%s", err)
  assert.Equal(t, 200, resp.StatusCode, "Status code should be 200")
 
@@ -69,7 +70,7 @@ func TestPromServerObserve(t *testing.T) {
  r.Error = "REQUEST FAILED"
  pm.Observe(r)
 
- resp, err = http.Get("http://localhost:8880/metrics")
+ resp, err = http.Get("http://localhost:8880")
  assert.Nil(t, err, "Error calling prometheus metrics. err=%s", err)
  assert.Equal(t, 200, resp.StatusCode, "Status code should be 200")