NBS-4789: [Disk Manager] refactor sre tools (#177)

* refactor sre tools * resolve issues * fix build * fix issue
ydb-platform · Jan 19, 2024 · a2fe687 · a2fe687
1 parent 093ecde
commit a2fe687
Show file tree

Hide file tree

Showing 7 changed files with 59 additions and 33 deletions.
diff --git a/cloud/disk_manager/pkg/admin/tasks.go b/cloud/disk_manager/pkg/admin/tasks.go
@@ -225,7 +225,7 @@ func newListReadyToRunCmd(
 		serverConfig,
 		"ready_to_run",
 		func(ctx context.Context, storage tasks_storage.Storage, limit uint64) ([]string, error) {
-			taskInfos, err := storage.ListTasksReadyToRun(
+			tasks, err := storage.ListTasksReadyToRun(
 				ctx,
 				limit,
 				nil, // taskTypeWhitelist
@@ -234,7 +234,7 @@ func newListReadyToRunCmd(
 				return []string{}, err
 			}
 
-			return getTaskIDs(taskInfos), nil
+			return getTaskIDs(tasks), nil
 		},
 	)
 }
@@ -249,7 +249,7 @@ func newListReadyToCancelCmd(
 		serverConfig,
 		"ready_to_cancel",
 		func(ctx context.Context, storage tasks_storage.Storage, limit uint64) ([]string, error) {
-			taskInfos, err := storage.ListTasksReadyToCancel(
+			tasks, err := storage.ListTasksReadyToCancel(
 				ctx,
 				limit,
 				nil, // taskTypeWhitelist
@@ -258,7 +258,7 @@ func newListReadyToCancelCmd(
 				return []string{}, err
 			}
 
-			return getTaskIDs(taskInfos), nil
+			return getTaskIDs(tasks), nil
 		},
 	)
 }
@@ -273,7 +273,12 @@ func newListRunningCmd(
 		serverConfig,
 		"running",
 		func(ctx context.Context, storage tasks_storage.Storage, limit uint64) ([]string, error) {
-			return storage.ListTasksRunning(ctx, limit)
+			tasks, err := storage.ListTasksRunning(ctx, limit)
+			if err != nil {
+				return []string{}, err
+			}
+
+			return getTaskIDs(tasks), nil
 		},
 	)
 }
@@ -288,7 +293,12 @@ func newListCancellingCmd(
 		serverConfig,
 		"cancelling",
 		func(ctx context.Context, storage tasks_storage.Storage, limit uint64) ([]string, error) {
-			return storage.ListTasksCancelling(ctx, limit)
+			tasks, err := storage.ListTasksCancelling(ctx, limit)
+			if err != nil {
+				return []string{}, err
+			}
+
+			return getTaskIDs(tasks), nil
 		},
 	)
 }

diff --git a/cloud/tasks/storage/compound_storage.go b/cloud/tasks/storage/compound_storage.go
@@ -226,9 +226,9 @@ func (s *compoundStorage) ListTasksStallingWhileCancelling(
 func (s *compoundStorage) ListTasksRunning(
 	ctx context.Context,
 	limit uint64,
-) ([]string, error) {
+) ([]TaskInfo, error) {
 
-	tasks := []string{}
+	tasks := []TaskInfo{}
 	err := s.visit(ctx, func(storage Storage) error {
 		values, err := storage.ListTasksRunning(ctx, limit)
 		tasks = append(tasks, values...)
@@ -240,9 +240,9 @@ func (s *compoundStorage) ListTasksRunning(
 func (s *compoundStorage) ListTasksCancelling(
 	ctx context.Context,
 	limit uint64,
-) ([]string, error) {
+) ([]TaskInfo, error) {
 
-	tasks := []string{}
+	tasks := []TaskInfo{}
 	err := s.visit(ctx, func(storage Storage) error {
 		values, err := storage.ListTasksCancelling(ctx, limit)
 		tasks = append(tasks, values...)

diff --git a/cloud/tasks/storage/mocks/storage_mock.go b/cloud/tasks/storage/mocks/storage_mock.go
@@ -103,19 +103,19 @@ func (s *StorageMock) ListTasksStallingWhileCancelling(
 func (s *StorageMock) ListTasksRunning(
 	ctx context.Context,
 	limit uint64,
-) ([]string, error) {
+) ([]tasks_storage.TaskInfo, error) {
 
 	args := s.Called(ctx, limit)
-	res, _ := args.Get(0).([]string)
+	res, _ := args.Get(0).([]tasks_storage.TaskInfo)
 	return res, args.Error(1)
 }
 func (s *StorageMock) ListTasksCancelling(
 	ctx context.Context,
 	limit uint64,
-) ([]string, error) {
+) ([]tasks_storage.TaskInfo, error) {
 
 	args := s.Called(ctx, limit)
-	res, _ := args.Get(0).([]string)
+	res, _ := args.Get(0).([]tasks_storage.TaskInfo)
 	return res, args.Error(1)
 }
 

diff --git a/cloud/tasks/storage/storage.go b/cloud/tasks/storage/storage.go
@@ -313,8 +313,8 @@ type Storage interface {
 	) ([]TaskInfo, error)
 
 	// Used for SRE tools.
-	ListTasksRunning(ctx context.Context, limit uint64) ([]string, error)
-	ListTasksCancelling(ctx context.Context, limit uint64) ([]string, error)
+	ListTasksRunning(ctx context.Context, limit uint64) ([]TaskInfo, error)
+	ListTasksCancelling(ctx context.Context, limit uint64) ([]TaskInfo, error)
 	ListFailedTasks(ctx context.Context, since time.Time) ([]string, error)
 	ListSlowTasks(ctx context.Context, since time.Time, estimateMiss time.Duration) ([]string, error)
 

diff --git a/cloud/tasks/storage/storage_ydb.go b/cloud/tasks/storage/storage_ydb.go
@@ -208,15 +208,15 @@ func (s *storageYDB) ListTasksStallingWhileCancelling(
 func (s *storageYDB) ListTasksRunning(
 	ctx context.Context,
 	limit uint64,
-) ([]string, error) {
+) ([]TaskInfo, error) {
 
-	var tasks []string
+	var tasks []TaskInfo
 
 	err := s.db.Execute(
 		ctx,
 		func(ctx context.Context, session *persistence.Session) error {
 			var err error
-			tasks, err = s.listTaskIDs(
+			tasks, err = s.listTasks(
 				ctx,
 				session,
 				"running",
@@ -232,15 +232,15 @@ func (s *storageYDB) ListTasksRunning(
 func (s *storageYDB) ListTasksCancelling(
 	ctx context.Context,
 	limit uint64,
-) ([]string, error) {
+) ([]TaskInfo, error) {
 
-	var tasks []string
+	var tasks []TaskInfo
 
 	err := s.db.Execute(
 		ctx,
 		func(ctx context.Context, session *persistence.Session) error {
 			var err error
-			tasks, err = s.listTaskIDs(
+			tasks, err = s.listTasks(
 				ctx,
 				session,
 				"cancelling",

diff --git a/cloud/tasks/storage/storage_ydb_test.go b/cloud/tasks/storage/storage_ydb_test.go
@@ -1114,14 +1114,22 @@ func TestStorageYDBListTasksRunning(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	expectedTasks := []string{
-		taskIDRunning,
-		taskIDStallingWhileRunning,
+	expectedTasks := []TaskInfo{
+		TaskInfo{
+			ID:           taskIDRunning,
+			GenerationID: generationID,
+			TaskType:     "task1",
+		},
+		TaskInfo{
+			ID:           taskIDStallingWhileRunning,
+			GenerationID: generationID,
+			TaskType:     "task1",
+		},
 	}
 
-	taskIDs, err := storage.ListTasksRunning(ctx, 100500)
+	tasks, err := storage.ListTasksRunning(ctx, 100500)
 	require.NoError(t, err)
-	require.ElementsMatch(t, expectedTasks, taskIDs)
+	require.ElementsMatch(t, expectedTasks, tasks)
 }
 
 func TestStorageYDBListTasksCancelling(t *testing.T) {
@@ -1284,14 +1292,22 @@ func TestStorageYDBListTasksCancelling(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	expectedTasks := []string{
-		taskIDCancelling,
-		taskIDStallingWhileCancelling,
+	expectedTasks := []TaskInfo{
+		TaskInfo{
+			ID:           taskIDCancelling,
+			GenerationID: generationID,
+			TaskType:     "task1",
+		},
+		TaskInfo{
+			ID:           taskIDStallingWhileCancelling,
+			GenerationID: generationID,
+			TaskType:     "task1",
+		},
 	}
 
-	taskIDs, err := storage.ListTasksCancelling(ctx, 100500)
+	tasks, err := storage.ListTasksCancelling(ctx, 100500)
 	require.NoError(t, err)
-	require.ElementsMatch(t, expectedTasks, taskIDs)
+	require.ElementsMatch(t, expectedTasks, tasks)
 }
 
 func TestStorageYDBListFailedTasks(t *testing.T) {

diff --git a/cloud/tasks/tasks_tests/tasks_test.go b/cloud/tasks/tasks_tests/tasks_test.go
@@ -689,7 +689,7 @@ func TestTasksRunningLimit(t *testing.T) {
 
 			runningLongTaskCount := 0
 			for _, task := range runningTasks {
-				taskState, err := s.storage.GetTask(ctx, task)
+				taskState, err := s.storage.GetTask(ctx, task.ID)
 				require.NoError(t, err)
 
 				if taskState.TaskType == "long" {