优化缓存管理

internal/nodes/api_stream.go

@@ -159,26 +159,38 @@ func (this *APIStream) handleWriteCache(message *pb.NodeStreamMessage) error {
 	}
 
 	expiredAt := time.Now().Unix() + msg.LifeSeconds
-	writer, err := storage.Open(msg.Key, expiredAt)
+	writer, err := storage.OpenWriter(msg.Key, expiredAt, 200)
 	if err != nil {
 		this.replyFail(message.RequestId, "prepare writing failed: "+err.Error())
 		return err
 	}
 
-	_, err = writer.Write(msg.Value)
+	// 写入一个空的Header
+	_, err = writer.WriteHeader([]byte(":"))
 	if err != nil {
-		_ = writer.Discard()
 		this.replyFail(message.RequestId, "write failed: "+err.Error())
 		return err
 	}
-	err = writer.Close()
-	if err == nil {
-		storage.AddToList(&caches.Item{
-			Key:       msg.Key,
-			ExpiredAt: expiredAt,
-		})
+
+	// 写入数据
+	_, err = writer.Write(msg.Value)
+	if err != nil {
+		this.replyFail(message.RequestId, "write failed: "+err.Error())
+		return err
 	}
 
+	err = writer.Close()
+	if err != nil {
+		this.replyFail(message.RequestId, "write failed: "+err.Error())
+		return err
+	}
+	storage.AddToList(&caches.Item{
+		Key:        msg.Key,
+		ExpiredAt:  expiredAt,
+		HeaderSize: writer.HeaderSize(),
+		BodySize:   writer.BodySize(),
+	})
+
 	this.replyOk(message.RequestId, "write ok")
 
 	return nil
@@ -203,21 +215,20 @@ func (this *APIStream) handleReadCache(message *pb.NodeStreamMessage) error {
 		}()
 	}
 
-	buf := make([]byte, 1024)
-	size := 0
-	err = storage.Read(msg.Key, buf, func(data []byte, valueSize int64, expiredAt int64, isEOF bool) {
-		size += len(data)
-	})
+	reader, err := storage.OpenReader(msg.Key)
 	if err != nil {
 		if err == caches.ErrNotFound {
 			this.replyFail(message.RequestId, "key not found")
 			return nil
 		}
 		this.replyFail(message.RequestId, "read key failed: "+err.Error())
-		return err
+		return nil
 	}
+	defer func() {
+		_ = reader.Close()
+	}()
 
-	this.replyOk(message.RequestId, "value "+strconv.Itoa(size)+" bytes")
+	this.replyOk(message.RequestId, "value "+strconv.FormatInt(reader.BodySize(), 10)+" bytes")
 
 	return nil
 }
@@ -373,6 +384,13 @@ func (this *APIStream) handlePreheatCache(message *pb.NodeStreamMessage) error {
 				return
 			}
 
+			if resp.StatusCode != 200 {
+				locker.Lock()
+				errorMessages = append(errorMessages, "request failed: "+key+": status code '"+strconv.Itoa(resp.StatusCode)+"'")
+				locker.Unlock()
+				return
+			}
+
 			defer func() {
 				_ = resp.Body.Close()
 			}()
@@ -388,7 +406,7 @@ func (this *APIStream) handlePreheatCache(message *pb.NodeStreamMessage) error {
 			}
 
 			expiredAt := time.Now().Unix() + 8600
