[stats]: rewrite s3 stats collector to return shutdown func

weed/command/master.go

@@ -231,6 +231,8 @@ func startMaster(masterOption MasterOptions, masterWhiteList []string) {
 	}
 	go grpcS.Serve(grpcL)
 
+	rootCtx, rootCancel := context.WithCancel(context.Background())
+
 	timeSleep := 1500 * time.Millisecond
 	if !*masterOption.raftHashicorp {
 		go func() {
@@ -245,8 +247,8 @@ func startMaster(masterOption MasterOptions, masterWhiteList []string) {
 		}()
 	}
 
-	go ms.MasterClient.KeepConnectedToMaster(context.Background())
-	go weed_server.CollectCollectionsStats(context.Background(), ms, *masterOption.intervalToCollectStats)
+	go ms.MasterClient.KeepConnectedToMaster(rootCtx)
+	statsShutdown := weed_server.CollectCollectionsStats(rootCtx, ms, *masterOption.intervalToCollectStats)
 
 	// start http server
 	var (
@@ -286,6 +288,15 @@ func startMaster(masterOption MasterOptions, masterWhiteList []string) {
 
 	grace.OnInterrupt(ms.Shutdown)
 	grace.OnInterrupt(grpcS.Stop)
+	grace.OnInterrupt(func() {
+		rootCancel()
+
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+		if err := statsShutdown(ctx); err != nil {
+			glog.Errorf("graceful shutdown of collections stats: %v", err)
+		}
+	})
 	grace.OnReload(func() {
 		if ms.Topo.HashicorpRaft != nil && ms.Topo.HashicorpRaft.State() == hashicorpRaft.Leader {
 			ms.Topo.HashicorpRaft.LeadershipTransfer()

weed/server/master_server.go

@@ -35,6 +35,7 @@ import (
 	util_http "github.com/seaweedfs/seaweedfs/weed/util/http"
 	"github.com/seaweedfs/seaweedfs/weed/util/version"
 	"github.com/seaweedfs/seaweedfs/weed/wdclient"
+	"math"
 )
 
 const (
@@ -436,18 +437,28 @@ func (ms *MasterServer) Reload() {
 	)
 }
 
-func CollectCollectionsStats(ctx context.Context, ms *MasterServer, intervalToCollect time.Duration) {
-	ticker := time.NewTicker(intervalToCollect)
+func CollectCollectionsStats(parent context.Context, ms *MasterServer, interval time.Duration) (shutdown func(ctx context.Context) error) {
+	ctx, cancel := context.WithCancel(parent)
+	done := make(chan struct{})
+
+	go func() {
+		defer close(done)
+		ticker := time.NewTicker(interval)
 		defer ticker.Stop()
 
+		req := &master_pb.VolumeListRequest{}
 		for {
 			select {
 			case <-ctx.Done():
 				return
 			case <-ticker.C:
-			volumeList, err := ms.VolumeList(context.Background(), &master_pb.VolumeListRequest{})
+				volumeList, err := ms.VolumeList(ctx, req)
 				if err != nil {
 					glog.Errorf("collect volume list: %v", err)
+					continue
+				}
+				if volumeList == nil || volumeList.TopologyInfo == nil {
+					continue
 				}
 
 				collectionInfos := make(map[string]*shell.CollectionInfo)
@@ -457,8 +468,8 @@ func CollectCollectionsStats(ctx context.Context, ms *MasterServer, intervalToCo
 							for _, diskInfo := range dn.DiskInfos {
 								for _, vif := range diskInfo.VolumeInfos {
 									c := vif.Collection
-								cif, found := collectionInfos[c]
-								if !found {
+									cif := collectionInfos[c]
+									if cif == nil {
 										cif = &shell.CollectionInfo{}
 										collectionInfos[c] = cif
 									}
@@ -475,10 +486,21 @@ func CollectCollectionsStats(ctx context.Context, ms *MasterServer, intervalToCo
 					}
 				}
 
-			for k, v := range collectionInfos {
-				stats.S3BucketFileCount.WithLabelValues(k).Set(v.FileCount - v.DeleteCount)
-				stats.S3BucketSize.WithLabelValues(k).Set(v.Size - v.DeletedByteCount)
+				for bucket, v := range collectionInfos {
+					stats.S3BucketFileCount.WithLabelValues(bucket).Set(math.Round(v.FileCount - v.DeleteCount))
+					stats.S3BucketSize.WithLabelValues(bucket).Set(math.Round(v.Size - v.DeletedByteCount))
+				}
 			}
 		}
+	}()
+
+	return func(ctx context.Context) error {
+		cancel()
+		select {
+		case <-done:
+			return nil
+		case <-ctx.Done():
+			return ctx.Err()
+		}
 	}
 }