Add interface for custom source code loaders

stacktrace.go

@@ -143,7 +143,7 @@ func NewStacktraceFrame(pc uintptr, file string, line, context int, appPackagePr
 	}
 
 	if context > 0 {
-		contextLines, lineIdx := fileContext(file, line, context)
+		contextLines, lineIdx := sourceCodeLoader.Load(file, line, context)
 		if len(contextLines) > 0 {
 			for i, line := range contextLines {
 				switch {
@@ -157,7 +157,7 @@ func NewStacktraceFrame(pc uintptr, file string, line, context int, appPackagePr
 			}
 		}
 	} else if context == -1 {
-		contextLine, _ := fileContext(file, line, 0)
+		contextLine, _ := sourceCodeLoader.Load(file, line, 0)
 		if len(contextLine) > 0 {
 			frame.ContextLine = string(contextLine[0])
 		}
@@ -191,24 +191,36 @@ func splitFunctionName(name string) (string, string) {
 	return pack, name
 }
 
-var fileCacheLock sync.Mutex
-var fileCache = make(map[string][][]byte)
+type SourceCodeLoader interface {
+	Load(filename string, line, context int) ([][]byte, int)
+}
+
+var sourceCodeLoader SourceCodeLoader = fsLoader{cache: make(map[string][][]byte)}
+
+func SetSourceCodeLoader(loader SourceCodeLoader) {
+	sourceCodeLoader = loader
+}
+
+type fsLoader struct {
+	mu    sync.Mutex
+	cache map[string][][]byte
+}
 
-func fileContext(filename string, line, context int) ([][]byte, int) {
-	fileCacheLock.Lock()
-	defer fileCacheLock.Unlock()
-	lines, ok := fileCache[filename]
+func (fs fsLoader) Load(filename string, line, context int) ([][]byte, int) {
+	fs.mu.Lock()
+	defer fs.mu.Unlock()
+	lines, ok := fs.cache[filename]
 	if !ok {
 		data, err := ioutil.ReadFile(filename)
 		if err != nil {
 			// cache errors as nil slice: code below handles it correctly
 			// otherwise when missing the source or running as a different user, we try
 			// reading the file on each error which is unnecessary
-			fileCache[filename] = nil
+			fs.cache[filename] = nil
 			return nil, 0
 		}
 		lines = bytes.Split(data, []byte{'\n'})
-		fileCache[filename] = lines
+		fs.cache[filename] = lines
 	}
 
 	if lines == nil {

stacktrace_test.go

@@ -48,8 +48,8 @@ func TestFunctionName(t *testing.T) {
 }
 
 func TestSplitFunctionName(t *testing.T) {
-	tests := []struct{
-		in string
+	tests := []struct {
+		in         string
 		pack, name string
 	}{
 		{"", "", ""},
@@ -169,7 +169,7 @@ func TestNewStacktrace_noFrames(t *testing.T) {
 
 func TestFileContext(t *testing.T) {
 	// reset the cache
-	fileCache = make(map[string][][]byte)
+	sourceCodeLoader = fsLoader{cache: make(map[string][][]byte)}
 
 	tempdir, err := ioutil.TempDir("", "")
 	if err != nil {
@@ -200,13 +200,14 @@ func TestFileContext(t *testing.T) {
 		{noPermissionPath, 0, 0},
 	}
 	for i, test := range tests {
-		lines, index := fileContext(test.path, 1, 0)
+		lines, index := sourceCodeLoader.Load(test.path, 1, 0)
 		if !(len(lines) == test.expectedLines && index == test.expectedIndex) {
 			t.Errorf("%d: fileContext(%#v, 1, 0) = %v, %v; expected len()=%d, %d",
 				i, test.path, lines, index, test.expectedLines, test.expectedIndex)
 		}
-		if len(fileCache) != i+1 {
-			t.Errorf("%d: result was not cached; len(fileCached)=%d", i, len(fileCache))
+		cacheLen := len(sourceCodeLoader.(fsLoader).cache)
+		if cacheLen != i+1 {
+			t.Errorf("%d: result was not cached; len=%d", i, cacheLen)
 		}
 	}
 }