be/src/load/routine_load/consumer_helpers.h

// 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.

#pragma once

#include <chrono>
#include <string>
#include <unordered_map>

namespace doris {

// Helper class for tracking metrics in data consumers
class ConsumerMetrics {
public:
    ConsumerMetrics() = default;

    // Track a single message fetch operation
    void track_get_msg(int64_t latency_us);

    // Track consumed bytes
    void track_consume_bytes(int64_t bytes);

    // Track consumed rows
    void track_consume_rows(int64_t rows);

    // Get accumulated metrics
    int64_t get_msg_count() const { return _get_msg_count; }
    int64_t consume_bytes() const { return _consume_bytes; }
    int64_t consume_rows() const { return _consume_rows; }

private:
    int64_t _get_msg_count = 0;
    int64_t _consume_bytes = 0;
    int64_t _consume_rows = 0;
};

// Helper class for retry logic with configurable backoff
class RetryPolicy {
public:
    explicit RetryPolicy(int max_retries = 3, int backoff_ms = 200)
            : _max_retries(max_retries), _backoff_ms(backoff_ms), _retry_count(0) {}

    // Check if should retry
    bool should_retry() const { return _retry_count < _max_retries; }

    // Increment retry count and sleep with backoff
    void retry_with_backoff();

    // Reset retry counter (call on success)
    void reset() { _retry_count = 0; }

    // Get current retry count
    int retry_count() const { return _retry_count; }

private:
    int _max_retries;
    int _backoff_ms;
    int _retry_count;
};

// Helper class for exponential backoff (used for throttling)
class ThrottleBackoff {
public:
    explicit ThrottleBackoff(int initial_backoff_ms = 1000, int max_backoff_ms = 10000)
            : _initial_backoff_ms(initial_backoff_ms),
              _max_backoff_ms(max_backoff_ms),
              _throttle_count(0) {}

    // Increment throttle count and sleep with exponential backoff
    void backoff_and_sleep();

    // Reset throttle counter (call on success)
    void reset() { _throttle_count = 0; }

    // Get current throttle count
    int throttle_count() const { return _throttle_count; }

private:
    int _initial_backoff_ms;
    int _max_backoff_ms;
    int _throttle_count;
};

// Helper class for comparing custom properties between consumers
class PropertyMatcher {
public:
    // Check if two property maps match exactly
    template <typename MapType1, typename MapType2>
    static bool properties_match(const MapType1& props1, const MapType2& props2) {
        if (props1.size() != props2.size()) {
            return false;
        }

        for (const auto& [key, value] : props1) {
            auto it = props2.find(key);
            if (it == props2.end() || it->second != value) {
                return false;
            }
        }

        return true;
    }
};

} // namespace doris