From 4a90895458196055e3abbd8c2068b3086acdf8ce Mon Sep 17 00:00:00 2001
From: Aleksey Skvortsov <46933547+Rattysed@users.noreply.github.com>
Date: Mon, 7 Oct 2024 21:50:37 +0300
Subject: [PATCH] [Disk Manager] Speeding up takeBaseDisksToSchedule (#2221)

* fix sql query for getPoolConfigs

* add getPoolConfigsForActiveImages

* add more duration buckets

* update parallelism config

* rename method, edit config and make exponentioal duration buckets

* rename method
---
 .../pkg/services/pools/config/config.proto    |  2 +-
 .../services/pools/storage/metrics/metrics.go |  6 +--
 .../pools/storage/storage_ydb_impl.go         | 40 ++++++++++++++++++-
 3 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/cloud/disk_manager/internal/pkg/services/pools/config/config.proto b/cloud/disk_manager/internal/pkg/services/pools/config/config.proto
index 1c75022b7e5..0048dcb67a2 100644
--- a/cloud/disk_manager/internal/pkg/services/pools/config/config.proto
+++ b/cloud/disk_manager/internal/pkg/services/pools/config/config.proto
@@ -12,7 +12,7 @@ message PoolsConfig {
     // Maximum number of base disks simultaneously creating from one image.
     optional uint32 MaxBaseDisksInflight = 2 [default = 5];
     optional uint32 MaxBaseDiskUnits = 3 [default = 640];
-    optional uint32 TakeBaseDisksToScheduleParallelism = 4 [default = 20];
+    optional uint32 TakeBaseDisksToScheduleParallelism = 4 [default = 100];
     optional string ScheduleBaseDisksTaskScheduleInterval = 5 [default = "1m"];
     optional string DeleteBaseDisksTaskScheduleInterval = 6 [default = "1m"];
     optional string StorageFolder = 7 [default = "pools"];
diff --git a/cloud/disk_manager/internal/pkg/services/pools/storage/metrics/metrics.go b/cloud/disk_manager/internal/pkg/services/pools/storage/metrics/metrics.go
index 767db9ec383..bf1fe778093 100644
--- a/cloud/disk_manager/internal/pkg/services/pools/storage/metrics/metrics.go
+++ b/cloud/disk_manager/internal/pkg/services/pools/storage/metrics/metrics.go
@@ -43,10 +43,8 @@ func (s *callStats) onError() {
 ////////////////////////////////////////////////////////////////////////////////
 
 func callDurationBuckets() common_metrics.DurationBuckets {
-	return common_metrics.NewDurationBuckets(
-		10*time.Millisecond, 20*time.Millisecond, 50*time.Millisecond,
-		100*time.Millisecond, 200*time.Millisecond, 500*time.Millisecond,
-		1*time.Second, 2*time.Second, 5*time.Second,
+	return common_metrics.NewExponentialDurationBuckets(
+		10*time.Millisecond, 2, 14,
 	)
 }
 
diff --git a/cloud/disk_manager/internal/pkg/services/pools/storage/storage_ydb_impl.go b/cloud/disk_manager/internal/pkg/services/pools/storage/storage_ydb_impl.go
index cc984203529..76a33682c8e 100644
--- a/cloud/disk_manager/internal/pkg/services/pools/storage/storage_ydb_impl.go
+++ b/cloud/disk_manager/internal/pkg/services/pools/storage/storage_ydb_impl.go
@@ -1979,6 +1979,44 @@ func (s *storageYDB) getPoolConfigs(
 	return configs, nil
 }
 
+func (s *storageYDB) getPoolConfigsWithNonZeroCapacity(
+	ctx context.Context,
+	session *persistence.Session,
+) (configs []poolConfig, err error) {
+
+	res, err := session.StreamExecuteRO(ctx, fmt.Sprintf(`
+		--!syntax_v1
+		pragma TablePathPrefix = "%v";
+
+		select *
+		from configs 
+		where capacity > 0
+	`, s.tablesPath,
+	))
+	if err != nil {
+		return nil, err
+	}
+	defer res.Close()
+
+	for res.NextResultSet(ctx) {
+		for res.NextRow() {
+			config, err := scanPoolConfig(res)
+			if err != nil {
+				return nil, err
+			}
+			configs = append(configs, config)
+		}
+	}
+
+	// NOTE: always check stream query result after iteration.
+	err = res.Err()
+	if err != nil {
+		return nil, errors.NewRetriableError(err)
+	}
+
+	return configs, nil
+}
+
 func (s *storageYDB) getBaseDisksScheduling(
 	ctx context.Context,
 	session *persistence.Session,
@@ -2181,7 +2219,7 @@ func (s *storageYDB) takeBaseDisksToSchedule(
 
 	defer s.metrics.StatCall("takeBaseDisksToSchedule")(&err)
 
-	configs, err := s.getPoolConfigs(ctx, session)
+	configs, err := s.getPoolConfigsWithNonZeroCapacity(ctx, session)
 	if err != nil {
 		return nil, err
 	}