feat(a2agrpc): add GRPCConnectionPool for gRPC client connection reuse (fixes #109)

a2agrpc/v1/pool.go

@@ -0,0 +1,151 @@
+// Copyright 2026 The A2A Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package a2agrpc
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"github.com/a2aproject/a2a-go/v2/a2a"
+	"github.com/a2aproject/a2a-go/v2/a2aclient"
+	a2apb "github.com/a2aproject/a2a-go/v2/a2apb/v1"
+	"google.golang.org/grpc"
+)
+
+// GRPCConnectionPool manages reusable gRPC connections keyed by URL.
+// gRPC ClientConns are long-lived and expensive to recreate — a pool
+// avoids the hidden performance cost of creating a new connection on
+// every client creation.
+type GRPCConnectionPool interface {
+	// Acquire returns an existing or newly created connection for url.
+	Acquire(ctx context.Context, url string, opts ...grpc.DialOption) (*grpc.ClientConn, error)
+
+	// Release returns a connection to the pool. The caller must not
+	// close the connection after releasing it.
+	Release(*grpc.ClientConn) error
+}
+
+// DefaultGRPCConnectionPool is a simple in-memory pool that reuses
+// gRPC connections by URL. Connections idle longer than ttl are
+// closed and removed on the next Acquire. A zero or negative ttl
+// disables eviction (connections live forever).
+type DefaultGRPCConnectionPool struct {
+	mu    sync.Mutex
+	conns map[string]*pooledConn
+	ttl   time.Duration
+}
+
+type pooledConn struct {
+	conn     *grpc.ClientConn
+	lastUsed time.Time
+}
+
+// NewDefaultGRPCConnectionPool creates a pool with the given idle TTL.
+func NewDefaultGRPCConnectionPool(ttl time.Duration) *DefaultGRPCConnectionPool {
+	return &DefaultGRPCConnectionPool{
+		conns: make(map[string]*pooledConn),
+		ttl:   ttl,
+	}
+}
+
+// Acquire returns a cached connection for url or creates one with the
+// provided dial options. Expired connections are evicted before lookup.
+func (p *DefaultGRPCConnectionPool) Acquire(ctx context.Context, url string, opts ...grpc.DialOption) (*grpc.ClientConn, error) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	p.evictExpired()
+
+	if pc, ok := p.conns[url]; ok {
+		pc.lastUsed = time.Now()
+		return pc.conn, nil
+	}
+
+	conn, err := grpc.NewClient(url, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	p.conns[url] = &pooledConn{conn: conn, lastUsed: time.Now()}
+	return conn, nil
+}
+
+// Release returns a connection to the pool. The connection is NOT closed.
+func (p *DefaultGRPCConnectionPool) Release(conn *grpc.ClientConn) error {
+	// Connections are kept in the pool until evicted by TTL.
+	// On release we just update the timestamp so the TTL clock resets.
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	for _, pc := range p.conns {
+		if pc.conn == conn {
+			pc.lastUsed = time.Now()
+			return nil
+		}
+	}
+	return nil
+}
+
+// Close closes all pooled connections.
+func (p *DefaultGRPCConnectionPool) Close() error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	for url, pc := range p.conns {
+		_ = pc.conn.Close()
+		delete(p.conns, url)
+	}
+	return nil
+}
+
+// Len returns the current number of pooled connections.
+func (p *DefaultGRPCConnectionPool) Len() int {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return len(p.conns)
+}
+
+func (p *DefaultGRPCConnectionPool) evictExpired() {
+	if p.ttl <= 0 {
+		return
+	}
+	now := time.Now()
+	for url, pc := range p.conns {
+		if now.Sub(pc.lastUsed) > p.ttl {
+			_ = pc.conn.Close()
+			delete(p.conns, url)
+		}
+	}
+}
+
+// WithPooledGRPCTransport creates a gRPC transport backed by a connection pool.
+// Each client creation acquires a connection from the pool, and releasing the
+// transport returns the connection to the pool instead of closing it.
+func WithPooledGRPCTransport(pool GRPCConnectionPool, opts ...grpc.DialOption) a2aclient.FactoryOption {
+	return a2aclient.WithTransport(
+		a2a.TransportProtocolGRPC,
+		a2aclient.TransportFactoryFn(func(ctx context.Context, card *a2a.AgentCard, iface *a2a.AgentInterface) (a2aclient.Transport, error) {
+			conn, err := pool.Acquire(ctx, iface.URL, opts...)
+			if err != nil {
+				return nil, err
+			}
+			return &grpcTransport{
+				client:      a2apb.NewA2AServiceClient(conn),
+				closeConnFn: func() error { return pool.Release(conn) },
+			}, nil
+		}),
+	)
+}

a2agrpc/v1/pool_test.go

