From 47ed047f7e300581e68de252b854748f480b26a8 Mon Sep 17 00:00:00 2001
From: Arne Maier <Arne_Maier@gmx.de>
Date: Mon, 11 Nov 2024 21:43:12 +0100
Subject: [PATCH] Add throttle to context

---
 main.go                 | 29 ++++++++++++++---------------
 throttling/rateLimit.go |  3 ---
 2 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/main.go b/main.go
index e316da7..205d633 100644
--- a/main.go
+++ b/main.go
@@ -13,22 +13,26 @@ const burstLimit = 2
 
 func main() {
 	ctx := context.Background()
-	throttles := make(map[string]<-chan time.Time)
+	throttles := make(map[string]context.Context)
 	calledLastHour := make([]string, 0)
 
 	http.HandleFunc("/", serveIndexFile)
 	http.HandleFunc("/input.txt", serveInputFile)
 	http.HandleFunc("/result", func(w http.ResponseWriter, r *http.Request) {
 
-		ctx, cancel := context.WithCancel(ctx)
-		throttle := throttling.CreateThrottle(ctx, burstLimit)
-		a := throttleWithCancel{
-			throttle: throttle,
-			cancel:   cancel,
+		cctxV, ok := throttles[r.RemoteAddr]
+		if !ok {
+			cctx, cancel := context.WithCancel(ctx)
+			throttle := throttling.CreateThrottle(cctx, burstLimit)
+			a := throttleWithCancel{
+				throttle: throttle,
+				cancel:   cancel,
+			}
+			cctxV = context.WithValue(ctx, "throttle", a)
+			throttles[r.RemoteAddr] = cctxV
 		}
-		con := context.WithValue(ctx, "throttle", a)
 
-		tryResult(ctx, w, r, throttles, calledLastHour)
+		tryResult(cctxV, w, r, calledLastHour)
 	})
 
 	err := http.ListenAndServe(":3333", nil)
@@ -51,19 +55,14 @@ func serveInputFile(w http.ResponseWriter, r *http.Request) {
 	http.ServeFile(w, r, "static/input.txt")
 }
 
-func tryResult(ctx context.Context, w http.ResponseWriter, r *http.Request, throttles map[string]<-chan time.Time, calledLastHour []string) {
+func tryResult(ctx context.Context, w http.ResponseWriter, r *http.Request, calledLastHour []string) {
 	clientIP := r.RemoteAddr
 	clientResult := r.URL.Query().Get("result")
 	fmt.Println(clientIP, clientResult)
 	if !slices.Contains(calledLastHour, clientIP) {
 		calledLastHour = append(calledLastHour, clientIP)
 	}
-	throttle, ok := throttles[clientIP]
-	if !ok {
-		fmt.Println("Creating new throttle")
-		throttle = throttling.CreateThrottle(ctx, burstLimit)
-		throttles[clientIP] = throttle
-	}
+	throttle := ctx.Value("throttle").(throttleWithCancel).throttle
 	payload := throttling.Payload{
 		R:            r,
 		W:            w,
diff --git a/throttling/rateLimit.go b/throttling/rateLimit.go
index 3e9c610..2ba2b43 100644
--- a/throttling/rateLimit.go
+++ b/throttling/rateLimit.go
@@ -24,9 +24,6 @@ type Payload struct {
 // CallFunction allows burst rate limiting client calls with the
 // payloads.
 func CallFunction(ctx context.Context, client Client, payload *Payload, throttle <-chan time.Time) {
-	//ctx, cancel := context.WithCancel(ctx)
-	// defer cancel()
-
 	<-throttle // rate limit our client calls
 	client.Call(payload)
 }