go pprof server

go-pprof-server.go

// a server fetch remote go pprof file and then reverse proxy to `go tool pprof --http` web 
package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"
)

const (
	downloadDir = "/tmp"
	port        = 18889
)

type PprofServer struct {
	Port       int
	Process    *os.Process
	LastAccess time.Time
	FileName   string
}

var (
	pprofServers sync.Map // Stores mapping of file to PprofServer
	cleanupMutex sync.Mutex
)

func main() {
	http.HandleFunc("/pprof", handlePprof)
	http.HandleFunc("/view", handleView)
	http.HandleFunc("/proxy/", handleProxy) // Add proxy route

	// Start cleanup goroutine
	go cleanupRoutine()

	log.Println("Server starting on port ", port)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", port), nil))
}

// GET /pprof handles profile requests
func handlePprof(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// Get pprof URL from query string
	pprofURL := r.URL.Query().Get("url")
	if pprofURL == "" {
		http.Error(w, "usage: /pprof?url=http://host:port/debug/pprof/heap", http.StatusBadRequest)
		return
	}

	// Ensure URL includes protocol
	if !strings.HasPrefix(pprofURL, "http://") && !strings.HasPrefix(pprofURL, "https://") {
		pprofURL = "http://" + pprofURL
	}

	log.Printf("Downloading pprof from: %s", pprofURL)

	// Download pprof file
	filename, err := downloadPprof(pprofURL)
	if err != nil {
		log.Printf("Failed to download pprof: %v", err)
		http.Error(w, fmt.Sprintf("Failed to download pprof: %v", err), http.StatusInternalServerError)
		return
	}

	log.Printf("Pprof file saved as: %s", filename)

	// Redirect to /view with filename parameter
	redirectURL := fmt.Sprintf("/view?file=%s", filepath.Base(filename))
	http.Redirect(w, r, redirectURL, http.StatusSeeOther)
}

// Proxy handler forwards requests to pprof server
func handleProxy(w http.ResponseWriter, r *http.Request) {
	// Parse path to get filename
	path := strings.TrimPrefix(r.URL.Path, "/proxy/")
	parts := strings.SplitN(path, "/", 2)
	if len(parts) < 1 {
		http.Error(w, "Invalid proxy path", http.StatusBadRequest)
		return
	}

	fileName := parts[0]
	if fileName == "" {
		http.Error(w, "Missing file name in proxy path", http.StatusBadRequest)
		return
	}

	// Get corresponding server info
	serverInterface, exists := pprofServers.Load(fileName)
	if !exists {
		http.Error(w, "Pprof server not found for this file", http.StatusNotFound)
		return
	}

	server := serverInterface.(*PprofServer)

	// Check if server is still running
	if !checkPprofServer(server.Port) {
		pprofServers.Delete(fileName)
		http.Error(w, "Pprof server is not running", http.StatusServiceUnavailable)
		return
	}

	// Update access time
	server.LastAccess = time.Now()

	// Create reverse proxy
	target, err := url.Parse(fmt.Sprintf("http://localhost:%d", server.Port))
	if err != nil {
		http.Error(w, "Failed to parse target URL", http.StatusInternalServerError)
		return
	}

	// Modify request path, remove proxy prefix
	originalPath := r.URL.Path
	if len(parts) > 1 {
		r.URL.Path = "/" + parts[1]
	} else {
		r.URL.Path = "/"
	}

	// Create reverse proxy
	proxy := httputil.NewSingleHostReverseProxy(target)

	// Modify response to handle relative paths
	proxy.ModifyResponse = func(resp *http.Response) error {
		// Modify Location header (redirect)
		if location := resp.Header.Get("Location"); location != "" {
			if strings.HasPrefix(location, "/") {
				resp.Header.Set("Location", fmt.Sprintf("/proxy/%s%s", fileName, location))
			}
		}

		// If response is HTML, modify links
		if strings.Contains(resp.Header.Get("Content-Type"), "text/html") {
			// Add logic to modify HTML content, converting relative links to proxy links
			// Simplified for now, not processing HTML content modification
		}

		return nil
	}

	log.Printf("Proxying request from %s to http://localhost:%d%s", originalPath, server.Port, r.URL.Path)
	proxy.ServeHTTP(w, r)
}

// Cleanup goroutine periodically checks and removes expired pprof servers
func cleanupRoutine() {
	ticker := time.NewTicker(10 * time.Minute) // Check every 10 minutes
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			cleanupExpiredServers()
		}
	}
}

