lislgosms/gateway/internal/upstream/pool.go

package upstream

import (
	"cmpp-platform/gateway/internal/protocollog"
	"cmpp-platform/gateway/internal/queue"
	"context"
	"net/http"
	"sync"
	"time"
)

type connectionPool struct {
	channelID            string
	connectionID         string
	config               queue.UpstreamConfig
	apiBaseURL           string
	httpClient           *http.Client
	reporter             func(context.Context, ConnectionState) error
	protocolLogPublisher interface {
		Publish(context.Context, protocollog.Event) error
	}
	gatewayInstanceID string
	eventPublisher    interface {
		PublishReceipt(context.Context, queue.ReceiptEvent) error
		PublishUplink(context.Context, queue.UplinkEvent) error
	}

	mu                     sync.Mutex
	connectMu              sync.Mutex
	conns                  []*connection
	next                   int
	reconnectSignal        chan struct{}
	stopCh                 chan struct{}
	stopOnce               sync.Once
	supervisorOnce         sync.Once
	reconnectCount         int
	lastReconnectAttemptAt time.Time
	nextReconnectAt        time.Time
	lastErrorCategory      string
}

func (p *connectionPool) matches(config queue.UpstreamConfig) bool {
	return p.config == normalizeUpstreamConfig(config)
}

func (p *connectionPool) sameEndpoint(config queue.UpstreamConfig) bool {
	next := normalizeUpstreamConfig(config)
	current := p.config
	return current.GatewayHost == next.GatewayHost && current.GatewayPort == next.GatewayPort &&
		current.Account == next.Account && current.PasswordCipher == next.PasswordCipher && current.CMPPVersion == next.CMPPVersion
}

// reconfigure changes only runtime capacity on the existing pool. Existing
// sequence mappings remain attached to their physical connection. Scale down
// marks surplus connections draining before closing them; scale up is serialized.
func (p *connectionPool) reconfigure(config queue.UpstreamConfig) {
	next := normalizeUpstreamConfig(config)
	p.mu.Lock()
	p.config = next
	for _, conn := range p.conns {
		conn.mu.Lock()
		conn.config = next
		conn.windowLimit = next.WindowSize
		conn.mu.Unlock()
	}
	p.mu.Unlock()
	go p.reconcileCapacity()
}

func (p *connectionPool) reconcileCapacity() {
	p.connectMu.Lock()
	defer p.connectMu.Unlock()
	if p.stopped() {
		return
	}

	p.mu.Lock()
	desired := max(1, min(maximumConnections, p.config.DesiredConnections))
	if len(p.conns) > desired {
		retiring := append([]*connection(nil), p.conns[desired:]...)
		p.conns = p.conns[:desired]
		for _, conn := range retiring {
			conn.mu.Lock()
			conn.draining = true
			conn.mu.Unlock()
		}
		p.mu.Unlock()
		deadline := time.Now().Add(time.Duration(p.config.ConnectionDrainSeconds) * time.Second)
		for _, conn := range retiring {
			for len(conn.window) > 0 && time.Now().Before(deadline) {
				time.Sleep(10 * time.Millisecond)
			}
			conn.retire()
		}
		_ = p.reportState(context.Background(), "connected", nil)
		return
	}
	p.mu.Unlock()

	for {
		p.mu.Lock()
		if len(p.conns) >= desired || p.stopped() {
			p.mu.Unlock()
			break
		}
		index := len(p.conns)
		config := p.config
		p.mu.Unlock()
		conn := p.newConnection(index, config)
		if _, err := conn.ensureConnected(); err != nil {
			p.scheduleReconnect(err)
			_ = p.reportState(context.Background(), "failed", err)
			return
		}
		p.mu.Lock()
		p.conns = append(p.conns, conn)
		p.mu.Unlock()
		_ = p.reportState(context.Background(), "connected", nil)
		if len(p.conns) < desired && config.ConnectionWarmupSeconds > 0 {
			timer := time.NewTimer(time.Duration(config.ConnectionWarmupSeconds) * time.Second)
			select {
			case <-p.stopCh:
				timer.Stop()
				return
			case <-timer.C:
			}
		}
	}
}

func (p *connectionPool) newConnection(index int, config queue.UpstreamConfig) *connection {
	return &connection{
		channelID: p.channelID, config: config, index: index, pool: p,
		apiBaseURL: p.apiBaseURL, httpClient: p.httpClient,
		protocolLogPublisher: p.protocolLogPublisher, gatewayInstanceID: p.gatewayInstanceID,
		eventPublisher: p.eventPublisher,
		window:         make(chan struct{}, maximumWindowSize), windowLimit: config.WindowSize,
		pending: make(map[uint32]chan submitPartResponse), tracker: make(map[uint64]queue.SubmitCommand),
		longUplink: make(map[string]*longUplinkAssembly), heartbeatPending: make(map[uint32]time.Time),
	}
}

func (p *connectionPool) ensureConnected() error {
	p.connectMu.Lock()
	defer p.connectMu.Unlock()
	desired := p.config.DesiredConnections
	if desired <= 0 {
		desired = 1
	}
	connectedAny := false
	for {
		p.mu.Lock()
		active := p.conns[:0]
		for _, existing := range p.conns {
			existing.mu.Lock()
			usable := existing.client != nil && !existing.closed
			existing.mu.Unlock()
			if usable {
				active = append(active, existing)
			}
		}
		p.conns = active
		if len(p.conns) >= desired {
			p.mu.Unlock()
			break
		}
		if connectedAny {
			p.mu.Unlock()
			break
		}
		index := len(p.conns)
		conn := p.newConnection(index, p.config)
		p.mu.Unlock()

		connected, err := conn.ensureConnected()
		if err != nil {
			return err
		}

		p.mu.Lock()
		p.conns = append(p.conns, conn)
		p.mu.Unlock()
		connectedAny = connectedAny || connected
	}
	if connectedAny {
		_ = p.reportState(context.Background(), "connected", nil)
	}
	go p.reconcileCapacity()
	return nil
}

func (p *connectionPool) close() {
	p.connectMu.Lock()
	defer p.connectMu.Unlock()
	p.stopOnce.Do(func() {
		close(p.stopCh)
	})
	p.mu.Lock()
	conns := p.conns
	p.conns = nil
	p.mu.Unlock()
	for _, conn := range conns {
		conn.close()
	}
}

func (p *connectionPool) countActiveConnections() int {
	p.mu.Lock()
	defer p.mu.Unlock()
	count := 0
	for _, conn := range p.conns {
		conn.mu.Lock()
		active := conn.client != nil && !conn.closed
		conn.mu.Unlock()
		if active {
			count += 1
		}
	}
	return count
}