chore: add more tests (#4930)

core/service/serviceconf.go

@@ -39,7 +39,7 @@ type (
 		Telemetry  trace.Config      `json:",optional"`
 		DevServer  DevServerConfig   `json:",optional"`
 		Shutdown   proc.ShutdownConf `json:",optional"`
-		// Profiling is the configuration for profiling.
+		// Profiling is the configuration for continuous profiling.
 		Profiling profiling.Config `json:",optional"`
 	}
 )
@@ -73,8 +73,8 @@ func (sc ServiceConf) SetUp() error {
 	if len(sc.MetricsUrl) > 0 {
 		stat.SetReportWriter(stat.NewRemoteWriter(sc.MetricsUrl))
 	}
-	devserver.StartAgent(sc.DevServer)
 
+	devserver.StartAgent(sc.DevServer)
 	profiling.Start(sc.Profiling)
 
 	return nil

internal/profiling/profiling.go

@@ -6,27 +6,32 @@ import (
 	"time"
 
 	"github.com/grafana/pyroscope-go"
-
 	"github.com/zeromicro/go-zero/core/logx"
 	"github.com/zeromicro/go-zero/core/proc"
 	"github.com/zeromicro/go-zero/core/stat"
 	"github.com/zeromicro/go-zero/core/threading"
 )
 
+const (
+	defaultCheckInterval     = time.Second * 10
+	defaultProfilingDuration = time.Minute * 2
+	defaultUploadRate        = time.Second * 15
+)
+
 type (
 	Config struct {
 		// Name is the name of the application.
 		Name string `json:",optional,inherit"`
-		// ServerAddress is the address of the profiling server.
+		// ServerAddr is the address of the profiling server.
-		ServerAddress string
+		ServerAddr string
 		// AuthUser is the username for basic authentication.
 		AuthUser string `json:",optional"`
 		// AuthPassword is the password for basic authentication.
 		AuthPassword string `json:",optional"`
-		// UploadDuration is the duration for which profiling data is uploaded.
+		// UploadRate is the duration for which profiling data is uploaded.
-		UploadDuration time.Duration `json:",default=15s"`
+		UploadRate time.Duration `json:",default=15s"`
-		// IntervalDuration is the interval for which profiling data is collected.
+		// CheckInterval is the interval to check if profiling should start.
-		IntervalDuration time.Duration `json:",default=10s"`
+		CheckInterval time.Duration `json:",default=10s"`
 		// ProfilingDuration is the duration for which profiling data is collected.
 		ProfilingDuration time.Duration `json:",default=2m"`
 		// CpuThreshold the collection is allowed only when the current service cpu < CpuThreshold
@@ -56,7 +61,7 @@ type (
 		Stop() error
 	}
 
-	pyProfiler struct {
+	pyroscopeProfiler struct {
 		c        Config
 		profiler *pyroscope.Profiler
 	}
@@ -66,63 +71,58 @@ var (
 	once sync.Once
 
 	newProfiler = func(c Config) profiler {
-		return newPyProfiler(c)
+		return newPyroscopeProfiler(c)
 	}
 )
 
 // Start initializes the pyroscope profiler with the given configuration.
 func Start(c Config) {
 	// check if the profiling is enabled
-	if c.ServerAddress == "" {
+	if len(c.ServerAddr) == 0 {
 		return
 	}
 
 	// set default values for the configuration
 	if c.ProfilingDuration <= 0 {
-		c.ProfilingDuration = time.Minute * 2
+		c.ProfilingDuration = defaultProfilingDuration
 	}
 
 	// set default values for the configuration
-	if c.IntervalDuration <= 0 {
+	if c.CheckInterval <= 0 {
-		c.IntervalDuration = time.Second * 10
+		c.CheckInterval = defaultCheckInterval
 	}
 
-	if c.UploadDuration <= 0 {
+	if c.UploadRate <= 0 {
-		c.UploadDuration = time.Second * 15
+		c.UploadRate = defaultUploadRate
 	}
 
 	once.Do(func() {
 		logx.Info("continuous profiling started")
 
-		var done = make(chan struct{})
-		proc.AddShutdownListener(func() {
-			done <- struct{}{}
-			close(done)
-		})
-
 		threading.GoSafe(func() {
-			startPyroScope(c, done)
+			startPyroscope(c, proc.Done())
 		})
 	})
 }
 
-// startPyroScope starts the pyroscope profiler with the given configuration.
+// startPyroscope starts the pyroscope profiler with the given configuration.
-func startPyroScope(c Config, done <-chan struct{}) {
+func startPyroscope(c Config, done <-chan struct{}) {
 	var (
-		intervalTicker  = time.NewTicker(c.IntervalDuration)
+		pr                  profiler
-		profilingTicker = time.NewTicker(c.ProfilingDuration)
+		err                 error
-
-		pr  profiler
-		err error
-
 		latestProfilingTime time.Time
+		intervalTicker      = time.NewTicker(c.CheckInterval)
+		profilingTicker     = time.NewTicker(c.ProfilingDuration)
 	)
 
+	defer profilingTicker.Stop()
+	defer intervalTicker.Stop()
+
 	for {
 		select {
 		case <-intervalTicker.C:
 			// Check if the machine is overloaded and if the profiler is not running
-			if pr == nil && checkMachinePerformance(c) {
+			if pr == nil && isCpuOverloaded(c) {
 				pr = newProfiler(c)
 				if err := pr.Start(); err != nil {
 					logx.Errorf("failed to start profiler: %v", err)
@@ -131,7 +131,7 @@ func startPyroScope(c Config, done <-chan struct{}) {
 
 				// record the latest profiling time
 				latestProfilingTime = time.Now()
-				logx.Infof("pyroScope profiler started.")
+				logx.Infof("pyroscope profiler started.")
 			}
 		case <-profilingTicker.C:
 			// check if the profiling duration has passed
@@ -144,30 +144,26 @@ func startPyroScope(c Config, done <-chan struct{}) {
 				if err = pr.Stop(); err != nil {
 					logx.Errorf("failed to stop profiler: %v", err)
 				}
-				logx.Infof("pyroScope profiler stopped.")
+				logx.Infof("pyroscope profiler stopped.")
 				pr = nil
 			}
 		case <-done:
 			logx.Infof("continuous profiling stopped.")
-			intervalTicker.Stop()
-			profilingTicker.Stop()
 			return
 		}
 	}
 }
 
-// genPyroScopeConf generates the pyroscope configuration based on the given config.
+// genPyroscopeConf generates the pyroscope configuration based on the given config.
-func genPyroScopeConf(c Config) pyroscope.Config {
+func genPyroscopeConf(c Config) pyroscope.Config {
 	pConf := pyroscope.Config{
-		UploadRate:        c.UploadDuration,
+		UploadRate:        c.UploadRate,
 		ApplicationName:   c.Name,
 		BasicAuthUser:     c.AuthUser,     // http basic auth user
 		BasicAuthPassword: c.AuthPassword, // http basic auth password
-		ServerAddress:     c.ServerAddress,
+		ServerAddress:     c.ServerAddr,
 		Logger:            nil,
-
+		HTTPHeaders:       map[string]string{},
-		HTTPHeaders: map[string]string{},
-
 		// you can provide static tags via a map:
 		Tags: map[string]string{
 			"name": c.Name,
@@ -200,46 +196,46 @@ func genPyroScopeConf(c Config) pyroscope.Config {
 	return pConf
 }
 
-// checkMachinePerformance checks the machine performance based on the given configuration.
+// isCpuOverloaded checks the machine performance based on the given configuration.
-func checkMachinePerformance(c Config) bool {
+func isCpuOverloaded(c Config) bool {
 	currentValue := stat.CpuUsage()
 	if currentValue >= c.CpuThreshold {
-		logx.Infof("continuous profiling cpu overload, cpu:%d", currentValue)
+		logx.Infof("continuous profiling cpu overload, cpu: %d", currentValue)
 		return true
 	}
 
 	return false
 }
 
-func newPyProfiler(c Config) profiler {
+func newPyroscopeProfiler(c Config) profiler {
-	return &pyProfiler{
+	return &pyroscopeProfiler{
 		c: c,
 	}
 }
 
-func (p *pyProfiler) Start() error {
+func (p *pyroscopeProfiler) Start() error {
-	pConf := genPyroScopeConf(p.c)
+	pConf := genPyroscopeConf(p.c)
 	// set mutex and block profile rate
 	setFraction(p.c)
-	profiler, err := pyroscope.Start(pConf)
+	prof, err := pyroscope.Start(pConf)
 	if err != nil {
 		resetFraction(p.c)
 		return err
 	}
 
-	p.profiler = profiler
+	p.profiler = prof
 	return nil
 }
 
-func (p *pyProfiler) Stop() error {
+func (p *pyroscopeProfiler) Stop() error {
 	if p.profiler == nil {
 		return nil
 	}
 
-	err := p.profiler.Stop()
+	if err := p.profiler.Stop(); err != nil {
-	if err != nil {
 		return err
 	}
+
 	resetFraction(p.c)
 	p.profiler = nil
 

internal/profiling/profiling_test.go

@@ -7,77 +7,111 @@ import (
 
 	"github.com/grafana/pyroscope-go"
 	"github.com/stretchr/testify/assert"
+	"github.com/zeromicro/go-zero/core/conf"
+	"github.com/zeromicro/go-zero/core/syncx"
 )
 
 func TestStart(t *testing.T) {
+	t.Run("profiling", func(t *testing.T) {
+		var c Config
+		assert.NoError(t, conf.FillDefault(&c))
+		c.Name = "test"
+		p := newProfiler(c)
+		assert.NotNil(t, p)
+		assert.NoError(t, p.Start())
+		assert.NoError(t, p.Stop())
+	})
+
 	t.Run("invalid config", func(t *testing.T) {
-		var mockProfiler = &mockProfiler{}
+		mp := &mockProfiler{}
 		newProfiler = func(c Config) profiler {
-			return mockProfiler
+			return mp
 		}
 
 		Start(Config{})
 
 		Start(Config{
-			ServerAddress: "localhost:4040",
+			ServerAddr: "localhost:4040",
 		})
 	})
 
 	t.Run("test start profiler", func(t *testing.T) {
-		var mockProfiler = &mockProfiler{}
+		mp := &mockProfiler{}
 		newProfiler = func(c Config) profiler {
-			return mockProfiler
+			return mp
 		}
 
 		c := Config{
 			Name:              "test",
-			ServerAddress:     "localhost:4040",
+			ServerAddr:        "localhost:4040",
-			IntervalDuration:  time.Millisecond,
+			CheckInterval:     time.Millisecond,
 			ProfilingDuration: time.Millisecond * 10,
 			CpuThreshold:      0,
 		}
 		var done = make(chan struct{})
-		go startPyroScope(c, done)
+		go startPyroscope(c, done)
 
 		time.Sleep(time.Millisecond * 50)
-		done <- struct{}{}
+		close(done)
 
-		assert.True(t, mockProfiler.started)
+		assert.True(t, mp.started.True())
-		assert.True(t, mockProfiler.stopped)
+		assert.True(t, mp.stopped.True())
+	})
+
+	t.Run("test start profiler with cpu overloaded", func(t *testing.T) {
+		mp := &mockProfiler{}
+		newProfiler = func(c Config) profiler {
+			return mp
+		}
+
+		c := Config{
+			Name:              "test",
+			ServerAddr:        "localhost:4040",
+			CheckInterval:     time.Millisecond,
+			ProfilingDuration: time.Millisecond * 10,
+			CpuThreshold:      900,
+		}
+		var done = make(chan struct{})
+		go startPyroscope(c, done)
+
+		time.Sleep(time.Millisecond * 50)
+		close(done)
+
+		assert.False(t, mp.started.True())
 	})
 
 	t.Run("start/stop err", func(t *testing.T) {
-		var mockProfiler = &mockProfiler{
+		mp := &mockProfiler{
 			err: assert.AnError,
 		}
 		newProfiler = func(c Config) profiler {
-			return mockProfiler
+			return mp
 		}
 
 		c := Config{
 			Name:              "test",
-			ServerAddress:     "localhost:4040",
+			ServerAddr:        "localhost:4040",
-			IntervalDuration:  time.Millisecond,
+			CheckInterval:     time.Millisecond,
 			ProfilingDuration: time.Millisecond * 10,
 			CpuThreshold:      0,
 		}
 		var done = make(chan struct{})
-		go startPyroScope(c, done)
+		go startPyroscope(c, done)
 
 		time.Sleep(time.Millisecond * 50)
-		done <- struct{}{}
+		close(done)
 
-		assert.False(t, mockProfiler.started)
+		assert.False(t, mp.started.True())
-		assert.False(t, mockProfiler.stopped)
+		assert.False(t, mp.stopped.True())
 	})
 }
 
-func TestGenPyroScopeConf(t *testing.T) {
+func TestGenPyroscopeConf(t *testing.T) {
 	c := Config{
-		Name:          "",
+		Name:         "",
-		ServerAddress: "localhost:4040",
+		ServerAddr:   "localhost:4040",
-		AuthUser:      "user",
+		AuthUser:     "user",
-		AuthPassword:  "password",
+		AuthPassword: "password",
 		ProfileType: ProfileType{
 			Logger:     true,
 			CPU:        true,
@@ -88,30 +122,30 @@ func TestGenPyroScopeConf(t *testing.T) {
 		},
 	}
 
-	conf := genPyroScopeConf(c)
+	pyroscopeConf := genPyroscopeConf(c)
-	assert.Equal(t, c.ServerAddress, conf.ServerAddress)
+	assert.Equal(t, c.ServerAddr, pyroscopeConf.ServerAddress)
-	assert.Equal(t, c.AuthUser, conf.BasicAuthUser)
+	assert.Equal(t, c.AuthUser, pyroscopeConf.BasicAuthUser)
-	assert.Equal(t, c.AuthPassword, conf.BasicAuthPassword)
+	assert.Equal(t, c.AuthPassword, pyroscopeConf.BasicAuthPassword)
-	assert.Equal(t, c.Name, conf.ApplicationName)
+	assert.Equal(t, c.Name, pyroscopeConf.ApplicationName)
-	assert.Contains(t, conf.ProfileTypes, pyroscope.ProfileCPU)
+	assert.Contains(t, pyroscopeConf.ProfileTypes, pyroscope.ProfileCPU)
-	assert.Contains(t, conf.ProfileTypes, pyroscope.ProfileGoroutines)
+	assert.Contains(t, pyroscopeConf.ProfileTypes, pyroscope.ProfileGoroutines)
-	assert.Contains(t, conf.ProfileTypes, pyroscope.ProfileAllocObjects)
+	assert.Contains(t, pyroscopeConf.ProfileTypes, pyroscope.ProfileAllocObjects)
-	assert.Contains(t, conf.ProfileTypes, pyroscope.ProfileAllocSpace)
+	assert.Contains(t, pyroscopeConf.ProfileTypes, pyroscope.ProfileAllocSpace)
-	assert.Contains(t, conf.ProfileTypes, pyroscope.ProfileInuseObjects)
+	assert.Contains(t, pyroscopeConf.ProfileTypes, pyroscope.ProfileInuseObjects)
-	assert.Contains(t, conf.ProfileTypes, pyroscope.ProfileInuseSpace)
+	assert.Contains(t, pyroscopeConf.ProfileTypes, pyroscope.ProfileInuseSpace)
-	assert.Contains(t, conf.ProfileTypes, pyroscope.ProfileMutexCount)
+	assert.Contains(t, pyroscopeConf.ProfileTypes, pyroscope.ProfileMutexCount)
-	assert.Contains(t, conf.ProfileTypes, pyroscope.ProfileMutexDuration)
+	assert.Contains(t, pyroscopeConf.ProfileTypes, pyroscope.ProfileMutexDuration)
-	assert.Contains(t, conf.ProfileTypes, pyroscope.ProfileBlockCount)
+	assert.Contains(t, pyroscopeConf.ProfileTypes, pyroscope.ProfileBlockCount)
-	assert.Contains(t, conf.ProfileTypes, pyroscope.ProfileBlockDuration)
+	assert.Contains(t, pyroscopeConf.ProfileTypes, pyroscope.ProfileBlockDuration)
 
 	setFraction(c)
 	resetFraction(c)
 
-	newPyProfiler(c)
+	newPyroscopeProfiler(c)
 }
 
-func TestNewPyProfiler(t *testing.T) {
+func TestNewPyroscopeProfiler(t *testing.T) {
-	p := newPyProfiler(Config{})
+	p := newPyroscopeProfiler(Config{})
 
 	assert.Error(t, p.Start())
 	assert.NoError(t, p.Stop())
@@ -119,15 +153,15 @@ func TestNewPyProfiler(t *testing.T) {
 
 type mockProfiler struct {
 	mutex   sync.Mutex
-	started bool
+	started syncx.AtomicBool
-	stopped bool
+	stopped syncx.AtomicBool
 	err     error
 }
 
 func (m *mockProfiler) Start() error {
 	m.mutex.Lock()
 	if m.err == nil {
-		m.started = true
+		m.started.Set(true)
 	}
 	m.mutex.Unlock()
 	return m.err
@@ -136,7 +170,7 @@ func (m *mockProfiler) Start() error {
 func (m *mockProfiler) Stop() error {
 	m.mutex.Lock()
 	if m.err == nil {
-		m.stopped = true
+		m.stopped.Set(true)
 	}
 	m.mutex.Unlock()
 	return m.err

tools/goctl/api/swagger/swagger.go

@@ -6,7 +6,6 @@ import (
 	"time"
 
 	"github.com/go-openapi/spec"
-
 	apiSpec "github.com/zeromicro/go-zero/tools/goctl/api/spec"
 	"github.com/zeromicro/go-zero/tools/goctl/internal/version"
 )