Various improvements

Author: KillianMeersman

go.mod

@@ -1,6 +1,7 @@
 module github.com/KillianMeersman/chaperone
 
-go 1.22.0
+go 1.23.0
+
 toolchain go1.23.7
 
 require (

pkg/proxy/niceclient.go

@@ -127,6 +127,7 @@ func (c *NiceClient) RoundTripWithOptions(req *http.Request, options *RequestOpt
 	// Start span for retry loop
 	ctx, span := telemetry.Tracer.Start(ctx, "HTTP request retry loop", trace.WithAttributes(attribute.String("http.method", req.Method), attribute.String("http.url", req.URL.String())))
 	defer span.End()
+	req = req.WithContext(ctx)
 
 	for {
 		logger.Debug("waiting to make request")
@@ -139,7 +140,7 @@ func (c *NiceClient) RoundTripWithOptions(req *http.Request, options *RequestOpt
 		logger.Debug("Making request")
 		ctx, span := telemetry.Tracer.Start(ctx, req.Method, trace.WithSpanKind(trace.SpanKindClient), trace.WithAttributes(attribute.String("http.method", req.Method), attribute.String("http.url", req.URL.String())))
 		defer span.End()
-		req = req.WithContext(ctx)
+		req := req.WithContext(ctx)
 
 		res, err := c.roundtripper.RoundTrip(req)
 		switch err {
@@ -160,7 +161,7 @@ func (c *NiceClient) RoundTripWithOptions(req *http.Request, options *RequestOpt
 		}
 
 		span.SetAttributes(attribute.Int("http.status_code", res.StatusCode))
-		logger = logger.With("status_code", fmt.Sprint(res.StatusCode))
+		logger := logger.With("status_code", fmt.Sprint(res.StatusCode))
 		logger.Debug("got response", "status_code", fmt.Sprint(res.StatusCode))
 
 		switch res.StatusCode {
@@ -180,6 +181,8 @@ func (c *NiceClient) RoundTripWithOptions(req *http.Request, options *RequestOpt
 			jitterDuration := time.Duration(jitterMilliseconds) * time.Millisecond
 			blockDuration := time.Duration(waitTime + jitterDuration)
 			span.AddEvent("blocking throttle", trace.WithAttributes(attribute.Int("duration_seconds", int(blockDuration.Seconds()))))
+
+			logger.Warning("blocking throttle", "seconds", fmt.Sprint(blockDuration.Seconds()))
 			c.throttle.Block(res.Request, blockDuration)
 		case 301, 302, 307, 308:
 			// Handle redirects.

pkg/proxy/niceclient_test.go

@@ -10,11 +10,11 @@ import (
 	"time"
 )
 
-type mockRountTripper struct {
+type mockRoundTripper struct {
 	Response []byte
 }
 
-func (m *mockRountTripper) RoundTrip(r *http.Request) (*http.Response, error) {
+func (m *mockRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
 	bodyReader := bytes.NewReader(m.Response)
 	bodyReadCloser := io.NopCloser(bodyReader)
 
@@ -29,7 +29,7 @@ func TestNiceClient(t *testing.T) {
 	throttle := NewMemoryHTTPThrottle()
 	cache := NewMemoryHTTPCache(context.Background(), 1000)
 
-	roundTripper := &mockRountTripper{
+	roundTripper := &mockRoundTripper{
 		Response: []byte("test"),
 	}
 
@@ -47,12 +47,67 @@ func TestNiceClient(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+
 	throttle.Block(&http.Request{
 		URL:    url,
 		Method: "GET",
 	}, 3*time.Second)
+
 	client.RoundTrip(req)
-	if time.Since(start) < 4*time.Second {
+
+	if time.Since(start) < 3*time.Second {
 		t.FailNow()
 	}
 }
+
+func TestNiceClientThrottle(t *testing.T) {
+	cache := NewMemoryHTTPCache(context.Background(), 1000)
+	throttle := NewMemoryHTTPThrottle()
+
+	throttleDuration := 20 * time.Millisecond
+
+	url, err := url.Parse("http://example.com")
+	if err != nil {
+		panic(err)
+	}
+	throttle.SetThrottle(&http.Request{
+		Method: "GET",
+		URL:    url,
+	}, throttleDuration)
+	throttle.SetThrottle(&http.Request{
+		Method: "POST",
+		URL:    url,
+	}, throttleDuration*100)
+
+	url, err = url.Parse("http://example.com/test")
+	if err != nil {
+		panic(err)
+	}
+	throttle.SetThrottle(&http.Request{
+		Method: "GET",
+		URL:    url,
+	}, throttleDuration*100)
+
+	roundTripper := &mockRoundTripper{
+		Response: []byte("test"),
+	}
+
+	client := NewNiceClient(context.Background(), roundTripper, throttle, cache)
+
+	for i := 0; i < 100; i++ {
+		req, _ := http.NewRequest("GET", "http://example.com", nil)
+		start := time.Now()
+		client.RoundTrip(req)
+
+		timeTaken := time.Since(start)
+
+		if timeTaken < throttleDuration-5*time.Millisecond {
+			t.Fatal("not enough time since last request")
+		}
+
+		if timeTaken > throttleDuration+5*time.Millisecond {
+			t.Fatal("too much time since last request")
+		}
+	}
+
+}