Supercharge Your Go Web Service: Building a Custom Profiler

Introduction

As Go developers, we often reach for built-in profiling tools when optimizing our applications. But what if we could create a profiler that speaks our application's language? In this guide, we'll construct a custom profiler for a Go web service, focusing on request handling, database operations, and memory usage.

The Case for Custom Profiling

While Go's standard profiler is powerful, it might not capture everything specific to your web service:

Patterns in web request handling across different endpoints
Database query performance for various operations
Memory usage fluctuations during peak loads

Let's build a profiler that addresses these exact needs.

Our Sample Web Service

First, let's set up a basic web service to profile:

package main

import (
    "database/sql"
    "encoding/json"
    "log"
    "net/http"

    _ "github.com/lib/pq"
)

type User struct {
    ID   int    `json:"id"`
    Name string `json:"name"`
}

var db *sql.DB

func main() {
    // Initialize database connection
    var err error
    db, err = sql.Open("postgres", "postgres://username:password@localhost/database?sslmode=disable")
    if err != nil {
        log.Fatal(err)
    }
    defer db.Close()

    // Set up routes
    http.HandleFunc("/user", handleUser)

    // Start the server
    log.Println("Server starting on :8080")
    log.Fatal(http.ListenAndServe(":8080", nil))
}

func handleUser(w http.ResponseWriter, r *http.Request) {
    // Handle GET and POST requests for users
    // Implementation omitted for brevity
}

Now, let's build our custom profiler to gain deep insights into this service.

Custom Profiler Implementation

1. Request Duration Tracking

We'll start by measuring how long each request takes:

import (
    "time"
    "sync"
)

var (
    requestDurations = make(map[string]time.Duration)
    requestMutex     sync.RWMutex
)

func trackRequestDuration(handler http.HandlerFunc) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        start := time.Now()
        handler(w, r)
        duration := time.Since(start)

        requestMutex.Lock()
        requestDurations[r.URL.Path] += duration
        requestMutex.Unlock()
    }
}

// In main(), wrap your handlers:
http.HandleFunc("/user", trackRequestDuration(handleUser))

2. Database Query Profiling

Next, let's keep tabs on our database performance:

type QueryStats struct {
    Count    int
    Duration time.Duration
}

var (
    queryStats = make(map[string]QueryStats)
    queryMutex sync.RWMutex
)

func trackQuery(query string, duration time.Duration) {
    queryMutex.Lock()
    defer queryMutex.Unlock()

    stats := queryStats[query]
    stats.Count++
    stats.Duration += duration
    queryStats[query] = stats
}

// Use this function to wrap your database queries:
func profiledQuery(query string, args ...interface{}) (*sql.Rows, error) {
    start := time.Now()
    rows, err := db.Query(query, args...)
    duration := time.Since(start)
    trackQuery(query, duration)
    return rows, err
}

3. Memory Usage Tracking

Let's add memory usage tracking to complete our profiler:

import "runtime"

func getMemStats() runtime.MemStats {
    var m runtime.MemStats
    runtime.ReadMemStats(&m)
    return m
}

func logMemStats() {
    stats := getMemStats()
    log.Printf("Alloc = %v MiB", bToMb(stats.Alloc))
    log.Printf("TotalAlloc = %v MiB", bToMb(stats.TotalAlloc))
    log.Printf("Sys = %v MiB", bToMb(stats.Sys))
    log.Printf("NumGC = %v", stats.NumGC)
}

func bToMb(b uint64) uint64 {
    return b / 1024 / 1024
}

// Call this periodically in a goroutine:
go func() {
    ticker := time.NewTicker(1 * time.Minute)
    for range ticker.C {
        logMemStats()
    }
}()

4. Profiler API Endpoint

Finally, let's create an endpoint to expose our profiling data:

func handleProfile(w http.ResponseWriter, r *http.Request) {
    requestMutex.RLock()
    queryMutex.RLock()
    defer requestMutex.RUnlock()
    defer queryMutex.RUnlock()

    profile := map[string]interface{}{
        "requestDurations": requestDurations,
        "queryStats":       queryStats,
        "memStats":         getMemStats(),
    }

    w.Header().Set("Content-Type", "application/json")
    json.NewEncoder(w).Encode(profile)
}

// In main():
http.HandleFunc("/debug/profile", handleProfile)

Putting It All Together

Now that we have our profiler components, let's integrate them into our main application:

func main() {
    // ... (previous database initialization code) ...

    // Set up profiled routes
    http.HandleFunc("/user", trackRequestDuration(handleUser))
    http.HandleFunc("/debug/profile", handleProfile)

    // Start memory stats logging
    go func() {
        ticker := time.NewTicker(1 * time.Minute)
        for range ticker.C {
            logMemStats()
        }
    }()

    // Start the server
    log.Println("Server starting on :8080")
    log.Fatal(http.ListenAndServe(":8080", nil))
}

Using Our Custom Profiler

To gain insights into your web service:

Run your web service as usual.
Generate some traffic to your /user endpoint.
Visit http://localhost:8080/debug/profile to view the profiling data.

Analyzing the Results

With this custom profiler, you can now:

Identify your slowest endpoints (check requestDurations).
Pinpoint problematic database queries (examine queryStats).
Monitor memory usage trends over time (review memStats).

Pro Tips

Sampling: For high-traffic services, consider sampling your requests to reduce overhead.
Alerting: Set up alerts based on your profiling data to catch performance issues early.
Visualization: Use tools like Grafana to create dashboards from your profiling data.
Continuous Profiling: Implement a system to continuously collect and analyze profiling data in production.

Conclusion

We've built a custom profiler tailored to our Go web service needs, allowing us to gather specific insights that generic profilers might miss. This targeted approach empowers you to make informed optimizations and deliver a faster, more efficient application.

Remember, while custom profiling is powerful, it does add some overhead. Use it judiciously, especially in production environments. Start with development and staging environments, and gradually roll out to production as you refine your profiling strategy.

By understanding the unique performance characteristics of your Go web service, you're now equipped to take your optimization game to the next level. Happy profiling!

How did you like this deep dive into custom Go profiling? Let me know in the comments, and don't forget to share your own profiling tips and tricks!

Blog