Stats: Fix GetOrRegister*() races (#6468)

Fixes https://github.com/XTLS/Xray-core/pull/6467#issue-4856851059

---------

Co-authored-by: Joshua Su <i@joshua.su>
This commit is contained in:
风扇滑翔翼 2026-07-11 04:34:27 +08:00 committed by GitHub
parent d5bc58dc6b
commit 0bafca9486
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 84 additions and 45 deletions

View file

@ -162,7 +162,7 @@ func (d *DefaultDispatcher) getLink(ctx context.Context) (*transport.Link, *tran
p := d.policy.ForLevel(user.Level)
if p.Stats.UserUplink {
name := "user>>>" + user.Email + ">>>traffic>>>uplink"
if c, _ := stats.GetOrRegisterCounter(d.stats, name); c != nil {
if c, _ := d.stats.GetOrRegisterCounter(name); c != nil {
inboundLink.Writer = &SizeStatWriter{
Counter: c,
Writer: inboundLink.Writer,
@ -171,7 +171,7 @@ func (d *DefaultDispatcher) getLink(ctx context.Context) (*transport.Link, *tran
}
if p.Stats.UserDownlink {
name := "user>>>" + user.Email + ">>>traffic>>>downlink"
if c, _ := stats.GetOrRegisterCounter(d.stats, name); c != nil {
if c, _ := d.stats.GetOrRegisterCounter(name); c != nil {
outboundLink.Writer = &SizeStatWriter{
Counter: c,
Writer: outboundLink.Writer,
@ -200,13 +200,13 @@ func WrapLink(ctx context.Context, policyManager policy.Manager, statsManager st
p := policyManager.ForLevel(user.Level)
if p.Stats.UserUplink {
name := "user>>>" + user.Email + ">>>traffic>>>uplink"
if c, _ := stats.GetOrRegisterCounter(statsManager, name); c != nil {
if c, _ := statsManager.GetOrRegisterCounter(name); c != nil {
link.Reader.(*buf.TimeoutWrapperReader).Counter = c
}
}
if p.Stats.UserDownlink {
name := "user>>>" + user.Email + ">>>traffic>>>downlink"
if c, _ := stats.GetOrRegisterCounter(statsManager, name); c != nil {
if c, _ := statsManager.GetOrRegisterCounter(name); c != nil {
link.Writer = &SizeStatWriter{
Counter: c,
Writer: link.Writer,
@ -223,7 +223,7 @@ func WrapLink(ctx context.Context, policyManager policy.Manager, statsManager st
func trackOnlineIP(ctx context.Context, sm stats.Manager, email, ip string) {
name := "user>>>" + email + ">>>online"
if om, _ := stats.GetOrRegisterOnlineMap(sm, name); om != nil {
if om, _ := sm.GetOrRegisterOnlineMap(name); om != nil {
om.AddIP(ip)
context.AfterFunc(ctx, func() { om.RemoveIP(ip) })
}

View file

@ -26,7 +26,7 @@ func getStatCounter(v *core.Instance, tag string) (stats.Counter, stats.Counter)
if len(tag) > 0 && policy.ForSystem().Stats.InboundUplink {
statsManager := v.GetFeature(stats.ManagerType()).(stats.Manager)
name := "inbound>>>" + tag + ">>>traffic>>>uplink"
c, _ := stats.GetOrRegisterCounter(statsManager, name)
c, _ := statsManager.GetOrRegisterCounter(name)
if c != nil {
uplinkCounter = c
}
@ -34,7 +34,7 @@ func getStatCounter(v *core.Instance, tag string) (stats.Counter, stats.Counter)
if len(tag) > 0 && policy.ForSystem().Stats.InboundDownlink {
statsManager := v.GetFeature(stats.ManagerType()).(stats.Manager)
name := "inbound>>>" + tag + ">>>traffic>>>downlink"
c, _ := stats.GetOrRegisterCounter(statsManager, name)
c, _ := statsManager.GetOrRegisterCounter(name)
if c != nil {
downlinkCounter = c
}

View file

@ -39,7 +39,7 @@ func getStatCounter(v *core.Instance, tag string) (stats.Counter, stats.Counter)
if len(tag) > 0 && policy.ForSystem().Stats.OutboundUplink {
statsManager := v.GetFeature(stats.ManagerType()).(stats.Manager)
name := "outbound>>>" + tag + ">>>traffic>>>uplink"
c, _ := stats.GetOrRegisterCounter(statsManager, name)
c, _ := statsManager.GetOrRegisterCounter(name)
if c != nil {
uplinkCounter = c
}
@ -47,7 +47,7 @@ func getStatCounter(v *core.Instance, tag string) (stats.Counter, stats.Counter)
if len(tag) > 0 && policy.ForSystem().Stats.OutboundDownlink {
statsManager := v.GetFeature(stats.ManagerType()).(stats.Manager)
name := "outbound>>>" + tag + ">>>traffic>>>downlink"
c, _ := stats.GetOrRegisterCounter(statsManager, name)
c, _ := statsManager.GetOrRegisterCounter(name)
if c != nil {
downlinkCounter = c
}

View file

@ -48,6 +48,20 @@ func (m *Manager) RegisterCounter(name string) (stats.Counter, error) {
return c, nil
}
// GetOrRegisterCounter implements stats.Manager.
func (m *Manager) GetOrRegisterCounter(name string) (stats.Counter, error) {
m.access.Lock()
defer m.access.Unlock()
if c, found := m.counters[name]; found {
return c, nil
}
errors.LogDebug(context.Background(), "create new counter ", name)
c := new(Counter)
m.counters[name] = c
return c, nil
}
// UnregisterCounter implements stats.Manager.
func (m *Manager) UnregisterCounter(name string) error {
m.access.Lock()
@ -97,6 +111,20 @@ func (m *Manager) RegisterOnlineMap(name string) (stats.OnlineMap, error) {
return om, nil
}
// GetOrRegisterOnlineMap implements stats.Manager.
func (m *Manager) GetOrRegisterOnlineMap(name string) (stats.OnlineMap, error) {
m.access.Lock()
defer m.access.Unlock()
if om, found := m.onlineMaps[name]; found {
return om, nil
}
errors.LogDebug(context.Background(), "create new OnlineMap ", name)
om := NewOnlineMap()
m.onlineMaps[name] = om
return om, nil
}
// UnregisterOnlineMap implements stats.Manager.
func (m *Manager) UnregisterOnlineMap(name string) error {
m.access.Lock()
@ -149,6 +177,26 @@ func (m *Manager) RegisterChannel(name string) (stats.Channel, error) {
return c, nil
}
// GetOrRegisterChannel implements stats.Manager.
func (m *Manager) GetOrRegisterChannel(name string) (stats.Channel, error) {
m.access.Lock()
defer m.access.Unlock()
if c, found := m.channels[name]; found {
return c, nil
}
errors.LogDebug(context.Background(), "create new channel ", name)
c := NewChannel(&ChannelConfig{BufferSize: 64, Blocking: false})
if m.running {
// Start before publishing so no goroutine can observe an unstarted channel.
if err := c.Start(); err != nil {
return nil, err
}
}
m.channels[name] = c
return c, nil
}
// UnregisterChannel implements stats.Manager.
func (m *Manager) UnregisterChannel(name string) error {
m.access.Lock()

View file

@ -81,6 +81,8 @@ type Manager interface {
// RegisterCounter registers a new counter to the manager. The identifier string must not be empty, and unique among other counters.
RegisterCounter(string) (Counter, error)
// GetOrRegisterCounter returns the counter by its identifier, atomically creating and registering it if absent.
GetOrRegisterCounter(string) (Counter, error)
// UnregisterCounter unregisters a counter from the manager by its identifier.
UnregisterCounter(string) error
// GetCounter returns a counter by its identifier.
@ -91,6 +93,8 @@ type Manager interface {
// RegisterOnlineMap registers a new OnlineMap to the manager. The identifier string must not be empty, and unique among other OnlineMaps.
RegisterOnlineMap(string) (OnlineMap, error)
// GetOrRegisterOnlineMap returns the OnlineMap by its identifier, atomically creating and registering it if absent.
GetOrRegisterOnlineMap(string) (OnlineMap, error)
// UnregisterOnlineMap unregisters an OnlineMap from the manager by its identifier.
UnregisterOnlineMap(string) error
// GetOnlineMap returns an OnlineMap by its identifier.
@ -101,6 +105,8 @@ type Manager interface {
// RegisterChannel registers a new channel to the manager. The identifier string must not be empty, and unique among other channels.
RegisterChannel(string) (Channel, error)
// GetOrRegisterChannel returns the channel by its identifier, atomically creating and registering it if absent.
GetOrRegisterChannel(string) (Channel, error)
// UnregisterChannel unregisters a channel from the manager by its identifier.
UnregisterChannel(string) error
// GetChannel returns a channel by its identifier.
@ -110,36 +116,6 @@ type Manager interface {
GetAllOnlineUsers() []string
}
// GetOrRegisterCounter tries to get the StatCounter first. If not exist, it then tries to create a new counter.
func GetOrRegisterCounter(m Manager, name string) (Counter, error) {
counter := m.GetCounter(name)
if counter != nil {
return counter, nil
}
return m.RegisterCounter(name)
}
// GetOrRegisterOnlineMap tries to get the OnlineMap first. If not exist, it then tries to create a new OnlineMap.
func GetOrRegisterOnlineMap(m Manager, name string) (OnlineMap, error) {
onlineMap := m.GetOnlineMap(name)
if onlineMap != nil {
return onlineMap, nil
}
return m.RegisterOnlineMap(name)
}
// GetOrRegisterChannel tries to get the StatChannel first. If not exist, it then tries to create a new channel.
func GetOrRegisterChannel(m Manager, name string) (Channel, error) {
channel := m.GetChannel(name)
if channel != nil {
return channel, nil
}
return m.RegisterChannel(name)
}
// ManagerType returns the type of Manager interface. Can be used to implement common.HasType.
//
// xray:api:stable
@ -160,6 +136,11 @@ func (NoopManager) RegisterCounter(string) (Counter, error) {
return nil, errors.New("not implemented")
}
// GetOrRegisterCounter implements Manager.
func (NoopManager) GetOrRegisterCounter(string) (Counter, error) {
return nil, errors.New("not implemented")
}
// UnregisterCounter implements Manager.
func (NoopManager) UnregisterCounter(string) error {
return nil
@ -178,6 +159,11 @@ func (NoopManager) RegisterOnlineMap(string) (OnlineMap, error) {
return nil, errors.New("not implemented")
}
// GetOrRegisterOnlineMap implements Manager.
func (NoopManager) GetOrRegisterOnlineMap(string) (OnlineMap, error) {
return nil, errors.New("not implemented")
}
// UnregisterOnlineMap implements Manager.
func (NoopManager) UnregisterOnlineMap(string) error {
return nil
@ -196,6 +182,11 @@ func (NoopManager) RegisterChannel(string) (Channel, error) {
return nil, errors.New("not implemented")
}
// GetOrRegisterChannel implements Manager.
func (NoopManager) GetOrRegisterChannel(string) (Channel, error) {
return nil, errors.New("not implemented")
}
// UnregisterChannel implements Manager.
func (NoopManager) UnregisterChannel(string) error {
return nil

View file

@ -65,7 +65,7 @@ func (t *Handler) Init(ctx context.Context, pm policy.Manager, dispatcher routin
if len(t.tag) > 0 && pm.ForSystem().Stats.InboundUplink {
statsManager := core.MustFromContext(ctx).GetFeature(stats.ManagerType()).(stats.Manager)
name := "inbound>>>" + t.tag + ">>>traffic>>>uplink"
c, _ := stats.GetOrRegisterCounter(statsManager, name)
c, _ := statsManager.GetOrRegisterCounter(name)
if c != nil {
t.uplinkCounter = c
}
@ -73,7 +73,7 @@ func (t *Handler) Init(ctx context.Context, pm policy.Manager, dispatcher routin
if len(t.tag) > 0 && pm.ForSystem().Stats.InboundDownlink {
statsManager := core.MustFromContext(ctx).GetFeature(stats.ManagerType()).(stats.Manager)
name := "inbound>>>" + t.tag + ">>>traffic>>>downlink"
c, _ := stats.GetOrRegisterCounter(statsManager, name)
c, _ := statsManager.GetOrRegisterCounter(name)
if c != nil {
t.downlinkCounter = c
}

View file

@ -57,7 +57,7 @@ func NewClient(ctx context.Context, conf *DeviceConfig) (*Handler, error) {
if len(tag) > 0 && p.ForSystem().Stats.OutboundUplink {
statsManager := v.GetFeature(stats.ManagerType()).(stats.Manager)
name := "outbound>>>" + tag + ">>>traffic>>>uplink"
c, _ := stats.GetOrRegisterCounter(statsManager, name)
c, _ := statsManager.GetOrRegisterCounter(name)
if c != nil {
uplinkCounter = c
}
@ -65,7 +65,7 @@ func NewClient(ctx context.Context, conf *DeviceConfig) (*Handler, error) {
if len(tag) > 0 && p.ForSystem().Stats.OutboundDownlink {
statsManager := v.GetFeature(stats.ManagerType()).(stats.Manager)
name := "outbound>>>" + tag + ">>>traffic>>>downlink"
c, _ := stats.GetOrRegisterCounter(statsManager, name)
c, _ := statsManager.GetOrRegisterCounter(name)
if c != nil {
downlinkCounter = c
}

View file

@ -66,7 +66,7 @@ func NewServer(ctx context.Context, conf *DeviceConfig) (*Server, error) {
if len(tag) > 0 && p.ForSystem().Stats.InboundUplink {
statsManager := v.GetFeature(stats.ManagerType()).(stats.Manager)
name := "inbound>>>" + tag + ">>>traffic>>>uplink"
c, _ := stats.GetOrRegisterCounter(statsManager, name)
c, _ := statsManager.GetOrRegisterCounter(name)
if c != nil {
uplinkCounter = c
}
@ -74,7 +74,7 @@ func NewServer(ctx context.Context, conf *DeviceConfig) (*Server, error) {
if len(tag) > 0 && p.ForSystem().Stats.InboundDownlink {
statsManager := v.GetFeature(stats.ManagerType()).(stats.Manager)
name := "inbound>>>" + tag + ">>>traffic>>>downlink"
c, _ := stats.GetOrRegisterCounter(statsManager, name)
c, _ := statsManager.GetOrRegisterCounter(name)
if c != nil {
downlinkCounter = c
}