From e7fb706d8ae4d2adfe263bc175c080d852ad6b89 Mon Sep 17 00:00:00 2001
From: GoEdgeLab <goedgecloud@gmail.com>
Date: Wed, 27 Dec 2023 20:55:12 +0800
Subject: [PATCH] =?UTF-8?q?=E5=86=99=E7=BC=93=E5=AD=98=E5=A4=B1=E8=B4=A5?=
 =?UTF-8?q?=E6=97=B6=EF=BC=8C=E5=85=81=E8=AE=B8=E7=BB=A7=E7=BB=AD=E8=AF=BB?=
 =?UTF-8?q?=E5=8F=96=E6=BA=90=E7=AB=99=E5=86=85=E5=AE=B9?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/caches/writer_memory.go            |  2 --
 internal/nodes/http_writer.go               |  2 +-
 internal/utils/readers/reader_closer_tee.go | 19 +++++++++++++++----
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/internal/caches/writer_memory.go b/internal/caches/writer_memory.go
index 66c165a..4e7092d 100644
--- a/internal/caches/writer_memory.go
+++ b/internal/caches/writer_memory.go
@@ -125,7 +125,6 @@ func (this *MemoryWriter) Close() error {
 	// 需要在Locker之外
 	defer this.once.Do(func() {
 		this.endFunc(this.item)
-		this.item = nil // free memory
 	})
 
 	if this.item == nil {
@@ -164,7 +163,6 @@ func (this *MemoryWriter) Discard() error {
 	// 需要在Locker之外
 	defer this.once.Do(func() {
 		this.endFunc(this.item)
-		this.item = nil // free memory
 	})
 
 	this.storage.locker.Lock()
diff --git a/internal/nodes/http_writer.go b/internal/nodes/http_writer.go
index bdd5e7d..510810c 100644
--- a/internal/nodes/http_writer.go
+++ b/internal/nodes/http_writer.go
@@ -494,7 +494,7 @@ func (this *HTTPWriter) PrepareCache(resp *http.Response, size int64) {
 		return
 	}
 
-	var cacheReader = readers.NewTeeReaderCloser(resp.Body, this.cacheWriter)
+	var cacheReader = readers.NewTeeReaderCloser(resp.Body, this.cacheWriter, false)
 	resp.Body = cacheReader
 	this.rawReader = cacheReader
 
diff --git a/internal/utils/readers/reader_closer_tee.go b/internal/utils/readers/reader_closer_tee.go
index b57146e..4be6cc2 100644
--- a/internal/utils/readers/reader_closer_tee.go
+++ b/internal/utils/readers/reader_closer_tee.go
@@ -12,12 +12,17 @@ type TeeReaderCloser struct {
 
 	onFail func(err error)
 	onEOF  func()
+
+	mustWrite bool
 }
 
-func NewTeeReaderCloser(reader io.Reader, writer io.Writer) *TeeReaderCloser {
+// NewTeeReaderCloser
+// mustWrite - ensure writing MUST be successfully
+func NewTeeReaderCloser(reader io.Reader, writer io.Writer, mustWrite bool) *TeeReaderCloser {
 	return &TeeReaderCloser{
-		r: reader,
-		w: writer,
+		r:         reader,
+		w:         writer,
+		mustWrite: mustWrite,
 	}
 }
 
@@ -26,7 +31,13 @@ func (this *TeeReaderCloser) Read(p []byte) (n int, err error) {
 	if n > 0 {
 		_, wErr := this.w.Write(p[:n])
 		if (err == nil || err == io.EOF) && wErr != nil {
-			err = wErr
+			if this.mustWrite {
+				err = wErr
+			} else {
+				if this.onFail != nil {
+					this.onFail(wErr)
+				}
+			}
 		}
 	}
 	if err != nil {