@@ -0,0 +1,200 @@
+// Copyright 2026 The A2A Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package a2agrpc
+
+import (
+	"context"
+	"net"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/test/bufconn"
+)
+
+func newBufconnServer(t *testing.T) (*grpc.Server, *bufconn.Listener, func()) {
+	t.Helper()
+	lis := bufconn.Listen(1024 * 1024)
+	srv := grpc.NewServer()
+	go func() { _ = srv.Serve(lis) }()
+	return srv, lis, func() { srv.Stop() }
+}
+
+func bufconnDialer(lis *bufconn.Listener) func(context.Context, string) (net.Conn, error) {
+	return func(context.Context, string) (net.Conn, error) {
+		return lis.Dial()
+	}
+}
+
+func bufconnOpts(lis *bufconn.Listener) []grpc.DialOption {
+	return []grpc.DialOption{
+		grpc.WithContextDialer(bufconnDialer(lis)),
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+	}
+}
+
+func TestDefaultPool_AcquireCreatesConnection(t *testing.T) {
+	pool := NewDefaultGRPCConnectionPool(1 * time.Minute)
+	defer func() { _ = pool.Close() }()
+
+	_, lis, stop := newBufconnServer(t)
+	defer stop()
+
+	conn, err := pool.Acquire(context.Background(), "bufnet", bufconnOpts(lis)...)
+	if err != nil {
+		t.Fatalf("Acquire failed: %v", err)
+	}
+	if conn == nil {
+		t.Fatal("expected non-nil connection")
+	}
+
+	if n := pool.Len(); n != 1 {
+		t.Fatalf("expected 1 pooled connection, got %d", n)
+	}
+}
+
+func TestDefaultPool_ReusesConnection(t *testing.T) {
+	pool := NewDefaultGRPCConnectionPool(1 * time.Minute)
+	defer func() { _ = pool.Close() }()
+
+	_, lis, stop := newBufconnServer(t)
+	defer stop()
+
+	conn1, err := pool.Acquire(context.Background(), "bufnet", bufconnOpts(lis)...)
+	if err != nil {
+		t.Fatalf("Acquire 1 failed: %v", err)
+	}
+	_ = pool.Release(conn1)
+
+	conn2, err := pool.Acquire(context.Background(), "bufnet", bufconnOpts(lis)...)
+	if err != nil {
+		t.Fatalf("Acquire 2 failed: %v", err)
+	}
+
+	if conn1 != conn2 {
+		t.Fatal("expected same connection to be reused")
+	}
+
+	if n := pool.Len(); n != 1 {
+		t.Fatalf("expected 1 pooled connection, got %d", n)
+	}
+}
+
+func TestDefaultPool_DifferentURLs(t *testing.T) {
+	pool := NewDefaultGRPCConnectionPool(1 * time.Minute)
+	defer func() { _ = pool.Close() }()
+
+	_, lis, stop := newBufconnServer(t)
+	defer stop()
+	opts := bufconnOpts(lis)
+
+	conn1, err := pool.Acquire(context.Background(), "agent-a", opts...)
+	if err != nil {
+		t.Fatalf("Acquire agent-a failed: %v", err)
+	}
+	defer func() { _ = pool.Release(conn1) }()
+
+	conn2, err := pool.Acquire(context.Background(), "agent-b", opts...)
+	if err != nil {
+		t.Fatalf("Acquire agent-b failed: %v", err)
+	}
+	defer func() { _ = pool.Release(conn2) }()
+
+	if conn1 == conn2 {
+		t.Fatal("different URLs should not reuse the same connection")
+	}
+
+	if n := pool.Len(); n != 2 {
+		t.Fatalf("expected 2 pooled connections, got %d", n)
+	}
+}
+
+func TestDefaultPool_TTLEviction(t *testing.T) {
+	pool := NewDefaultGRPCConnectionPool(50 * time.Millisecond)
+	defer func() { _ = pool.Close() }()
+
+	_, lis, stop := newBufconnServer(t)
+	defer stop()
+	opts := bufconnOpts(lis)
+
+	conn1, err := pool.Acquire(context.Background(), "bufnet", opts...)
+	if err != nil {
+		t.Fatalf("Acquire failed: %v", err)
+	}
+	_ = pool.Release(conn1)
+
+	if n := pool.Len(); n != 1 {
+		t.Fatalf("expected 1 pooled connection before TTL, got %d", n)
+	}
+
+	time.Sleep(100 * time.Millisecond)
+
+	conn2, err := pool.Acquire(context.Background(), "bufnet", opts...)
+	if err != nil {
+		t.Fatalf("Acquire after TTL failed: %v", err)
+	}
+
+	if conn1 == conn2 {
+		t.Fatal("expired connection should have been evicted")
+	}
+
+	if n := pool.Len(); n != 1 {
+		t.Fatalf("expected 1 pooled connection after eviction, got %d", n)
+	}
+}
+
+func TestDefaultPool_ZeroTTLNeverEvicts(t *testing.T) {
+	pool := NewDefaultGRPCConnectionPool(0)
+	defer func() { _ = pool.Close() }()
+
+	_, lis, stop := newBufconnServer(t)
+	defer stop()
+	opts := bufconnOpts(lis)
+
+	conn1, err := pool.Acquire(context.Background(), "bufnet", opts...)
+	if err != nil {
+		t.Fatalf("Acquire failed: %v", err)
+	}
+	_ = pool.Release(conn1)
+
+	time.Sleep(50 * time.Millisecond)
+
+	conn2, err := pool.Acquire(context.Background(), "bufnet", opts...)
+	if err != nil {
+		t.Fatalf("Acquire 2 failed: %v", err)
+	}
+
+	if conn1 != conn2 {
+		t.Fatal("zero TTL should never evict, expected same connection")
+	}
+}
+
+func TestDefaultPool_Release(t *testing.T) {
+	pool := NewDefaultGRPCConnectionPool(1 * time.Minute)
+	defer func() { _ = pool.Close() }()
+
+	_, lis, stop := newBufconnServer(t)
+	defer stop()
+
+	conn, err := pool.Acquire(context.Background(), "bufnet", bufconnOpts(lis)...)
+	if err != nil {
+		t.Fatalf("Acquire failed: %v", err)
+	}
+
+	if err := pool.Release(conn); err != nil {
+		t.Fatalf("Release failed: %v", err)
+	}
+}