Hein
106 lines
2.9 KiB
Go
106 lines
2.9 KiB
Go
package pool
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
"sync"
|
|
"sync/atomic"
|
|
)
|
|
|
|
type PubSubStats struct {
|
|
Created uint32
|
|
Untracked uint32
|
|
Active uint32
|
|
}
|
|
|
|
// PubSubPool manages a pool of PubSub connections.
|
|
type PubSubPool struct {
|
|
opt *Options
|
|
netDialer func(ctx context.Context, network, addr string) (net.Conn, error)
|
|
|
|
// Map to track active PubSub connections
|
|
activeConns sync.Map // map[uint64]*Conn (connID -> conn)
|
|
closed atomic.Bool
|
|
stats PubSubStats
|
|
}
|
|
|
|
// NewPubSubPool implements a pool for PubSub connections.
|
|
// It intentionally does not implement the Pooler interface
|
|
func NewPubSubPool(opt *Options, netDialer func(ctx context.Context, network, addr string) (net.Conn, error)) *PubSubPool {
|
|
return &PubSubPool{
|
|
opt: opt,
|
|
netDialer: netDialer,
|
|
}
|
|
}
|
|
|
|
func (p *PubSubPool) NewConn(ctx context.Context, network string, addr string, channels []string) (*Conn, error) {
|
|
if p.closed.Load() {
|
|
return nil, ErrClosed
|
|
}
|
|
|
|
netConn, err := p.netDialer(ctx, network, addr)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
cn := NewConnWithBufferSize(netConn, p.opt.ReadBufferSize, p.opt.WriteBufferSize)
|
|
cn.pubsub = true
|
|
// Set pool name for metrics
|
|
cn.SetPoolName(p.opt.Name)
|
|
atomic.AddUint32(&p.stats.Created, 1)
|
|
return cn, nil
|
|
}
|
|
|
|
func (p *PubSubPool) TrackConn(cn *Conn) {
|
|
atomic.AddUint32(&p.stats.Active, 1)
|
|
p.activeConns.Store(cn.GetID(), cn)
|
|
// Emit +1 used for PubSub connection
|
|
if cb := getMetricConnectionCountCallback(); cb != nil {
|
|
cb(context.Background(), 1, cn, "used", true)
|
|
}
|
|
}
|
|
|
|
func (p *PubSubPool) UntrackConn(cn *Conn) {
|
|
// LoadAndDelete ensures each connection is only decremented once,
|
|
// guarding against double-decrement if Close() already untracked it.
|
|
if _, loaded := p.activeConns.LoadAndDelete(cn.GetID()); !loaded {
|
|
return
|
|
}
|
|
atomic.AddUint32(&p.stats.Active, ^uint32(0))
|
|
atomic.AddUint32(&p.stats.Untracked, 1)
|
|
// Emit -1 used for PubSub connection
|
|
if cb := getMetricConnectionCountCallback(); cb != nil {
|
|
cb(context.Background(), -1, cn, "used", true)
|
|
}
|
|
}
|
|
|
|
func (p *PubSubPool) Close() error {
|
|
p.closed.Store(true)
|
|
cb := getMetricConnectionCountCallback()
|
|
p.activeConns.Range(func(key, value interface{}) bool {
|
|
cn := value.(*Conn)
|
|
// Use LoadAndDelete to atomically claim ownership of this entry.
|
|
// If a concurrent UntrackConn already removed it, skip to avoid double-decrement.
|
|
if _, loaded := p.activeConns.LoadAndDelete(key); !loaded {
|
|
return true
|
|
}
|
|
atomic.AddUint32(&p.stats.Active, ^uint32(0))
|
|
atomic.AddUint32(&p.stats.Untracked, 1)
|
|
// Emit -1 used for each PubSub connection being closed
|
|
if cb != nil {
|
|
cb(context.Background(), -1, cn, "used", true)
|
|
}
|
|
_ = cn.Close()
|
|
return true
|
|
})
|
|
return nil
|
|
}
|
|
|
|
func (p *PubSubPool) Stats() *PubSubStats {
|
|
// load stats atomically
|
|
return &PubSubStats{
|
|
Created: atomic.LoadUint32(&p.stats.Created),
|
|
Untracked: atomic.LoadUint32(&p.stats.Untracked),
|
|
Active: atomic.LoadUint32(&p.stats.Active),
|
|
}
|
|
}
|