1// Copyright 2018 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package main
6
7import (
8	"bytes"
9	"fmt"
10	"runtime"
11	"strings"
12)
13
14func init() {
15	register("TracebackAncestors", TracebackAncestors)
16}
17
18const numGoroutines = 3
19const numFrames = 2
20
21func TracebackAncestors() {
22	w := make(chan struct{})
23	recurseThenCallGo(w, numGoroutines, numFrames, true)
24	<-w
25	printStack()
26	close(w)
27}
28
29var ignoreGoroutines = make(map[string]bool)
30
31func printStack() {
32	buf := make([]byte, 1024)
33	for {
34		n := runtime.Stack(buf, true)
35		if n < len(buf) {
36			all := string(buf[:n])
37			var saved string
38
39			// Delete any ignored goroutines, if present.
40			for all != "" {
41				var g string
42				g, all, _ = strings.Cut(all, "\n\n")
43
44				if strings.HasPrefix(g, "goroutine ") {
45					id, _, _ := strings.Cut(strings.TrimPrefix(g, "goroutine "), " ")
46					if ignoreGoroutines[id] {
47						continue
48					}
49				}
50				if saved != "" {
51					saved += "\n\n"
52				}
53				saved += g
54			}
55
56			fmt.Print(saved)
57			return
58		}
59		buf = make([]byte, 2*len(buf))
60	}
61}
62
63func recurseThenCallGo(w chan struct{}, frames int, goroutines int, main bool) {
64	if frames == 0 {
65		// Signal to TracebackAncestors that we are done recursing and starting goroutines.
66		w <- struct{}{}
67		<-w
68		return
69	}
70	if goroutines == 0 {
71		// Record which goroutine this is so we can ignore it
72		// in the traceback if it hasn't finished exiting by
73		// the time we printStack.
74		if !main {
75			ignoreGoroutines[goroutineID()] = true
76		}
77
78		// Start the next goroutine now that there are no more recursions left
79		// for this current goroutine.
80		go recurseThenCallGo(w, frames-1, numFrames, false)
81		return
82	}
83	recurseThenCallGo(w, frames, goroutines-1, main)
84}
85
86func goroutineID() string {
87	buf := make([]byte, 128)
88	runtime.Stack(buf, false)
89	prefix := []byte("goroutine ")
90	var found bool
91	if buf, found = bytes.CutPrefix(buf, prefix); !found {
92		panic(fmt.Sprintf("expected %q at beginning of traceback:\n%s", prefix, buf))
93	}
94	id, _, _ := bytes.Cut(buf, []byte(" "))
95	return string(id)
96}
97