improve collect lister metrics task

cloud/tasks/collect_lister_metrics_task.go

@@ -13,11 +13,16 @@ import (
 
 ////////////////////////////////////////////////////////////////////////////////
 
+const totalHangingTaskCountGaugeName = "totalHangingTaskCount"
+
+////////////////////////////////////////////////////////////////////////////////
+
 type collectListerMetricsTask struct {
 	registry                  metrics.Registry
 	storage                   storage.Storage
 	metricsCollectionInterval time.Duration
 
+	taskTypes                 []string
 	hangingTaskGaugesByID     map[string]metrics.Gauge
 	maxHangingTaskIDsToReport int64
 }
@@ -134,6 +139,9 @@ func (c *collectListerMetricsTask) collectTasksMetrics(
 ) error {
 
 	tasksByType := make(map[string]uint64)
+	for _, taskType := range c.taskTypes {
+		tasksByType[taskType] = 0
+	}
 
 	taskInfos, err := getTaskInfos(ctx)
 	if err != nil {
@@ -158,14 +166,22 @@ func (c *collectListerMetricsTask) collectHangingTasksMetrics(
 	ctx context.Context,
 ) error {
 
-	totalHangingTasksGauge := c.registry.Gauge("totalHangingTaskCount")
-
 	taskInfos, err := c.storage.ListHangingTasks(ctx, ^uint64(0))
 	if err != nil {
 		return err
 	}
 
-	totalHangingTasksGauge.Set(float64(len(taskInfos)))
+	err = c.collectTasksMetrics(
+		ctx,
+		func(context.Context) ([]storage.TaskInfo, error) {
+			return taskInfos, nil
+		},
+		totalHangingTaskCountGaugeName,
+	)
+	if err != nil {
+		return err
+	}
+
 	taskInfoByID := make(map[string]storage.TaskInfo)
 	for _, taskInfo := range taskInfos {
 		taskInfoByID[taskInfo.ID] = taskInfo

cloud/tasks/scheduler_impl.go

@@ -620,6 +620,7 @@ func NewScheduler(
 				storage:                   storage,
 				metricsCollectionInterval: listerMetricsCollectionInterval,
 
+				taskTypes:                 registry.TaskTypesForExecution(),
 				hangingTaskGaugesByID:     make(map[string]metrics.Gauge),
 				maxHangingTaskIDsToReport: config.GetMaxHangingTaskIDsToReport(),
 			}

cloud/tasks/tasks_tests/tasks_test.go

@@ -1252,11 +1252,15 @@ func TestHangingTasksMetrics(t *testing.T) {
 	gaugeSetWg := sync.WaitGroup{}
 	gaugeUnsetWg := sync.WaitGroup{}
 
-	registry.GetGauge("totalHangingTaskCount", map[string]string{}).On(
+	registry.GetGauge(
+		"totalHangingTaskCount",
+		map[string]string{"type": "tasks.hanging"},
+	).On(
 		"Set",
-		mock.Anything,
+		float64(1),
 	).Return(mock.Anything)
-	gaugeSetCall := registry.GetGauge(
+
+	hangingTasksGaugeSetCall := registry.GetGauge(
 		"hangingTasks",
 		map[string]string{"type": "tasks.hanging", "id": taskID},
 	).On("Set", float64(1)).Return(mock.Anything).Run(
@@ -1273,7 +1277,7 @@ func TestHangingTasksMetrics(t *testing.T) {
 		"Set",
 		float64(0),
 	).NotBefore(
-		gaugeSetCall,
+		hangingTasksGaugeSetCall,
 	).Return(mock.Anything).Run(
 		func(args mock.Arguments) {
 			gaugeUnsetWg.Done()