diff --git a/core/load/adaptiveshedder.go b/core/load/adaptiveshedder.go index 8d6737be..c0c14481 100644 --- a/core/load/adaptiveshedder.go +++ b/core/load/adaptiveshedder.go @@ -31,6 +31,8 @@ var ( // default to be enabled enabled = syncx.ForAtomicBool(true) + // default to be enabled + logEnabled = syncx.ForAtomicBool(true) // make it a variable for unit test systemOverloadChecker = func(cpuThreshold int64) bool { return stat.CpuUsage() >= cpuThreshold @@ -80,6 +82,11 @@ func Disable() { enabled.Set(false) } +// DisableLog disables the stat logs for load shedding. +func DisableLog() { + logEnabled.Set(false) +} + // NewAdaptiveShedder returns an adaptive shedder. // opts can be used to customize the Shedder. func NewAdaptiveShedder(opts ...ShedderOption) Shedder { diff --git a/core/load/adaptiveshedder_test.go b/core/load/adaptiveshedder_test.go index 1f9e4b3a..212bcb6f 100644 --- a/core/load/adaptiveshedder_test.go +++ b/core/load/adaptiveshedder_test.go @@ -25,6 +25,7 @@ func init() { } func TestAdaptiveShedder(t *testing.T) { + DisableLog() shedder := NewAdaptiveShedder(WithWindow(bucketDuration), WithBuckets(buckets), WithCpuThreshold(100)) var wg sync.WaitGroup var drop int64 diff --git a/core/load/sheddingstat.go b/core/load/sheddingstat.go index 2a12d969..09ec462d 100644 --- a/core/load/sheddingstat.go +++ b/core/load/sheddingstat.go @@ -48,6 +48,25 @@ func (s *SheddingStat) IncrementDrop() { atomic.AddInt64(&s.drop, 1) } +func (s *SheddingStat) loop(c <-chan time.Time) { + for range c { + st := s.reset() + + if !logEnabled.True() { + continue + } + + c := stat.CpuUsage() + if st.Drop == 0 { + logx.Statf("(%s) shedding_stat [1m], cpu: %d, total: %d, pass: %d, drop: %d", + s.name, c, st.Total, st.Pass, st.Drop) + } else { + logx.Statf("(%s) shedding_stat_drop [1m], cpu: %d, total: %d, pass: %d, drop: %d", + s.name, c, st.Total, st.Pass, st.Drop) + } + } +} + func (s *SheddingStat) reset() snapshot { return snapshot{ Total: atomic.SwapInt64(&s.total, 0), @@ -59,15 +78,6 @@ func (s *SheddingStat) reset() snapshot { func (s *SheddingStat) run() { ticker := time.NewTicker(time.Minute) defer ticker.Stop() - for range ticker.C { - c := stat.CpuUsage() - st := s.reset() - if st.Drop == 0 { - logx.Statf("(%s) shedding_stat [1m], cpu: %d, total: %d, pass: %d, drop: %d", - s.name, c, st.Total, st.Pass, st.Drop) - } else { - logx.Statf("(%s) shedding_stat_drop [1m], cpu: %d, total: %d, pass: %d, drop: %d", - s.name, c, st.Total, st.Pass, st.Drop) - } - } + + s.loop(ticker.C) } diff --git a/core/load/sheddingstat_test.go b/core/load/sheddingstat_test.go index 351eefc0..498341f4 100644 --- a/core/load/sheddingstat_test.go +++ b/core/load/sheddingstat_test.go @@ -2,6 +2,7 @@ package load import ( "testing" + "time" "github.com/stretchr/testify/assert" ) @@ -22,3 +23,32 @@ func TestSheddingStat(t *testing.T) { assert.Equal(t, int64(5), result.Pass) assert.Equal(t, int64(7), result.Drop) } + +func TestLoopTrue(t *testing.T) { + ch := make(chan time.Time, 1) + ch <- time.Now() + close(ch) + st := new(SheddingStat) + logEnabled.Set(true) + st.loop(ch) +} + +func TestLoopTrueAndDrop(t *testing.T) { + ch := make(chan time.Time, 1) + ch <- time.Now() + close(ch) + st := new(SheddingStat) + st.IncrementDrop() + logEnabled.Set(true) + st.loop(ch) +} + +func TestLoopFalseAndDrop(t *testing.T) { + ch := make(chan time.Time, 1) + ch <- time.Now() + close(ch) + st := new(SheddingStat) + st.IncrementDrop() + logEnabled.Set(false) + st.loop(ch) +} diff --git a/core/stat/usage.go b/core/stat/usage.go index b6a5e3df..251cd4da 100644 --- a/core/stat/usage.go +++ b/core/stat/usage.go @@ -38,7 +38,9 @@ func init() { atomic.StoreInt64(&cpuUsage, usage) }) case <-allTicker.C: - printUsage() + if logEnabled.True() { + printUsage() + } } } }()