// Clean up pprof servers not accessed for over 1 hour
func cleanupExpiredServers() {
	cleanupMutex.Lock()
	defer cleanupMutex.Unlock()

	now := time.Now()
	var toDelete []string

	pprofServers.Range(func(key, value interface{}) bool {
		fileName := key.(string)
		server := value.(*PprofServer)

		// Check if not accessed for over 1 hour
		if now.Sub(server.LastAccess) > time.Hour {
			toDelete = append(toDelete, fileName)
			log.Printf("Marking pprof server for %s for cleanup (last access: %v)", fileName, server.LastAccess)
		}

		return true
	})

	// Delete expired servers
	for _, fileName := range toDelete {
		if serverInterface, exists := pprofServers.Load(fileName); exists {
			server := serverInterface.(*PprofServer)

			// Stop pprof process
			if server.Process != nil {
				log.Printf("Stopping pprof server process for %s (PID: %d)", fileName, server.Process.Pid)
				if err := server.Process.Signal(syscall.SIGTERM); err != nil {
					log.Printf("Failed to send SIGTERM to process %d: %v", server.Process.Pid, err)
					// If SIGTERM fails, try SIGKILL
					if err := server.Process.Kill(); err != nil {
						log.Printf("Failed to kill process %d: %v", server.Process.Pid, err)
					}
				} else {
					// Wait for process to exit
					go func(proc *os.Process, name string) {
						if _, err := proc.Wait(); err != nil {
							log.Printf("Process %d for %s exited with error: %v", proc.Pid, name, err)
						} else {
							log.Printf("Process %d for %s exited cleanly", proc.Pid, name)
						}
					}(server.Process, fileName)
				}
			}

			// Remove from map
			pprofServers.Delete(fileName)
			log.Printf("Cleaned up pprof server for %s", fileName)
		}
	}

	if len(toDelete) > 0 {
		log.Printf("Cleanup completed, removed %d expired pprof servers", len(toDelete))
	}
}

// GET /view starts go tool pprof -http
func handleView(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// Get filename from query string
	fileName := r.URL.Query().Get("file")
	if fileName == "" {
		http.Error(w, "Missing 'file' parameter", http.StatusBadRequest)
		return
	}

	// Construct full file path
	pprofFile := filepath.Join(downloadDir, fileName)

	// Check if file exists
	if _, err := os.Stat(pprofFile); os.IsNotExist(err) {
		http.Error(w, "Pprof file not found", http.StatusNotFound)
		return
	}

	// Check if a pprof server is already running for this file
	if serverInterface, exists := pprofServers.Load(fileName); exists {
		server := serverInterface.(*PprofServer)
		if checkPprofServer(server.Port) {
			// Update access time
			server.LastAccess = time.Now()
			// Server is already running, redirect to proxy
			proxyURL := fmt.Sprintf("/proxy/%s/", fileName)
			http.Redirect(w, r, proxyURL, http.StatusSeeOther)
			return
		} else {
			// Server stopped, remove record
			pprofServers.Delete(fileName)
		}
	}

	// Find an available random port
	port, err := findAvailablePort()
	if err != nil {
		http.Error(w, "Failed to find available port", http.StatusInternalServerError)
		return
	}

	// Start go tool pprof -http
	var process *os.Process
	go func() {
		cmd := exec.Command("go", "tool", "pprof", "-http", fmt.Sprintf(":%d", port), pprofFile)
		log.Printf("Starting pprof server for %s on port %d: %s", fileName, port, cmd.String())

		if err := cmd.Start(); err != nil {
			log.Printf("Failed to start pprof server for %s: %v", fileName, err)
			pprofServers.Delete(fileName)
			return
		}

		process = cmd.Process

		if err := cmd.Wait(); err != nil {
			log.Printf("Pprof server process for %s exited: %v", fileName, err)
			pprofServers.Delete(fileName)
		}
	}()

	// Wait for pprof server to start
	time.Sleep(2 * time.Second)

	// Check if pprof server started successfully
	if !checkPprofServer(port) {
		http.Error(w, "Failed to start pprof server", http.StatusInternalServerError)
		return
	}

	// Store server information
	server := &PprofServer{
		Port:       port,
		Process:    process,
		LastAccess: time.Now(),
		FileName:   fileName,
	}
	pprofServers.Store(fileName, server)

	// Redirect to proxy path
	proxyURL := fmt.Sprintf("/proxy/%s/", fileName)
	log.Printf("Redirecting to proxy interface: %s", proxyURL)
	http.Redirect(w, r, proxyURL, http.StatusSeeOther)
}

// Download pprof file
func downloadPprof(url string) (string, error) {
	// Create HTTP client
	client := &http.Client{
		Timeout: 60 * time.Second,
	}

	resp, err := client.Get(url)
	if err != nil {
		return "", fmt.Errorf("failed to download: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("HTTP error: %d %s", resp.StatusCode, resp.Status)
	}

	// Generate filename
	timestamp := time.Now().Format("20060102_150405")
	filename := fmt.Sprintf("pprof_%s.pb.gz", timestamp)
	filepath := filepath.Join(downloadDir, filename)

	// Create file
	file, err := os.Create(filepath)
	if err != nil {
		return "", fmt.Errorf("failed to create file: %w", err)
	}
	defer file.Close()

	// Copy data
	_, err = io.Copy(file, resp.Body)
	if err != nil {
		return "", fmt.Errorf("failed to save file: %w", err)
	}

	return filepath, nil
}

// Find available port
func findAvailablePort() (int, error) {
	// Try ports starting from 8080
	for port := 8080; port < 9000; port++ {
		if isPortAvailable(port) {
			return port, nil
		}
	}
	return 0, fmt.Errorf("no available port found")
}

// Check if port is available
func isPortAvailable(port int) bool {
	conn, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		return false
	}
	conn.Close()
	return true
}

// Check if pprof server is running
func checkPprofServer(port int) bool {
	url := fmt.Sprintf("http://localhost:%d", port)
	client := &http.Client{Timeout: 5 * time.Second}

	// Retry a few times
	for i := 0; i < 10; i++ {
		resp, err := client.Get(url)
		if err == nil && resp.StatusCode == http.StatusOK {
			resp.Body.Close()
			return true
		}
		if resp != nil {
			resp.Body.Close()
		}
		time.Sleep(500 * time.Millisecond)
	}
	return false
}