From 0bafca94868a4ec38cd3dba0f2ba7ec6f44bbf18 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=A3=8E=E6=89=87=E6=BB=91=E7=BF=94=E7=BF=BC?= Date: Sat, 11 Jul 2026 04:34:27 +0800 Subject: [PATCH] Stats: Fix GetOrRegister*() races (#6468) Fixes https://github.com/XTLS/Xray-core/pull/6467#issue-4856851059 --------- Co-authored-by: Joshua Su --- app/dispatcher/default.go | 10 +++---- app/proxyman/inbound/always.go | 4 +-- app/proxyman/outbound/handler.go | 4 +-- app/stats/stats.go | 48 ++++++++++++++++++++++++++++++ features/stats/stats.go | 51 +++++++++++++------------------- proxy/tun/handler.go | 4 +-- proxy/wireguard/client.go | 4 +-- proxy/wireguard/server.go | 4 +-- 8 files changed, 84 insertions(+), 45 deletions(-) diff --git a/app/dispatcher/default.go b/app/dispatcher/default.go index f6cfd76eb..2717df5fe 100644 --- a/app/dispatcher/default.go +++ b/app/dispatcher/default.go @@ -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) }) } diff --git a/app/proxyman/inbound/always.go b/app/proxyman/inbound/always.go index 462f894f5..e0d6c0925 100644 --- a/app/proxyman/inbound/always.go +++ b/app/proxyman/inbound/always.go @@ -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 } diff --git a/app/proxyman/outbound/handler.go b/app/proxyman/outbound/handler.go index 2fe0fec55..5f4833e9b 100644 --- a/app/proxyman/outbound/handler.go +++ b/app/proxyman/outbound/handler.go @@ -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 } diff --git a/app/stats/stats.go b/app/stats/stats.go index 68fb9be18..6df8c499e 100644 --- a/app/stats/stats.go +++ b/app/stats/stats.go @@ -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() diff --git a/features/stats/stats.go b/features/stats/stats.go index 1d490c10b..e1333c32a 100644 --- a/features/stats/stats.go +++ b/features/stats/stats.go @@ -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 diff --git a/proxy/tun/handler.go b/proxy/tun/handler.go index 5bf8f231f..56fe5653a 100644 --- a/proxy/tun/handler.go +++ b/proxy/tun/handler.go @@ -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 } diff --git a/proxy/wireguard/client.go b/proxy/wireguard/client.go index c50dd995f..cae846077 100644 --- a/proxy/wireguard/client.go +++ b/proxy/wireguard/client.go @@ -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 } diff --git a/proxy/wireguard/server.go b/proxy/wireguard/server.go index e26479ad6..51416c435 100644 --- a/proxy/wireguard/server.go +++ b/proxy/wireguard/server.go @@ -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 }