MCP Gateway Design: Scaling Sessions and Backpressure

package gateway

import (
	"context"
	"fmt"
	"net/http"
	"sync"
	"time"
)

type ConnState int

const (
	ConnIdle ConnState = iota
	ConnActive
	ConnDraining
)

type SSEConn struct {
	ID        string
	ServerURL string
	State     ConnState
	CreatedAt time.Time
	LastUsed  time.Time
	mu        sync.Mutex
	client    *http.Client
	eventCh   chan []byte
	closeCh   chan struct{}
}

type ConnPool struct {
	mu          sync.RWMutex
	conns       map[string][]*SSEConn // serverURL -> connections
	maxPerHost  int
	idleTimeout time.Duration
	maxLifetime time.Duration
}

func NewConnPool(maxPerHost int, idleTimeout, maxLifetime time.Duration) *ConnPool {
	pool := &ConnPool{
		conns:       make(map[string][]*SSEConn),
		maxPerHost:  maxPerHost,
		idleTimeout: idleTimeout,
		maxLifetime: maxLifetime,
	}
	go pool.evictLoop()
	return pool
}

func (p *ConnPool) Acquire(ctx context.Context, serverURL string) (*SSEConn, error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	conns := p.conns[serverURL]
	for _, conn := range conns {
		conn.mu.Lock()
		if conn.State == ConnIdle && time.Since(conn.CreatedAt) < p.maxLifetime {
			conn.State = ConnActive
			conn.LastUsed = time.Now()
			conn.mu.Unlock()
			return conn, nil
		}
		conn.mu.Unlock()
	}

	if len(conns) >= p.maxPerHost {
		return nil, fmt.Errorf("connection pool exhausted for %s", serverURL)
	}

	conn, err := p.dial(ctx, serverURL)
	if err != nil {
		return nil, err
	}
	conn.State = ConnActive
	p.conns[serverURL] = append(p.conns[serverURL], conn)
	return conn, nil
}

func (p *ConnPool) Release(conn *SSEConn) {
	conn.mu.Lock()
	defer conn.mu.Unlock()
	conn.State = ConnIdle
	conn.LastUsed = time.Now()
}

func (p *ConnPool) evictLoop() {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for range ticker.C {
		p.mu.Lock()
		for url, conns := range p.conns {
			alive := conns[:0]
			for _, conn := range conns {
				conn.mu.Lock()
				expired := conn.State == ConnIdle &&
					(time.Since(conn.LastUsed) > p.idleTimeout ||
						time.Since(conn.CreatedAt) > p.maxLifetime)
				if expired {
					close(conn.closeCh)
					conn.mu.Unlock()
					continue
				}
				conn.mu.Unlock()
				alive = append(alive, conn)
			}
			p.conns[url] = alive
		}
		p.mu.Unlock()
	}
}

func (p *ConnPool) dial(ctx context.Context, serverURL string) (*SSEConn, error) {
	conn := &SSEConn{
		ID:        fmt.Sprintf("conn-%d", time.Now().UnixNano()),
		ServerURL: serverURL,
		CreatedAt: time.Now(),
		LastUsed:  time.Now(),
		client:    &http.Client{Timeout: 0},
		eventCh:   make(chan []byte, 256),
		closeCh:   make(chan struct{}),
	}

	req, err := http.NewRequestWithContext(ctx, "GET", serverURL+"/sse", nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "text/event-stream")

	go conn.readLoop(req)
	return conn, nil
}

func (c *SSEConn) readLoop(req *http.Request) {
	resp, err := c.client.Do(req)
	if err != nil {
		return
	}
	defer resp.Body.Close()

	buf := make([]byte, 4096)
	for {
		select {
		case <-c.closeCh:
			return
		default:
			n, err := resp.Body.Read(buf)
			if err != nil {
				return
			}
			if n > 0 {
				data := make([]byte, n)
				copy(data, buf[:n])
				select {
				case c.eventCh <- data:
				case <-c.closeCh:
					return
				}
			}
		}
	}
}

package gateway

import (
	"fmt"
	"hash/crc32"
	"sort"
	"sync"
)

type ServerInfo struct {
	URL     string
	Weight  int
	Tools   []string
	Healthy bool
}

type ToolRouter struct {
	mu       sync.RWMutex
	toolMap  map[string][]*ServerInfo // toolName -> servers
	hashRing *ConsistentHash
}

type ConsistentHash struct {
	ring     map[uint32]*ServerInfo
	keys     []uint32
	replicas int
}

func NewConsistentHash(replicas int) *ConsistentHash {
	return &ConsistentHash{
		ring:     make(map[uint32]*ServerInfo),
		replicas: replicas,
	}
}

func (ch *ConsistentHash) Add(server *ServerInfo) {
	for i := 0; i < ch.replicas; i++ {
		key := crc32.ChecksumIEEE([]byte(fmt.Sprintf("%s-%d", server.URL, i)))
		ch.ring[key] = server
		ch.keys = append(ch.keys, key)
	}
	sort.Slice(ch.keys, func(i, j int) bool { return ch.keys[i] < ch.keys[j] })
}

