mirror of
https://github.com/lukaszraczylo/traefikoidc.git
synced 2026-10-13 04:07:37 +00:00
Complete rebuild of the plugin
* Fix bug affecting Azure OIDC authentication ( and most likely others ) * Fixes issue #51 * Ensure that appended roles are unique. Update the documentation. * Improvements targetting possible memory usage spikes. * Additional fixes and cleanup * Refactoring code to fix the issues identified by the users. * Modernize run * Fieldalignment * Multiple changes to improve performance and reduce complexity. - Optimise the errors and recovery. - Deduplicate code in metadata cache. - Remove unused performance monitoring code. - Simplify session management and settings handling. * Fix claims issue. * Add ability to overwrite the default scopes in the settings file * Well.. that escalated quickly. Completely forgot that Traefik uses outdated Yaegi and requires compatibility with 1.20 ( pre-generic Go code ). * Bugfix #51: Ensures that user provided scopes overrides work. * fixup! Bugfix #51: Ensures that user provided scopes overrides work. * fixup! fixup! Bugfix #51: Ensures that user provided scopes overrides work. * Abstract the provider logic into a separate package. * Additional micro fixes and cleanups. * Simplify all the things. * fixup! Simplify all the things. * fixup! fixup! Simplify all the things. * fixup! fixup! fixup! Simplify all the things. * fixup! fixup! fixup! fixup! Simplify all the things. * ... * Cleanup tests. * fixup! Cleanup tests. * fixup! fixup! fixup! Cleanup tests. * fixup! fixup! fixup! fixup! Cleanup tests. * fixup! fixup! fixup! fixup! fixup! Cleanup tests. * Issue #53: Fix CSRF token handling in reverse proxy 1. ✅ HTTPS Detection Fixed (session.go:723) - Now uses X-Forwarded-Proto header instead of r.URL.Scheme - Properly detects HTTPS in reverse proxy environments 2. ✅ SameSite Cookie Attribute Fixed - Removed automatic SameSiteStrictMode for HTTPS (would break OAuth) - Keeps SameSiteLaxMode to allow OAuth callbacks from external domains - Only uses Strict for AJAX requests which don't involve OAuth redirects 3. ✅ Cookie Domain Handling Fixed - Now respects X-Forwarded-Host header for cookie domain - Ensures cookies are set for the public domain, not internal proxy domain 4. ✅ EnhanceSessionSecurity Properly Integrated - Function is now actually called during session save - Applies security enhancements without breaking OAuth flow Why Issue #53 Failed Before: 1. Cookies were not marked Secure in HTTPS environments (browser wouldn't send them back) 2. If they had been Secure with SameSite=Strict, Azure callbacks would still fail 3. Cookie domain might have been wrong (internal vs public domain) Why It Works Now: 1. Cookies are properly marked Secure for HTTPS 2. Uses SameSite=Lax to allow OAuth provider callbacks 3. Cookie domain uses public domain from X-Forwarded-Host 4. CSRF token persists through the entire OAuth flow * Next set of enhancements together with memory usage improvements. * Memory leak fixes and optimisations. * CSRF and Cookie Domain fixes * fixup! CSRF and Cookie Domain fixes * Metadata cache leak fix + profiling * fixup! Metadata cache leak fix + profiling * Memory leaks hunting, part 1337. * Further pursue of perfection. * fixup! Further pursue of perfection. * fixup! fixup! Further pursue of perfection. * fixup! fixup! fixup! Further pursue of perfection. * fixup! fixup! fixup! fixup! Further pursue of perfection. * fixup! fixup! fixup! fixup! fixup! Further pursue of perfection. * fixup! fixup! fixup! fixup! fixup! fixup! Further pursue of perfection. * fixup! fixup! fixup! fixup! fixup! fixup! fixup! Further pursue of perfection. * fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! Further pursue of perfection. * fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! fixup! Further pursue of perfection. * Clear race conditions * fixup! Clear race conditions * Weekend fun with memory leaks * Splitting code into multiple files with reasonable testing coverage. ``` ok github.com/lukaszraczylo/traefikoidc 117.017s coverage: 72.6% of statements ok github.com/lukaszraczylo/traefikoidc/auth 0.505s coverage: 87.1% of statements ok github.com/lukaszraczylo/traefikoidc/circuit_breaker 0.283s coverage: 99.0% of statements github.com/lukaszraczylo/traefikoidc/config coverage: 0.0% of statements ok github.com/lukaszraczylo/traefikoidc/handlers 0.349s coverage: 98.2% of statements ok github.com/lukaszraczylo/traefikoidc/internal/providers (cached) coverage: 94.3% of statements ok github.com/lukaszraczylo/traefikoidc/middleware 0.808s coverage: 78.0% of statements ok github.com/lukaszraczylo/traefikoidc/recovery 0.653s coverage: 100.0% of statements ok github.com/lukaszraczylo/traefikoidc/session/chunking (cached) coverage: 87.8% of statements ok github.com/lukaszraczylo/traefikoidc/session/core (cached) coverage: 85.6% of statements ok github.com/lukaszraczylo/traefikoidc/session/crypto (cached) coverage: 81.8% of statements ok github.com/lukaszraczylo/traefikoidc/session/storage (cached) coverage: 93.5% of statements ok github.com/lukaszraczylo/traefikoidc/session/validators (cached) coverage: 98.8% of statements ```` * fixup! Splitting code into multiple files with reasonable testing coverage. * fixup! fixup! Splitting code into multiple files with reasonable testing coverage. * Weekend fun with further optimisations. * fixup! Weekend fun with further optimisations. * fixup! fixup! Weekend fun with further optimisations. * fixup! fixup! fixup! Weekend fun with further optimisations. * fixup! fixup! fixup! fixup! Weekend fun with further optimisations. * fixup! fixup! fixup! fixup! fixup! Weekend fun with further optimisations. * Pre-release cleanup. * Enhance test coverage. * fixup! Enhance test coverage. * fixup! fixup! Enhance test coverage. * fixup! fixup! fixup! Enhance test coverage.
This commit is contained in:
@@ -0,0 +1,473 @@
|
||||
// Package pool provides a unified, centralized memory pool management system
|
||||
// for the entire application. It consolidates all duplicate pool implementations
|
||||
// into a single, efficient, and thread-safe package.
|
||||
package pool
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// Manager is the centralized pool manager that consolidates all memory pools
|
||||
// used throughout the application. It provides a single entry point for
|
||||
// all pooling operations, reducing duplicate code and improving maintainability.
|
||||
type Manager struct {
|
||||
// Buffer pools
|
||||
smallBufferPool *sync.Pool // 1KB buffers
|
||||
mediumBufferPool *sync.Pool // 4KB buffers
|
||||
largeBufferPool *sync.Pool // 8KB buffers
|
||||
xlBufferPool *sync.Pool // 16KB buffers
|
||||
|
||||
// Compression pools
|
||||
gzipWriterPool *sync.Pool
|
||||
gzipReaderPool *sync.Pool
|
||||
|
||||
// String builder pool
|
||||
stringBuilderPool *sync.Pool
|
||||
|
||||
// JWT parsing buffers
|
||||
jwtBufferPool *sync.Pool
|
||||
|
||||
// HTTP response buffers
|
||||
httpResponsePool *sync.Pool
|
||||
|
||||
// Byte slice pools for various sizes
|
||||
byteSlicePools map[int]*sync.Pool
|
||||
poolMu sync.RWMutex
|
||||
|
||||
// Statistics
|
||||
stats PoolStats
|
||||
}
|
||||
|
||||
// PoolStats tracks pool usage statistics
|
||||
type PoolStats struct {
|
||||
BufferGets uint64
|
||||
BufferPuts uint64
|
||||
GzipGets uint64
|
||||
GzipPuts uint64
|
||||
StringGets uint64
|
||||
StringPuts uint64
|
||||
JWTGets uint64
|
||||
JWTPuts uint64
|
||||
HTTPGets uint64
|
||||
HTTPPuts uint64
|
||||
OversizedRejects uint64
|
||||
}
|
||||
|
||||
// JWTBuffer provides pre-allocated buffers for JWT parsing
|
||||
type JWTBuffer struct {
|
||||
Header []byte
|
||||
Payload []byte
|
||||
Signature []byte
|
||||
}
|
||||
|
||||
var (
|
||||
// globalManager is the singleton pool manager instance
|
||||
globalManager *Manager
|
||||
// managerOnce ensures single initialization
|
||||
managerOnce sync.Once
|
||||
)
|
||||
|
||||
// Get returns the global pool manager instance
|
||||
func Get() *Manager {
|
||||
managerOnce.Do(func() {
|
||||
globalManager = newManager()
|
||||
})
|
||||
return globalManager
|
||||
}
|
||||
|
||||
// newManager creates a new pool manager with all pools initialized
|
||||
func newManager() *Manager {
|
||||
m := &Manager{
|
||||
byteSlicePools: make(map[int]*sync.Pool),
|
||||
}
|
||||
|
||||
// Initialize buffer pools with different sizes
|
||||
m.smallBufferPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
return bytes.NewBuffer(make([]byte, 0, 1024))
|
||||
},
|
||||
}
|
||||
|
||||
m.mediumBufferPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
return bytes.NewBuffer(make([]byte, 0, 4096))
|
||||
},
|
||||
}
|
||||
|
||||
m.largeBufferPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
return bytes.NewBuffer(make([]byte, 0, 8192))
|
||||
},
|
||||
}
|
||||
|
||||
m.xlBufferPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
return bytes.NewBuffer(make([]byte, 0, 16384))
|
||||
},
|
||||
}
|
||||
|
||||
// Initialize compression pools
|
||||
m.gzipWriterPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
w, _ := gzip.NewWriterLevel(nil, gzip.BestSpeed)
|
||||
return w
|
||||
},
|
||||
}
|
||||
|
||||
m.gzipReaderPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
return (*gzip.Reader)(nil)
|
||||
},
|
||||
}
|
||||
|
||||
// Initialize string builder pool
|
||||
m.stringBuilderPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
sb := &strings.Builder{}
|
||||
sb.Grow(1024)
|
||||
return sb
|
||||
},
|
||||
}
|
||||
|
||||
// Initialize JWT buffer pool
|
||||
m.jwtBufferPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
return &JWTBuffer{
|
||||
Header: make([]byte, 0, 512),
|
||||
Payload: make([]byte, 0, 2048),
|
||||
Signature: make([]byte, 0, 512),
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
// Initialize HTTP response buffer pool
|
||||
m.httpResponsePool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
buf := make([]byte, 0, 8192)
|
||||
return &buf
|
||||
},
|
||||
}
|
||||
|
||||
// Initialize common byte slice pools
|
||||
for _, size := range []int{256, 512, 1024, 2048, 4096, 8192, 16384} {
|
||||
size := size // capture for closure
|
||||
m.byteSlicePools[size] = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
b := make([]byte, size)
|
||||
return &b
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
return m
|
||||
}
|
||||
|
||||
// GetBuffer returns a buffer from the appropriate pool based on size hint
|
||||
func (m *Manager) GetBuffer(sizeHint int) *bytes.Buffer {
|
||||
atomic.AddUint64(&m.stats.BufferGets, 1)
|
||||
|
||||
switch {
|
||||
case sizeHint <= 1024:
|
||||
return m.smallBufferPool.Get().(*bytes.Buffer)
|
||||
case sizeHint <= 4096:
|
||||
return m.mediumBufferPool.Get().(*bytes.Buffer)
|
||||
case sizeHint <= 8192:
|
||||
return m.largeBufferPool.Get().(*bytes.Buffer)
|
||||
case sizeHint <= 16384:
|
||||
return m.xlBufferPool.Get().(*bytes.Buffer)
|
||||
default:
|
||||
// For very large buffers, create new ones
|
||||
return bytes.NewBuffer(make([]byte, 0, sizeHint))
|
||||
}
|
||||
}
|
||||
|
||||
// PutBuffer returns a buffer to the appropriate pool
|
||||
func (m *Manager) PutBuffer(buf *bytes.Buffer) {
|
||||
if buf == nil {
|
||||
return
|
||||
}
|
||||
|
||||
atomic.AddUint64(&m.stats.BufferPuts, 1)
|
||||
|
||||
// Reset buffer before returning to pool
|
||||
capacity := buf.Cap()
|
||||
buf.Reset()
|
||||
|
||||
// Reject oversized buffers to prevent memory bloat
|
||||
if capacity > 32768 {
|
||||
atomic.AddUint64(&m.stats.OversizedRejects, 1)
|
||||
return
|
||||
}
|
||||
|
||||
// Return to appropriate pool based on capacity
|
||||
switch {
|
||||
case capacity <= 1024:
|
||||
m.smallBufferPool.Put(buf)
|
||||
case capacity <= 4096:
|
||||
m.mediumBufferPool.Put(buf)
|
||||
case capacity <= 8192:
|
||||
m.largeBufferPool.Put(buf)
|
||||
case capacity <= 16384:
|
||||
m.xlBufferPool.Put(buf)
|
||||
}
|
||||
}
|
||||
|
||||
// GetGzipWriter returns a gzip writer from the pool
|
||||
func (m *Manager) GetGzipWriter() *gzip.Writer {
|
||||
atomic.AddUint64(&m.stats.GzipGets, 1)
|
||||
return m.gzipWriterPool.Get().(*gzip.Writer)
|
||||
}
|
||||
|
||||
// PutGzipWriter returns a gzip writer to the pool
|
||||
func (m *Manager) PutGzipWriter(w *gzip.Writer) {
|
||||
if w == nil {
|
||||
return
|
||||
}
|
||||
atomic.AddUint64(&m.stats.GzipPuts, 1)
|
||||
w.Reset(nil)
|
||||
m.gzipWriterPool.Put(w)
|
||||
}
|
||||
|
||||
// GetGzipReader returns a gzip reader from the pool
|
||||
func (m *Manager) GetGzipReader() *gzip.Reader {
|
||||
atomic.AddUint64(&m.stats.GzipGets, 1)
|
||||
r := m.gzipReaderPool.Get()
|
||||
if r == nil {
|
||||
return nil
|
||||
}
|
||||
return r.(*gzip.Reader)
|
||||
}
|
||||
|
||||
// PutGzipReader returns a gzip reader to the pool
|
||||
func (m *Manager) PutGzipReader(r *gzip.Reader) {
|
||||
if r == nil {
|
||||
return
|
||||
}
|
||||
atomic.AddUint64(&m.stats.GzipPuts, 1)
|
||||
r.Reset(nil)
|
||||
m.gzipReaderPool.Put(r)
|
||||
}
|
||||
|
||||
// GetStringBuilder returns a string builder from the pool
|
||||
func (m *Manager) GetStringBuilder() *strings.Builder {
|
||||
atomic.AddUint64(&m.stats.StringGets, 1)
|
||||
sb := m.stringBuilderPool.Get().(*strings.Builder)
|
||||
sb.Reset()
|
||||
return sb
|
||||
}
|
||||
|
||||
// PutStringBuilder returns a string builder to the pool
|
||||
func (m *Manager) PutStringBuilder(sb *strings.Builder) {
|
||||
if sb == nil {
|
||||
return
|
||||
}
|
||||
|
||||
atomic.AddUint64(&m.stats.StringPuts, 1)
|
||||
|
||||
// Reject oversized builders
|
||||
if sb.Cap() > 16384 {
|
||||
atomic.AddUint64(&m.stats.OversizedRejects, 1)
|
||||
return
|
||||
}
|
||||
|
||||
sb.Reset()
|
||||
m.stringBuilderPool.Put(sb)
|
||||
}
|
||||
|
||||
// GetJWTBuffer returns JWT parsing buffers from the pool
|
||||
func (m *Manager) GetJWTBuffer() *JWTBuffer {
|
||||
atomic.AddUint64(&m.stats.JWTGets, 1)
|
||||
return m.jwtBufferPool.Get().(*JWTBuffer)
|
||||
}
|
||||
|
||||
// PutJWTBuffer returns JWT parsing buffers to the pool
|
||||
func (m *Manager) PutJWTBuffer(buf *JWTBuffer) {
|
||||
if buf == nil {
|
||||
return
|
||||
}
|
||||
|
||||
atomic.AddUint64(&m.stats.JWTPuts, 1)
|
||||
|
||||
// Check for oversized buffers
|
||||
if cap(buf.Header) > 2048 || cap(buf.Payload) > 8192 || cap(buf.Signature) > 2048 {
|
||||
atomic.AddUint64(&m.stats.OversizedRejects, 1)
|
||||
return
|
||||
}
|
||||
|
||||
// Reset slices to zero length
|
||||
buf.Header = buf.Header[:0]
|
||||
buf.Payload = buf.Payload[:0]
|
||||
buf.Signature = buf.Signature[:0]
|
||||
m.jwtBufferPool.Put(buf)
|
||||
}
|
||||
|
||||
// GetHTTPResponseBuffer returns an HTTP response buffer from the pool
|
||||
func (m *Manager) GetHTTPResponseBuffer() []byte {
|
||||
atomic.AddUint64(&m.stats.HTTPGets, 1)
|
||||
return *m.httpResponsePool.Get().(*[]byte)
|
||||
}
|
||||
|
||||
// PutHTTPResponseBuffer returns an HTTP response buffer to the pool
|
||||
func (m *Manager) PutHTTPResponseBuffer(buf []byte) {
|
||||
if buf == nil {
|
||||
return
|
||||
}
|
||||
|
||||
atomic.AddUint64(&m.stats.HTTPPuts, 1)
|
||||
|
||||
// Reject oversized buffers
|
||||
if cap(buf) > 32768 {
|
||||
atomic.AddUint64(&m.stats.OversizedRejects, 1)
|
||||
return
|
||||
}
|
||||
|
||||
buf = buf[:0]
|
||||
m.httpResponsePool.Put(&buf)
|
||||
}
|
||||
|
||||
// GetByteSlice returns a byte slice of the specified size from the pool
|
||||
func (m *Manager) GetByteSlice(size int) []byte {
|
||||
m.poolMu.RLock()
|
||||
pool, exists := m.byteSlicePools[size]
|
||||
m.poolMu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
// Round up to nearest power of 2
|
||||
poolSize := 1
|
||||
for poolSize < size {
|
||||
poolSize *= 2
|
||||
}
|
||||
|
||||
m.poolMu.Lock()
|
||||
// Double-check after acquiring write lock
|
||||
pool, exists = m.byteSlicePools[poolSize]
|
||||
if !exists {
|
||||
pool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
b := make([]byte, poolSize)
|
||||
return &b
|
||||
},
|
||||
}
|
||||
m.byteSlicePools[poolSize] = pool
|
||||
}
|
||||
m.poolMu.Unlock()
|
||||
}
|
||||
|
||||
b := pool.Get().(*[]byte)
|
||||
return (*b)[:size]
|
||||
}
|
||||
|
||||
// PutByteSlice returns a byte slice to the pool
|
||||
func (m *Manager) PutByteSlice(b []byte) {
|
||||
if b == nil || cap(b) > 65536 { // Don't pool very large slices
|
||||
return
|
||||
}
|
||||
|
||||
size := cap(b)
|
||||
m.poolMu.RLock()
|
||||
pool, exists := m.byteSlicePools[size]
|
||||
m.poolMu.RUnlock()
|
||||
|
||||
if exists {
|
||||
b = b[:0]
|
||||
pool.Put(&b)
|
||||
}
|
||||
}
|
||||
|
||||
// GetStats returns current pool statistics
|
||||
func (m *Manager) GetStats() PoolStats {
|
||||
return PoolStats{
|
||||
BufferGets: atomic.LoadUint64(&m.stats.BufferGets),
|
||||
BufferPuts: atomic.LoadUint64(&m.stats.BufferPuts),
|
||||
GzipGets: atomic.LoadUint64(&m.stats.GzipGets),
|
||||
GzipPuts: atomic.LoadUint64(&m.stats.GzipPuts),
|
||||
StringGets: atomic.LoadUint64(&m.stats.StringGets),
|
||||
StringPuts: atomic.LoadUint64(&m.stats.StringPuts),
|
||||
JWTGets: atomic.LoadUint64(&m.stats.JWTGets),
|
||||
JWTPuts: atomic.LoadUint64(&m.stats.JWTPuts),
|
||||
HTTPGets: atomic.LoadUint64(&m.stats.HTTPGets),
|
||||
HTTPPuts: atomic.LoadUint64(&m.stats.HTTPPuts),
|
||||
OversizedRejects: atomic.LoadUint64(&m.stats.OversizedRejects),
|
||||
}
|
||||
}
|
||||
|
||||
// ResetStats resets all statistics counters
|
||||
func (m *Manager) ResetStats() {
|
||||
atomic.StoreUint64(&m.stats.BufferGets, 0)
|
||||
atomic.StoreUint64(&m.stats.BufferPuts, 0)
|
||||
atomic.StoreUint64(&m.stats.GzipGets, 0)
|
||||
atomic.StoreUint64(&m.stats.GzipPuts, 0)
|
||||
atomic.StoreUint64(&m.stats.StringGets, 0)
|
||||
atomic.StoreUint64(&m.stats.StringPuts, 0)
|
||||
atomic.StoreUint64(&m.stats.JWTGets, 0)
|
||||
atomic.StoreUint64(&m.stats.JWTPuts, 0)
|
||||
atomic.StoreUint64(&m.stats.HTTPGets, 0)
|
||||
atomic.StoreUint64(&m.stats.HTTPPuts, 0)
|
||||
atomic.StoreUint64(&m.stats.OversizedRejects, 0)
|
||||
}
|
||||
|
||||
// Global convenience functions
|
||||
|
||||
// Buffer returns a buffer from the global pool
|
||||
func Buffer(sizeHint int) *bytes.Buffer {
|
||||
return Get().GetBuffer(sizeHint)
|
||||
}
|
||||
|
||||
// ReturnBuffer returns a buffer to the global pool
|
||||
func ReturnBuffer(buf *bytes.Buffer) {
|
||||
Get().PutBuffer(buf)
|
||||
}
|
||||
|
||||
// GzipWriter returns a gzip writer from the global pool
|
||||
func GzipWriter() *gzip.Writer {
|
||||
return Get().GetGzipWriter()
|
||||
}
|
||||
|
||||
// ReturnGzipWriter returns a gzip writer to the global pool
|
||||
func ReturnGzipWriter(w *gzip.Writer) {
|
||||
Get().PutGzipWriter(w)
|
||||
}
|
||||
|
||||
// StringBuilder returns a string builder from the global pool
|
||||
func StringBuilder() *strings.Builder {
|
||||
return Get().GetStringBuilder()
|
||||
}
|
||||
|
||||
// ReturnStringBuilder returns a string builder to the global pool
|
||||
func ReturnStringBuilder(sb *strings.Builder) {
|
||||
Get().PutStringBuilder(sb)
|
||||
}
|
||||
|
||||
// JWTBuffers returns JWT parsing buffers from the global pool
|
||||
func JWTBuffers() *JWTBuffer {
|
||||
return Get().GetJWTBuffer()
|
||||
}
|
||||
|
||||
// ReturnJWTBuffers returns JWT parsing buffers to the global pool
|
||||
func ReturnJWTBuffers(buf *JWTBuffer) {
|
||||
Get().PutJWTBuffer(buf)
|
||||
}
|
||||
|
||||
// HTTPBuffer returns an HTTP response buffer from the global pool
|
||||
func HTTPBuffer() []byte {
|
||||
return Get().GetHTTPResponseBuffer()
|
||||
}
|
||||
|
||||
// ReturnHTTPBuffer returns an HTTP response buffer to the global pool
|
||||
func ReturnHTTPBuffer(buf []byte) {
|
||||
Get().PutHTTPResponseBuffer(buf)
|
||||
}
|
||||
|
||||
// ByteSlice returns a byte slice from the global pool
|
||||
func ByteSlice(size int) []byte {
|
||||
return Get().GetByteSlice(size)
|
||||
}
|
||||
|
||||
// ReturnByteSlice returns a byte slice to the global pool
|
||||
func ReturnByteSlice(b []byte) {
|
||||
Get().PutByteSlice(b)
|
||||
}
|
||||
@@ -0,0 +1,586 @@
|
||||
package pool
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestManager_Singleton tests that Get() returns the same instance
|
||||
func TestManager_Singleton(t *testing.T) {
|
||||
manager1 := Get()
|
||||
manager2 := Get()
|
||||
|
||||
if manager1 != manager2 {
|
||||
t.Error("Get() should return the same instance (singleton)")
|
||||
}
|
||||
|
||||
if manager1 == nil {
|
||||
t.Error("Get() should not return nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_BufferPools tests buffer pool operations
|
||||
func TestManager_BufferPools(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
sizeHint int
|
||||
expected int // expected capacity range
|
||||
}{
|
||||
{"small buffer", 512, 1024},
|
||||
{"medium buffer", 2048, 4096},
|
||||
{"large buffer", 6144, 8192},
|
||||
{"xl buffer", 12288, 16384},
|
||||
{"oversized buffer", 32768, 32768}, // Should create new buffer
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
buf := manager.GetBuffer(test.sizeHint)
|
||||
if buf == nil {
|
||||
t.Error("GetBuffer should not return nil")
|
||||
}
|
||||
|
||||
if buf.Cap() < test.sizeHint {
|
||||
t.Errorf("Buffer capacity %d is less than size hint %d", buf.Cap(), test.sizeHint)
|
||||
}
|
||||
|
||||
// Write some data
|
||||
buf.WriteString("test data")
|
||||
if buf.String() != "test data" {
|
||||
t.Error("Buffer should contain written data")
|
||||
}
|
||||
|
||||
// Return to pool
|
||||
manager.PutBuffer(buf)
|
||||
|
||||
// Buffer should be reset when returned to pool
|
||||
buf2 := manager.GetBuffer(test.sizeHint)
|
||||
if buf2.Len() != 0 {
|
||||
t.Error("Buffer from pool should be reset")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_PutBuffer_Nil tests putting nil buffer
|
||||
func TestManager_PutBuffer_Nil(t *testing.T) {
|
||||
manager := Get()
|
||||
// Should not panic
|
||||
manager.PutBuffer(nil)
|
||||
}
|
||||
|
||||
// TestManager_PutBuffer_Oversized tests rejection of oversized buffers
|
||||
func TestManager_PutBuffer_Oversized(t *testing.T) {
|
||||
manager := Get()
|
||||
manager.ResetStats()
|
||||
|
||||
// Create oversized buffer
|
||||
buf := bytes.NewBuffer(make([]byte, 0, 40000))
|
||||
manager.PutBuffer(buf)
|
||||
|
||||
stats := manager.GetStats()
|
||||
if stats.OversizedRejects == 0 {
|
||||
t.Error("Oversized buffer should be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_GzipPools tests gzip writer and reader pools
|
||||
func TestManager_GzipPools(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
// Test gzip writer
|
||||
writer := manager.GetGzipWriter()
|
||||
if writer == nil {
|
||||
t.Error("GetGzipWriter should not return nil")
|
||||
}
|
||||
|
||||
// Test that we can use it
|
||||
var buf bytes.Buffer
|
||||
writer.Reset(&buf)
|
||||
writer.Write([]byte("test data"))
|
||||
writer.Close()
|
||||
|
||||
if buf.Len() == 0 {
|
||||
t.Error("Gzip writer should have written compressed data")
|
||||
}
|
||||
|
||||
// Return to pool
|
||||
manager.PutGzipWriter(writer)
|
||||
|
||||
// Test gzip reader
|
||||
reader := manager.GetGzipReader()
|
||||
// Reader might be nil from pool initially
|
||||
if reader != nil {
|
||||
manager.PutGzipReader(reader)
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_GzipPools_Nil tests putting nil gzip objects
|
||||
func TestManager_GzipPools_Nil(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
// Should not panic
|
||||
manager.PutGzipWriter(nil)
|
||||
manager.PutGzipReader(nil)
|
||||
}
|
||||
|
||||
// TestManager_StringBuilderPool tests string builder pool
|
||||
func TestManager_StringBuilderPool(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
sb := manager.GetStringBuilder()
|
||||
if sb == nil {
|
||||
t.Error("GetStringBuilder should not return nil")
|
||||
}
|
||||
|
||||
// Should be reset
|
||||
if sb.Len() != 0 {
|
||||
t.Error("String builder from pool should be reset")
|
||||
}
|
||||
|
||||
// Test writing
|
||||
sb.WriteString("test")
|
||||
sb.WriteString(" data")
|
||||
if sb.String() != "test data" {
|
||||
t.Error("String builder should contain written data")
|
||||
}
|
||||
|
||||
// Return to pool
|
||||
manager.PutStringBuilder(sb)
|
||||
|
||||
// Get another one - should be reset
|
||||
sb2 := manager.GetStringBuilder()
|
||||
if sb2.Len() != 0 {
|
||||
t.Error("String builder from pool should be reset")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_StringBuilderPool_Nil tests putting nil string builder
|
||||
func TestManager_StringBuilderPool_Nil(t *testing.T) {
|
||||
manager := Get()
|
||||
// Should not panic
|
||||
manager.PutStringBuilder(nil)
|
||||
}
|
||||
|
||||
// TestManager_StringBuilderPool_Oversized tests rejection of oversized string builders
|
||||
func TestManager_StringBuilderPool_Oversized(t *testing.T) {
|
||||
manager := Get()
|
||||
manager.ResetStats()
|
||||
|
||||
// Create oversized string builder
|
||||
sb := &strings.Builder{}
|
||||
sb.Grow(20000)
|
||||
sb.WriteString("test")
|
||||
|
||||
manager.PutStringBuilder(sb)
|
||||
|
||||
stats := manager.GetStats()
|
||||
if stats.OversizedRejects == 0 {
|
||||
t.Error("Oversized string builder should be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_JWTBufferPool tests JWT buffer pool
|
||||
func TestManager_JWTBufferPool(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
jwtBuf := manager.GetJWTBuffer()
|
||||
if jwtBuf == nil {
|
||||
t.Error("GetJWTBuffer should not return nil")
|
||||
return
|
||||
}
|
||||
|
||||
// Check structure
|
||||
if jwtBuf.Header == nil || jwtBuf.Payload == nil || jwtBuf.Signature == nil {
|
||||
t.Error("JWT buffer should have all fields initialized")
|
||||
}
|
||||
|
||||
// Should be empty initially
|
||||
if len(jwtBuf.Header) != 0 || len(jwtBuf.Payload) != 0 || len(jwtBuf.Signature) != 0 {
|
||||
t.Error("JWT buffer from pool should be reset")
|
||||
}
|
||||
|
||||
// Use the buffer
|
||||
jwtBuf.Header = append(jwtBuf.Header, []byte("header")...)
|
||||
jwtBuf.Payload = append(jwtBuf.Payload, []byte("payload")...)
|
||||
jwtBuf.Signature = append(jwtBuf.Signature, []byte("signature")...)
|
||||
|
||||
// Return to pool
|
||||
manager.PutJWTBuffer(jwtBuf)
|
||||
|
||||
// Get another one - should be reset
|
||||
jwtBuf2 := manager.GetJWTBuffer()
|
||||
if len(jwtBuf2.Header) != 0 || len(jwtBuf2.Payload) != 0 || len(jwtBuf2.Signature) != 0 {
|
||||
t.Error("JWT buffer from pool should be reset")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_JWTBufferPool_Nil tests putting nil JWT buffer
|
||||
func TestManager_JWTBufferPool_Nil(t *testing.T) {
|
||||
manager := Get()
|
||||
// Should not panic
|
||||
manager.PutJWTBuffer(nil)
|
||||
}
|
||||
|
||||
// TestManager_JWTBufferPool_Oversized tests rejection of oversized JWT buffers
|
||||
func TestManager_JWTBufferPool_Oversized(t *testing.T) {
|
||||
manager := Get()
|
||||
manager.ResetStats()
|
||||
|
||||
// Create oversized JWT buffer
|
||||
jwtBuf := &JWTBuffer{
|
||||
Header: make([]byte, 0, 3000), // Over 2048 limit
|
||||
Payload: make([]byte, 0, 10000), // Over 8192 limit
|
||||
Signature: make([]byte, 0, 3000), // Over 2048 limit
|
||||
}
|
||||
|
||||
manager.PutJWTBuffer(jwtBuf)
|
||||
|
||||
stats := manager.GetStats()
|
||||
if stats.OversizedRejects == 0 {
|
||||
t.Error("Oversized JWT buffer should be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_HTTPResponsePool tests HTTP response buffer pool
|
||||
func TestManager_HTTPResponsePool(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
buf := manager.GetHTTPResponseBuffer()
|
||||
if buf == nil {
|
||||
t.Error("GetHTTPResponseBuffer should not return nil")
|
||||
}
|
||||
|
||||
// Should be empty initially
|
||||
if len(buf) != 0 {
|
||||
t.Error("HTTP buffer from pool should be empty")
|
||||
}
|
||||
|
||||
// Use the buffer
|
||||
buf = append(buf, []byte("HTTP response data")...)
|
||||
|
||||
// Return to pool
|
||||
manager.PutHTTPResponseBuffer(buf)
|
||||
|
||||
// Get another one - should be reset
|
||||
buf2 := manager.GetHTTPResponseBuffer()
|
||||
if len(buf2) != 0 {
|
||||
t.Error("HTTP buffer from pool should be reset")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_HTTPResponsePool_Nil tests putting nil HTTP buffer
|
||||
func TestManager_HTTPResponsePool_Nil(t *testing.T) {
|
||||
manager := Get()
|
||||
// Should not panic
|
||||
manager.PutHTTPResponseBuffer(nil)
|
||||
}
|
||||
|
||||
// TestManager_HTTPResponsePool_Oversized tests rejection of oversized HTTP buffers
|
||||
func TestManager_HTTPResponsePool_Oversized(t *testing.T) {
|
||||
manager := Get()
|
||||
manager.ResetStats()
|
||||
|
||||
// Create oversized buffer
|
||||
buf := make([]byte, 0, 40000)
|
||||
manager.PutHTTPResponseBuffer(buf)
|
||||
|
||||
stats := manager.GetStats()
|
||||
if stats.OversizedRejects == 0 {
|
||||
t.Error("Oversized HTTP buffer should be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_ByteSlicePool tests byte slice pool with dynamic sizing
|
||||
func TestManager_ByteSlicePool(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
tests := []int{256, 512, 1024, 2048, 4096, 8192, 16384}
|
||||
|
||||
for _, size := range tests {
|
||||
t.Run(strings.Join([]string{"size", string(rune(size))}, "_"), func(t *testing.T) {
|
||||
slice := manager.GetByteSlice(size)
|
||||
if slice == nil {
|
||||
t.Error("GetByteSlice should not return nil")
|
||||
}
|
||||
|
||||
if len(slice) != size {
|
||||
t.Errorf("Byte slice length %d != requested size %d", len(slice), size)
|
||||
}
|
||||
|
||||
if cap(slice) < size {
|
||||
t.Errorf("Byte slice capacity %d < requested size %d", cap(slice), size)
|
||||
}
|
||||
|
||||
// Use the slice
|
||||
copy(slice, []byte("test data"))
|
||||
|
||||
// Return to pool
|
||||
manager.PutByteSlice(slice)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_ByteSlicePool_CustomSize tests byte slice pool with non-standard sizes
|
||||
func TestManager_ByteSlicePool_CustomSize(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
// Test custom size (should round up to power of 2)
|
||||
slice := manager.GetByteSlice(300)
|
||||
if slice == nil {
|
||||
t.Error("GetByteSlice should not return nil")
|
||||
}
|
||||
|
||||
if len(slice) != 300 {
|
||||
t.Errorf("Byte slice length %d != requested size 300", len(slice))
|
||||
}
|
||||
|
||||
// Capacity should be >= 300 (likely 512 as next power of 2)
|
||||
if cap(slice) < 300 {
|
||||
t.Error("Byte slice capacity should be at least 300")
|
||||
}
|
||||
|
||||
manager.PutByteSlice(slice)
|
||||
}
|
||||
|
||||
// TestManager_ByteSlicePool_Nil tests putting nil byte slice
|
||||
func TestManager_ByteSlicePool_Nil(t *testing.T) {
|
||||
manager := Get()
|
||||
// Should not panic
|
||||
manager.PutByteSlice(nil)
|
||||
}
|
||||
|
||||
// TestManager_ByteSlicePool_Oversized tests rejection of oversized byte slices
|
||||
func TestManager_ByteSlicePool_Oversized(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
// Create oversized slice
|
||||
slice := make([]byte, 100000)
|
||||
|
||||
// Should not panic and should not be pooled
|
||||
manager.PutByteSlice(slice)
|
||||
}
|
||||
|
||||
// TestManager_Stats tests statistics tracking
|
||||
func TestManager_Stats(t *testing.T) {
|
||||
manager := Get()
|
||||
manager.ResetStats()
|
||||
|
||||
initialStats := manager.GetStats()
|
||||
if initialStats.BufferGets != 0 || initialStats.BufferPuts != 0 {
|
||||
t.Error("Stats should be zero after reset")
|
||||
}
|
||||
|
||||
// Perform operations
|
||||
buf := manager.GetBuffer(1024)
|
||||
manager.PutBuffer(buf)
|
||||
|
||||
writer := manager.GetGzipWriter()
|
||||
manager.PutGzipWriter(writer)
|
||||
|
||||
sb := manager.GetStringBuilder()
|
||||
manager.PutStringBuilder(sb)
|
||||
|
||||
jwtBuf := manager.GetJWTBuffer()
|
||||
manager.PutJWTBuffer(jwtBuf)
|
||||
|
||||
httpBuf := manager.GetHTTPResponseBuffer()
|
||||
manager.PutHTTPResponseBuffer(httpBuf)
|
||||
|
||||
// Check stats
|
||||
stats := manager.GetStats()
|
||||
if stats.BufferGets == 0 || stats.BufferPuts == 0 {
|
||||
t.Error("Buffer stats should be incremented")
|
||||
}
|
||||
if stats.GzipGets == 0 || stats.GzipPuts == 0 {
|
||||
t.Error("Gzip stats should be incremented")
|
||||
}
|
||||
if stats.StringGets == 0 || stats.StringPuts == 0 {
|
||||
t.Error("String stats should be incremented")
|
||||
}
|
||||
if stats.JWTGets == 0 || stats.JWTPuts == 0 {
|
||||
t.Error("JWT stats should be incremented")
|
||||
}
|
||||
if stats.HTTPGets == 0 || stats.HTTPPuts == 0 {
|
||||
t.Error("HTTP stats should be incremented")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_ResetStats tests statistics reset
|
||||
func TestManager_ResetStats(t *testing.T) {
|
||||
manager := Get()
|
||||
|
||||
// Perform some operations
|
||||
buf := manager.GetBuffer(1024)
|
||||
manager.PutBuffer(buf)
|
||||
|
||||
// Check that stats are non-zero
|
||||
stats := manager.GetStats()
|
||||
if stats.BufferGets == 0 {
|
||||
t.Error("Stats should be non-zero before reset")
|
||||
}
|
||||
|
||||
// Reset stats
|
||||
manager.ResetStats()
|
||||
|
||||
// Check that stats are zero
|
||||
resetStats := manager.GetStats()
|
||||
if resetStats.BufferGets != 0 || resetStats.BufferPuts != 0 {
|
||||
t.Error("Stats should be zero after reset")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManager_ConcurrentAccess tests concurrent access to pools
|
||||
func TestManager_ConcurrentAccess(t *testing.T) {
|
||||
manager := Get()
|
||||
manager.ResetStats()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
numGoroutines := 50
|
||||
operationsPerGoroutine := 10
|
||||
|
||||
wg.Add(numGoroutines)
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < operationsPerGoroutine; j++ {
|
||||
// Test buffer pool
|
||||
buf := manager.GetBuffer(1024)
|
||||
buf.WriteString("test")
|
||||
manager.PutBuffer(buf)
|
||||
|
||||
// Test string builder pool
|
||||
sb := manager.GetStringBuilder()
|
||||
sb.WriteString("test")
|
||||
manager.PutStringBuilder(sb)
|
||||
|
||||
// Test JWT buffer pool
|
||||
jwtBuf := manager.GetJWTBuffer()
|
||||
jwtBuf.Header = append(jwtBuf.Header, byte(j))
|
||||
manager.PutJWTBuffer(jwtBuf)
|
||||
|
||||
// Test byte slice pool
|
||||
slice := manager.GetByteSlice(256)
|
||||
slice[0] = byte(j)
|
||||
manager.PutByteSlice(slice)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
// Check that operations completed without panic
|
||||
stats := manager.GetStats()
|
||||
expectedOps := uint64(numGoroutines * operationsPerGoroutine)
|
||||
if stats.BufferGets < expectedOps || stats.StringGets < expectedOps || stats.JWTGets < expectedOps {
|
||||
t.Error("Some operations may have failed during concurrent access")
|
||||
}
|
||||
}
|
||||
|
||||
// TestGlobalConvenienceFunctions tests the global convenience functions
|
||||
func TestGlobalConvenienceFunctions(t *testing.T) {
|
||||
// Test buffer functions
|
||||
buf := Buffer(1024)
|
||||
if buf == nil {
|
||||
t.Error("Buffer() should not return nil")
|
||||
}
|
||||
buf.WriteString("test")
|
||||
ReturnBuffer(buf)
|
||||
|
||||
// Test gzip functions
|
||||
writer := GzipWriter()
|
||||
if writer == nil {
|
||||
t.Error("GzipWriter() should not return nil")
|
||||
}
|
||||
ReturnGzipWriter(writer)
|
||||
|
||||
// Test string builder functions
|
||||
sb := StringBuilder()
|
||||
if sb == nil {
|
||||
t.Error("StringBuilder() should not return nil")
|
||||
}
|
||||
sb.WriteString("test")
|
||||
ReturnStringBuilder(sb)
|
||||
|
||||
// Test JWT buffer functions
|
||||
jwtBuf := JWTBuffers()
|
||||
if jwtBuf == nil {
|
||||
t.Error("JWTBuffers() should not return nil")
|
||||
}
|
||||
ReturnJWTBuffers(jwtBuf)
|
||||
|
||||
// Test HTTP buffer functions
|
||||
httpBuf := HTTPBuffer()
|
||||
if httpBuf == nil {
|
||||
t.Error("HTTPBuffer() should not return nil")
|
||||
}
|
||||
ReturnHTTPBuffer(httpBuf)
|
||||
|
||||
// Test byte slice functions
|
||||
slice := ByteSlice(256)
|
||||
if slice == nil {
|
||||
t.Error("ByteSlice() should not return nil")
|
||||
}
|
||||
if len(slice) != 256 {
|
||||
t.Error("ByteSlice() should return correct size")
|
||||
}
|
||||
ReturnByteSlice(slice)
|
||||
}
|
||||
|
||||
// Benchmark tests for performance verification
|
||||
func BenchmarkManager_GetBuffer(b *testing.B) {
|
||||
manager := Get()
|
||||
b.ResetTimer()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
buf := manager.GetBuffer(1024)
|
||||
manager.PutBuffer(buf)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkManager_GetStringBuilder(b *testing.B) {
|
||||
manager := Get()
|
||||
b.ResetTimer()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
sb := manager.GetStringBuilder()
|
||||
manager.PutStringBuilder(sb)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkManager_GetJWTBuffer(b *testing.B) {
|
||||
manager := Get()
|
||||
b.ResetTimer()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
jwtBuf := manager.GetJWTBuffer()
|
||||
manager.PutJWTBuffer(jwtBuf)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkManager_GetByteSlice(b *testing.B) {
|
||||
manager := Get()
|
||||
b.ResetTimer()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
slice := manager.GetByteSlice(1024)
|
||||
manager.PutByteSlice(slice)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkManager_ConcurrentAccess(b *testing.B) {
|
||||
manager := Get()
|
||||
b.ResetTimer()
|
||||
|
||||
b.RunParallel(func(pb *testing.PB) {
|
||||
for pb.Next() {
|
||||
buf := manager.GetBuffer(1024)
|
||||
buf.WriteString("test")
|
||||
manager.PutBuffer(buf)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,370 @@
|
||||
package pool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"net"
|
||||
"net/http"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TransportPool manages a pool of shared HTTP transports to prevent connection exhaustion
|
||||
// and resource leaks. It provides centralized management of HTTP client transports with
|
||||
// proper lifecycle management and security controls.
|
||||
type TransportPool struct {
|
||||
mu sync.RWMutex
|
||||
transports map[string]*sharedTransport
|
||||
maxConns int
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
clientCount int32 // Track total HTTP clients
|
||||
maxClients int32 // Limit total clients
|
||||
}
|
||||
|
||||
// sharedTransport wraps an HTTP transport with reference counting
|
||||
type sharedTransport struct {
|
||||
transport *http.Transport
|
||||
refCount int32
|
||||
lastUsed time.Time
|
||||
config TransportConfig
|
||||
}
|
||||
|
||||
// TransportConfig defines configuration for HTTP transports
|
||||
type TransportConfig struct {
|
||||
// Timeouts
|
||||
DialTimeout time.Duration
|
||||
TLSHandshakeTimeout time.Duration
|
||||
ResponseHeaderTimeout time.Duration
|
||||
ExpectContinueTimeout time.Duration
|
||||
IdleConnTimeout time.Duration
|
||||
KeepAlive time.Duration
|
||||
|
||||
// Connection limits
|
||||
MaxIdleConns int
|
||||
MaxIdleConnsPerHost int
|
||||
MaxConnsPerHost int
|
||||
|
||||
// Features
|
||||
ForceHTTP2 bool
|
||||
DisableKeepAlives bool
|
||||
DisableCompression bool
|
||||
|
||||
// Buffer sizes
|
||||
WriteBufferSize int
|
||||
ReadBufferSize int
|
||||
|
||||
// TLS
|
||||
InsecureSkipVerify bool
|
||||
MinTLSVersion uint16
|
||||
}
|
||||
|
||||
var (
|
||||
// globalTransportPool is the singleton transport pool instance
|
||||
globalTransportPool *TransportPool
|
||||
// transportPoolOnce ensures single initialization
|
||||
transportPoolOnce sync.Once
|
||||
)
|
||||
|
||||
// GetTransportPool returns the global transport pool instance
|
||||
func GetTransportPool() *TransportPool {
|
||||
transportPoolOnce.Do(func() {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
globalTransportPool = &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
clientCount: 0,
|
||||
maxClients: 5,
|
||||
}
|
||||
go globalTransportPool.cleanupRoutine(ctx)
|
||||
})
|
||||
return globalTransportPool
|
||||
}
|
||||
|
||||
// DefaultTransportConfig returns a secure default configuration
|
||||
func DefaultTransportConfig() TransportConfig {
|
||||
return TransportConfig{
|
||||
DialTimeout: 30 * time.Second,
|
||||
TLSHandshakeTimeout: 10 * time.Second,
|
||||
ResponseHeaderTimeout: 10 * time.Second,
|
||||
ExpectContinueTimeout: 1 * time.Second,
|
||||
IdleConnTimeout: 30 * time.Second,
|
||||
KeepAlive: 30 * time.Second,
|
||||
MaxIdleConns: 10,
|
||||
MaxIdleConnsPerHost: 2,
|
||||
MaxConnsPerHost: 5,
|
||||
ForceHTTP2: true,
|
||||
DisableKeepAlives: false,
|
||||
DisableCompression: false,
|
||||
WriteBufferSize: 4096,
|
||||
ReadBufferSize: 4096,
|
||||
InsecureSkipVerify: false,
|
||||
MinTLSVersion: tls.VersionTLS12,
|
||||
}
|
||||
}
|
||||
|
||||
// GetTransport gets or creates a shared transport with the given config
|
||||
func (p *TransportPool) GetTransport(config TransportConfig) *http.Transport {
|
||||
// Check client limit
|
||||
if atomic.LoadInt32(&p.clientCount) >= p.maxClients {
|
||||
return p.getExistingTransport()
|
||||
}
|
||||
|
||||
key := p.configKey(config)
|
||||
|
||||
// Fast path: check with read lock
|
||||
p.mu.RLock()
|
||||
if shared, exists := p.transports[key]; exists {
|
||||
atomic.AddInt32(&shared.refCount, 1)
|
||||
shared.lastUsed = time.Now()
|
||||
p.mu.RUnlock()
|
||||
return shared.transport
|
||||
}
|
||||
p.mu.RUnlock()
|
||||
|
||||
// Slow path: create new transport
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
// Double-check after acquiring write lock
|
||||
if shared, exists := p.transports[key]; exists {
|
||||
atomic.AddInt32(&shared.refCount, 1)
|
||||
shared.lastUsed = time.Now()
|
||||
return shared.transport
|
||||
}
|
||||
|
||||
// Create new transport
|
||||
transport := p.createTransport(config)
|
||||
shared := &sharedTransport{
|
||||
transport: transport,
|
||||
refCount: 1,
|
||||
lastUsed: time.Now(),
|
||||
config: config,
|
||||
}
|
||||
|
||||
p.transports[key] = shared
|
||||
atomic.AddInt32(&p.clientCount, 1)
|
||||
|
||||
return transport
|
||||
}
|
||||
|
||||
// ReleaseTransport decrements the reference count for a transport
|
||||
func (p *TransportPool) ReleaseTransport(transport *http.Transport) {
|
||||
if transport == nil {
|
||||
return
|
||||
}
|
||||
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
|
||||
for _, shared := range p.transports {
|
||||
if shared.transport == transport {
|
||||
count := atomic.AddInt32(&shared.refCount, -1)
|
||||
if count <= 0 {
|
||||
shared.lastUsed = time.Now()
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// getExistingTransport returns any available transport when limit is reached
|
||||
func (p *TransportPool) getExistingTransport() *http.Transport {
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
|
||||
for _, shared := range p.transports {
|
||||
if shared != nil && shared.transport != nil {
|
||||
atomic.AddInt32(&shared.refCount, 1)
|
||||
shared.lastUsed = time.Now()
|
||||
return shared.transport
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// createTransport creates a new HTTP transport with the given config
|
||||
func (p *TransportPool) createTransport(config TransportConfig) *http.Transport {
|
||||
// Set secure defaults
|
||||
if config.MinTLSVersion == 0 {
|
||||
config.MinTLSVersion = tls.VersionTLS12
|
||||
}
|
||||
|
||||
tlsConfig := &tls.Config{
|
||||
MinVersion: config.MinTLSVersion,
|
||||
MaxVersion: tls.VersionTLS13,
|
||||
CipherSuites: []uint16{
|
||||
tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
|
||||
tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
|
||||
tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
|
||||
tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
|
||||
},
|
||||
PreferServerCipherSuites: true,
|
||||
InsecureSkipVerify: config.InsecureSkipVerify,
|
||||
}
|
||||
|
||||
return &http.Transport{
|
||||
Proxy: http.ProxyFromEnvironment,
|
||||
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
|
||||
dialer := &net.Dialer{
|
||||
Timeout: config.DialTimeout,
|
||||
KeepAlive: config.KeepAlive,
|
||||
}
|
||||
return dialer.DialContext(ctx, network, addr)
|
||||
},
|
||||
TLSClientConfig: tlsConfig,
|
||||
ForceAttemptHTTP2: config.ForceHTTP2,
|
||||
TLSHandshakeTimeout: config.TLSHandshakeTimeout,
|
||||
ExpectContinueTimeout: config.ExpectContinueTimeout,
|
||||
MaxIdleConns: config.MaxIdleConns,
|
||||
MaxIdleConnsPerHost: config.MaxIdleConnsPerHost,
|
||||
IdleConnTimeout: config.IdleConnTimeout,
|
||||
DisableKeepAlives: config.DisableKeepAlives,
|
||||
MaxConnsPerHost: config.MaxConnsPerHost,
|
||||
ResponseHeaderTimeout: config.ResponseHeaderTimeout,
|
||||
DisableCompression: config.DisableCompression,
|
||||
WriteBufferSize: config.WriteBufferSize,
|
||||
ReadBufferSize: config.ReadBufferSize,
|
||||
}
|
||||
}
|
||||
|
||||
// configKey generates a unique key for a transport config
|
||||
func (p *TransportPool) configKey(config TransportConfig) string {
|
||||
// Create a simple key based on critical parameters
|
||||
sb := Get().GetStringBuilder()
|
||||
defer Get().PutStringBuilder(sb)
|
||||
|
||||
sb.WriteByte(byte(config.MaxConnsPerHost))
|
||||
sb.WriteByte(byte(config.MaxIdleConnsPerHost))
|
||||
sb.WriteByte(byte(config.MaxIdleConns))
|
||||
if config.ForceHTTP2 {
|
||||
sb.WriteByte(1)
|
||||
} else {
|
||||
sb.WriteByte(0)
|
||||
}
|
||||
if config.DisableKeepAlives {
|
||||
sb.WriteByte(1)
|
||||
} else {
|
||||
sb.WriteByte(0)
|
||||
}
|
||||
if config.DisableCompression {
|
||||
sb.WriteByte(1)
|
||||
} else {
|
||||
sb.WriteByte(0)
|
||||
}
|
||||
if config.InsecureSkipVerify {
|
||||
sb.WriteByte(1)
|
||||
} else {
|
||||
sb.WriteByte(0)
|
||||
}
|
||||
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// cleanupRoutine periodically cleans up unused transports
|
||||
func (p *TransportPool) cleanupRoutine(ctx context.Context) {
|
||||
ticker := time.NewTicker(1 * time.Minute)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
p.cleanup()
|
||||
return
|
||||
case <-ticker.C:
|
||||
p.cleanupIdle()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cleanupIdle removes idle transports
|
||||
func (p *TransportPool) cleanupIdle() {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
for key, shared := range p.transports {
|
||||
refCount := atomic.LoadInt32(&shared.refCount)
|
||||
if refCount <= 0 && now.Sub(shared.lastUsed) > 2*time.Minute {
|
||||
shared.transport.CloseIdleConnections()
|
||||
delete(p.transports, key)
|
||||
atomic.AddInt32(&p.clientCount, -1)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cleanup closes all transports
|
||||
func (p *TransportPool) cleanup() {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
for _, shared := range p.transports {
|
||||
shared.transport.CloseIdleConnections()
|
||||
}
|
||||
p.transports = make(map[string]*sharedTransport)
|
||||
atomic.StoreInt32(&p.clientCount, 0)
|
||||
}
|
||||
|
||||
// Shutdown gracefully shuts down the transport pool
|
||||
func (p *TransportPool) Shutdown() {
|
||||
if p.cancel != nil {
|
||||
p.cancel()
|
||||
}
|
||||
}
|
||||
|
||||
// Stats returns transport pool statistics
|
||||
type TransportPoolStats struct {
|
||||
ActiveTransports int
|
||||
TotalClients int32
|
||||
MaxClients int32
|
||||
}
|
||||
|
||||
// GetStats returns current pool statistics
|
||||
func (p *TransportPool) GetStats() TransportPoolStats {
|
||||
p.mu.RLock()
|
||||
defer p.mu.RUnlock()
|
||||
|
||||
activeCount := 0
|
||||
for _, shared := range p.transports {
|
||||
if atomic.LoadInt32(&shared.refCount) > 0 {
|
||||
activeCount++
|
||||
}
|
||||
}
|
||||
|
||||
return TransportPoolStats{
|
||||
ActiveTransports: activeCount,
|
||||
TotalClients: atomic.LoadInt32(&p.clientCount),
|
||||
MaxClients: p.maxClients,
|
||||
}
|
||||
}
|
||||
|
||||
// CreateHTTPClient creates an HTTP client using the transport pool
|
||||
func CreateHTTPClient(config TransportConfig, timeout time.Duration) *http.Client {
|
||||
pool := GetTransportPool()
|
||||
transport := pool.GetTransport(config)
|
||||
|
||||
if transport == nil {
|
||||
// Fallback to a basic client if pool is exhausted
|
||||
return &http.Client{
|
||||
Timeout: timeout,
|
||||
}
|
||||
}
|
||||
|
||||
client := &http.Client{
|
||||
Transport: transport,
|
||||
Timeout: timeout,
|
||||
}
|
||||
|
||||
// Configure redirect policy
|
||||
client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
|
||||
if len(via) >= 10 {
|
||||
return http.ErrUseLastResponse
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
return client
|
||||
}
|
||||
@@ -0,0 +1,593 @@
|
||||
package pool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"net/http"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestGetTransportPool_Singleton tests that GetTransportPool returns the same instance
|
||||
func TestGetTransportPool_Singleton(t *testing.T) {
|
||||
pool1 := GetTransportPool()
|
||||
pool2 := GetTransportPool()
|
||||
|
||||
if pool1 != pool2 {
|
||||
t.Error("GetTransportPool() should return the same instance (singleton)")
|
||||
}
|
||||
|
||||
if pool1 == nil {
|
||||
t.Error("GetTransportPool() should not return nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestDefaultTransportConfig tests the default transport configuration
|
||||
func TestDefaultTransportConfig(t *testing.T) {
|
||||
config := DefaultTransportConfig()
|
||||
|
||||
// Verify security defaults
|
||||
if config.MinTLSVersion != tls.VersionTLS12 {
|
||||
t.Errorf("Default MinTLSVersion should be TLS 1.2, got %d", config.MinTLSVersion)
|
||||
}
|
||||
|
||||
if config.InsecureSkipVerify {
|
||||
t.Error("Default should not skip TLS verification")
|
||||
}
|
||||
|
||||
if !config.ForceHTTP2 {
|
||||
t.Error("Default should force HTTP/2")
|
||||
}
|
||||
|
||||
// Verify reasonable timeouts
|
||||
if config.DialTimeout <= 0 {
|
||||
t.Error("DialTimeout should be positive")
|
||||
}
|
||||
|
||||
if config.TLSHandshakeTimeout <= 0 {
|
||||
t.Error("TLSHandshakeTimeout should be positive")
|
||||
}
|
||||
|
||||
if config.ResponseHeaderTimeout <= 0 {
|
||||
t.Error("ResponseHeaderTimeout should be positive")
|
||||
}
|
||||
|
||||
// Verify connection limits
|
||||
if config.MaxIdleConns <= 0 {
|
||||
t.Error("MaxIdleConns should be positive")
|
||||
}
|
||||
|
||||
if config.MaxIdleConnsPerHost <= 0 {
|
||||
t.Error("MaxIdleConnsPerHost should be positive")
|
||||
}
|
||||
|
||||
if config.MaxConnsPerHost <= 0 {
|
||||
t.Error("MaxConnsPerHost should be positive")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_GetTransport tests transport creation and reuse
|
||||
func TestTransportPool_GetTransport(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
|
||||
// First call should create new transport
|
||||
transport1 := pool.GetTransport(config)
|
||||
if transport1 == nil {
|
||||
t.Error("GetTransport should not return nil")
|
||||
}
|
||||
|
||||
// Second call with same config should return same transport
|
||||
transport2 := pool.GetTransport(config)
|
||||
if transport2 == nil {
|
||||
t.Error("GetTransport should not return nil")
|
||||
}
|
||||
|
||||
if transport1 != transport2 {
|
||||
t.Error("GetTransport should return same transport for same config")
|
||||
}
|
||||
|
||||
// Verify reference counting
|
||||
pool.mu.RLock()
|
||||
key := pool.configKey(config)
|
||||
shared := pool.transports[key]
|
||||
refCount := shared.refCount
|
||||
pool.mu.RUnlock()
|
||||
|
||||
if refCount != 2 {
|
||||
t.Errorf("Reference count should be 2, got %d", refCount)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_GetTransport_DifferentConfigs tests transport creation with different configs
|
||||
func TestTransportPool_GetTransport_DifferentConfigs(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
}
|
||||
|
||||
config1 := DefaultTransportConfig()
|
||||
config2 := DefaultTransportConfig()
|
||||
config2.MaxConnsPerHost = 10 // Different from default
|
||||
|
||||
transport1 := pool.GetTransport(config1)
|
||||
transport2 := pool.GetTransport(config2)
|
||||
|
||||
if transport1 == transport2 {
|
||||
t.Error("Different configs should produce different transports")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_GetTransport_ClientLimit tests client limit enforcement
|
||||
func TestTransportPool_GetTransport_ClientLimit(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 2, // Low limit for testing
|
||||
clientCount: 2, // Already at limit
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
|
||||
// Should return existing transport when limit reached
|
||||
transport := pool.GetTransport(config)
|
||||
// Transport might be nil if no existing transports
|
||||
if transport != nil && pool.clientCount > pool.maxClients {
|
||||
t.Error("Should not exceed client limit")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_ReleaseTransport tests transport reference counting
|
||||
func TestTransportPool_ReleaseTransport(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
|
||||
// Get transport
|
||||
transport := pool.GetTransport(config)
|
||||
if transport == nil {
|
||||
t.Error("GetTransport should not return nil")
|
||||
}
|
||||
|
||||
// Release transport
|
||||
pool.ReleaseTransport(transport)
|
||||
|
||||
// Verify reference count decreased
|
||||
pool.mu.RLock()
|
||||
key := pool.configKey(config)
|
||||
shared := pool.transports[key]
|
||||
refCount := shared.refCount
|
||||
pool.mu.RUnlock()
|
||||
|
||||
if refCount != 0 {
|
||||
t.Errorf("Reference count should be 0 after release, got %d", refCount)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_ReleaseTransport_Nil tests releasing nil transport
|
||||
func TestTransportPool_ReleaseTransport_Nil(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
}
|
||||
|
||||
// Should not panic
|
||||
pool.ReleaseTransport(nil)
|
||||
}
|
||||
|
||||
// TestTransportPool_ReleaseTransport_Unknown tests releasing unknown transport
|
||||
func TestTransportPool_ReleaseTransport_Unknown(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
}
|
||||
|
||||
// Create a transport not from the pool
|
||||
transport := &http.Transport{}
|
||||
|
||||
// Should not panic
|
||||
pool.ReleaseTransport(transport)
|
||||
}
|
||||
|
||||
// TestTransportPool_createTransport tests transport creation with different configs
|
||||
func TestTransportPool_createTransport(t *testing.T) {
|
||||
pool := &TransportPool{}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
config TransportConfig
|
||||
}{
|
||||
{
|
||||
"default config",
|
||||
DefaultTransportConfig(),
|
||||
},
|
||||
{
|
||||
"custom timeouts",
|
||||
TransportConfig{
|
||||
DialTimeout: 10 * time.Second,
|
||||
TLSHandshakeTimeout: 5 * time.Second,
|
||||
MinTLSVersion: tls.VersionTLS13,
|
||||
},
|
||||
},
|
||||
{
|
||||
"insecure config",
|
||||
TransportConfig{
|
||||
InsecureSkipVerify: true,
|
||||
MinTLSVersion: tls.VersionTLS10,
|
||||
},
|
||||
},
|
||||
{
|
||||
"no HTTP/2",
|
||||
TransportConfig{
|
||||
ForceHTTP2: false,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
transport := pool.createTransport(test.config)
|
||||
|
||||
if transport == nil {
|
||||
t.Error("createTransport should not return nil")
|
||||
return
|
||||
}
|
||||
|
||||
// Verify TLS config
|
||||
if transport.TLSClientConfig == nil {
|
||||
t.Error("Transport should have TLS config")
|
||||
return
|
||||
}
|
||||
|
||||
// Verify minimum TLS version
|
||||
expectedMinVersion := test.config.MinTLSVersion
|
||||
if expectedMinVersion == 0 {
|
||||
expectedMinVersion = tls.VersionTLS12 // Default
|
||||
}
|
||||
if transport.TLSClientConfig.MinVersion != expectedMinVersion {
|
||||
t.Errorf("TLS MinVersion should be %d, got %d", expectedMinVersion, transport.TLSClientConfig.MinVersion)
|
||||
}
|
||||
|
||||
// Verify max TLS version
|
||||
if transport.TLSClientConfig.MaxVersion != tls.VersionTLS13 {
|
||||
t.Errorf("TLS MaxVersion should be %d, got %d", tls.VersionTLS13, transport.TLSClientConfig.MaxVersion)
|
||||
}
|
||||
|
||||
// Verify InsecureSkipVerify
|
||||
if transport.TLSClientConfig.InsecureSkipVerify != test.config.InsecureSkipVerify {
|
||||
t.Errorf("InsecureSkipVerify should be %v, got %v", test.config.InsecureSkipVerify, transport.TLSClientConfig.InsecureSkipVerify)
|
||||
}
|
||||
|
||||
// Verify HTTP/2
|
||||
if transport.ForceAttemptHTTP2 != test.config.ForceHTTP2 {
|
||||
t.Errorf("ForceAttemptHTTP2 should be %v, got %v", test.config.ForceHTTP2, transport.ForceAttemptHTTP2)
|
||||
}
|
||||
|
||||
// Verify timeouts
|
||||
if test.config.TLSHandshakeTimeout > 0 && transport.TLSHandshakeTimeout != test.config.TLSHandshakeTimeout {
|
||||
t.Errorf("TLSHandshakeTimeout should be %v, got %v", test.config.TLSHandshakeTimeout, transport.TLSHandshakeTimeout)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_configKey tests configuration key generation
|
||||
func TestTransportPool_configKey(t *testing.T) {
|
||||
pool := &TransportPool{}
|
||||
|
||||
config1 := DefaultTransportConfig()
|
||||
config2 := DefaultTransportConfig()
|
||||
|
||||
key1 := pool.configKey(config1)
|
||||
key2 := pool.configKey(config2)
|
||||
|
||||
if key1 != key2 {
|
||||
t.Error("Same configs should generate same key")
|
||||
}
|
||||
|
||||
// Different config
|
||||
config3 := config1
|
||||
config3.MaxConnsPerHost = 999
|
||||
key3 := pool.configKey(config3)
|
||||
|
||||
if key1 == key3 {
|
||||
t.Error("Different configs should generate different keys")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_cleanupIdle tests idle transport cleanup
|
||||
func TestTransportPool_cleanupIdle(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
transport := pool.createTransport(config)
|
||||
|
||||
// Add transport to pool with old timestamp
|
||||
shared := &sharedTransport{
|
||||
transport: transport,
|
||||
refCount: 0,
|
||||
lastUsed: time.Now().Add(-5 * time.Minute), // Old
|
||||
config: config,
|
||||
}
|
||||
|
||||
key := pool.configKey(config)
|
||||
pool.transports[key] = shared
|
||||
|
||||
// Run cleanup
|
||||
pool.cleanupIdle()
|
||||
|
||||
// Transport should be removed
|
||||
if _, exists := pool.transports[key]; exists {
|
||||
t.Error("Old idle transport should be cleaned up")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_cleanup tests full cleanup
|
||||
func TestTransportPool_cleanup(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
clientCount: 3,
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
transport := pool.createTransport(config)
|
||||
|
||||
// Add transport to pool
|
||||
shared := &sharedTransport{
|
||||
transport: transport,
|
||||
refCount: 1,
|
||||
lastUsed: time.Now(),
|
||||
config: config,
|
||||
}
|
||||
|
||||
key := pool.configKey(config)
|
||||
pool.transports[key] = shared
|
||||
|
||||
// Run cleanup
|
||||
pool.cleanup()
|
||||
|
||||
// All transports should be removed
|
||||
if len(pool.transports) != 0 {
|
||||
t.Error("All transports should be cleaned up")
|
||||
}
|
||||
|
||||
// Client count should be reset
|
||||
if pool.clientCount != 0 {
|
||||
t.Error("Client count should be reset")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_Shutdown tests graceful shutdown
|
||||
func TestTransportPool_Shutdown(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
}
|
||||
|
||||
// Should not panic
|
||||
pool.Shutdown()
|
||||
}
|
||||
|
||||
// TestTransportPool_GetStats tests statistics
|
||||
func TestTransportPool_GetStats(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 5,
|
||||
clientCount: 3,
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
|
||||
// Add some transports
|
||||
for i := 0; i < 3; i++ {
|
||||
transport := pool.createTransport(config)
|
||||
shared := &sharedTransport{
|
||||
transport: transport,
|
||||
refCount: int32(i % 2), // Some active, some idle
|
||||
lastUsed: time.Now(),
|
||||
config: config,
|
||||
}
|
||||
pool.transports[string(rune(i))] = shared
|
||||
}
|
||||
|
||||
stats := pool.GetStats()
|
||||
|
||||
if stats.TotalClients != 3 {
|
||||
t.Errorf("TotalClients should be 3, got %d", stats.TotalClients)
|
||||
}
|
||||
|
||||
if stats.MaxClients != 5 {
|
||||
t.Errorf("MaxClients should be 5, got %d", stats.MaxClients)
|
||||
}
|
||||
|
||||
if stats.ActiveTransports < 0 || stats.ActiveTransports > 3 {
|
||||
t.Errorf("ActiveTransports should be between 0 and 3, got %d", stats.ActiveTransports)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCreateHTTPClient tests HTTP client creation
|
||||
func TestCreateHTTPClient(t *testing.T) {
|
||||
config := DefaultTransportConfig()
|
||||
timeout := 30 * time.Second
|
||||
|
||||
client := CreateHTTPClient(config, timeout)
|
||||
|
||||
if client == nil {
|
||||
t.Error("CreateHTTPClient should not return nil")
|
||||
return
|
||||
}
|
||||
|
||||
if client.Timeout != timeout {
|
||||
t.Errorf("Client timeout should be %v, got %v", timeout, client.Timeout)
|
||||
}
|
||||
|
||||
if client.Transport == nil {
|
||||
t.Error("Client should have transport")
|
||||
}
|
||||
|
||||
if client.CheckRedirect == nil {
|
||||
t.Error("Client should have redirect policy")
|
||||
}
|
||||
|
||||
// Test redirect policy
|
||||
req := &http.Request{}
|
||||
var via []*http.Request
|
||||
|
||||
// Should allow up to 9 redirects (10 total requests)
|
||||
for i := 0; i < 9; i++ {
|
||||
via = append(via, &http.Request{})
|
||||
err := client.CheckRedirect(req, via)
|
||||
if err != nil {
|
||||
t.Errorf("Should allow %d redirects, got error: %v", i+1, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Should reject 10th redirect (11th total request)
|
||||
via = append(via, &http.Request{})
|
||||
err := client.CheckRedirect(req, via)
|
||||
if err != http.ErrUseLastResponse {
|
||||
t.Error("Should reject too many redirects")
|
||||
}
|
||||
}
|
||||
|
||||
// TestCreateHTTPClient_Fallback tests fallback when pool is exhausted
|
||||
func TestCreateHTTPClient_Fallback(t *testing.T) {
|
||||
// Override global pool with limited one
|
||||
originalPool := globalTransportPool
|
||||
defer func() {
|
||||
globalTransportPool = originalPool
|
||||
}()
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
globalTransportPool = &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
clientCount: 10,
|
||||
maxClients: 1, // Very low limit
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
timeout := 30 * time.Second
|
||||
|
||||
client := CreateHTTPClient(config, timeout)
|
||||
|
||||
if client == nil {
|
||||
t.Error("CreateHTTPClient should not return nil even when pool is exhausted")
|
||||
return
|
||||
}
|
||||
|
||||
if client.Timeout != timeout {
|
||||
t.Errorf("Client timeout should be %v, got %v", timeout, client.Timeout)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportPool_ConcurrentAccess tests concurrent access to transport pool
|
||||
func TestTransportPool_ConcurrentAccess(t *testing.T) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 50, // High limit for concurrent test
|
||||
}
|
||||
|
||||
// Use different configs to reduce contention on single transport
|
||||
baseConfig := DefaultTransportConfig()
|
||||
configs := make([]TransportConfig, 10)
|
||||
for i := range configs {
|
||||
configs[i] = baseConfig
|
||||
configs[i].MaxConnsPerHost = 5 + i // Make each config unique
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
numGoroutines := 10
|
||||
operationsPerGoroutine := 3
|
||||
|
||||
wg.Add(numGoroutines)
|
||||
for i := 0; i < numGoroutines; i++ {
|
||||
go func(goroutineID int) {
|
||||
defer wg.Done()
|
||||
config := configs[goroutineID%len(configs)]
|
||||
for j := 0; j < operationsPerGoroutine; j++ {
|
||||
transport := pool.GetTransport(config)
|
||||
if transport == nil {
|
||||
continue
|
||||
}
|
||||
// Use transport briefly
|
||||
time.Sleep(time.Millisecond)
|
||||
pool.ReleaseTransport(transport)
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
// Should not panic and should have reasonable stats
|
||||
stats := pool.GetStats()
|
||||
if stats.TotalClients < 0 || stats.TotalClients > int32(numGoroutines) {
|
||||
t.Errorf("Unexpected client count: %d", stats.TotalClients)
|
||||
}
|
||||
}
|
||||
|
||||
// Benchmark tests for performance verification
|
||||
func BenchmarkTransportPool_GetTransport(b *testing.B) {
|
||||
pool := &TransportPool{
|
||||
transports: make(map[string]*sharedTransport),
|
||||
maxConns: 20,
|
||||
maxClients: 100,
|
||||
}
|
||||
|
||||
config := DefaultTransportConfig()
|
||||
b.ResetTimer()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
transport := pool.GetTransport(config)
|
||||
pool.ReleaseTransport(transport)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkCreateHTTPClient(b *testing.B) {
|
||||
config := DefaultTransportConfig()
|
||||
timeout := 30 * time.Second
|
||||
b.ResetTimer()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
CreateHTTPClient(config, timeout)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkTransportPool_configKey(b *testing.B) {
|
||||
pool := &TransportPool{}
|
||||
config := DefaultTransportConfig()
|
||||
b.ResetTimer()
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
pool.configKey(config)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user