improve rebalances

Author: niksaveliev

ydb/core/kafka_proxy/actors/kafka_balance_actor_sql.cpp

@@ -9,9 +9,8 @@ const TString INSERT_NEW_GROUP = R"sql(
     DECLARE $State AS Uint64;
     DECLARE $Database AS Utf8;
     DECLARE $Master AS Utf8;
-    DECLARE $LastHeartbeat AS Datetime;
+    DECLARE $LastMasterHeartbeat AS Datetime;
     DECLARE $ProtocolType AS Utf8;
-    DECLARE $RebalanceTimeoutMs AS Uint32;
 
     INSERT INTO `%s`
     (
@@ -21,19 +20,17 @@ const TString INSERT_NEW_GROUP = R"sql(
         database,
         last_heartbeat_time,
         master,
-        protocol_type,
-        rebalance_timeout_ms
+        protocol_type
     )
     VALUES
     (
         $ConsumerGroup,
         $Generation,
         $State,
         $Database,
-        $LastHeartbeat,
+        $LastMasterHeartbeat,
         $Master,
-        $ProtocolType,
-        $RebalanceTimeoutMs
+        $ProtocolType
     );
 )sql";
 
@@ -44,16 +41,14 @@ const TString UPDATE_GROUP = R"sql(
     DECLARE $Generation AS Uint64;
     DECLARE $Database AS Utf8;
     DECLARE $Master AS Utf8;
-    DECLARE $LastHeartbeat AS Datetime;
-    DECLARE $RebalanceTimeoutMs AS Uint32;
+    DECLARE $LastMasterHeartbeat AS Datetime;
 
     UPDATE `%s`
     SET
         state = $State,
         generation = $Generation,
-        last_heartbeat_time = $LastHeartbeat,
-        master = $Master,
-        rebalance_timeout_ms = $RebalanceTimeoutMs
+        last_heartbeat_time = $LastMasterHeartbeat,
+        master = $Master
     WHERE database = $Database
       AND consumer_group = $ConsumerGroup;
 )sql";
@@ -77,12 +72,12 @@ const TString UPDATE_GROUP_STATE_AND_PROTOCOL = R"sql(
     DECLARE $State AS Uint64;
     DECLARE $Database AS Utf8;
     DECLARE $Protocol AS Utf8;
-    DECLARE $LastHeartbeat AS Datetime;
+    DECLARE $LastMasterHeartbeat AS Datetime;
 
     UPDATE `%s`
     SET
         state = $State,
-        last_heartbeat_time = $LastHeartbeat,
+        last_heartbeat_time = $LastMasterHeartbeat,
         protocol = $Protocol
     WHERE database = $Database
     AND consumer_group = $ConsumerGroup;
@@ -93,6 +88,7 @@ const TString INSERT_MEMBER = R"sql(
     DECLARE $ConsumerGroup       AS Utf8;
     DECLARE $Generation          AS Uint64;
     DECLARE $MemberId            AS Utf8;
+    DECLARE $InstanceId          AS Utf8;
     DECLARE $WorkerStateProto    AS String;
     DECLARE $Database            AS Utf8;
     DECLARE $HeartbeatDeadline AS Datetime;
@@ -103,6 +99,7 @@ const TString INSERT_MEMBER = R"sql(
         consumer_group,
         generation,
         member_id,
+        instance_id,
         heartbeat_deadline,
         worker_state_proto,
         database,
@@ -112,6 +109,7 @@ const TString INSERT_MEMBER = R"sql(
         $ConsumerGroup,
         $Generation,
         $MemberId,
+        $InstanceId,
         $HeartbeatDeadline,
         $WorkerStateProto,
         $Database,
@@ -126,7 +124,7 @@ const TString UPSERT_ASSIGNMENTS_AND_SET_WORKING_STATE = R"sql(
     DECLARE $Database AS Utf8;
     DECLARE $Generation AS Uint64;
     DECLARE $State AS Uint64;
-    DECLARE $LastHeartbeat AS Datetime;
+    DECLARE $LastMasterHeartbeat AS Datetime;
 
     UPSERT INTO `%s`
     SELECT