func (ch *ConsistentHash) Get(key string) *ServerInfo {
	if len(ch.keys) == 0 {
		return nil
	}
	hash := crc32.ChecksumIEEE([]byte(key))
	idx := sort.Search(len(ch.keys), func(i int) bool { return ch.keys[i] >= hash })
	if idx >= len(ch.keys) {
		idx = 0
	}
	return ch.ring[ch.keys[idx]]
}

func NewToolRouter() *ToolRouter {
	return &ToolRouter{
		toolMap:  make(map[string][]*ServerInfo),
		hashRing: NewConsistentHash(150),
	}
}

func (r *ToolRouter) Register(server *ServerInfo) {
	r.mu.Lock()
	defer r.mu.Unlock()

	for _, tool := range server.Tools {
		r.toolMap[tool] = append(r.toolMap[tool], server)
	}
	r.hashRing.Add(server)
}

func (r *ToolRouter) Route(toolName, sessionID string) (*ServerInfo, error) {
	r.mu.RLock()
	defer r.mu.RUnlock()

	servers, ok := r.toolMap[toolName]
	if !ok || len(servers) == 0 {
		return nil, fmt.Errorf("no server registered for tool: %s", toolName)
	}

	healthy := make([]*ServerInfo, 0, len(servers))
	for _, s := range servers {
		if s.Healthy {
			healthy = append(healthy, s)
		}
	}
	if len(healthy) == 0 {
		return nil, fmt.Errorf("all servers for tool %s are unhealthy", toolName)
	}

	if len(healthy) == 1 {
		return healthy[0], nil
	}

	target := r.hashRing.Get(sessionID + ":" + toolName)
	if target != nil && target.Healthy {
		return target, nil
	}
	return healthy[0], nil
}

package gateway

import (
	"context"
	"sync"
	"time"
)

type RateLimiter struct {
	tokens     chan struct{}
	maxTokens  int
	refillRate time.Duration
	stopCh     chan struct{}
}

func NewRateLimiter(maxTokens int, refillRate time.Duration) *RateLimiter {
	rl := &RateLimiter{
		tokens:     make(chan struct{}, maxTokens),
		maxTokens:  maxTokens,
		refillRate: refillRate,
		stopCh:     make(chan struct{}),
	}
	for i := 0; i < maxTokens; i++ {
		rl.tokens <- struct{}{}
	}
	go rl.refill()
	return rl
}

func (rl *RateLimiter) refill() {
	ticker := time.NewTicker(rl.refillRate)
	defer ticker.Stop()

	for {
		select {
		case <-rl.stopCh:
			return
		case <-ticker.C:
			select {
			case rl.tokens <- struct{}{}:
			default:
			}
		}
	}
}

func (rl *RateLimiter) Allow(ctx context.Context) bool {
	select {
	case <-rl.tokens:
		return true
	case <-ctx.Done():
		return false
	}
}

type BackpressureController struct {
	semaphore   chan struct{}
	rateLimiter *RateLimiter
	queueSize   int64
	mu          sync.Mutex
	metrics     *BackpressureMetrics
}

type BackpressureMetrics struct {
	Accepted int64
	Rejected int64
	Queued   int64
}

func NewBackpressureController(maxConcurrent, maxRPS int) *BackpressureController {
	return &BackpressureController{
		semaphore:   make(chan struct{}, maxConcurrent),
		rateLimiter: NewRateLimiter(maxRPS, time.Second/time.Duration(maxRPS)),
		metrics:     &BackpressureMetrics{},
	}
}

func (bp *BackpressureController) Execute(
	ctx context.Context,
	fn func(context.Context) (any, error),
) (any, error) {
	if !bp.rateLimiter.Allow(ctx) {
		bp.mu.Lock()
		bp.metrics.Rejected++
		bp.mu.Unlock()
		return nil, fmt.Errorf("rate limit exceeded")
	}

	select {
	case bp.semaphore <- struct{}{}:
		defer func() { <-bp.semaphore }()
	case <-ctx.Done():
		bp.mu.Lock()
		bp.metrics.Rejected++
		bp.mu.Unlock()
		return nil, ctx.Err()
	}

	bp.mu.Lock()
	bp.metrics.Accepted++
	bp.mu.Unlock()

	return fn(ctx)
}

package gateway

import (
	"context"
	"encoding/json"
	"fmt"
	"time"
)

type SessionData struct {
	SessionID    string            `json:"session_id"`
	UserID       string            `json:"user_id"`
	GatewayNode  string            `json:"gateway_node"`
	ServerURL    string            `json:"server_url"`
	Capabilities map[string]any    `json:"capabilities"`
	Metadata     map[string]string `json:"metadata"`
	CreatedAt    time.Time         `json:"created_at"`
	LastActiveAt time.Time         `json:"last_active_at"`
	TTL          time.Duration     `json:"ttl"`
}

