be/src/io/fs/s3_file_system.cpp

Source
// 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 "io/fs/s3_file_system.h"

#include <fmt/format.h>

#include <cstddef>

#include "common/compiler_util.h" // IWYU pragma: keep
// IWYU pragma: no_include <bits/chrono.h>
#include <aws/core/utils/threading/Executor.h>
#include <aws/s3/S3Client.h>

#include <chrono> // IWYU pragma: keep
#include <filesystem>
#include <fstream> // IWYU pragma: keep
#include <future>
#include <memory>
#include <variant>

#include "common/check.h"
#include "common/config.h"
#include "common/logging.h"
#include "common/status.h"
#include "cpp/obj-client/obj_storage_client.h"
#include "cpp/obj-client/s3_common.h"
#include "cpp/sync_point.h"
#include "io/fs/err_utils.h"
#include "io/fs/file_system.h"
#include "io/fs/file_writer.h"
#include "io/fs/local_file_system.h"
#include "io/fs/remote_file_system.h"
#include "io/fs/s3_file_reader.h"
#include "io/fs/s3_file_writer.h"
#include "io/obj_storage/obj_client_factory.h"
#include "io/obj_storage/obj_storage_config_resolver.h"
#include "io/obj_storage/obj_storage_uri.h"
#include "runtime/exec_env.h"
#include "runtime/thread_context.h"
#include "util/s3_uri.h"
#include "util/s3_util.h"

