package transport import ( "context" "errors" "sync" "time" "github.com/jeremytregunna/kevo/pkg/transport" ) var ( ErrPoolClosed = errors.New("connection pool is closed") ErrPoolFull = errors.New("connection pool is full") ErrPoolEmptyNoWait = errors.New("connection pool is empty and wait is disabled") ) // ConnectionPool manages a pool of gRPC connections type ConnectionPool struct { manager *GRPCTransportManager address string maxIdle int maxActive int idlePool chan *GRPCConnection activePool chan struct{} mu sync.Mutex closed bool idleTime time.Duration } // NewConnectionPool creates a new connection pool func NewConnectionPool(manager *GRPCTransportManager, address string, maxIdle, maxActive int, idleTime time.Duration) *ConnectionPool { if maxIdle <= 0 { maxIdle = 2 } if maxActive <= 0 { maxActive = 10 } if idleTime <= 0 { idleTime = 5 * time.Minute } pool := &ConnectionPool{ manager: manager, address: address, maxIdle: maxIdle, maxActive: maxActive, idlePool: make(chan *GRPCConnection, maxIdle), activePool: make(chan struct{}, maxActive), idleTime: idleTime, } return pool } // Get retrieves a connection from the pool or creates a new one func (p *ConnectionPool) Get(ctx context.Context, wait bool) (*GRPCConnection, error) { p.mu.Lock() if p.closed { p.mu.Unlock() return nil, ErrPoolClosed } p.mu.Unlock() // Try to get an idle connection select { case conn := <-p.idlePool: return conn, nil default: // No idle connections available } // Check if we can create a new connection select { case p.activePool <- struct{}{}: // We acquired a slot to create a new connection conn, err := p.createConnection(ctx) if err != nil { // If connection creation fails, release the active slot <-p.activePool return nil, err } return conn, nil default: // Pool is full, check if we should wait if !wait { return nil, ErrPoolEmptyNoWait } } // Wait for a connection to become available or context to expire select { case conn := <-p.idlePool: // Got an idle connection return conn, nil case p.activePool <- struct{}{}: // Got permission to create a new connection conn, err := p.createConnection(ctx) if err != nil { <-p.activePool return nil, err } return conn, nil case <-ctx.Done(): // Context deadline exceeded or canceled return nil, ctx.Err() } } // createConnection creates a new connection to the server func (p *ConnectionPool) createConnection(ctx context.Context) (*GRPCConnection, error) { conn, err := p.manager.Connect(ctx, p.address) if err != nil { return nil, err } // Type assert to our concrete connection type grpcConn, ok := conn.(*GRPCConnection) if !ok { conn.Close() return nil, errors.New("received incorrect connection type from transport manager") } return grpcConn, nil } // Put returns a connection to the pool func (p *ConnectionPool) Put(conn *GRPCConnection) error { p.mu.Lock() if p.closed { p.mu.Unlock() conn.Close() // Close the connection since the pool is closed <-p.activePool return nil } p.mu.Unlock() // Try to add to the idle pool select { case p.idlePool <- conn: // Successfully returned to idle pool return nil default: // Idle pool full, close the connection conn.Close() <-p.activePool return nil } } // Close closes the connection pool and all idle connections func (p *ConnectionPool) Close() error { p.mu.Lock() defer p.mu.Unlock() if p.closed { return ErrPoolClosed } p.closed = true // Close all idle connections close(p.idlePool) for conn := range p.idlePool { conn.Close() <-p.activePool } return nil } // ConnectionPoolManager manages multiple connection pools type ConnectionPoolManager struct { manager *GRPCTransportManager pools sync.Map // map[string]*ConnectionPool defaultMaxIdle int defaultMaxActive int defaultIdleTime time.Duration } // NewConnectionPoolManager creates a new connection pool manager func NewConnectionPoolManager(manager *GRPCTransportManager, maxIdle, maxActive int, idleTime time.Duration) *ConnectionPoolManager { return &ConnectionPoolManager{ manager: manager, defaultMaxIdle: maxIdle, defaultMaxActive: maxActive, defaultIdleTime: idleTime, } } // GetPool gets or creates a connection pool for the given address func (m *ConnectionPoolManager) GetPool(address string) *ConnectionPool { // Check if pool exists if pool, ok := m.pools.Load(address); ok { return pool.(*ConnectionPool) } // Create new pool pool := NewConnectionPool(m.manager, address, m.defaultMaxIdle, m.defaultMaxActive, m.defaultIdleTime) m.pools.Store(address, pool) return pool } // CloseAll closes all connection pools func (m *ConnectionPoolManager) CloseAll() { m.pools.Range(func(key, value interface{}) bool { pool := value.(*ConnectionPool) pool.Close() m.pools.Delete(key) return true }) } // Client is a wrapper around the gRPC client that supports connection pooling and retries type Client struct { pool *ConnectionPool maxRetries int retryDelay time.Duration } // NewClient creates a new client with the given connection pool func NewClient(pool *ConnectionPool, maxRetries int, retryDelay time.Duration) *Client { if maxRetries <= 0 { maxRetries = 3 } if retryDelay <= 0 { retryDelay = 100 * time.Millisecond } return &Client{ pool: pool, maxRetries: maxRetries, retryDelay: retryDelay, } } // Execute executes the given function with a connection from the pool func (c *Client) Execute(ctx context.Context, fn func(ctx context.Context, client interface{}) (interface{}, error)) (interface{}, error) { var conn *GRPCConnection var err error var result interface{} // Get a connection from the pool for i := 0; i < c.maxRetries; i++ { if i > 0 { // Wait before retrying select { case <-time.After(c.retryDelay): // Continue with retry case <-ctx.Done(): return nil, ctx.Err() } } conn, err = c.pool.Get(ctx, true) if err != nil { continue // Try to get another connection } // Execute the function with the connection client := conn.GetClient() result, err = fn(ctx, client) // Return connection to the pool regardless of error putErr := c.pool.Put(conn) if putErr != nil { // Log the error but continue with the original error transport.LogError("Failed to return connection to pool: %v", putErr) } if err == nil { // Success, return the result return result, nil } } return nil, err }