@@ -140,11 +138,32 @@ const TString UPSERT_ASSIGNMENTS_AND_SET_WORKING_STATE = R"sql(
     UPDATE `%s`
     SET
         state = $State,
-        last_heartbeat_time = $LastHeartbeat
+        last_heartbeat_time = $LastMasterHeartbeat,
+        last_success_generation = $Generation
     WHERE database = $Database
       AND consumer_group = $ConsumerGroup;
 )sql";
 
+const TString SELECT_ALIVE_MEMBERS = R"sql(
+    --!syntax_v1
+    DECLARE $ConsumerGroup AS Utf8;
+    DECLARE $Generation AS Uint64;
+    DECLARE $Database AS Utf8;
+    DECLARE $PaginationMemberId AS Utf8;
+    DECLARE $Limit AS Uint64;
+
+    SELECT member_id, instance_id, rebalance_timeout_ms
+    FROM `%s`
+    VIEW PRIMARY KEY
+    WHERE database = $Database
+      AND consumer_group = $ConsumerGroup
+      AND generation = $Generation
+      AND member_id > $PaginationMemberId
+      AND (leaved IS NULL OR leaved = False)
+    ORDER BY member_id
+    LIMIT $Limit;
+)sql";
+
 const TString SELECT_WORKER_STATES = R"sql(
     --!syntax_v1
     DECLARE $ConsumerGroup AS Utf8;
@@ -153,7 +172,7 @@ const TString SELECT_WORKER_STATES = R"sql(
     DECLARE $PaginationMemberId AS Utf8;
     DECLARE $Limit AS Uint64;
 
-    SELECT worker_state_proto, member_id
+    SELECT worker_state_proto, member_id, instance_id
     FROM `%s`
     VIEW PRIMARY KEY
     WHERE database = $Database
@@ -169,7 +188,7 @@ const TString CHECK_GROUP_STATE = R"sql(
     DECLARE $ConsumerGroup AS Utf8;
     DECLARE $Database AS Utf8;
 
-    SELECT state, generation, master, last_heartbeat_time, consumer_group, database, protocol, protocol_type, rebalance_timeout_ms
+    SELECT state, generation, master, last_heartbeat_time, consumer_group, database, protocol, protocol_type, last_success_generation
     FROM `%s`
     VIEW PRIMARY KEY
     WHERE database = $Database
@@ -218,20 +237,21 @@ const TString CHECK_DEAD_MEMBERS = R"sql(
 
 )sql";
 
-const TString UPDATE_LASTHEARTBEATS = R"sql(
+const TString UPDATE_LAST_HEARTBEATS = R"sql(
     --!syntax_v1
     DECLARE $ConsumerGroup AS Utf8;
     DECLARE $Generation AS Uint64;
     DECLARE $MemberId AS Utf8;
     DECLARE $Database AS Utf8;
-    DECLARE $LastHeartbeat AS Datetime;
+    DECLARE $LastMasterHeartbeat AS Datetime;
     DECLARE $HeartbeatDeadline AS Datetime;
     DECLARE $UpdateGroupHeartbeat AS Bool;
 
     UPDATE `%s`
-    SET last_heartbeat_time = $LastHeartbeat
+    SET last_heartbeat_time = $LastMasterHeartbeat
     WHERE consumer_group = $ConsumerGroup
         AND database = $Database
+        AND generation = $Generation
         AND $UpdateGroupHeartbeat = True;
 
     UPDATE `%s`
@@ -248,10 +268,11 @@ const TString UPDATE_LASTHEARTBEAT_TO_LEAVE_GROUP = R"sql(
     DECLARE $ConsumerGroup AS Utf8;
     DECLARE $MemberId AS Utf8;
     DECLARE $Database AS Utf8;
-    DECLARE $LastHeartbeat AS Datetime;
+    DECLARE $LastMasterHeartbeat AS Datetime;
 
     UPDATE `%s`
-    SET heartbeat_deadline = $LastHeartbeat
+    SET heartbeat_deadline = $LastMasterHeartbeat,
+        leaved = True
     WHERE consumer_group = $ConsumerGroup
     AND member_id = $MemberId
     AND database = $Database;