perf: Serialize concurrent certificate generation (#210)

* perf: Serialize concurrent certificate generation

* fix: Code review fixes

* fix: Code review fixes

go.mod

@@ -19,6 +19,7 @@ require (
 	github.com/spf13/pflag v1.0.10
 	github.com/spf13/viper v1.21.0
 	github.com/stretchr/testify v1.11.1
+	golang.org/x/sync v0.20.0
 	golang.org/x/term v0.39.0
 	google.golang.org/grpc v1.75.0
 	gopkg.in/yaml.v3 v3.0.1

go.sum

@@ -240,6 +240,8 @@ golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
+golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

proxy/certmanager/certmanager_test.go

@@ -3,13 +3,16 @@ package certmanager
 import (
 	"crypto/tls"
 	"crypto/x509"
+	"fmt"
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestGenerateCA(t *testing.T) {
@@ -202,6 +205,45 @@ func TestNewCertificateManagerWithNilCA(t *testing.T) {
 	}
 }
 
+func TestConcurrentCertGeneration(t *testing.T) {
+	config := DefaultCertManagerConfig()
+	ca, err := GenerateCA(config)
+	require.NoError(t, err)
+
+	cm, err := NewCertificateManagerWithCA(ca, config)
+	require.NoError(t, err)
+
+	hostname := "registry.npmjs.org"
+	const goroutines = 50
+
+	certs := make([]*Certificate, goroutines)
+	errs := make([]error, goroutines)
+
+	var wg sync.WaitGroup
+	wg.Add(goroutines)
+
+	// Launch many goroutines requesting the same hostname concurrently.
+	// Singleflight should ensure only one RSA key generation happens.
+	for i := 0; i < goroutines; i++ {
+		go func(idx int) {
+			defer wg.Done()
+			certs[idx], errs[idx] = cm.GenerateCertForHost(hostname)
+		}(i)
+	}
+	wg.Wait()
+
+	for i := 0; i < goroutines; i++ {
+		require.NoError(t, errs[i], "goroutine %d should not error", i)
+		require.NotNil(t, certs[i], "goroutine %d should get a certificate", i)
+	}
+
+	// All goroutines should receive the same cached certificate.
+	for i := 1; i < goroutines; i++ {
+		assert.Equal(t, certs[0].X509Cert.SerialNumber, certs[i].X509Cert.SerialNumber,
+			"goroutine %d should get the same certificate as goroutine 0", i)
+	}
+}
+
 func BenchmarkGenerateCA(b *testing.B) {
 	config := DefaultCertManagerConfig()
 
@@ -227,6 +269,50 @@ func BenchmarkGenerateHostCert(b *testing.B) {
 	}
 }
 
+func BenchmarkGenerateHostCertUncached(b *testing.B) {
+	config := DefaultCertManagerConfig()
+	ca, err := GenerateCA(config)
+	require.NoError(b, err, "Failed to generate CA")
+
+	cm, err := NewCertificateManagerWithCA(ca, config)
+	require.NoError(b, err, "Failed to create certificate manager")
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		hostname := fmt.Sprintf("host-%d.example.com", i)
+		_, err := cm.GenerateCertForHost(hostname)
+		assert.NoError(b, err, "Failed to generate host certificate")
+	}
+}
+
+func BenchmarkGetTLSConfig(b *testing.B) {
+	config := DefaultCertManagerConfig()
+	ca, err := GenerateCA(config)
+	require.NoError(b, err, "Failed to generate CA")
+
+	cm, err := NewCertificateManagerWithCA(ca, config)
+	require.NoError(b, err, "Failed to create certificate manager")
+
+	b.Run("Uncached", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			hostname := fmt.Sprintf("host-%d.example.com", i)
+			_, err := cm.GetTLSConfig(hostname)
+			assert.NoError(b, err, "Failed to get TLS config")
+		}
+	})
+
+	b.Run("Cached", func(b *testing.B) {
+		_, err := cm.GetTLSConfig("cached.example.com")
+		assert.NoError(b, err)
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			_, err := cm.GetTLSConfig("cached.example.com")
+			assert.NoError(b, err, "Failed to get TLS config")
+		}
+	})
+}
+
 func BenchmarkCacheOperations(b *testing.B) {
 	cache := NewInMemoryCache()
 	cert := &Certificate{

proxy/certmanager/generator.go

@@ -15,10 +15,12 @@ import (
 	"time"
 
 	"github.com/safedep/dry/log"
+	"golang.org/x/sync/singleflight"
 )
 
 const (
 	maxSystemCABundleBytes int64         = 10 * 1024 * 1024
+	certExpiryThreshold    time.Duration = 1 * time.Hour
 	goosDarwin                           = "darwin"
 	goosLinux                            = "linux"
 	goosWindows                          = "windows"
@@ -29,6 +31,7 @@ type certManager struct {
 	ca     *Certificate
 	cache  CertificateCache
 	config CertManagerConfig
+	group  singleflight.Group
 }
 
 // NewCertificateManagerWithCA creates a new certificate manager with an existing CA certificate
@@ -51,7 +54,7 @@ func NewCertificateManagerWithCA(ca *Certificate, config CertManagerConfig) (Cer
 		ca = parsedCA
 	}
 
-	if ca.IsExpired(1 * time.Hour) {
+	if ca.IsExpired(certExpiryThreshold) {
 		return nil, fmt.Errorf("CA certificate is expired")
 	}
 
@@ -67,11 +70,22 @@ func (cm *certManager) GetCA() (*Certificate, error) {
 	return cm.ca, nil
 }
 
-// GenerateCertForHost creates a certificate for the given hostname
+// GenerateCertForHost creates a certificate for the given hostname.
-// Uses caching to avoid regeneration
+// Uses caching and singleflight to ensure only one goroutine generates
+// a certificate for a given hostname at a time, preventing CPU starvation
+// when many concurrent CONNECT requests arrive for the same host.
 func (cm *certManager) GenerateCertForHost(hostname string) (*Certificate, error) {
 	if cached, found := cm.cache.Get(hostname); found {
-		if !cached.IsExpired(1 * time.Hour) {
+		if !cached.IsExpired(certExpiryThreshold) {
+			return cached, nil
+		}
+	}
+
+	result, err, _ := cm.group.Do(hostname, func() (interface{}, error) {
+		// Re-check cache: another goroutine in a previous singleflight
+		// group may have populated it while we were waiting.
+		if cached, found := cm.cache.Get(hostname); found {
+			if !cached.IsExpired(certExpiryThreshold) {
 				return cached, nil
 			}
 		}
@@ -82,8 +96,13 @@ func (cm *certManager) GenerateCertForHost(hostname string) (*Certificate, error
 		}
 
 		cm.cache.Set(hostname, cert)
-
 		return cert, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return result.(*Certificate), nil
 }
 
 // GetTLSConfig returns a tls.Config for the given hostname