// 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 "util/disk_info.h"

// IWYU pragma: no_include <bthread/errno.h>

#include <absl/strings/str_split.h>

#include <errno.h> // IWYU pragma: keep

#include <stdio.h>

#include <stdlib.h>

#include <string.h>

#include <sys/stat.h>

#include <sys/sysmacros.h>

#include <sys/types.h>

#include <algorithm>

#include <boost/algorithm/string/classification.hpp>

#include <boost/algorithm/string/detail/classification.hpp>

#include <boost/algorithm/string/trim.hpp>

#include <fstream>

#include <iterator>

#include <memory>

#include <utility>

#include "common/cast_set.h"

#include "io/fs/local_file_system.h"

namespace doris {

#include "common/compile_check_begin.h"

bool DiskInfo::_s_initialized;

std::vector<DiskInfo::Disk> DiskInfo::_s_disks;

std::map<dev_t, int> DiskInfo::_s_device_id_to_disk_id;

std::map<std::string, int> DiskInfo::_s_disk_name_to_disk_id;

int DiskInfo::_s_num_datanode_dirs;

// Parses /proc/partitions to get the number of disks.  A bit of looking around

// seems to indicate this as the best way to do this.

// TODO: is there not something better than this?

void DiskInfo::get_device_names() {

    // Format of this file is:

    //    major, minor, #blocks, name

    // We are only interesting in name which is formatted as device_name<partition #>

    // The same device will show up multiple times for each partition (e.g. sda1, sda2).

    std::ifstream partitions("/proc/partitions", std::ios::in);

    while (partitions.good() && !partitions.eof()) {

        std::string line;

        getline(partitions, line);

        boost::trim(line);

        std::vector<std::string> fields = absl::StrSplit(line, " ", absl::SkipWhitespace());

        if (fields.size() != 4) {

            continue;

        }

        std::string name = fields[3];

        if (name == "name") {

            continue;

        }

        // Remove the partition# from the name.  e.g. sda2 --> sda

        boost::trim_right_if(name, boost::is_any_of("0123456789"));

        // Create a mapping of all device ids (one per partition) to the disk id.

        int major_dev_id = atoi(fields[0].c_str());

        int minor_dev_id = atoi(fields[1].c_str());

        dev_t dev = makedev(major_dev_id, minor_dev_id);

        DCHECK(_s_device_id_to_disk_id.find(dev) == _s_device_id_to_disk_id.end());

        int disk_id = -1;

        std::map<std::string, int>::iterator it = _s_disk_name_to_disk_id.find(name);

        if (it == _s_disk_name_to_disk_id.end()) {

            // First time seeing this disk

            disk_id = cast_set<int>(_s_disks.size());

            _s_disks.push_back(Disk(name, disk_id));

            _s_disk_name_to_disk_id[name] = disk_id;

        } else {

            disk_id = it->second;

        }

        _s_device_id_to_disk_id[dev] = disk_id;

    }

    if (partitions.is_open()) {

        partitions.close();

    }

    if (_s_disks.empty()) {

        // If all else fails, return 1

        LOG(WARNING) << "Could not determine number of disks on this machine.";

        _s_disks.push_back(Disk("sda", 0));

    }

    // Determine if the disk is rotational or not.

    for (int i = 0; i < _s_disks.size(); ++i) {

        // We can check if it is rotational by reading:

        // /sys/block/<device>/queue/rotational

        // If the file is missing or has unexpected data, default to rotational.

        std::stringstream ss;

        ss << "/sys/block/" << _s_disks[i].name << "/queue/rotational";

        std::ifstream rotational(ss.str().c_str(), std::ios::in);

        if (rotational.good()) {

            std::string line;

            getline(rotational, line);

            if (line == "0") {

                _s_disks[i].is_rotational = false;

            }

        }

        if (rotational.is_open()) {

            rotational.close();

        }

    }

}

void DiskInfo::init() {

    get_device_names();

    _s_initialized = true;

}

int DiskInfo::disk_id(const char* path) {

    struct stat s;

    stat(path, &s);

    std::map<dev_t, int>::iterator it = _s_device_id_to_disk_id.find(s.st_dev);

    if (it == _s_device_id_to_disk_id.end()) {

        return -1;

    }

    return it->second;

}

std::string DiskInfo::debug_string() {

    DCHECK(_s_initialized);

    std::stringstream stream;

    stream << "Disk Info: " << std::endl;

    stream << "  Num disks " << num_disks() << ": ";

    for (int i = 0; i < _s_disks.size(); ++i) {

        stream << _s_disks[i].name;

        if (i < num_disks() - 1) {

            stream << ", ";

        }

    }

    stream << std::endl;

    return stream.str();

}

Status DiskInfo::get_disk_devices(const std::vector<std::string>& paths,

                                  std::set<std::string>* devices) {

    std::vector<std::string> real_paths;

    for (auto& path : paths) {

        std::string p;

        Status st = io::global_local_filesystem()->canonicalize(path, &p);

        if (!st.ok()) {

            LOG(WARNING) << "skip disk monitoring of path. " << st;

            continue;

        }

        real_paths.emplace_back(std::move(p));

    }

    FILE* fp = fopen("/proc/mounts", "r");

    if (fp == nullptr) {

        std::stringstream ss;

        char buf[64];

        ss << "open /proc/mounts failed, errno:" << errno

           << ", message:" << strerror_r(errno, buf, 64);

        LOG(WARNING) << ss.str();

        return Status::InternalError(ss.str());

    }

    Status status;

    char* line_ptr = 0;

    size_t line_buf_size = 0;

    for (auto& path : real_paths) {

        size_t max_mount_size = 0;

        std::string match_dev;

        rewind(fp);

        while (getline(&line_ptr, &line_buf_size, fp) > 0) {

            char dev_path[4096];

            char mount_path[4096];

            int num = sscanf(line_ptr, "%4095s %4095s", dev_path, mount_path);

            if (num < 2) {

                continue;

            }

            size_t mount_size = strlen(mount_path);

            if (mount_size < max_mount_size || path.size() < mount_size ||

                strncmp(path.c_str(), mount_path, mount_size) != 0) {

                continue;

            }

            std::string dev(basename(dev_path));

            boost::trim_right_if(dev, boost::is_any_of("0123456789"));

            if (_s_disk_name_to_disk_id.find(dev) != std::end(_s_disk_name_to_disk_id)) {

                max_mount_size = mount_size;

                match_dev = dev;

            }

        }

        if (ferror(fp) != 0) {

            std::stringstream ss;

            char buf[64];

            ss << "open /proc/mounts failed, errno:" << errno

               << ", message:" << strerror_r(errno, buf, 64);

            LOG(WARNING) << ss.str();

            status = Status::InternalError(ss.str());

            break;

        }

        if (max_mount_size > 0) {

            devices->emplace(match_dev);

        }

    }

    if (line_ptr != nullptr) {

        free(line_ptr);

    }

    fclose(fp);

    return status;

}

#include "common/compile_check_end.h"

} // namespace doris