type RedisSessionStore struct {
	client     RedisClient
	keyPrefix  string
	defaultTTL time.Duration
}

type RedisClient interface {
	Get(ctx context.Context, key string) (string, error)
	Set(ctx context.Context, key string, value any, ttl time.Duration) error
	Del(ctx context.Context, keys ...string) error
	Publish(ctx context.Context, channel string, message any) error
}

func NewRedisSessionStore(client RedisClient, prefix string, ttl time.Duration) *RedisSessionStore {
	return &RedisSessionStore{
		client:     client,
		keyPrefix:  prefix,
		defaultTTL: ttl,
	}
}

func (s *RedisSessionStore) Save(ctx context.Context, session *SessionData) error {
	key := fmt.Sprintf("%s:session:%s", s.keyPrefix, session.SessionID)
	data, err := json.Marshal(session)
	if err != nil {
		return fmt.Errorf("marshal session: %w", err)
	}

	ttl := session.TTL
	if ttl == 0 {
		ttl = s.defaultTTL
	}

	return s.client.Set(ctx, key, data, ttl)
}

func (s *RedisSessionStore) Load(ctx context.Context, sessionID string) (*SessionData, error) {
	key := fmt.Sprintf("%s:session:%s", s.keyPrefix, sessionID)
	val, err := s.client.Get(ctx, key)
	if err != nil {
		return nil, fmt.Errorf("load session %s: %w", sessionID, err)
	}

	var session SessionData
	if err := json.Unmarshal([]byte(val), &session); err != nil {
		return nil, fmt.Errorf("unmarshal session: %w", err)
	}
	return &session, nil
}

func (s *RedisSessionStore) UpdateActivity(ctx context.Context, sessionID string) error {
	session, err := s.Load(ctx, sessionID)
	if err != nil {
		return err
	}
	session.LastActiveAt = time.Now()
	return s.Save(ctx, session)
}

func (s *RedisSessionStore) NotifyNode(ctx context.Context, targetNode, sessionID, message string) error {
	channel := fmt.Sprintf("%s:node:%s", s.keyPrefix, targetNode)
	payload := map[string]string{
		"session_id": sessionID,
		"message":    message,
	}
	data, _ := json.Marshal(payload)
	return s.client.Publish(ctx, channel, data)
}

package gateway

import (
	"fmt"
	"sync"
	"time"
)

type CircuitState int

const (
	StateClosed   CircuitState = iota
	StateOpen
	StateHalfOpen
)

type CircuitBreaker struct {
	mu               sync.Mutex
	state            CircuitState
	failureCount     int
	successCount     int
	failureThreshold int
	successThreshold int
	timeout          time.Duration
	lastFailureTime  time.Time
	onStateChange    func(from, to CircuitState)
}

type CircuitBreakerConfig struct {
	FailureThreshold int
	SuccessThreshold int
	Timeout          time.Duration
	OnStateChange    func(from, to CircuitState)
}

func NewCircuitBreaker(cfg CircuitBreakerConfig) *CircuitBreaker {
	return &CircuitBreaker{
		state:            StateClosed,
		failureThreshold: cfg.FailureThreshold,
		successThreshold: cfg.SuccessThreshold,
		timeout:          cfg.Timeout,
		onStateChange:    cfg.OnStateChange,
	}
}

func (cb *CircuitBreaker) Allow() (bool, error) {
	cb.mu.Lock()
	defer cb.mu.Unlock()

	switch cb.state {
	case StateClosed:
		return true, nil
	case StateOpen:
		if time.Since(cb.lastFailureTime) > cb.timeout {
			cb.transitionTo(StateHalfOpen)
			return true, nil
		}
		return false, fmt.Errorf("circuit breaker is open")
	case StateHalfOpen:
		return true, nil
	}
	return false, fmt.Errorf("unknown circuit state")
}

func (cb *CircuitBreaker) RecordSuccess() {
	cb.mu.Lock()
	defer cb.mu.Unlock()

	switch cb.state {
	case StateClosed:
		cb.failureCount = 0
	case StateHalfOpen:
		cb.successCount++
		if cb.successCount >= cb.successThreshold {
			cb.transitionTo(StateClosed)
		}
	}
}

func (cb *CircuitBreaker) RecordFailure() {
	cb.mu.Lock()
	defer cb.mu.Unlock()

	cb.lastFailureTime = time.Now()

	switch cb.state {
	case StateClosed:
		cb.failureCount++
		if cb.failureCount >= cb.failureThreshold {
			cb.transitionTo(StateOpen)
		}
	case StateHalfOpen:
		cb.transitionTo(StateOpen)
	}
}

func (cb *CircuitBreaker) transitionTo(newState CircuitState) {
	oldState := cb.state
	cb.state = newState
	cb.failureCount = 0
	cb.successCount = 0

	if cb.onStateChange != nil {
		cb.onStateChange(oldState, newState)
	}
}

func (cb *CircuitBreaker) State() CircuitState {
	cb.mu.Lock()
	defer cb.mu.Unlock()
	return cb.state
}