From cb576fb1789b5aaabbc8849e03503c98451a1f6d Mon Sep 17 00:00:00 2001
From: Benson Wong <mostlygeek@gmail.com>
Date: Thu, 3 Oct 2024 20:33:55 -0700
Subject: [PATCH] replace io.Copy to improve performance sending data to client

---
 llama-proxy.go | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/llama-proxy.go b/llama-proxy.go
index b855f7a..b3e79de 100644
--- a/llama-proxy.go
+++ b/llama-proxy.go
@@ -181,7 +181,27 @@ func proxyRequest(w http.ResponseWriter, r *http.Request, config *Config, state
 	}
 	w.WriteHeader(resp.StatusCode)
 
-	io.Copy(w, resp.Body)
+	buf := make([]byte, 32*1024) // Buffer size set to 32KB
+	for {
+		n, err := resp.Body.Read(buf)
+		if n > 0 {
+			if _, writeErr := w.Write(buf[:n]); writeErr != nil {
+				http.Error(w, writeErr.Error(), http.StatusInternalServerError)
+				return
+			}
+			// Flush the buffer to the client
+			if flusher, ok := w.(http.Flusher); ok {
+				flusher.Flush()
+			}
+		}
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusBadGateway)
+			return
+		}
+	}
 }
 
 func main() {