feat: add gouroutine-leak profiling

1 files changed, 27 insertions, 1 deletions

diff --git a/main.go b/main.go
index 27d3e27..db64348 100644
--- a/main.go
+++ b/main.go
@@ -7,7 +7,10 @@ import (
 	"fmt"
 	"log"
 	"net/url"
+	"runtime/pprof"
+	"strings"
 	"sync"
+	"time"
 )
 
 func main() {
@@ -38,7 +41,9 @@ func main() {
 		log.Fatal(err)
 	}
 
-	pool(*startURL, *maxConcurrency, *maxURLs)
+	getLeakProfile(func() {
+		pool(*startURL, *maxConcurrency, *maxURLs)
+	})
 }
 
 func pool(startURL url.URL, maxConcurrency, maxURLs int) {
@@ -117,3 +122,24 @@ func fanIn(chans ...<-chan []url.URL) <-chan []url.URL {
 
 	return out
 }
+
+// getLeakProfile runs a leaky program snippet, extracts the goroutine leak profile,
+// and writes it to stdout.
+func getLeakProfile(leakySnippet func()) {
+	prof := pprof.Lookup("goroutineleak")
+	defer func() {
+		time.Sleep(2 * time.Second)
+		var content strings.Builder
+
+		prof.WriteTo(&content, 2)
+		// Ignore non leaked goroutines
+		leaks := strings.Split(content.String(), "\n\n")
+		for _, leak := range leaks {
+			if strings.Contains(leak, "(leaked)") {
+				fmt.Println(leak + "\n")
+			}
+		}
+	}()
+
+	leakySnippet()
+}