be/src/exec/common/memory.cpp

// 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.
// This file is copied from
// https://github.com/ClickHouse/ClickHouse/blob/master/src/Common/formatIPv6.cpp
// and modified by Doris

#include "exec/common/memory.h"

namespace doris {

MemShareArbitrator::MemShareArbitrator(const TUniqueId& qid, int64_t query_mem_limit,
                                       double max_scan_ratio)
        : query_id(qid),
          query_mem_limit(query_mem_limit),
          mem_limit(std::max<int64_t>(
                  1, static_cast<int64_t>(static_cast<double>(query_mem_limit) * max_scan_ratio))) {
}

void MemShareArbitrator::register_scan_node() {
    total_mem_bytes.fetch_add(DEFAULT_SCANNER_MEM_BYTES);
}

int64_t MemShareArbitrator::update_mem_bytes(int64_t old_value, int64_t new_value) {
    int64_t diff = new_value - old_value;
    int64_t total = total_mem_bytes.fetch_add(diff) + diff;
    if (new_value == 0) return 0;
    if (total <= 0) return mem_limit;
    // Proportional sharing: allocate based on this context's share of total usage
    double ratio = static_cast<double>(new_value) / static_cast<double>(std::max(total, new_value));
    return static_cast<int64_t>(static_cast<double>(mem_limit) * ratio);
}

// ==================== MemLimiter ====================
int MemLimiter::available_scanner_count(int ins_idx) const {
    int64_t mem_limit_value = mem_limit.load();
    int64_t running_tasks_count_value = running_tasks_count.load();
    int64_t estimated_block_mem_bytes_value = get_estimated_block_mem_bytes();
    DCHECK_GT(estimated_block_mem_bytes_value, 0);

    int64_t max_count = std::max(1L, mem_limit_value / estimated_block_mem_bytes_value);
    int64_t avail_count = max_count;
    int64_t per_count = avail_count / parallelism;
    if (serial_operator) {
        per_count += (avail_count - per_count * parallelism);
    } else if (ins_idx < avail_count - per_count * parallelism) {
        per_count += 1;
    }

    VLOG_DEBUG << "available_scanner_count. max_count=" << max_count << "("
               << running_tasks_count_value << "/" << estimated_block_mem_bytes_value
               << "), operator_mem_limit = " << operator_mem_limit
               << ", running_tasks_count = " << running_tasks_count_value
               << ", parallelism = " << parallelism << ", avail_count = " << avail_count
               << ", ins_id = " << ins_idx << ", per_count = " << per_count
               << " debug_string: " << debug_string();

    return cast_set<int>(per_count);
}

void MemLimiter::reestimated_block_mem_bytes(int64_t value) {
    if (value == 0) return;
    value = std::min(value, operator_mem_limit);

    std::lock_guard<std::mutex> L(lock);
    auto old_value = estimated_block_mem_bytes.load();
    int64_t total =
            get_estimated_block_mem_bytes() * estimated_block_mem_bytes_update_count + value;
    estimated_block_mem_bytes_update_count += 1;
    estimated_block_mem_bytes = total / estimated_block_mem_bytes_update_count;
    VLOG_DEBUG << fmt::format(
            "reestimated_block_mem_bytes. MemLimiter = {}, estimated_block_mem_bytes = {}, "
            "old_value = {}, value: {}",
            debug_string(), estimated_block_mem_bytes, old_value, value);
}

} // namespace doris