fix: align channel copy and gateway connection state

This commit is contained in:
hectorzhao
2026-07-09 14:57:55 +08:00
parent 2a65b41c4e
commit 06e6be3c39
12 changed files with 309 additions and 157 deletions
@@ -1,14 +1,26 @@
package upstream
import (
"context"
"errors"
"testing"
)
func TestHandleConnectionLossNotifiesPendingSubmitters(t *testing.T) {
var reported ConnectionState
conn := &connection{
pending: make(map[uint32]chan submitPartResponse),
pool: &connectionPool{
channelID: "channel-1",
connectionID: "channel-1:primary",
config: normalizeUpstreamConfig(queueUpstreamConfigForTest()),
reporter: func(_ context.Context, state ConnectionState) error {
reported = state
return nil
},
},
}
conn.pool.conns = []*connection{conn}
waiter := make(chan submitPartResponse, 1)
conn.pending[7] = waiter
@@ -30,6 +42,9 @@ func TestHandleConnectionLossNotifiesPendingSubmitters(t *testing.T) {
if len(conn.pending) != 0 {
t.Fatalf("expected pending map to be reset, got %d entries", len(conn.pending))
}
if reported.Status != "disconnected" || reported.CurrentConnections != 0 {
t.Fatalf("unexpected reported state: %+v", reported)
}
}
func TestTemporaryReadTimeoutDetection(t *testing.T) {
+185 -34
View File
@@ -33,6 +33,19 @@ type Manager struct {
conns map[string]*connectionPool
}
type ConnectionState struct {
ChannelID string `json:"channelId"`
ConnectionID string `json:"connectionId"`
Status string `json:"status"`
DesiredConnections int `json:"desiredConnections"`
CurrentConnections int `json:"currentConnections"`
LastConnectedAt string `json:"lastConnectedAt,omitempty"`
LastDisconnectedAt string `json:"lastDisconnectedAt,omitempty"`
LastHeartbeatAt string `json:"lastHeartbeatAt,omitempty"`
ReconnectCount int `json:"reconnectCount,omitempty"`
LastError string `json:"lastError,omitempty"`
}
func (m *Manager) Submit(ctx context.Context, cmd queue.SubmitCommand) (queue.SubmitResult, error) {
if err := validateSubmitCommand(cmd); err != nil {
result := submitResult(cmd, 0, "", "rejected", "INVALID_COMMAND", err.Error())
@@ -64,28 +77,54 @@ func (m *Manager) Submit(ctx context.Context, cmd queue.SubmitCommand) (queue.Su
return result, nil
}
func (m *Manager) ConnectChannel(ctx context.Context, command queue.ConnectChannelCommand) (ConnectionState, error) {
if err := validateConnectChannelCommand(command); err != nil {
return ConnectionState{}, err
}
m.mu.Lock()
m.ensureDefaultsLocked()
pool := m.conns[command.ChannelID]
config := normalizeUpstreamConfig(queue.UpstreamConfig{
GatewayHost: command.Channel.GatewayHost,
GatewayPort: command.Channel.GatewayPort,
Account: command.Channel.Account,
PasswordCipher: command.Channel.PasswordCipher,
CMPPVersion: command.Channel.CMPPVersion,
DesiredConnections: command.DesiredConnections,
WindowSize: 16,
})
if pool == nil || !pool.matches(config) {
if pool != nil {
pool.close()
}
pool = m.newConnectionPool(command.ChannelID, command.ConnectionID, config)
m.conns[command.ChannelID] = pool
}
m.mu.Unlock()
if err := pool.ensureConnected(); err != nil {
_ = pool.reportState(ctx, "failed", err)
m.mu.Lock()
delete(m.conns, command.ChannelID)
m.mu.Unlock()
return pool.snapshotState("failed", err), nil
}
return pool.snapshotState("connected", nil), nil
}
func (m *Manager) connectionFor(cmd queue.SubmitCommand) (*connectionPool, error) {
m.mu.Lock()
defer m.mu.Unlock()
if m.HTTPClient == nil {
m.HTTPClient = &http.Client{Timeout: defaultHTTPTimeout}
}
if m.conns == nil {
m.conns = make(map[string]*connectionPool)
}
m.ensureDefaultsLocked()
pool := m.conns[cmd.ChannelID]
if pool == nil || !pool.matches(cmd.Upstream) {
if pool != nil {
pool.close()
}
pool = &connectionPool{
channelID: cmd.ChannelID,
config: normalizeUpstreamConfig(cmd.Upstream),
apiBaseURL: m.APIBaseURL,
httpClient: m.HTTPClient,
}
pool = m.newConnectionPool(cmd.ChannelID, defaultChannelConnectionID(cmd.ChannelID), normalizeUpstreamConfig(cmd.Upstream))
m.conns[cmd.ChannelID] = pool
}
if err := pool.ensureConnected(); err != nil {
@@ -103,11 +142,35 @@ func (m *Manager) post(ctx context.Context, path string, payload any) error {
return postJSON(ctx, client, m.APIBaseURL, path, payload)
}
func (m *Manager) ensureDefaultsLocked() {
if m.HTTPClient == nil {
m.HTTPClient = &http.Client{Timeout: defaultHTTPTimeout}
}
if m.conns == nil {
m.conns = make(map[string]*connectionPool)
}
}
func (m *Manager) newConnectionPool(channelID string, connectionID string, config queue.UpstreamConfig) *connectionPool {
return &connectionPool{
channelID: channelID,
connectionID: connectionID,
config: config,
apiBaseURL: m.APIBaseURL,
httpClient: m.HTTPClient,
reporter: func(ctx context.Context, state ConnectionState) error {
return m.post(ctx, "/admin/gateway/connections", state)
},
}
}
type connectionPool struct {
channelID string
config queue.UpstreamConfig
apiBaseURL string
httpClient *http.Client
channelID string
connectionID string
config queue.UpstreamConfig
apiBaseURL string
httpClient *http.Client
reporter func(context.Context, ConnectionState) error
mu sync.Mutex
conns []*connection
@@ -119,19 +182,23 @@ func (p *connectionPool) matches(config queue.UpstreamConfig) bool {
}
func (p *connectionPool) ensureConnected() error {
p.mu.Lock()
defer p.mu.Unlock()
desired := p.config.DesiredConnections
if desired <= 0 {
desired = 1
}
for len(p.conns) < desired {
connectedAny := false
for {
p.mu.Lock()
if len(p.conns) >= desired {
p.mu.Unlock()
break
}
index := len(p.conns)
conn := &connection{
channelID: p.channelID,
config: p.config,
index: index,
pool: p,
apiBaseURL: p.apiBaseURL,
httpClient: p.httpClient,
window: make(chan struct{}, p.config.WindowSize),
@@ -139,12 +206,22 @@ func (p *connectionPool) ensureConnected() error {
tracker: make(map[uint64]queue.SubmitCommand),
longUplink: make(map[string]*longUplinkAssembly),
}
if err := conn.ensureConnected(); err != nil {
p.mu.Unlock()
connected, err := conn.ensureConnected()
if err != nil {
conn.close()
p.closeLocked()
p.close()
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)
}
return nil
}
@@ -198,7 +275,7 @@ func (p *connectionPool) acquireConnection(ctx context.Context) (*connection, fu
for {
if conn, release := p.tryAcquireConnection(); conn != nil {
if err := conn.ensureConnected(); err != nil {
if connected, err := conn.ensureConnected(); err != nil {
release()
select {
case <-waitCtx.Done():
@@ -206,6 +283,8 @@ func (p *connectionPool) acquireConnection(ctx context.Context) (*connection, fu
case <-ticker.C:
continue
}
} else if connected {
_ = p.reportState(context.Background(), "connected", nil)
}
return conn, release, nil
}
@@ -236,21 +315,66 @@ func (p *connectionPool) tryAcquireConnection() (*connection, func()) {
func (p *connectionPool) close() {
p.mu.Lock()
defer p.mu.Unlock()
p.closeLocked()
}
func (p *connectionPool) closeLocked() {
for _, conn := range p.conns {
conns := p.conns
p.conns = nil
p.mu.Unlock()
for _, conn := range conns {
conn.close()
}
p.conns = nil
}
func (p *connectionPool) reportState(ctx context.Context, status string, stateErr error) error {
if p.reporter == nil {
return nil
}
return p.reporter(ctx, p.snapshotState(status, stateErr))
}
func (p *connectionPool) snapshotState(status string, stateErr error) ConnectionState {
now := time.Now().UTC().Format(time.RFC3339Nano)
state := ConnectionState{
ChannelID: p.channelID,
ConnectionID: p.connectionID,
Status: status,
DesiredConnections: p.config.DesiredConnections,
CurrentConnections: p.countActiveConnections(),
}
if state.DesiredConnections <= 0 {
state.DesiredConnections = 1
}
switch status {
case "connected":
state.LastConnectedAt = now
state.LastHeartbeatAt = now
case "disconnected", "failed":
state.LastDisconnectedAt = now
}
if stateErr != nil {
state.LastError = stateErr.Error()
}
return state
}
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
}
type connection struct {
channelID string
config queue.UpstreamConfig
index int
pool *connectionPool
apiBaseURL string
httpClient *http.Client
@@ -274,23 +398,23 @@ func (c *connection) matches(config queue.UpstreamConfig) bool {
return c.config == normalizeUpstreamConfig(config)
}
func (c *connection) ensureConnected() error {
func (c *connection) ensureConnected() (bool, error) {
c.mu.Lock()
defer c.mu.Unlock()
if c.client != nil && !c.closed {
return nil
return false, nil
}
client := cmpp.NewClient(protocolVersion(c.config.CMPPVersion))
addr := fmt.Sprintf("%s:%d", c.config.GatewayHost, c.config.GatewayPort)
if err := client.Connect(addr, c.config.Account, c.config.PasswordCipher, defaultConnectTimeout); err != nil {
client.Disconnect()
return err
return false, err
}
c.client = client
c.closed = false
go c.readLoop()
return nil
return true, nil
}
func (c *connection) submitPart(ctx context.Context, cmd queue.SubmitCommand, part submitPart) (uint32, string, queue.SubmitResult, error) {
@@ -540,6 +664,13 @@ func (c *connection) handleConnectionLoss(err error) {
default:
}
}
if c.pool != nil {
status := "disconnected"
if c.pool.countActiveConnections() > 0 {
status = "connected"
}
_ = c.pool.reportState(context.Background(), status, err)
}
}
func submitResult(cmd queue.SubmitCommand, sequenceID uint32, gatewayMessageID string, status string, code string, message string) queue.SubmitResult {
@@ -599,6 +730,26 @@ func validateSubmitCommand(cmd queue.SubmitCommand) error {
return nil
}
func validateConnectChannelCommand(command queue.ConnectChannelCommand) error {
if command.MessageType != queue.MessageTypeConnectChannel {
return fmt.Errorf("unsupported messageType %q", command.MessageType)
}
if command.ChannelID == "" || command.ConnectionID == "" {
return fmt.Errorf("channelId and connectionId are required")
}
if command.Channel.GatewayHost == "" || command.Channel.GatewayPort <= 0 {
return fmt.Errorf("gatewayHost and gatewayPort are required")
}
if command.Channel.Account == "" || command.Channel.PasswordCipher == "" {
return fmt.Errorf("account and passwordCipher are required")
}
return nil
}
func defaultChannelConnectionID(channelID string) string {
return fmt.Sprintf("%s:primary", channelID)
}
func normalizeUpstreamConfig(config queue.UpstreamConfig) queue.UpstreamConfig {
if config.DesiredConnections <= 0 {
config.DesiredConnections = 1