From 9d79c9693b436041cc5cd1b50b73bd68d8fc565d Mon Sep 17 00:00:00 2001
From: perennial <mail@perennialte.ch>
Date: Sat, 5 Oct 2024 21:16:22 +1000
Subject: [PATCH] Implement worker pool for proxy requests

This commit introduces worker pooling for handling media proxy requests,
allowing for actual concurrency in this area.

Key changes:
- Add ProxyJob struct to represent worker pool jobs
- Create a channel for ProxyJob and initialise workers in init function
- Modify ProxyRequest to use the worker pool
- Add proxyWorker and processProxyJob functions for job processing

The number of workers and job queue size are currently hardcoded but can
be made configurable in the future.
---
 core/requests.go | 75 +++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 62 insertions(+), 13 deletions(-)

diff --git a/core/requests.go b/core/requests.go
index 53edd6f..d6d7cef 100644
--- a/core/requests.go
+++ b/core/requests.go
@@ -26,12 +26,28 @@ type SimpleHTTPResponse struct {
 
 var retryClient *retryablehttp.Client
 
+// ProxyJob represents a job for the proxy worker pool
+type ProxyJob struct {
+	Request  *http.Request
+	Response http.ResponseWriter
+	Done     chan error
+}
+
+var proxyJobQueue chan ProxyJob
+var numWorkers = 10 // TODO: allow number of workers to be configured
+
 func init() {
 	retryClient = retryablehttp.NewClient()
 	retryClient.RetryMax = config.GlobalConfig.APIMaxRetries
 	retryClient.RetryWaitMin = config.GlobalConfig.APIBaseTimeout
 	retryClient.RetryWaitMax = config.GlobalConfig.APIMaxBackoffTime
 	retryClient.HTTPClient = utils.HttpClient
+
+	// Initialize the proxy worker pool
+	proxyJobQueue = make(chan ProxyJob, 100) // TODO: allow buffer size to be configured
+	for i := 0; i < numWorkers; i++ {
+		go proxyWorker(proxyJobQueue)
+	}
 }
 
 // retryRequest performs a request with automatic retries and token management
@@ -67,17 +83,16 @@ func retryRequest(ctx context.Context, reqFunc func(context.Context, string) (*r
 			return nil, err
 		}
 
-
 		start := time.Now()
 		resp, err := retryClient.Do(req)
 		end := time.Now()
 
-                // Unwrap the body here so that we could log stuff correctly
-                defer resp.Body.Close()
-                body, err := io.ReadAll(resp.Body)
-                if err != nil {
-                        return nil, err
-                }
+		// Unwrap the body here so that we could log stuff correctly
+		defer resp.Body.Close()
+		body, err := io.ReadAll(resp.Body)
+		if err != nil {
+			return nil, err
+		}
 
 		audit.LogAPIRoundTrip(audit.APIRequestSpan{
 			RequestId: request_context.GetFromContext(ctx).RequestId,
@@ -193,18 +208,52 @@ func API_POST(ctx context.Context, url, payload, userToken, csrf string, isJSON
 }
 
 // ProxyRequest forwards an HTTP request to the target server and copies the response back
-func ProxyRequest(w http.ResponseWriter, req *http.Request) error {
-	resp, err := utils.HttpClient.Do(req)
+func ProxyRequest(w http.ResponseWriter, req *http.Request) {
+	done := make(chan error, 1)
+	job := ProxyJob{
+		Request:  req,
+		Response: w,
+		Done:     done,
+	}
+
+	proxyJobQueue <- job
+
+	err := <-done
 	if err != nil {
-		return err
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+	}
+}
+
+// proxyWorker processes jobs from the proxyJobQueue
+func proxyWorker(jobs <-chan ProxyJob) {
+	for job := range jobs {
+		err := processProxyJob(job)
+		job.Done <- err
+	}
+}
+
+// processProxyJob handles the actual proxying of the request
+func processProxyJob(job ProxyJob) error {
+	resp, err := utils.HttpClient.Do(job.Request)
+	if err != nil {
+		return fmt.Errorf("failed to process request: %w", err)
 	}
 	defer resp.Body.Close()
 
-	header := w.Header()
+	// Copy response headers
+	header := job.Response.Header()
 	for k, v := range resp.Header {
 		header[k] = v
 	}
 
-	_, err = io.Copy(w, resp.Body)
-	return err
+	// Set the status code
+	job.Response.WriteHeader(resp.StatusCode)
+
+	// Copy the body from the response to the original writer
+	_, err = io.Copy(job.Response, resp.Body)
+	if err != nil {
+		return fmt.Errorf("failed to copy response body: %w", err)
+	}
+
+	return nil
 }