namespace doris::io {
namespace {
constexpr std::string_view OSS_PRIVATE_ENDPOINT_SUFFIX = "-internal.aliyuncs.com";
constexpr int LEN_OF_OSS_PRIVATE_SUFFIX = 9; // length of "-internal"

#ifndef CHECK_S3_CLIENT
#define CHECK_S3_CLIENT(client)                                 \
    if (!client) {                                              \
        return Status::InvalidArgument("init s3 client error"); \
    }
#endif

} // namespace

ObjClientHolder::ObjClientHolder(S3ClientConf conf)
        : _state(std::make_shared<ObjClientState>(ObjClientState {
                  .config = ObjStorageConfigResolver::from_client_conf(conf),
                  .legacy_conf = std::move(conf),
                  .client = nullptr,
                  .generation = 0,
          })) {}

ObjClientHolder::ObjClientHolder(ObjClientConfig config)
        : _state(std::make_shared<ObjClientState>(ObjClientState {
                  .config = std::move(config),
                  .legacy_conf = std::nullopt,
                  .client = nullptr,
                  .generation = 0,
          })) {}

ObjClientHolder::~ObjClientHolder() = default;

Status ObjClientHolder::init() {
    auto current = state();
    std::shared_ptr<ObjStorageClient> client;
    if (current->legacy_conf.has_value()) {
        // A legacy conf still goes through S3ClientFactory, which applies the legacy
        // validation rules (e.g. a non-empty region) and the test client hook.
        client = DORIS_TRY(S3ClientFactory::instance().create(*current->legacy_conf));
    } else {
        client = DORIS_TRY(ObjClientFactory::instance().get_or_create(current->config));
    }
    return _publish(current, current->config, current->legacy_conf, std::move(client));
}

S3ClientConf ObjClientHolder::s3_client_conf() const {
    auto current = state();
    DORIS_CHECK(current->legacy_conf.has_value());
    return *current->legacy_conf;
}

Status ObjClientHolder::reset(const S3ClientConf& conf) {
    auto current = state();
    DORIS_CHECK(current->legacy_conf.has_value());
    S3ClientConf reset_conf = *current->legacy_conf;
    reset_conf.ak = conf.ak;
    reset_conf.sk = conf.sk;
    reset_conf.token = conf.token;
    reset_conf.bucket = conf.bucket;
    reset_conf.connect_timeout_ms = conf.connect_timeout_ms;
    reset_conf.max_connections = conf.max_connections;
    reset_conf.request_timeout_ms = conf.request_timeout_ms;
    reset_conf.use_virtual_addressing = conf.use_virtual_addressing;
    reset_conf.is_internal_bucket = conf.is_internal_bucket;

    reset_conf.role_arn = conf.role_arn;
    reset_conf.external_id = conf.external_id;
    reset_conf.cred_provider_type = conf.cred_provider_type;

    // Compare full-field equality of the merged conf, not get_hash(): the hash is
    // an XOR of crc32s and distinct configurations can collide, which would skip a
    // required client rebuild (e.g. a credential update).
    if (reset_conf == *current->legacy_conf) {
        return Status::OK(); // Same conf
    }

    auto client = DORIS_TRY(S3ClientFactory::instance().create(reset_conf));

    LOG(WARNING) << "reset s3 client with new conf: " << conf.to_string();

    auto config = ObjStorageConfigResolver::from_client_conf(reset_conf);
    return _publish(current, std::move(config), std::move(reset_conf), std::move(client));
}

Status ObjClientHolder::_publish(const std::shared_ptr<const ObjClientState>& expected,
                                 ObjClientConfig config, std::optional<S3ClientConf> legacy_conf,
                                 std::shared_ptr<ObjStorageClient> client) {
    // Declared before the lock: after the swap below it holds the replaced state, and
    // possibly the last reference to its client, which is then released after unlocking.
    std::shared_ptr<const ObjClientState> next = std::make_shared<ObjClientState>(ObjClientState {
            .config = std::move(config),
            .legacy_conf = std::move(legacy_conf),
            .client = std::move(client),
            .generation = expected->generation + 1,
    });
    {
        std::lock_guard lock(_mtx);
        // Two resets that both start from generation N merge their conf onto N. Publishing
        // the second one on top of N + 1 would silently drop the first update, so it is
        // rejected and the caller retries from the new state.
        if (_state->generation != expected->generation) {
            return Status::Aborted(
                    "s3 client was updated concurrently, expected generation {}, current {}",
                    expected->generation, _state->generation);
        }
        _state.swap(next);
    }
    return Status::OK();
}

Result<int64_t> ObjClientHolder::object_file_size(const std::string& bucket,
                                                  const std::string& key) const {
    auto client = get();
    if (!client) {
        return ResultError(Status::InvalidArgument("init s3 client error"));
    }

    auto resp = client->head_object({
            .bucket = bucket,
            .key = key,
    });

    if (resp.resp.status.code != ErrorCode::OK) {
        return ResultError(std::move(Status(resp.resp.status.code, std::move(resp.resp.status.msg))
                                             .append(fmt::format("failed to head s3 file {}",
                                                                 full_s3_path(bucket, key)))));
    }

    return resp.file_size;
}

std::string ObjClientHolder::full_s3_path(std::string_view bucket, std::string_view key) const {
    return fmt::format("{}/{}/{}", state()->config.endpoint(), bucket, key);
}

std::string S3FileSystem::full_s3_path(std::string_view key) const {
    return _client->full_s3_path(_bucket, key);
}

Result<std::shared_ptr<S3FileSystem>> S3FileSystem::create(S3Conf s3_conf, std::string id) {
    std::shared_ptr<S3FileSystem> fs(new S3FileSystem(std::move(s3_conf), std::move(id)));
    RETURN_IF_ERROR_RESULT(fs->init());
    return fs;
}

Result<std::shared_ptr<S3FileSystem>> S3FileSystem::create(ObjClientConfig config,
                                                           std::string bucket, std::string id) {
    std::shared_ptr<S3FileSystem> fs(
            new S3FileSystem(std::move(config), std::move(bucket), std::move(id)));
    RETURN_IF_ERROR_RESULT(fs->init());
    return fs;
}

S3FileSystem::S3FileSystem(S3Conf s3_conf, std::string id)
        : RemoteFileSystem(s3_conf.prefix, std::move(id), FileSystemType::S3),
          _bucket(std::move(s3_conf.bucket)),
          _prefix(std::move(s3_conf.prefix)),
          _client(std::make_shared<ObjClientHolder>(std::move(s3_conf.client_conf))) {
    // FIXME(plat1ko): Normalize prefix
    // remove the first and last '/'
    if (!_prefix.empty()) {
        size_t start = _prefix.find_first_not_of('/');
        if (start == std::string::npos) {
            _prefix = "";
        } else {
            size_t end = _prefix.find_last_not_of('/');
            if (start > 0 || end < _prefix.size() - 1) {
                _prefix = _prefix.substr(start, end - start + 1);
            }
        }
    }
}

S3FileSystem::S3FileSystem(ObjClientConfig config, std::string bucket, std::string id)
        : RemoteFileSystem("", std::move(id), FileSystemType::S3),
          _bucket(std::move(bucket)),
          _native_azure_location(std::holds_alternative<AzureClientConfig>(config.backend)),
          _client(std::make_shared<ObjClientHolder>(std::move(config))) {}

Status S3FileSystem::init() {
    return _client->init();
}

Result<std::string> S3FileSystem::_get_key(const Path& full_path) const {
    if (_native_azure_location) {
        // Typed Azure access keeps the location FE got from the table, e.g.
        // abfss://c1@acct.dfs.core.windows.net/db/t/f.parquet, which S3URI cannot parse.
        auto uri = ObjStorageUri::parse_azure(full_path.native());
        if (!uri.has_value()) {
            return ResultError(std::move(uri).error());
        }
        return std::move(uri.value().key);
    }
    // FIXME(plat1ko): Check bucket in full path and support relative path
    S3URI uri(full_path.native());
    RETURN_IF_ERROR_RESULT(uri.parse());
    return uri.get_key();
}

S3FileSystem::~S3FileSystem() = default;

Status S3FileSystem::create_file_impl(const Path& file, FileWriterPtr* writer,
                                      const FileWriterOptions* opts) {
    auto client = _client->get();
    CHECK_S3_CLIENT(client);
    auto key = DORIS_TRY(_get_key(file));
    *writer = std::make_unique<S3FileWriter>(_client, _bucket, std::move(key), opts);
    return Status::OK();
}

Status S3FileSystem::open_file_internal(const Path& file, FileReaderSPtr* reader,
                                        const FileReaderOptions& opts) {
    TEST_SYNC_POINT_CALLBACK("S3FileSystem::open_file_internal", &file, &opts);
    auto key = DORIS_TRY(_get_key(file));
    *reader = DORIS_TRY(S3FileReader::create(_client, _bucket, key, opts.file_size, nullptr));
    return Status::OK();
}

Status S3FileSystem::create_directory_impl(const Path& dir, bool failed_if_exists) {
    return Status::OK();
}

Status S3FileSystem::delete_file_impl(const Path& file) {
    auto client = _client->get();
    CHECK_S3_CLIENT(client);

    auto key = DORIS_TRY(_get_key(file));

    auto resp = client->delete_object({.bucket = _bucket, .key = key});

    if (resp.status.code == ErrorCode::OK || resp.status.code == ErrorCode::NOT_FOUND) {
        return Status::OK();
    }
    return std::move(Status(resp.status.code, std::move(resp.status.msg))
                             .append(fmt::format("failed to delete file {}", full_s3_path(key))));
}

Status S3FileSystem::delete_directory_impl(const Path& dir) {
    auto client = _client->get();
    CHECK_S3_CLIENT(client);

    auto prefix = DORIS_TRY(_get_key(dir));
    if (!prefix.empty() && prefix.back() != '/') {
        prefix.push_back('/');
    }

    ObjStoragePath delete_path {
            .path = full_s3_path(prefix),
            .bucket = _bucket,
            .prefix = prefix,
    };
    auto resp = delete_objects_recursively(std::move(client), delete_path);
    return {resp.status.code, std::move(resp.status.msg)};
}

Status S3FileSystem::batch_delete_impl(const std::vector<Path>& remote_files) {
    auto client = _client->get();
    CHECK_S3_CLIENT(client);

    // `DeleteObjectsRequest` can only contain 1000 keys at most.
    constexpr size_t max_delete_batch = 1000;
    auto path_iter = remote_files.begin();

    do {
        std::vector<std::string> objects;
        auto path_begin = path_iter;
        for (; path_iter != remote_files.end() && (path_iter - path_begin < max_delete_batch);
             ++path_iter) {
            auto key = DORIS_TRY(_get_key(*path_iter));
            objects.emplace_back(std::move(key));
        }
        if (objects.empty()) {
            return Status::OK();
        }
        // clang-format off
        if (auto resp = client->delete_objects( {.bucket = _bucket,}, std::move(objects)); resp.status.code != ErrorCode::OK) {
            return {resp.status.code, std::move(resp.status.msg)};
        }
        // clang-format on
    } while (path_iter != remote_files.end());

    return Status::OK();
}

Status S3FileSystem::exists_impl(const Path& path, bool* res) const {
    auto client = _client->get();
    CHECK_S3_CLIENT(client);
    auto key = DORIS_TRY(_get_key(path));

    VLOG_DEBUG << "key:" << key << " path:" << path;

    auto resp = client->head_object({.bucket = _bucket, .key = key});

    if (resp.resp.status.code == ErrorCode::OK) {
        *res = true;
    } else if (resp.resp.status.code == ErrorCode::NOT_FOUND) {
        *res = false;
    } else {
        return std::move(
                Status(resp.resp.status.code, std::move(resp.resp.status.msg))
                        .append(fmt::format(" failed to check exists {}", full_s3_path(key))));
    }
    return Status::OK();
}

Status S3FileSystem::file_size_impl(const Path& file, int64_t* file_size) const {
    auto key = DORIS_TRY(_get_key(file));
    *file_size = DORIS_TRY(_client->object_file_size(_bucket, key));
    return Status::OK();
}

Status S3FileSystem::list_impl(const Path& dir, bool only_file, std::vector<FileInfo>* files,
                               bool* exists) {
    // For object storage, this path is always not exist.
    // So we ignore this property and set exists to true.
    *exists = true;
    auto client = _client->get();
    CHECK_S3_CLIENT(client);
    auto prefix = DORIS_TRY(_get_key(dir));
    if (!prefix.empty() && prefix.back() != '/') {
        prefix.push_back('/');
    }

    std::vector<ObjectMeta> objects;
    auto resp = client->list_objects({.bucket = _bucket, .prefix = prefix}, &objects);
    if (!resp.ok()) {
        files->clear();
        return {resp.status.code, std::move(resp.status.msg)};
    }
    files->reserve(files->size() + objects.size());
    for (auto& obj : objects) {
        obj.key.erase(0, prefix.size());
        bool is_dir = obj.key.empty() || obj.key.back() == '/';
        files->emplace_back(FileInfo {
                .file_name = std::move(obj.key), .file_size = obj.size, .is_file = !is_dir});
    }

    return Status::OK();
}

Status S3FileSystem::rename_impl(const Path& orig_name, const Path& new_name) {
    return Status::NotSupported("S3FileSystem::rename_impl");
}

Status S3FileSystem::upload_impl(const Path& local_file, const Path& remote_file) {
    auto client = _client->get();
    CHECK_S3_CLIENT(client);

    auto key = DORIS_TRY(_get_key(remote_file));
    auto start = std::chrono::steady_clock::now();
    FileWriterPtr obj_writer;
    RETURN_IF_ERROR(create_file_impl(key, &obj_writer, nullptr));
    FileReaderSPtr local_reader;
    RETURN_IF_ERROR(io::global_local_filesystem()->open_file(local_file, &local_reader));
    size_t local_buffer_size = config::s3_file_system_local_upload_buffer_size;
    std::unique_ptr<char[]> write_buffer =
            std::make_unique_for_overwrite<char[]>(local_buffer_size);
    size_t cur_read = 0;
    while (cur_read < local_reader->size()) {
        size_t bytes_read = 0;
        RETURN_IF_ERROR(local_reader->read_at(
                cur_read, Slice {write_buffer.get(), local_buffer_size}, &bytes_read));
        RETURN_IF_ERROR(obj_writer->append({write_buffer.get(), bytes_read}));
        cur_read += bytes_read;
    }
    RETURN_IF_ERROR(obj_writer->close());
    auto duration = std::chrono::duration<float>(std::chrono::steady_clock::now() - start);

    auto size = local_reader->size();
    LOG(INFO) << "Upload " << local_file.native() << " to " << full_s3_path(key)
              << ", duration=" << duration.count() << ", bytes=" << size;

    return Status::OK();
}

Status S3FileSystem::batch_upload_impl(const std::vector<Path>& local_files,
                                       const std::vector<Path>& remote_files) {
    auto client = _client->get();
    CHECK_S3_CLIENT(client);

    if (local_files.size() != remote_files.size()) {
        return Status::InvalidArgument("local_files.size({}) != remote_files.size({})",
                                       local_files.size(), remote_files.size());
    }

    std::vector<FileWriterPtr> obj_writers(local_files.size());

    auto upload_task = [&, this](size_t idx) {
        SCOPED_ATTACH_TASK(ExecEnv::GetInstance()->s3_file_buffer_tracker());
        const auto& local_file = local_files[idx];
        const auto& remote_file = remote_files[idx];
        auto& obj_writer = obj_writers[idx];
        auto key = DORIS_TRY(_get_key(remote_file));
        LOG(INFO) << "Start to upload " << local_file.native() << " to " << full_s3_path(key);
        RETURN_IF_ERROR(create_file_impl(key, &obj_writer, nullptr));
        FileReaderSPtr local_reader;
        RETURN_IF_ERROR(io::global_local_filesystem()->open_file(local_file, &local_reader));
        size_t local_buffer_size = config::s3_file_system_local_upload_buffer_size;
        std::unique_ptr<char[]> write_buffer =
                std::make_unique_for_overwrite<char[]>(local_buffer_size);
        size_t cur_read = 0;
        while (cur_read < local_reader->size()) {
            size_t bytes_read = 0;
            RETURN_IF_ERROR(local_reader->read_at(
                    cur_read, Slice {write_buffer.get(), local_buffer_size}, &bytes_read));
            RETURN_IF_ERROR((*obj_writer).append({write_buffer.get(), bytes_read}));
            cur_read += bytes_read;
        }
        RETURN_IF_ERROR((*obj_writer).close());
        return Status::OK();
    };

    Status s = Status::OK();
    std::vector<std::future<Status>> futures;
    for (int i = 0; i < local_files.size(); ++i) {
        auto task = std::make_shared<std::packaged_task<Status(size_t idx)>>(upload_task);
        futures.emplace_back(task->get_future());
        auto st = ExecEnv::GetInstance()->s3_file_system_thread_pool()->submit_func(
                [t = std::move(task), idx = i]() mutable { (*t)(idx); });
        // We shouldn't return immediately since the previous submitted tasks might still be running in the thread pool
        if (!st.ok()) {
            s = st;
            break;
        }
    }
    for (auto&& f : futures) {
        auto cur_s = f.get();
        if (!cur_s.ok()) {
            s = std::move(cur_s);
        }
    }
    return s;
}

Status S3FileSystem::download_impl(const Path& remote_file, const Path& local_file) {
    auto client = _client->get();
    CHECK_S3_CLIENT(client);
    auto key = DORIS_TRY(_get_key(remote_file));
    int64_t size;
    RETURN_IF_ERROR(file_size(remote_file, &size));
    std::unique_ptr<char[]> buf = std::make_unique_for_overwrite<char[]>(size);
    size_t bytes_read = 0;
    // clang-format off
    auto resp = client->get_object( {.bucket = _bucket, .key = key,},
            buf.get(), 0, size, &bytes_read);
    // clang-format on
    if (resp.status.code != ErrorCode::OK) {
        return {resp.status.code, std::move(resp.status.msg)};
    }
    Aws::OFStream local_file_s;
    local_file_s.open(local_file, std::ios::out | std::ios::binary);
    if (local_file_s.good()) {
        local_file_s << StringViewStream(buf.get(), size).rdbuf();
    } else {
        return localfs_error(errno, fmt::format("failed to write file {}", local_file.native()));
    }

    return Status::OK();
}

// oss has public endpoint and private endpoint, is_public_endpoint determines
// whether to return a public endpoint.
std::string S3FileSystem::generate_presigned_url(const Path& path, int64_t expiration_secs,
                                                 bool is_public_endpoint) const {
    std::string key = fmt::format("{}/{}", _prefix, path.native());
    std::shared_ptr<ObjStorageClient> client;
    // Read the conf and the client from one snapshot so a concurrent reset cannot pair the
    // endpoint of one generation with the client of another.
    auto state = _client->state();
    if (is_public_endpoint && state->config.endpoint().ends_with(OSS_PRIVATE_ENDPOINT_SUFFIX)) {
        // Only S3-protocol backends have an OSS private endpoint. Their file systems are
        // created from an S3Conf (the typed constructor is used for Azure only), so the holder
        // keeps the legacy conf that the public-endpoint client is derived from.
        DORIS_CHECK(state->legacy_conf.has_value());
        auto new_s3_conf = *state->legacy_conf;
        new_s3_conf.endpoint.erase(new_s3_conf.endpoint.size() - OSS_PRIVATE_ENDPOINT_SUFFIX.size(),
                                   LEN_OF_OSS_PRIVATE_SUFFIX);
        auto client_result = S3ClientFactory::instance().create(new_s3_conf);
        if (!client_result) {
            LOG(WARNING) << "failed to create S3 client for presigned URL: "
                         << client_result.error();
            return {};
        }
        client = std::move(client_result).value();
    } else {
        client = state->client;
    }
    return client->generate_presigned_url({.bucket = _bucket, .key = key, .prefix = ""},
                                          expiration_secs);
}

} // namespace doris::io

Coverage Report

Created: 2026-10-21 05:08