[feature](be) Support cross-AZ success quorum waiting

be/src/exec/sink/writer/vtablet_writer.cpp

@@ -490,6 +490,32 @@ bool IndexChannel::_quorum_success(const std::unordered_set<int64_t>& unfinished
         }
     }
 
+    const auto& table_sink = _parent->_t_sink.olap_table_sink;
+    if (table_sink.__isset.cross_az_succ_quorum) {
+        std::unordered_set<int64_t> finished_node_ids;
+        for (const auto& [node_id, node_channel] : _node_channels) {
+            if (!unfinished_node_channel_ids.contains(node_id) &&
+                node_channel->check_status().ok()) {
+                finished_node_ids.insert(node_id);
+            }
+        }
+        for (int64_t tablet_id : need_finish_tablets) {
+            const auto* tablet = _parent->_location->find_tablet(tablet_id);
+            if (tablet == nullptr) {
+                continue;
+            }
+            const auto gap_it = _parent->_tablet_version_gap_backends.find(tablet_id);
+            const auto* version_gap_node_ids = gap_it == _parent->_tablet_version_gap_backends.end()
+                                                       ? nullptr
+                                                       : &gap_it->second;
+            if (!_parent->_nodes_info->is_cross_az_quorum_success(
+                        table_sink.cross_az_succ_quorum, tablet->node_ids, finished_node_ids,
+                        version_gap_node_ids)) {
+                return false;
+            }
+        }
+    }
+
     return true;
 }
 

be/src/exec/sink/writer/vtablet_writer_v2.cpp

@@ -943,6 +943,23 @@ bool VTabletWriterV2::_quorum_success(
             return false;
         }
     }
+    const auto& table_sink = _t_sink.olap_table_sink;
+    if (table_sink.__isset.cross_az_succ_quorum) {
+        for (int64_t tablet_id : need_finish_tablets) {
+            const auto* tablet = _location->find_tablet(tablet_id);
+            if (tablet == nullptr) {
+                continue;
+            }
+            const auto gap_it = _tablet_version_gap_backends.find(tablet_id);
+            const auto* version_gap_node_ids =
+                    gap_it == _tablet_version_gap_backends.end() ? nullptr : &gap_it->second;
+            if (!_nodes_info->is_cross_az_quorum_success(table_sink.cross_az_succ_quorum,
+                                                         tablet->node_ids, finished_dst_ids,
+                                                         version_gap_node_ids)) {
+                return false;
+            }
+        }
+    }
     return true;
 }
 

be/src/storage/tablet_info.cpp