-			writer, err := storage.Open(key, expiredAt) // TODO 可以设置缓存过期时间
+			writer, err := storage.OpenWriter(key, expiredAt, 200) // TODO 可以设置缓存过期时间
 			if err != nil {
 				locker.Lock()
 				errorMessages = append(errorMessages, "open cache writer failed: "+key+": "+err.Error())
@@ -398,6 +416,21 @@ func (this *APIStream) handlePreheatCache(message *pb.NodeStreamMessage) error {
 
 			buf := make([]byte, 16*1024)
 			isClosed := false
+
+			// 写入Header
+			for k, v := range resp.Header {
+				for _, v1 := range v {
+					_, err = writer.WriteHeader([]byte(k + ":" + v1 + "\n"))
+					if err != nil {
+						locker.Lock()
+						errorMessages = append(errorMessages, "write failed: "+key+": "+err.Error())
+						locker.Unlock()
+						return
+					}
+				}
+			}
+
+			// 写入Body
 			for {
 				n, err := resp.Body.Read(buf)
 				if n > 0 {
@@ -411,6 +444,7 @@ func (this *APIStream) handlePreheatCache(message *pb.NodeStreamMessage) error {
 				}
 				if err != nil {
 					if err == io.EOF {
+
 						err = writer.Close()
 						if err == nil {
 							storage.AddToList(&caches.Item{
@@ -484,7 +518,7 @@ func (this *APIStream) handleCheckSystemdService(message *pb.NodeStreamMessage)
 	cmd.Add(systemctl, "is-enabled", shortName)
 	output, err := cmd.Run()
 	if err != nil {
-		this.replyFail(message.RequestId, "'systemctl' command error: " + err.Error())
+		this.replyFail(message.RequestId, "'systemctl' command error: "+err.Error())
 		return nil
 	}
 	if output == "enabled" {

internal/nodes/http_request_cache.go

@@ -2,9 +2,9 @@ package nodes
 
 import (
 	"bytes"
+	"errors"
 	"github.com/TeaOSLab/EdgeNode/internal/caches"
 	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
-	"github.com/iwind/TeaGo/types"
 	"net/http"
 	"strconv"
 )
@@ -66,74 +66,12 @@ func (this *HTTPRequest) doCacheRead() (shouldStop bool) {
 		return
 	}
 
-	isBroken := false
-	headerBuf := []byte{}
-	statusCode := http.StatusOK
-	statusFound := false
-	headerFound := false
-
 	buf := bytePool32k.Get()
-	err := storage.Read(key, buf, func(data []byte, valueSize int64, expiredAt int64, isEOF bool) {
-		if isBroken {
-			return
-		}
-
-		// 如果Header已发送完毕
-		if headerFound {
-			_, _ = this.writer.Write(data)
-			return
-		}
-
-		headerBuf = append(headerBuf, data...)
-
-		if !statusFound {
-			lineIndex := bytes.IndexByte(headerBuf, '\n')
-			if lineIndex < 0 {
-				return
-			}
-
-			pieces := bytes.Split(headerBuf[:lineIndex], []byte{' '})
-			if len(pieces) < 2 {
-				isBroken = true
-				return
-			}
-			statusCode = types.Int(string(pieces[1]))
-			statusFound = true
-			headerBuf = headerBuf[lineIndex+1:]
-
-			// cache相关变量
-			this.varMapping["cache.status"] = "HIT"
-		}
-
-		for {
-			lineIndex := bytes.IndexByte(headerBuf, '\n')
-			if lineIndex < 0 {
-				break
-			}
-			if lineIndex == 0 || lineIndex == 1 {
-				headerFound = true
-
-				this.processResponseHeaders(statusCode)
-				this.writer.WriteHeader(statusCode)
-
-				_, _ = this.writer.Write(headerBuf[lineIndex+1:])
-				headerBuf = nil
-				break
-			}
-
-			// 分解Header
-			line := headerBuf[:lineIndex]
-			colonIndex := bytes.IndexByte(line, ':')
-			if colonIndex <= 0 {
-				continue
-			}
-			this.writer.Header().Set(string(line[:colonIndex]), string(bytes.TrimSpace(line[colonIndex+1:])))
-			headerBuf = headerBuf[lineIndex+1:]
-		}
-	})
-
-	bytePool32k.Put(buf)
+	defer func() {
+		bytePool32k.Put(buf)
+	}()
 
+	reader, err := storage.OpenReader(key)
 	if err != nil {
 		if err == caches.ErrNotFound {
 			// cache相关变量
@@ -144,11 +82,56 @@ func (this *HTTPRequest) doCacheRead() (shouldStop bool) {
 		remotelogs.Error("REQUEST_CACHE", "read from cache failed: "+err.Error())
 		return
 	}
+	defer func() {
+		_ = reader.Close()
+	}()
 
-	if isBroken {
+	this.varMapping["cache.status"] = "HIT"
+
+	// 读取Header
+	headerBuf := []byte{}
+	err = reader.ReadHeader(buf, func(n int) (goNext bool, err error) {
+		headerBuf = append(headerBuf, buf[:n]...)
+		for {
+			nIndex := bytes.Index(headerBuf, []byte{'\n'})
+			if nIndex >= 0 {
+				row := headerBuf[:nIndex]
+				spaceIndex := bytes.Index(row, []byte{':'})
+				if spaceIndex <= 0 {
+					return false, errors.New("invalid header '" + string(row) + "'")
+				}
+
+				this.writer.Header().Set(string(row[:spaceIndex]), string(row[spaceIndex+1:]))
+				headerBuf = headerBuf[nIndex+1:]
+			} else {
+				break
+			}
+		}
+		return true, nil
+	})
+	if err != nil {
+		remotelogs.Error("REQUEST_CACHE", "read from cache failed: "+err.Error())
 		return
 	}
 
+	this.processResponseHeaders(reader.Status())
+	this.writer.WriteHeader(reader.Status())
+
+	// 输出Body
+	if this.RawReq.Method != http.MethodHead {
+		err = reader.ReadBody(buf, func(n int) (goNext bool, err error) {
+			_, err = this.writer.Write(buf[:n])
+			if err != nil {
+				return false, err
+			}
+			return true, nil
+		})
+		if err != nil {
+			remotelogs.Error("REQUEST_CACHE", "read from cache failed: "+err.Error())
+			return
+		}
+	}
+
 	this.cacheRef = nil // 终止读取不再往下传递
 	return true
 }

internal/nodes/http_writer.go

@@ -113,7 +113,6 @@ func (this *HTTPWriter) Write(data []byte) (n int, err error) {
 		if this.cacheWriter != nil {
 			_, err = this.cacheWriter.Write(data)
 			if err != nil {
-				_ = this.cacheWriter.Discard()
 				this.cacheWriter = nil
 				remotelogs.Error("REQUEST_WRITER", "write cache failed: "+err.Error())
 			}
@@ -216,9 +215,10 @@ func (this *HTTPWriter) Close() {
 		err := this.cacheWriter.Close()
 		if err == nil {
 			this.cacheStorage.AddToList(&caches.Item{
-				Key:       this.cacheWriter.Key(),
-				ExpiredAt: this.cacheWriter.ExpiredAt(),
-				ValueSize: this.cacheWriter.Size(),
+				Key:        this.cacheWriter.Key(),
+				ExpiredAt:  this.cacheWriter.ExpiredAt(),
+				HeaderSize: this.cacheWriter.HeaderSize(),
+				BodySize:   this.cacheWriter.BodySize(),
 			})
 		}
 	}
@@ -364,7 +364,7 @@ func (this *HTTPWriter) prepareCache(size int64) {
 		life = 60
 	}
 	expiredAt := utils.UnixTime() + life
-	cacheWriter, err := storage.Open(this.req.cacheKey, expiredAt)
+	cacheWriter, err := storage.OpenWriter(this.req.cacheKey, expiredAt, this.StatusCode())
 	if err != nil {
 		if err != caches.ErrFileIsWriting {
 			remotelogs.Error("REQUEST_WRITER", "write cache failed: "+err.Error())
@@ -377,12 +377,14 @@ func (this *HTTPWriter) prepareCache(size int64) {
 	}
 
 	// 写入Header
-	headerData := this.HeaderData()
-	_, err = cacheWriter.Write(headerData)
-	if err != nil {
-		remotelogs.Error("REQUEST_WRITER", "write cache failed: "+err.Error())
-		_ = this.cacheWriter.Discard()
-		this.cacheWriter = nil
-		return
+	for k, v := range this.Header() {
+		for _, v1 := range v {
+			_, err = cacheWriter.WriteHeader([]byte(k + ":" + v1 + "\n"))
+			if err != nil {
+				remotelogs.Error("REQUEST_WRITER", "write cache failed: "+err.Error())
+				this.cacheWriter = nil
+				return
+			}
+		}
 	}
 }