benchmark client enhancements

opencloud/pkg/command/benchmark.go

@@ -2,6 +2,7 @@ package command
 
 import (
 	"bytes"
+	"context"
 	"crypto/tls"
 	"encoding/base64"
 	"errors"
@@ -10,8 +11,11 @@ import (
 	"net/http"
 	"os"
 	"os/exec"
+	"os/signal"
 	"strconv"
 	"strings"
+	"sync"
+	"syscall"
 	"time"
 
 	"github.com/olekukonko/tablewriter"
@@ -64,8 +68,15 @@ func BenchmarkClientCommand(cfg *config.Config) *cli.Command {
 			&cli.StringFlag{
 				Name:    "data",
 				Aliases: []string{"d"},
-				Usage:   "Sends the specified data in a request to the HTTP server.",
-				// TODE support multiple data flags, support data-binary, data-raw
+				Usage:   "Sends the specified data in a POST request to the HTTP server, in the same way that a browser does when a user has filled in an HTML form and presses the submit button. If you start the data with the letter @, the rest should be a file name to read the data from, or - if you want to read the data from stdin. When -d, --data is told to read from a file like that, carriage returns and newlines are stripped out. If you do not want the @ character to have a special interpretation use --data-raw instead.",
+			},
+			&cli.StringFlag{
+				Name:  "data-binary",
+				Usage: "This posts data exactly as specified with no extra processing whatsoever. If you start the data with the letter @, the rest should be a file name to read the data from, or - if you want to read the data from stdin.",
+			},
+			&cli.StringFlag{
+				Name:  "data-raw",
+				Usage: "Sends the specified data in a request to the HTTP server.",
 			},
 			&cli.StringSliceFlag{
 				Name:    "header",
@@ -107,14 +118,97 @@ func BenchmarkClientCommand(cfg *config.Config) *cli.Command {
 		},
 		Category: "benchmark",
 		Action: func(c *cli.Context) error {
+			// Set up signal handling for Ctrl+C
+			ctx, cancel := context.WithCancel(c.Context)
+			defer cancel()
+
+			sigChan := make(chan os.Signal, 1)
+			signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
+			go func() {
+				<-sigChan
+				fmt.Println("\nReceived interrupt signal, shutting down...")
+				cancel()
+			}()
+
 			opt := clientOptions{
-				request:  c.String("request"),
 				url:      c.Args().First(),
 				insecure: c.Bool("insecure"),
 				jobs:     c.Int("jobs"),
 				headers:  make(map[string]string),
-				data:     []byte(c.String("data")),
 			}
+
+			if d := c.String("data-raw"); d != "" {
+				opt.request = "POST"
+				opt.headers["Content-Type"] = "application/x-www-form-urlencoded"
+				opt.data = []byte(d)
+			}
+
+			if d := c.String("data"); d != "" {
+				opt.request = "POST"
+				opt.headers["Content-Type"] = "application/x-www-form-urlencoded"
+				if strings.HasPrefix(d, "@") {
+					filePath := strings.TrimPrefix(d, "@")
+					var data []byte
+					var err error
+
+					// read from file or stdin and trim trailing newlines
+					if filePath == "-" {
+						data, err = os.ReadFile("/dev/stdin")
+					} else {
+						data, err = os.ReadFile(filePath)
+					}
+					if err != nil {
+						log.Fatal(errors.New("could not read data from file '" + filePath + "': " + err.Error()))
+					}
+
+					// clean byte array similar to curl's --data parameter
+					// It removes leading/trailing whitespace and converts line breaks to spaces
+
+					// Trim leading and trailing whitespace
+					data = bytes.TrimSpace(data)
+
+					// Replace newlines and carriage returns with spaces
+					data = bytes.ReplaceAll(data, []byte("\r\n"), []byte(" "))
+					data = bytes.ReplaceAll(data, []byte("\n"), []byte(" "))
+					data = bytes.ReplaceAll(data, []byte("\r"), []byte(" "))
+
+					// Replace multiple spaces with single space
+					for bytes.Contains(data, []byte("  ")) {
+						data = bytes.ReplaceAll(data, []byte("  "), []byte(" "))
+					}
+
+					opt.data = data
+				} else {
+					opt.data = []byte(d)
+				}
+			}
+
+			if d := c.String("data-binary"); d != "" {
+				opt.request = "POST"
+				opt.headers["Content-Type"] = "application/x-www-form-urlencoded"
+				if strings.HasPrefix(d, "@") {
+					filePath := strings.TrimPrefix(d, "@")
+					var data []byte
+					var err error
+					if filePath == "-" {
+						data, err = os.ReadFile("/dev/stdin")
+					} else {
+						data, err = os.ReadFile(filePath)
+					}
+					if err != nil {
+						log.Fatal(errors.New("could not read data from file '" + filePath + "': " + err.Error()))
+					}
+					opt.data = data
+				} else {
+					opt.data = []byte(d)
+				}
+			}
+
+			// override method if specified
+			if request := c.String("request"); request != "" {
+				opt.request = request
+			}
+
 			if opt.url == "" {
 				log.Fatal(errors.New("no URL specified"))
 			}
@@ -179,7 +273,7 @@ func BenchmarkClientCommand(cfg *config.Config) *cli.Command {
 				defer opt.ticker.Stop()
 			}
 
-			return client(opt)
+			return client(ctx, opt)
 
 		},
 	}
@@ -197,16 +291,19 @@ type clientOptions struct {
 	jobs      int
 }
 
-func client(o clientOptions) error {
-
+func client(ctx context.Context, o clientOptions) error {
 	type stat struct {
 		job      int
 		duration time.Duration
 		status   int
 	}
 	stats := make(chan stat)
+	var wg sync.WaitGroup
+
 	for i := 0; i < o.jobs; i++ {
+		wg.Add(1)
 		go func(i int) {
+			defer wg.Done()
 			tr := &http.Transport{
 				TLSClientConfig: &tls.Config{
 					MinVersion:         tls.VersionTLS12,
@@ -217,6 +314,13 @@ func client(o clientOptions) error {
 
 			cookies := map[string]*http.Cookie{}
 			for {
+				// Check if context is cancelled
+				select {
+				case <-ctx.Done():
+					return
+				default:
+				}
+
 				req, err := http.NewRequest(o.request, o.url, bytes.NewReader(o.data))
 				if err != nil {
 					log.Printf("client %d: could not create request: %s\n", i, err)
@@ -234,20 +338,35 @@ func client(o clientOptions) error {
 				res, err := client.Do(req)
 				duration := -time.Until(start)
 				if err != nil {
+					// Check if error is due to context cancellation
+					if ctx.Err() != nil {
+						return
+					}
 					log.Printf("client %d: could not create request: %s\n", i, err)
 					time.Sleep(time.Second)
 				} else {
 					res.Body.Close()
-					stats <- stat{
+					select {
+					case stats <- stat{
 						job:      i,
 						duration: duration,
 						status:   res.StatusCode,
+					}:
+					case <-ctx.Done():
+						return
 					}
 					for _, c := range res.Cookies() {
 						cookies[c.Name] = c
 					}
 				}
-				time.Sleep(o.rateDelay - duration)
+				// Sleep with context awareness
+				if o.rateDelay > duration {
+					select {
+					case <-time.After(o.rateDelay - duration):
+					case <-ctx.Done():
+						return
+					}
+				}
 			}
 		}(i)
 	}
@@ -256,9 +375,15 @@ func client(o clientOptions) error {
 	if o.ticker == nil {
 		// no ticker, just write every request
 		for {
-			stat := <-stats
-			numRequests++
-			fmt.Printf("req %d took %v and returned status %d\n", numRequests, stat.duration, stat.status)
+			select {
+			case stat := <-stats:
+				numRequests++
+				fmt.Printf("req %d took %v and returned status %d\n", numRequests, stat.duration, stat.status)
+			case <-ctx.Done():
+				fmt.Println("\nShutting down...")
+				wg.Wait()
+				return nil
+			}
 		}
 	}
 
@@ -274,6 +399,13 @@ func client(o clientOptions) error {
 				numRequests = 0
 				duration = 0
 			}
+		case <-ctx.Done():
+			if numRequests > 0 {
+				fmt.Printf("\n%d req at %v/req\n", numRequests, duration/time.Duration(numRequests))
+			}
+			fmt.Println("Shutting down...")
+			wg.Wait()
+			return nil
 		}
 	}