From 74161bf6859374a069eba6d9d5fb477c20c23cc2 Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Fri, 12 Apr 2024 08:41:14 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E4=BB=A3=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/nodes/http_writer.go      | 36 +++++++++++++-----------------
 internal/utils/buffer_pool_test.go | 26 +++++++++++++++++++++
 2 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/internal/nodes/http_writer.go b/internal/nodes/http_writer.go
index 16c57b1..700b81a 100644
--- a/internal/nodes/http_writer.go
+++ b/internal/nodes/http_writer.go
@@ -367,7 +367,7 @@ func (this *HTTPWriter) PrepareCache(resp *http.Response, size int64) {
 			if this.isPartial && k == "Content-Type" && strings.Contains(v1, "multipart/byteranges") {
 				continue
 			}
-			_, err = headerBuf.Write([]byte(k + ":" + v1 + "\n"))
+			_, err = headerBuf.WriteString(k + ":" + v1 + "\n")
 			if err != nil {
 				utils.SharedBufferPool.Put(headerBuf)
 
@@ -694,15 +694,15 @@ func (this *HTTPWriter) PrepareCompression(resp *http.Response, size int64) {
 		}
 
 		// 写入Header
-		var headerBuffer = utils.SharedBufferPool.Get()
+		var headerBuf = utils.SharedBufferPool.Get()
 		for k, v := range this.Header() {
 			if this.shouldIgnoreHeader(k) {
 				continue
 			}
 			for _, v1 := range v {
-				_, err = headerBuffer.Write([]byte(k + ":" + v1 + "\n"))
+				_, err = headerBuf.WriteString(k + ":" + v1 + "\n")
 				if err != nil {
-					utils.SharedBufferPool.Put(headerBuffer)
+					utils.SharedBufferPool.Put(headerBuf)
 					remotelogs.Error("HTTP_WRITER", "write compression cache failed: "+err.Error())
 					_ = compressionCacheWriter.Discard()
 					compressionCacheWriter = nil
@@ -710,28 +710,24 @@ func (this *HTTPWriter) PrepareCompression(resp *http.Response, size int64) {
 				}
 			}
 		}
-
-		_, err = compressionCacheWriter.WriteHeader(headerBuffer.Bytes())
-		utils.SharedBufferPool.Put(headerBuffer)
+		_, err = compressionCacheWriter.WriteHeader(headerBuf.Bytes())
+		utils.SharedBufferPool.Put(headerBuf)
 		if err != nil {
 			remotelogs.Error("HTTP_WRITER", "write compression cache failed: "+err.Error())
 			_ = compressionCacheWriter.Discard()
 			compressionCacheWriter = nil
 			return
 		}
-
-		if compressionCacheWriter != nil {
-			if this.compressionCacheWriter != nil {
-				_ = this.compressionCacheWriter.Close()
-			}
-			this.compressionCacheWriter = compressionCacheWriter
-			var teeWriter = writers.NewTeeWriterCloser(this.writer, compressionCacheWriter)
-			teeWriter.OnFail(func(err error) {
-				_ = compressionCacheWriter.Discard()
-				this.compressionCacheWriter = nil
-			})
-			this.writer = teeWriter
+		if this.compressionCacheWriter != nil {
+			_ = this.compressionCacheWriter.Close()
 		}
+		this.compressionCacheWriter = compressionCacheWriter
+		var teeWriter = writers.NewTeeWriterCloser(this.writer, compressionCacheWriter)
+		teeWriter.OnFail(func(err error) {
+			_ = compressionCacheWriter.Discard()
+			this.compressionCacheWriter = nil
+		})
+		this.writer = teeWriter
 	}
 
 	// compression writer
@@ -1294,6 +1290,6 @@ func (this *HTTPWriter) shouldIgnoreHeader(name string) bool {
 	case "Set-Cookie", "Strict-Transport-Security", "Alt-Svc", "Upgrade", "X-Cache":
 		return true
 	default:
-		return (this.isPartial && name == "Content-Range")
+		return this.isPartial && name == "Content-Range"
 	}
 }
diff --git a/internal/utils/buffer_pool_test.go b/internal/utils/buffer_pool_test.go
index 918424e..81a4627 100644
--- a/internal/utils/buffer_pool_test.go
+++ b/internal/utils/buffer_pool_test.go
@@ -45,3 +45,29 @@ func BenchmarkNewBufferPool2(b *testing.B) {
 		}
 	})
 }
+
+func BenchmarkNewBufferPool3(b *testing.B) {
+	var pool = utils.NewBufferPool()
+	var dataString = strings.Repeat("Hello", 1024)
+
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			var buffer = pool.Get()
+			buffer.Write([]byte(dataString))
+			pool.Put(buffer)
+		}
+	})
+}
+
+func BenchmarkNewBufferPool4(b *testing.B) {
+	var pool = utils.NewBufferPool()
+	var dataString = strings.Repeat("Hello", 1024)
+
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			var buffer = pool.Get()
+			buffer.WriteString(dataString)
+			pool.Put(buffer)
+		}
+	})
+}