139 lines
2.9 KiB
Go
139 lines
2.9 KiB
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
|
|
}
|
|
for i := 0; i < len(p.conns); i++ {
|
|
index := (p.next + i) % len(p.conns)
|
|
conn := p.conns[index]
|
|
if conn.tryAcquireWindow() {
|
|
p.next = (index + 1) % len(p.conns)
|
|
return conn, conn.releaseWindow
|
|
}
|
|
}
|
|
return nil, nil
|
|
}
|
|
|
|
func (c *connection) tryAcquireWindow() bool {
|
|
if c.window == nil {
|
|
c.window = make(chan struct{}, defaultWindowSize)
|
|
}
|
|
select {
|
|
case c.window <- struct{}{}:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
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.SendReqPkt(&cmpp.CmppActiveTestReqPkt{})
|
|
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()
|
|
}
|