From 0a0b537f00403004adc79523a56238d80fe1acab Mon Sep 17 00:00:00 2001
From: Flynn <flynn@buoyant.io>
Date: Thu, 18 Apr 2024 16:45:50 -0400
Subject: [PATCH] Support a simple load generation service

Signed-off-by: Flynn <flynn@buoyant.io>
---
 cmd/generic/main.go | 64 ++++++++++++++++++++++++++++++++++++++++-----
 go.mod              |  2 +-
 go.sum              |  4 +--
 3 files changed, 60 insertions(+), 10 deletions(-)

diff --git a/cmd/generic/main.go b/cmd/generic/main.go
index 3514fa5..4890740 100644
--- a/cmd/generic/main.go
+++ b/cmd/generic/main.go
@@ -18,8 +18,12 @@
 package main
 
 import (
+	"io"
 	"log"
+	"net/http"
 	"os"
+	"strconv"
+	"time"
 
 	"flag"
 	"fmt"
@@ -51,17 +55,63 @@ func main() {
 	case "color":
 		server = &faces.NewColorServer("ColorServer").BaseServer
 
+	case "load":
+		fmt.Printf("Running load generator")
+
 	default:
 		log.Fatalf("Unknown service: %s", service)
 	}
 
-	fmt.Printf("%s %s (generic), commit %s, built at %s\n", server.Name, version, commit, date)
+	if server != nil {
+		fmt.Printf("%s %s (generic), commit %s, built at %s\n", server.Name, version, commit, date)
+
+		// Define a command-line flag for the port number
+		port := flag.Int("port", 8000, "the port number to listen on")
+		flag.Parse()
+
+		// Use the port number from the command line flag
+		addr := fmt.Sprintf(":%d", *port)
+		server.ListenAndServe(addr)
+	} else {
+		target := os.Getenv("LOAD_TARGET")
+		rps := os.Getenv("LOAD_RPS")
+		debug, _ := strconv.ParseBool(os.Getenv("LOAD_DEBUG"))
+
+		// Convert rps to an integer
+		rpsInt, err := strconv.Atoi(rps)
+		if err != nil {
+			log.Fatalf("Failed to convert rps to an integer: %v", err)
+		}
+
+		// Create a ticker to control the rate of requests
+		ticker := time.NewTicker(time.Second / time.Duration(rpsInt))
+		defer ticker.Stop()
+		count := 0
 
-	// Define a command-line flag for the port number
-	port := flag.Int("port", 8000, "the port number to listen on")
-	flag.Parse()
+		// Start a goroutine to send requests
+		for range ticker.C {
+			go func() {
+				// Make a GET request to http://target/
+				resp, err := http.Get(fmt.Sprintf("http://%s/", target))
+				if err != nil {
+					log.Fatalf("Failed to make request: %v", err)
+				}
+				defer resp.Body.Close()
 
-	// Use the port number from the command line flag
-	addr := fmt.Sprintf(":%d", *port)
-	server.ListenAndServe(addr)
+				// Read the response body
+				body, _ := io.ReadAll(resp.Body)
+
+				if debug {
+					fmt.Printf("http://%s/ %d %s\n", target, resp.StatusCode, string(body))
+				}
+
+				count++
+
+				if count >= (rpsInt * 10) {
+					fmt.Printf("Sent %d requests\n", count)
+					count = 0
+				}
+			}()
+		}
+	}
 }
diff --git a/go.mod b/go.mod
index 83c4491..d9e6d43 100644
--- a/go.mod
+++ b/go.mod
@@ -4,4 +4,4 @@ go 1.20
 
 require github.com/warthog618/gpiod v0.8.2
 
-require golang.org/x/sys v0.10.0 // indirect
+require golang.org/x/sys v0.17.0 // indirect
diff --git a/go.sum b/go.sum
index a5d353f..0fa8bce 100644
--- a/go.sum
+++ b/go.sum
@@ -5,6 +5,6 @@ github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ
 github.com/warthog618/go-gpiosim v0.1.0 h1:2rTMTcKUVZxpUuvRKsagnKAbKpd3Bwffp87xywEDVGI=
 github.com/warthog618/gpiod v0.8.2 h1:2HgQ9pNowPp7W77sXhX5ut5Tqq1WoS3t7bXYDxtYvxc=
 github.com/warthog618/gpiod v0.8.2/go.mod h1:O7BNpHjCn/4YS5yFVmoFZAlY1LuYuQ8vhPf0iy/qdi4=
-golang.org/x/sys v0.10.0 h1:SqMFp9UcQJZa+pmYuAKjd9xq1f0j5rLcDIk0mj4qAsA=
-golang.org/x/sys v0.10.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.17.0 h1:25cE3gD+tdBA7lp7QfhuV+rJiE9YXTcS3VG1SqssI/Y=
+golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=