-
Notifications
You must be signed in to change notification settings - Fork 4
/
build.go
76 lines (70 loc) · 1.4 KB
/
build.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
package tensor
import (
"fmt"
"io"
"path/filepath"
"runtime"
"sort"
"strings"
"sync"
"sync/atomic"
)
var leaks struct {
sync.RWMutex
data map[uint64][]uintptr // don't ref to *Tensor
idx atomic.Uint64
}
func init() {
leaks.data = make(map[uint64][]uintptr)
}
func logBuildInfo(t *Tensor) {
idx := leaks.idx.Add(1)
t.idx = idx
pcs := make([]uintptr, 32)
n := runtime.Callers(2, pcs)
leaks.Lock()
leaks.data[idx] = pcs[:n]
leaks.Unlock()
}
func free(t *Tensor) {
leaks.Lock()
if _, ok := leaks.data[t.idx]; !ok {
panic("tensor: double free")
}
delete(leaks.data, t.idx)
leaks.Unlock()
}
func WriteLeaks(w io.Writer) {
raw := make(map[uint64][]uintptr, len(leaks.data))
ids := make([]uint64, 0, len(leaks.data))
leaks.RLock()
for idx, pcs := range leaks.data {
raw[idx] = pcs
ids = append(ids, idx)
}
leaks.RUnlock()
writeStack := func(pcs []uintptr) {
frames := runtime.CallersFrames(pcs)
for {
frame, more := frames.Next()
base := filepath.Base(frame.Function)
if strings.HasPrefix(base, "tensor.") {
continue
}
if strings.HasPrefix(base, "runtime.") {
continue
}
fmt.Fprintf(w, " - %s:%d => %s\n", frame.File, frame.Line, frame.Function)
if !more {
break
}
}
}
sort.Slice(ids, func(i, j int) bool {
return ids[i] < ids[j]
})
for _, id := range ids {
fmt.Fprintf(w, "tensor [>> tensor.%d <<] leaked:\n", id)
writeStack(raw[id])
}
}