be/src/pipeline/pipeline_tracing.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.

#include "pipeline_tracing.h"

#include <absl/time/clock.h>
#include <fcntl.h>
#include <sys/stat.h>

#include <boost/algorithm/string/predicate.hpp>
#include <chrono>
#include <cstdint>
#include <mutex>
#include <string>

#include "common/config.h"
#include "common/exception.h"
#include "common/status.h"
#include "io/fs/local_file_writer.h"
#include "util/time.h"

namespace doris::pipeline {

void PipelineTracerContext::record(ScheduleRecord record) {
    if (_dump_type == RecordType::None) [[unlikely]] {
        return;
    }

    auto map_ptr = std::atomic_load_explicit(&_data, std::memory_order_relaxed);
    auto it = map_ptr->find({record.query_id});
    if (it != map_ptr->end()) {
        it->second->enqueue(record);
    } else {
        _update([&](QueryTracesMap& new_map) {
            if (!new_map.contains({record.query_id})) {
                new_map[{record.query_id}].reset(new OneQueryTraces());
            }
            new_map[{record.query_id}]->enqueue(record);
        });
    }
}

void PipelineTracerContext::_update(std::function<void(QueryTracesMap&)>&& handler) {
    auto map_ptr = std::atomic_load_explicit(&_data, std::memory_order_relaxed);
    while (true) {
        auto new_map = std::make_shared<QueryTracesMap>(*map_ptr);
        handler(*new_map);
        if (std::atomic_compare_exchange_strong_explicit(&_data, &map_ptr, new_map,
                                                         std::memory_order_relaxed,
                                                         std::memory_order_relaxed)) {
            break;
        }
    }
}

void PipelineTracerContext::end_query(TUniqueId query_id, uint64_t workload_group) {
    {
        std::unique_lock<std::mutex> l(_tg_lock);
        _id_to_workload_group[query_id] = workload_group;
    }
    if (_dump_type == RecordType::PerQuery) {
        _dump_query(query_id);
    } else if (_dump_type == RecordType::Periodic) {
        auto now = MonotonicSeconds();
        auto interval = now - _last_dump_time;
        if (interval > _dump_interval_s) {
            _dump_timeslice();
        }
    }
}

Status PipelineTracerContext::change_record_params(
        const std::map<std::string, std::string>& params) {
    bool effective = false;
    if (auto it = params.find("type"); it != params.end()) {
        if (boost::iequals(it->second, "disable") || boost::iequals(it->second, "none")) {
            _dump_type = RecordType::None;
            effective = true;
        } else if (boost::iequals(it->second, "per_query") ||
                   boost::iequals(it->second, "perquery")) {
            _dump_type = RecordType::PerQuery;
            effective = true;
        } else if (boost::iequals(it->second, "periodic")) {
            _dump_type = RecordType::Periodic;
            _last_dump_time = MonotonicSeconds();
            effective = true;
        }
    }

    if (auto it = params.find("dump_interval"); it != params.end()) {
        _dump_interval_s = std::stoll(it->second); // s as unit
        effective = true;
    }

    return effective ? Status::OK()
                     : Status::InvalidArgument(
                               "No qualified param in changing tracing record method");
}

void PipelineTracerContext::_dump_query(TUniqueId query_id) {
    auto map_ptr = std::atomic_load_explicit(&_data, std::memory_order_relaxed);
    auto path = _log_dir / fmt::format("query{}", to_string(query_id));
    int fd = ::open(path.c_str(), O_CREAT | O_WRONLY | O_TRUNC,
                    S_ISGID | S_ISUID | S_IWUSR | S_IRUSR | S_IWGRP | S_IRGRP | S_IWOTH | S_IROTH);
    if (fd < 0) [[unlikely]] {
        throw Exception(Status::Error<ErrorCode::CREATE_FILE_ERROR>(
                "create tracing log file {} failed", path.c_str()));
    }
    auto writer = io::LocalFileWriter {path, fd};

    ScheduleRecord record;
    while ((*map_ptr)[QueryID {query_id}]->try_dequeue(record)) {
        uint64_t v = 0;
        {
            std::unique_lock<std::mutex> l(_tg_lock);
            v = _id_to_workload_group.at(query_id);
        }
        auto tmp_str = record.to_string(v);
        auto text = Slice {tmp_str};
        THROW_IF_ERROR(writer.appendv(&text, 1));
    }

    THROW_IF_ERROR(writer.close());

    _last_dump_time = MonotonicSeconds();

    _update([&](QueryTracesMap& new_map) { _data->erase(QueryID {query_id}); });

    {
        std::unique_lock<std::mutex> l(_tg_lock);
        _id_to_workload_group.erase(query_id);
    }
}

void PipelineTracerContext::_dump_timeslice() {
    auto new_map = std::make_shared<QueryTracesMap>();
    new_map.swap(_data);
    //TODO: if long time, per timeslice per file
    auto path = _log_dir /
                fmt::format("until{}", std::chrono::steady_clock::now().time_since_epoch().count());
    int fd = ::open(path.c_str(), O_CREAT | O_WRONLY | O_TRUNC,
                    S_ISGID | S_ISUID | S_IWUSR | S_IRUSR | S_IWGRP | S_IRGRP | S_IWOTH | S_IROTH);
    if (fd < 0) [[unlikely]] {
        throw Exception(Status::Error<ErrorCode::CREATE_FILE_ERROR>(
                "create tracing log file {} failed", path.c_str()));
    }
    auto writer = io::LocalFileWriter {path, fd};

    // dump all query traces in this time window to one file.
    for (auto& [query_id, trace] : (*new_map)) {
        ScheduleRecord record;
        while (trace->try_dequeue(record)) {
            uint64_t v = 0;
            {
                std::unique_lock<std::mutex> l(_tg_lock);
                v = _id_to_workload_group.at(query_id.query_id);
            }
            auto tmp_str = record.to_string(v);
            auto text = Slice {tmp_str};
            THROW_IF_ERROR(writer.appendv(&text, 1));
        }
    }
    THROW_IF_ERROR(writer.close());

    _last_dump_time = MonotonicSeconds();

    _id_to_workload_group.clear();
}
} // namespace doris::pipeline