refactor: use *http.ServeMux (#72)

* refactor: use `*http.ServeMux`

instead of default multiplexer.

also shutdown with current context

Signed-off-by: Dwi Siswanto <[email protected]>

* test: add TestStartMultipleTimes

Signed-off-by: Dwi Siswanto <[email protected]>

* chore: revert to avoid errcheck lint

Signed-off-by: Dwi Siswanto <[email protected]>

---------

Signed-off-by: Dwi Siswanto <[email protected]>

clistats.go

@@ -117,56 +117,59 @@ func NewWithOptions(ctx context.Context, options *Options) (*Statistics, error)
 	return statistics, nil
 }
 
+func (s *Statistics) metricsHandler(w http.ResponseWriter, req *http.Request) {
+	items := make(map[string]interface{})
+	for k, v := range s.counters {
+		items[k] = v.Load()
+	}
+	for k, v := range s.static {
+		items[k] = v
+	}
+	for k, v := range s.dynamic {
+		items[k] = v(s)
+	}
+
+	// Common fields
+	requests, hasRequests := s.GetCounter("requests")
+	startedAt, hasStartedAt := s.GetStatic("startedAt")
+	total, hasTotal := s.GetCounter("total")
+	var (
+		duration    time.Duration
+		hasDuration bool
+	)
+	// duration
+	if hasStartedAt {
+		if stAt, ok := startedAt.(time.Time); ok {
+			duration = time.Since(stAt)
+			items["duration"] = FmtDuration(duration)
+			hasDuration = true
+		}
+	}
+	// rps
+	if hasRequests && hasDuration {
+		items["rps"] = String(uint64(float64(requests) / duration.Seconds()))
+	}
+	// percent
+	if hasRequests && hasTotal {
+		percentData := (float64(requests) * float64(100)) / float64(total)
+		percent := String(uint64(percentData))
+		items["percent"] = percent
+	}
+
+	data, err := jsoniter.Marshal(items)
+	if err != nil {
+		w.WriteHeader(http.StatusInternalServerError)
+		_, _ = w.Write([]byte(fmt.Sprintf(`{"error":"%s"}`, err)))
+		return
+	}
+	_, _ = w.Write(data)
+}
+
 // Start starts the event loop of the stats client.
 func (s *Statistics) Start() error {
 	if s.Options.Web {
-		http.HandleFunc("/metrics", func(w http.ResponseWriter, req *http.Request) {
-			items := make(map[string]interface{})
-			for k, v := range s.counters {
-				items[k] = v.Load()
-			}
-			for k, v := range s.static {
-				items[k] = v
-			}
-			for k, v := range s.dynamic {
-				items[k] = v(s)
-			}
-
-			// Common fields
-			requests, hasRequests := s.GetCounter("requests")
-			startedAt, hasStartedAt := s.GetStatic("startedAt")
-			total, hasTotal := s.GetCounter("total")
-			var (
-				duration    time.Duration
-				hasDuration bool
-			)
-			// duration
-			if hasStartedAt {
-				if stAt, ok := startedAt.(time.Time); ok {
-					duration = time.Since(stAt)
-					items["duration"] = FmtDuration(duration)
-					hasDuration = true
-				}
-			}
-			// rps
-			if hasRequests && hasDuration {
-				items["rps"] = String(uint64(float64(requests) / duration.Seconds()))
-			}
-			// percent
-			if hasRequests && hasTotal {
-				percentData := (float64(requests) * float64(100)) / float64(total)
-				percent := String(uint64(percentData))
-				items["percent"] = percent
-			}
-
-			data, err := jsoniter.Marshal(items)
-			if err != nil {
-				w.WriteHeader(http.StatusInternalServerError)
-				_, _ = w.Write([]byte(fmt.Sprintf(`{"error":"%s"}`, err)))
-				return
-			}
-			_, _ = w.Write(data)
-		})
+		mux := http.NewServeMux()
+		mux.HandleFunc("/metrics", s.metricsHandler)
 
 		// check if the default port is available
 		port, err := freeport.GetPort(freeport.TCP, "127.0.0.1", s.Options.ListenPort)
@@ -181,7 +184,7 @@ func (s *Statistics) Start() error {
 
 		s.httpServer = &http.Server{
 			Addr:    fmt.Sprintf("%s:%d", port.Address, port.Port),
-			Handler: http.DefaultServeMux,
+			Handler: mux,
 		}
 
 		go func() {
@@ -227,11 +230,13 @@ func (s *Statistics) GetStatResponse(interval time.Duration, callback func(strin
 
 // Stop stops the event loop of the stats client
 func (s *Statistics) Stop() error {
-	s.cancel()
+	defer s.cancel()
+
 	if s.httpServer != nil {
-		if err := s.httpServer.Shutdown(context.Background()); err != nil {
+		if err := s.httpServer.Shutdown(s.ctx); err != nil {
 			return err
 		}
 	}
+
 	return nil
 }

clistats_test.go

@@ -38,3 +38,16 @@ func TestDynamicCallback_Elapsedtime(t *testing.T) {
 	elapsed := time.Since(startTime).Seconds()
 	require.True(t, elapsed > 0)
 }
+
+func TestStartMultipleTimes(t *testing.T) {
+	client, err := New()
+	require.Nil(t, err)
+
+	for i := 1; i <= 2; i++ {
+		err = client.Start()
+		require.Nil(t, err)
+
+		err = client.Stop()
+		require.Nil(t, err)
+	}
+}