@@ -60,6 +60,36 @@
 
 namespace doris {
 
+bool DorisNodesInfo::is_cross_az_quorum_success(
+        const std::map<std::string, int32_t>& cross_az_succ_quorum,
+        const std::vector<int64_t>& tablet_node_ids,
+        const std::unordered_set<int64_t>& finished_node_ids,
+        const std::unordered_set<int64_t>* version_gap_node_ids) const {
+    for (const auto& [az, configured_min] : cross_az_succ_quorum) {
+        int replica_num_in_az = 0;
+        int succ_in_az = 0;
+        for (int64_t node_id : tablet_node_ids) {
+            const auto* node = find_node(node_id);
+            if (node == nullptr || node->location != az) {
+                continue;
+            }
+            ++replica_num_in_az;
+            if (finished_node_ids.contains(node_id) &&
+                (version_gap_node_ids == nullptr || !version_gap_node_ids->contains(node_id))) {
+                ++succ_in_az;
+            }
+        }
+        const int required_in_az = std::min(configured_min, replica_num_in_az);
+        if (required_in_az == 0) {
+            continue;
+        }
+        if (succ_in_az < required_in_az) {
+            return false;
+        }
+    }
+    return true;
+}
+
 const OlapTableIndexSchema* OlapTableSchemaParam::row_binlog_index_schema(int64_t index_id) const {
     for (auto* schema : _row_binlog_index_schemas) {
         if (schema->index_id == index_id) {

be/src/storage/tablet_info.h

@@ -31,6 +31,7 @@
 #include <string>
 #include <tuple>
 #include <unordered_map>
+#include <unordered_set>
 #include <utility>
 #include <vector>
 
@@ -443,14 +444,16 @@ struct NodeInfo {
     int64_t option;
     std::string host;
     int32_t brpc_port;
+    std::string location;
 
     NodeInfo() = default;
 
     NodeInfo(const TNodeInfo& tnode)
             : id(tnode.id),
               option(tnode.option),
               host(tnode.host),
-              brpc_port(tnode.async_internal_port) {}
+              brpc_port(tnode.async_internal_port),
+              location(tnode.__isset.location ? tnode.location : "") {}
 };
 
 class DorisNodesInfo {
@@ -475,6 +478,11 @@ class DorisNodesInfo {
         return nullptr;
     }
 
+    bool is_cross_az_quorum_success(const std::map<std::string, int32_t>& cross_az_succ_quorum,
+                                    const std::vector<int64_t>& tablet_node_ids,
+                                    const std::unordered_set<int64_t>& finished_node_ids,
+                                    const std::unordered_set<int64_t>* version_gap_node_ids) const;
+
     void add_nodes(const std::vector<TNodeInfo>& t_nodes) {
         for (const auto& node : t_nodes) {
             const auto* node_info = find_node(node.id);

be/test/exec/sink/vtablet_writer_v2_test.cpp

@@ -520,4 +520,54 @@ TEST_F(TestVTabletWriterV2, quorum_excludes_streams_not_closing_in_current_stage
     ASSERT_TRUE(writer->_quorum_success(unfinished_streams, need_finish_tablets));
 }
 
+TEST_F(TestVTabletWriterV2, quorum_waits_for_cross_az_success) {
+    UniqueId load_id;
+    auto load_stream_map = std::make_shared<LoadStreamMap>(load_id, src_id, 1, 1, nullptr);
+    auto streams_1 = load_stream_map->get_or_create(1001);
+    auto streams_2 = load_stream_map->get_or_create(1002);
+    auto streams_3 = load_stream_map->get_or_create(1003);
+    for (const auto& streams : {streams_1, streams_2, streams_3}) {
+        streams->streams().front()->_is_closing.store(true);
+    }
+
+    TPaloNodesInfo t_nodes;
+    for (const auto& [node_id, location] : std::vector<std::pair<int64_t, std::string>> {
+                 {1001, "az1"}, {1002, "az1"}, {1003, "az2"}}) {
+        TNodeInfo node;
+        node.__set_id(node_id);
+        node.__set_location(location);
+        t_nodes.nodes.push_back(node);
+    }
+    DorisNodesInfo nodes_info(t_nodes);
+
+    TOlapTableLocationParam t_location;
+    TTabletLocation tablet;
+    tablet.__set_tablet_id(1);
+    tablet.__set_node_ids({1001, 1002, 1003});
+    t_location.tablets.push_back(tablet);
+    OlapTableLocationParam location(t_location);
+
+    auto writer = create_vtablet_writer();
+    writer->_load_stream_map = load_stream_map;
+    writer->_nodes_info = &nodes_info;
+    writer->_location = &location;
+    for (int64_t node_id : {1001, 1002, 1003}) {
+        writer->_tablets_by_node[node_id].insert(1);
+    }
+
+    std::unordered_set<std::shared_ptr<LoadStreamStub>> unfinished_streams {
+            streams_3->streams().front()};
+    std::unordered_set<int64_t> need_finish_tablets {1};
+    ASSERT_TRUE(writer->_quorum_success(unfinished_streams, need_finish_tablets));
+
+    writer->_t_sink.olap_table_sink.__set_cross_az_succ_quorum({{"az1", 2}, {"az2", 2}});
+    ASSERT_FALSE(writer->_quorum_success(unfinished_streams, need_finish_tablets));
+
+    unfinished_streams.clear();
+    ASSERT_TRUE(writer->_quorum_success(unfinished_streams, need_finish_tablets));
+
+    writer->_tablet_version_gap_backends[1].insert(1003);
+    ASSERT_FALSE(writer->_quorum_success(unfinished_streams, need_finish_tablets));
+}
+
 } // namespace doris

gensrc/thrift/DataSinks.thrift

@@ -319,6 +319,7 @@ struct TOlapTableSink {
     // initial partition list is empty, so auto-partition tables whose first partitions arrive at
     // runtime still enter the correct mode from the start.
     25: optional bool enable_adaptive_random_bucket
+    26: optional map<string, i32> cross_az_succ_quorum
 }
 
 struct THiveLocationParams {

gensrc/thrift/Descriptors.thrift

@@ -390,6 +390,7 @@ struct TNodeInfo {
     3: required string host
     // used to transfer data between nodes
     4: required i32 async_internal_port
+    5: optional string location
 }
 
 struct TPaloNodesInfo {

regression-test/suites/load_p0/cross_az_quorum/test_cross_az_succ_quorum.groovy

@@ -0,0 +1,127 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+import org.apache.doris.regression.suite.ClusterOptions
+import org.apache.doris.regression.util.NodeType
+import org.awaitility.Awaitility
+
+import static java.util.concurrent.TimeUnit.SECONDS
+
+suite('test_cross_az_succ_quorum', 'docker') {
+    def options = new ClusterOptions()
+    options.feConfigs += ['disable_tablet_scheduler=true']
+    options.beConfigs += ['quorum_success_min_wait_seconds=1']
+    options.enableDebugPoints()
+    options.cloudMode = false
+
+    docker(options) {
+        def backends = sql_return_maparray('SHOW BACKENDS')
+        assertEquals(3, backends.size())
+        sql """ALTER SYSTEM MODIFY BACKEND '${backends[0].BackendId}' SET ('tag.location' = 'az1')"""
+        sql """ALTER SYSTEM MODIFY BACKEND '${backends[1].BackendId}' SET ('tag.location' = 'az1')"""
+        sql """ALTER SYSTEM MODIFY BACKEND '${backends[2].BackendId}' SET ('tag.location' = 'az2')"""
+
+        // tag.location doubles as the compute group name, and WorkloadGroupChecker only creates the
+        // `normal` workload group for a newly seen compute group every workload_group_check_interval_ms
+        // (2s by default). Loading before that fails with "Can not find workload group normal in
+        // compute group az1", so wait for it rather than turning workload groups off -- production
+        // runs with them on.
+        Awaitility.await().atMost(60, SECONDS).pollInterval(1, SECONDS).until({
+            def normals = sql_return_maparray('SHOW WORKLOAD GROUPS')
+                    .findAll { it.Name == 'normal' }
+                    .collect { it.compute_group } as Set
+            normals.containsAll(['az1', 'az2'])
+        })
+
+        sql 'DROP TABLE IF EXISTS cross_az_quorum_table'
+        sql '''
+            CREATE TABLE cross_az_quorum_table (k INT)
+            DISTRIBUTED BY HASH(k) BUCKETS 1
+            PROPERTIES ('replication_allocation' = 'tag.location.az1: 2, tag.location.az2: 1')
+        '''
+
+        def injectName = 'TxnManager.prepare_txn.random_failed'
+        GetDebugPoint().enableDebugPoint(backends[0].Host, backends[0].HttpPort as int,
+                NodeType.BE, injectName, [percent: 1.0])
+
+        // The default empty config preserves the normal two-of-three quorum behavior.
+        sql 'INSERT INTO cross_az_quorum_table VALUES (1)'
+
+        setFeConfig('cross_az_succ_quorum', 'az1:2,az2:1')
+        test {
+            sql 'INSERT INTO cross_az_quorum_table VALUES (2)'
+            exception 'cross AZ success quorum failed for az1'
+        }
+
+        setFeConfig('cross_az_succ_quorum', '')
+        sql 'INSERT INTO cross_az_quorum_table VALUES (3)'
+        GetDebugPoint().disableDebugPoint(backends[0].Host, backends[0].HttpPort as int,
+                NodeType.BE, injectName)
+
+        // Losing all successful replicas in one AZ still leaves a normal two-of-three quorum.
+        sql 'DROP TABLE IF EXISTS cross_az_quorum_az2_table'
+        sql '''
+            CREATE TABLE cross_az_quorum_az2_table (k INT)
+            DISTRIBUTED BY HASH(k) BUCKETS 1
+            PROPERTIES ('replication_allocation' = 'tag.location.az1: 2, tag.location.az2: 1')
+        '''
+        GetDebugPoint().enableDebugPoint(backends[2].Host, backends[2].HttpPort as int,
+                NodeType.BE, injectName, [percent: 1.0])
+        sql 'INSERT INTO cross_az_quorum_az2_table VALUES (1)'
+        setFeConfig('cross_az_succ_quorum', 'az1:2,az2:1')
+        test {
+            sql 'INSERT INTO cross_az_quorum_az2_table VALUES (2)'
+            exception 'cross AZ success quorum failed for az2'
+        }
+        setFeConfig('cross_az_succ_quorum', '')
+        GetDebugPoint().disableDebugPoint(backends[2].Host, backends[2].HttpPort as int,
+                NodeType.BE, injectName)
+
+        // A slow remote AZ must remain in the first close-wait stage instead of being dropped
+        // after the two local replicas reach the ordinary quorum.
+        // Use a fresh table: the tables above already committed a version while their az2 replica
+        // was failing, so that replica carries a version gap and, with the tablet scheduler
+        // disabled, is never repaired -- it could never count as a success again.
+        sql 'DROP TABLE IF EXISTS cross_az_quorum_slow_table'
+        sql '''
+            CREATE TABLE cross_az_quorum_slow_table (k INT)
+            DISTRIBUTED BY HASH(k) BUCKETS 1
+            PROPERTIES ('replication_allocation' = 'tag.location.az1: 2, tag.location.az2: 1')
+        '''
+        setFeConfig('cross_az_succ_quorum', 'az1:2,az2:1')
+        GetDebugPoint().enableDebugPoint(backends[2].Host, backends[2].HttpPort as int,
+                NodeType.BE, 'TxnManager.prepare_txn.wait', [duration: 3000])
+        sql 'INSERT INTO cross_az_quorum_slow_table VALUES (1)'
+        GetDebugPoint().disableDebugPoint(backends[2].Host, backends[2].HttpPort as int,
+                NodeType.BE, 'TxnManager.prepare_txn.wait')
+        setFeConfig('cross_az_succ_quorum', '')
+
+        sql 'DROP TABLE IF EXISTS cross_az_quorum_clamp_table'
+        sql '''
+            CREATE TABLE cross_az_quorum_clamp_table (k INT)
+            DISTRIBUTED BY HASH(k) BUCKETS 1
+            PROPERTIES ('replication_allocation' = 'tag.location.az1: 1')
+        '''
+        setFeConfig('cross_az_succ_quorum', 'az1:2')
+        sql 'INSERT INTO cross_az_quorum_clamp_table VALUES (1)'
+
+        // Invalid entries are ignored and must never break the commit path.
+        setFeConfig('cross_az_succ_quorum', 'invalid,az1:not-a-number')
+        sql 'INSERT INTO cross_az_quorum_clamp_table VALUES (2)'
+        setFeConfig('cross_az_succ_quorum', '')
+    }
+}