Add go profiling flag, and enable on integration tests (#1382)

integration/hsic/hsic.go

@@ -16,6 +16,9 @@ import (
 	"net"
 	"net/http"
 	"net/url"
+	"os"
+	"path"
+	"strings"
 	"time"
 
 	"github.com/davecgh/go-spew/spew"
@@ -205,7 +208,10 @@ func New(
 		ContextDir: dockerContextPath,
 	}
 
-	env := []string{}
+	env := []string{
+		"HEADSCALE_PROFILING_ENABLED=1",
+		"HEADSCALE_PROFILING_PATH=/tmp/profile",
+	}
 	for key, value := range hsic.env {
 		env = append(env, fmt.Sprintf("%s=%s", key, value))
 	}
@@ -219,7 +225,7 @@ func New(
 		// Cmd:          []string{"headscale", "serve"},
 		// TODO(kradalby): Get rid of this hack, we currently need to give us some
 		// to inject the headscale configuration further down.
-		Entrypoint: []string{"/bin/bash", "-c", "/bin/sleep 3 ; headscale serve"},
+		Entrypoint: []string{"/bin/bash", "-c", "/bin/sleep 3 ; headscale serve ; /bin/sleep 30"},
 		Env:        env,
 	}
 
@@ -305,6 +311,33 @@ func (t *HeadscaleInContainer) hasTLS() bool {
 
 // Shutdown stops and cleans up the Headscale container.
 func (t *HeadscaleInContainer) Shutdown() error {
+	err := t.SaveLog("/tmp/control")
+	if err != nil {
+		log.Printf(
+			"Failed to save log from control: %s",
+			fmt.Errorf("failed to save log from control: %w", err),
+		)
+	}
+
+	// Send a interrupt signal to the "headscale" process inside the container
+	// allowing it to shut down gracefully and flush the profile to disk.
+	// The container will live for a bit longer due to the sleep at the end.
+	err = t.SendInterrupt()
+	if err != nil {
+		log.Printf(
+			"Failed to send graceful interrupt to control: %s",
+			fmt.Errorf("failed to send graceful interrupt to control: %w", err),
+		)
+	}
+
+	err = t.SaveProfile("/tmp/control")
+	if err != nil {
+		log.Printf(
+			"Failed to save profile from control: %s",
+			fmt.Errorf("failed to save profile from control: %w", err),
+		)
+	}
+
 	return t.pool.Purge(t.container)
 }
 
@@ -314,6 +347,24 @@ func (t *HeadscaleInContainer) SaveLog(path string) error {
 	return dockertestutil.SaveLog(t.pool, t.container, path)
 }
 
+func (t *HeadscaleInContainer) SaveProfile(savePath string) error {
+	tarFile, err := t.FetchPath("/tmp/profile")
+	if err != nil {
+		return err
+	}
+
+	err = os.WriteFile(
+		path.Join(savePath, t.hostname+".pprof.tar"),
+		tarFile,
+		os.ModePerm,
+	)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // Execute runs a command inside the Headscale container and returns the
 // result of stdout as a string.
 func (t *HeadscaleInContainer) Execute(
@@ -498,6 +549,26 @@ func (t *HeadscaleInContainer) WriteFile(path string, data []byte) error {
 	return integrationutil.WriteFileToContainer(t.pool, t.container, path, data)
 }
 
+// FetchPath gets a path from inside the Headscale container and returns a tar
+// file as byte array.
+func (t *HeadscaleInContainer) FetchPath(path string) ([]byte, error) {
+	return integrationutil.FetchPathFromContainer(t.pool, t.container, path)
+}
+
+func (t *HeadscaleInContainer) SendInterrupt() error {
+	pid, err := t.Execute([]string{"pidof", "headscale"})
+	if err != nil {
+		return err
+	}
+
+	_, err = t.Execute([]string{"kill", "-2", strings.Trim(pid, "'\n")})
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 // nolint
 func createCertificate(hostname string) ([]byte, []byte, error) {
 	// From: