diff --git a/app/dispatcher/default.go b/app/dispatcher/default.go index c1b4e3bf1..9acb5df02 100644 --- a/app/dispatcher/default.go +++ b/app/dispatcher/default.go @@ -15,12 +15,11 @@ import ( "v2ray.com/core/common/net" "v2ray.com/core/common/protocol" "v2ray.com/core/common/session" - "v2ray.com/core/common/stats" "v2ray.com/core/common/vio" "v2ray.com/core/features/outbound" "v2ray.com/core/features/policy" "v2ray.com/core/features/routing" - feature_stats "v2ray.com/core/features/stats" + "v2ray.com/core/features/stats" "v2ray.com/core/transport/pipe" ) @@ -87,7 +86,7 @@ type DefaultDispatcher struct { ohm outbound.Manager router routing.Router policy policy.Manager - stats feature_stats.Manager + stats stats.Manager } // NewDefaultDispatcher create a new DefaultDispatcher. @@ -144,8 +143,8 @@ func (d *DefaultDispatcher) getLink(ctx context.Context) (*vio.Link, *vio.Link) p := d.policy.ForLevel(user.Level) if p.Stats.UserUplink { name := "user>>>" + user.Email + ">>>traffic>>>uplink" - if c, _ := feature_stats.GetOrRegisterCounter(d.stats, name); c != nil { - inboundLink.Writer = &stats.SizeStatWriter{ + if c, _ := stats.GetOrRegisterCounter(d.stats, name); c != nil { + inboundLink.Writer = &vio.SizeStatWriter{ Counter: c, Writer: inboundLink.Writer, } @@ -153,8 +152,8 @@ func (d *DefaultDispatcher) getLink(ctx context.Context) (*vio.Link, *vio.Link) } if p.Stats.UserDownlink { name := "user>>>" + user.Email + ">>>traffic>>>downlink" - if c, _ := feature_stats.GetOrRegisterCounter(d.stats, name); c != nil { - outboundLink.Writer = &stats.SizeStatWriter{ + if c, _ := stats.GetOrRegisterCounter(d.stats, name); c != nil { + outboundLink.Writer = &vio.SizeStatWriter{ Counter: c, Writer: outboundLink.Writer, } diff --git a/common/stats/io.go b/common/vio/stats.go similarity index 97% rename from common/stats/io.go rename to common/vio/stats.go index 60ffabac6..4b6c81415 100644 --- a/common/stats/io.go +++ b/common/vio/stats.go @@ -1,4 +1,4 @@ -package stats +package vio import ( "v2ray.com/core/common" diff --git a/common/vio/stats_test.go b/common/vio/stats_test.go new file mode 100644 index 000000000..cab2bca01 --- /dev/null +++ b/common/vio/stats_test.go @@ -0,0 +1,46 @@ +package vio_test + +import ( + "testing" + + "v2ray.com/core/common" + "v2ray.com/core/common/buf" + . "v2ray.com/core/common/vio" +) + +type TestCounter int64 + +func (c *TestCounter) Value() int64 { + return int64(*c) +} + +func (c *TestCounter) Add(v int64) int64 { + x := int64(*c) + v + *c = TestCounter(x) + return x +} + +func (c *TestCounter) Set(v int64) int64 { + *c = TestCounter(v) + return v +} + +func TestStatsWriter(t *testing.T) { + var c TestCounter + writer := &SizeStatWriter{ + Counter: &c, + Writer: buf.Discard, + } + + var mb buf.MultiBuffer + common.Must2(mb.Write([]byte("abcd"))) + common.Must(writer.WriteMultiBuffer(mb)) + + mb.Release() + common.Must2(mb.Write([]byte("efg"))) + common.Must(writer.WriteMultiBuffer(mb)) + + if c.Value() != 7 { + t.Fatal("unexpected counter value. want 7, but got ", c.Value()) + } +}