lislgosms/gateway/internal/upstream/flow_control.go

package upstream

import (
	"context"
	"fmt"
	cmpp "github.com/bigwhite/gocmpp"
	"time"
)

// A window token belongs to one physical connection and must be released only
// after its Submit attempt completes, preserving per-connection CMPP flow control.

func (p *connectionPool) acquireConnection(ctx context.Context) (*connection, func(), error) {
	waitCtx, cancel := context.WithTimeout(ctx, defaultSubmitTimeout)
	defer cancel()
	ticker := time.NewTicker(10 * time.Millisecond)
	defer ticker.Stop()

	for {
		if conn, release := p.tryAcquireConnection(); conn != nil {
			if connected, err := conn.ensureConnected(); err != nil {
				release()
				select {
				case <-waitCtx.Done():
					return nil, nil, waitCtx.Err()
				case <-ticker.C:
					continue
				}
			} else if connected {
				_ = p.reportState(context.Background(), "connected", nil)
			}
			return conn, release, nil
		}
		select {
		case <-waitCtx.Done():
			return nil, nil, waitCtx.Err()
		case <-ticker.C:
		}
	}
}

func (p *connectionPool) tryAcquireConnection() (*connection, func()) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if len(p.conns) == 0 {
		return nil, nil
	}
	bestIndex := -1
	bestInFlight := int(^uint(0) >> 1)
	var bestRTT time.Duration
	for index, conn := range p.conns {
		inFlight, rtt, usable := conn.capacitySnapshot()
		if !usable || inFlight > bestInFlight || (inFlight == bestInFlight && bestIndex >= 0 && bestRTT > 0 && rtt >= bestRTT) {
			continue
		}
		bestIndex, bestInFlight, bestRTT = index, inFlight, rtt
	}
	if bestIndex >= 0 {
		conn := p.conns[bestIndex]
		if conn.tryAcquireWindow() {
			p.next = (bestIndex + 1) % len(p.conns)
			return conn, conn.releaseWindow
		}
	}
	return nil, nil
}

func (c *connection) tryAcquireWindow() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.window == nil {
		c.window = make(chan struct{}, maximumWindowSize)
	}
	limit := c.windowLimit
	if limit < 1 {
		limit = defaultWindowSize
	}
	if c.closed || c.draining || c.client == nil || len(c.window) >= limit {
		return false
	}
	select {
	case c.window <- struct{}{}:
		return true
	default:
		return false
	}
}

func (c *connection) capacitySnapshot() (int, time.Duration, bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	limit := c.windowLimit
	if limit < 1 {
		limit = defaultWindowSize
	}
	inFlight := len(c.window)
	return inFlight, c.lastSubmitRTT, !c.closed && !c.draining && c.client != nil && inFlight < limit && !time.Now().Before(c.cooldownUntil)
}

func (c *connection) markSubmitFailure() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.consecutiveFailures++
	if c.consecutiveFailures >= 3 {
		seconds := c.config.FailureCooldownSeconds
		if seconds <= 0 {
			seconds = 30
		}
		c.cooldownUntil = time.Now().Add(time.Duration(seconds) * time.Second)
	}
}

func (c *connection) markSubmitSuccess() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.consecutiveFailures = 0
	c.cooldownUntil = time.Time{}
}

func (c *connection) releaseWindow() {
	if c.window == nil {
		return
	}
	select {
	case <-c.window:
	default:
	}
}

func (c *connection) heartbeatLoop(ctx context.Context) {
	interval := time.Duration(c.config.HeartbeatIntervalSeconds) * time.Second
	if interval <= 0 {
		interval = defaultHeartbeatInterval
	}
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			if !c.sendHeartbeat() {
				return
			}
		}
	}
}

func (c *connection) sendHeartbeat() bool {
	threshold := c.config.HeartbeatMissThreshold
	if threshold <= 0 {
		threshold = defaultHeartbeatMissThreshold
	}
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return false
	}
	if len(c.heartbeatPending) >= threshold {
		c.mu.Unlock()
		c.handleConnectionLoss(fmt.Errorf("heartbeat timeout after %d unanswered ACTIVE_TEST requests", threshold))
		return false
	}
	if c.client == nil {
		c.mu.Unlock()
		return false
	}
	c.sendMu.Lock()
	seq, err := c.client.SendReqPktAvailable(&cmpp.CmppActiveTestReqPkt{}, c.sequenceAvailable)
	c.sendMu.Unlock()
	if err != nil {
		c.mu.Unlock()
		c.handleConnectionLoss(fmt.Errorf("send ACTIVE_TEST: %w", err))
		return false
	}
	if !c.closed {
		c.heartbeatPending[seq] = time.Now().UTC()
	}
	c.mu.Unlock()
	return true
}

func (c *connection) handleHeartbeatResponse(sequenceID uint32) {
	c.mu.Lock()
	delete(c.heartbeatPending, sequenceID)
	c.mu.Unlock()
}