优化结构

internal/logging/source.go

@@ -6,9 +6,12 @@ import (
 	"go/parser"
 	"go/token"
 	"os"
+	"path/filepath"
 	"regexp"
 	"strconv"
 	"strings"
+	"sync"
+	"time"
 )
 
 type stackFrame struct {
@@ -19,6 +22,23 @@ type stackFrame struct {
 
 var stackFileLinePattern = regexp.MustCompile(`\s*(.+\.go):(\d+)\s*$`)
 
+type sourceFunction struct {
+	name               string
+	source             string
+	startLine, endLine int
+}
+
+type sourceFileCache struct {
+	size    int64
+	modTime time.Time
+	items   []sourceFunction
+}
+
+var sourceCache = struct {
+	sync.RWMutex
+	files map[string]sourceFileCache
+}{files: make(map[string]sourceFileCache)}
+
 func finalApplicationCaller(stack string) (stackFrame, bool) {
 	if stack == "" {
 		return stackFrame{}, false
@@ -69,6 +89,17 @@ func skipStackFile(filename string) bool {
 }
 
 func functionSourceAt(filename string, line int) (name, source string, startLine, endLine int, err error) {
+	filename = filepath.Clean(filename)
+	info, err := os.Stat(filename)
+	if err != nil {
+		return "", "", 0, 0, fmt.Errorf("stat file failed: %w", err)
+	}
+	sourceCache.RLock()
+	cached, found := sourceCache.files[filename]
+	sourceCache.RUnlock()
+	if found && cached.size == info.Size() && cached.modTime.Equal(info.ModTime()) {
+		return sourceFunctionAtLine(cached.items, line, filename)
+	}
 	content, err := os.ReadFile(filename)
 	if err != nil {
 		return "", "", 0, 0, fmt.Errorf("read file failed: %w", err)
@@ -78,7 +109,15 @@ func functionSourceAt(filename string, line int) (name, source string, startLine
 	if err != nil {
 		return "", "", 0, 0, fmt.Errorf("parse file failed: %w", err)
 	}
-	var target *ast.FuncDecl
+	sourceCache.Lock()
+	sourceCache.files[filename] = sourceFileCache{size: info.Size(), modTime: info.ModTime(), items: allFunctions(parsed, files, content)}
+	cached = sourceCache.files[filename]
+	sourceCache.Unlock()
+	return sourceFunctionAtLine(cached.items, line, filename)
+}
+
+func allFunctions(parsed *ast.File, files *token.FileSet, content []byte) []sourceFunction {
+	items := make([]sourceFunction, 0)
 	ast.Inspect(parsed, func(node ast.Node) bool {
 		declaration, ok := node.(*ast.FuncDecl)
 		if !ok {
@@ -86,19 +125,21 @@ func functionSourceAt(filename string, line int) (name, source string, startLine
 		}
 		start := files.Position(declaration.Pos()).Line
 		end := files.Position(declaration.End()).Line
-		if line >= start && line <= end {
-			target, startLine, endLine = declaration, start, end
-			return false
+		startOffset := files.Position(declaration.Pos()).Offset
+		endOffset := files.Position(declaration.End()).Offset
+		if startOffset >= 0 && endOffset <= len(content) && startOffset < endOffset {
+			items = append(items, sourceFunction{name: declaration.Name.Name, source: string(content[startOffset:endOffset]), startLine: start, endLine: end})
 		}
 		return true
 	})
-	if target == nil {
-		return "", "", 0, 0, fmt.Errorf("no function encloses line %d in %s", line, filename)
+	return items
+}
+
+func sourceFunctionAtLine(items []sourceFunction, line int, filename string) (string, string, int, int, error) {
+	for _, item := range items {
+		if line >= item.startLine && line <= item.endLine {
+			return item.name, item.source, item.startLine, item.endLine, nil
+		}
 	}
-	start := files.Position(target.Pos()).Offset
-	end := files.Position(target.End()).Offset
-	if start < 0 || end > len(content) || start >= end {
-		return "", "", 0, 0, fmt.Errorf("invalid offsets for function: start=%d end=%d len=%d", start, end, len(content))
-	}
-	return target.Name.Name, string(content[start:end]), startLine, endLine, nil
+	return "", "", 0, 0, fmt.Errorf("no function encloses line %d in %s", line, filename)
 }