Copilot commented on code in PR #3677: URL: https://github.com/apache/celeborn/pull/3677#discussion_r3300784920
########## sbin/stop-lifecycle-manager.sh: ########## @@ -0,0 +1,142 @@ +#!/usr/bin/env bash +# +# 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. +# + +# Stops a Celeborn LifecycleManager daemon by PID file or port. +# Usage: +# stop-lifecycle-manager.sh --port <port> Stop the instance bound to <port> +# stop-lifecycle-manager.sh --all Stop all LifecycleManager instances + +set -euo pipefail + +if [ -z "${CELEBORN_HOME:-}" ]; then + export CELEBORN_HOME="$(cd "$(dirname "$0")/.."; pwd)" +fi + +LOG_DIR="${CELEBORN_HOME}/logs" +GRACEFUL_TIMEOUT=10 + +usage() { + cat <<EOF +Usage: $0 --port <port> | --all + + --port <port> Stop the LifecycleManager instance running on the specified port + --all Stop all LifecycleManager instances managed by this script +EOF + exit 1 +} + +# Stop a single instance by its PID file +stop_by_pid_file() { + local pid_file="$1" + local pid + + if [ ! -f "${pid_file}" ]; then + echo "PID file not found: ${pid_file}" >&2 + return 1 + fi + + pid=$(cat "${pid_file}") + if [ -z "${pid}" ]; then + echo "PID file is empty: ${pid_file}, removing." >&2 + rm -f "${pid_file}" + return 1 + fi + + if kill -0 "${pid}" 2>/dev/null; then + echo "Stopping LifecycleManager (PID: ${pid}) with SIGTERM ..." + kill -TERM "${pid}" + + # Wait for graceful shutdown + local waited=0 + while [ "${waited}" -lt "${GRACEFUL_TIMEOUT}" ]; do + if ! kill -0 "${pid}" 2>/dev/null; then + echo "LifecycleManager (PID: ${pid}) stopped." + rm -f "${pid_file}" + # Clean up the env file + rm -f "${LOG_DIR}/lifecyclemanager-${pid}.env" + return 0 + fi + sleep 1 + waited=$(( waited + 1 )) + done + + # Force kill if still alive + echo "LifecycleManager (PID: ${pid}) did not stop after ${GRACEFUL_TIMEOUT}s, sending SIGKILL ..." + kill -9 "${pid}" 2>/dev/null || true + sleep 1 + rm -f "${pid_file}" + rm -f "${LOG_DIR}/lifecyclemanager-${pid}.env" + echo "LifecycleManager (PID: ${pid}) killed." + else + echo "LifecycleManager (PID: ${pid}) is not running. Cleaning up stale PID file." + rm -f "${pid_file}" + rm -f "${LOG_DIR}/lifecyclemanager-${pid}.env" + fi +} + +# Parse arguments +if [ $# -eq 0 ]; then + usage +fi + +MODE="" +TARGET_PORT="" + +while [ $# -gt 0 ]; do + case "$1" in + --port) + MODE="port" + TARGET_PORT="${2:-}" + if [ -z "${TARGET_PORT}" ]; then + echo "Error: --port requires a port number." >&2 + usage + fi + shift 2 + ;; + --all) + MODE="all" + shift + ;; + *) + echo "Unknown option: $1" >&2 + usage + ;; + esac +done + +if [ -z "${MODE}" ]; then + usage +fi + +case "${MODE}" in + port) + PID_FILE="${LOG_DIR}/lifecyclemanager-${TARGET_PORT}.pid" + stop_by_pid_file "${PID_FILE}" + ;; + all) + found=0 + for pid_file in "${LOG_DIR}"/lifecyclemanager-*.pid; do + [ -f "${pid_file}" ] || continue + found=1 + stop_by_pid_file "${pid_file}" + done Review Comment: In `--all` mode, `stop_by_pid_file` can return non-zero (e.g., empty/stale PID file cleanup), and with `set -e` this aborts the whole script before stopping remaining instances. The loop should continue even if one PID file fails to stop/clean up. ########## sbin/start-lifecycle-manager.sh: ########## @@ -0,0 +1,175 @@ +#!/usr/bin/env bash +# +# 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. + +# Starts a standalone Celeborn LifecycleManager daemon in background. +# The RPC port is randomly selected from 30000~50000 if --port is not specified. +# After successful start, the port is exported as CELEBORN_LM_PORT. + +set -euo pipefail + +if [ -z "${CELEBORN_HOME:-}" ]; then + export CELEBORN_HOME="$(cd "$(dirname "$0")/.."; pwd)" +fi + +usage() { + cat <<EOF +Usage: $0 --app-id <id> --master-endpoints <ep1,ep2,...> [--port <port>] [--host <host>] [--properties-file <file>] + + --app-id REQUIRED unique application id + --master-endpoints REQUIRED comma-separated host:port of Celeborn Masters + --port OPTIONAL fixed RPC port to bind (default: random available port in 30000~50000) + --host OPTIONAL bind host (default: hostname) + --properties-file OPTIONAL path to celeborn-defaults.conf +EOF + exit 1 +} + +# Find a random available port in range [30000, 50000] +find_available_port() { + local port + local max_attempts=100 + local attempt=0 + while [ "$attempt" -lt "$max_attempts" ]; do + port=$(( RANDOM % 20001 + 30000 )) + # Check if the port is available (not in use) + if ! (echo >/dev/tcp/127.0.0.1/"$port") 2>/dev/null; then + echo "$port" + return 0 + fi + attempt=$(( attempt + 1 )) + done Review Comment: `find_available_port` only probes 127.0.0.1 via `/dev/tcp`. If an existing process is already listening on the chosen port but not on loopback (e.g., bound to a specific non-loopback interface), this check can miss the conflict and the daemon will fail to bind. Prefer a listener check that covers all interfaces when possible. ########## rust/resource/lib/README.md: ########## @@ -0,0 +1,51 @@ +# Prebuilt Celeborn native artifacts + +`celeborn-client-sys/build.rs` looks up the aggregated dylib here before +falling back to a from-source `cmake` build, so dropping the right file +in is all it takes for `cargo build` to work without any environment +variable. + +## Layout + +``` +resource/lib/ + <target-triple>/ + libceleborn_client.dylib # macOS + libceleborn_client.so # Linux +``` + +`<target-triple>` is the Cargo target identifier — the same string Cargo +prints from `rustc -vV` (`host:` line) or accepts in `cargo build --target`: + +| Platform | Triple | Filename | +|------------------|------------------------------|---------------------------| +| macOS Apple Si | `aarch64-apple-darwin` | `libceleborn_client.dylib`| +| macOS Intel | `x86_64-apple-darwin` | `libceleborn_client.dylib`| +| Linux x86_64 | `x86_64-unknown-linux-gnu` | `libceleborn_client.so` | Review Comment: Spelling: the platform table row says "macOS Apple Si" which looks truncated; it should be "macOS Apple Silicon". ########## cpp/celeborn/ffi/CelebornFfi.cc: ########## @@ -0,0 +1,281 @@ +#include "celeborn/ffi/CelebornFfi.h" + +#include <cstdlib> +#include <cstring> +#include <exception> +#include <memory> +#include <new> +#include <string> +#include <vector> + +#include "celeborn/client/ShuffleClient.h" +#include "celeborn/client/reader/CelebornInputStream.h" +#include "celeborn/conf/CelebornConf.h" + +namespace { + +struct ClientImpl { + std::shared_ptr<celeborn::conf::CelebornConf> conf; + std::shared_ptr<celeborn::client::ShuffleClientEndpoint> endpoint; + std::shared_ptr<celeborn::client::ShuffleClientImpl> client; + std::string app_id; + std::string lifecycle_manager_host; +}; + +// One open partition stream. The stream references the owning +// ShuffleClient, so callers must close every reader before shutting down +// the client. +struct PartitionReaderImpl { + std::unique_ptr<celeborn::client::CelebornInputStream> stream; +}; + +inline ClientImpl* as_impl(celeborn_ffi_handle* h) { + return reinterpret_cast<ClientImpl*>(h); +} + +inline PartitionReaderImpl* as_reader_impl(celeborn_ffi_partition_reader* r) { + return reinterpret_cast<PartitionReaderImpl*>(r); +} + +char* dup_message(const std::string& msg) { + char* out = static_cast<char*>(std::malloc(msg.size() + 1)); + if (!out) { + return nullptr; + } + std::memcpy(out, msg.data(), msg.size()); + out[msg.size()] = '\0'; + return out; +} + +void set_error(char** err_out, const std::string& msg) { + if (err_out) { + *err_out = dup_message("celeborn-ffi: " + msg); + } +} + +template <typename Fn> +celeborn_ffi_status guarded(char** err_out, Fn&& fn) { + try { + fn(); + return CELEBORN_FFI_OK; + } catch (const std::exception& e) { + set_error(err_out, e.what()); + return CELEBORN_FFI_ERROR; + } catch (...) { + set_error(err_out, "unknown C++ exception"); + return CELEBORN_FFI_ERROR; + } +} + +} // namespace + +extern "C" { + +void celeborn_ffi_free_error(char* err) { + std::free(err); +} + +void celeborn_ffi_free_buffer(uint8_t* data) { + delete[] data; +} + +celeborn_ffi_handle* celeborn_ffi_create_client( + const char* app_id, + size_t app_id_len, + int32_t push_buffer_max_size, + const char* codec, + size_t codec_len, + char** err_out) { + try { + auto impl = std::make_unique<ClientImpl>(); + impl->app_id.assign(app_id, app_id_len); + impl->conf = std::make_shared<celeborn::conf::CelebornConf>(); + + if (push_buffer_max_size > 0) { + impl->conf->registerProperty( + celeborn::conf::CelebornConf::kClientPushBufferMaxSize, + std::to_string(push_buffer_max_size) + "b"); + } + if (codec_len > 0) { + impl->conf->registerProperty( + celeborn::conf::CelebornConf::kShuffleCompressionCodec, + std::string(codec, codec_len)); + } Review Comment: `celeborn_ffi_create_client` dereferences `app_id`/`codec` via `std::string::assign` / `std::string(codec, ...)` without validating pointers. A null pointer from any non-Rust caller will segfault instead of returning `NULL` with an error message. ########## cpp/celeborn/ffi/CelebornFfi.cc: ########## @@ -0,0 +1,281 @@ +#include "celeborn/ffi/CelebornFfi.h" + +#include <cstdlib> +#include <cstring> +#include <exception> +#include <memory> +#include <new> +#include <string> +#include <vector> + +#include "celeborn/client/ShuffleClient.h" +#include "celeborn/client/reader/CelebornInputStream.h" +#include "celeborn/conf/CelebornConf.h" + +namespace { + +struct ClientImpl { + std::shared_ptr<celeborn::conf::CelebornConf> conf; + std::shared_ptr<celeborn::client::ShuffleClientEndpoint> endpoint; + std::shared_ptr<celeborn::client::ShuffleClientImpl> client; + std::string app_id; + std::string lifecycle_manager_host; +}; + +// One open partition stream. The stream references the owning +// ShuffleClient, so callers must close every reader before shutting down +// the client. +struct PartitionReaderImpl { + std::unique_ptr<celeborn::client::CelebornInputStream> stream; +}; + +inline ClientImpl* as_impl(celeborn_ffi_handle* h) { + return reinterpret_cast<ClientImpl*>(h); +} + +inline PartitionReaderImpl* as_reader_impl(celeborn_ffi_partition_reader* r) { + return reinterpret_cast<PartitionReaderImpl*>(r); +} + +char* dup_message(const std::string& msg) { + char* out = static_cast<char*>(std::malloc(msg.size() + 1)); + if (!out) { + return nullptr; + } + std::memcpy(out, msg.data(), msg.size()); + out[msg.size()] = '\0'; + return out; +} + +void set_error(char** err_out, const std::string& msg) { + if (err_out) { + *err_out = dup_message("celeborn-ffi: " + msg); + } +} + +template <typename Fn> +celeborn_ffi_status guarded(char** err_out, Fn&& fn) { + try { + fn(); + return CELEBORN_FFI_OK; + } catch (const std::exception& e) { + set_error(err_out, e.what()); + return CELEBORN_FFI_ERROR; + } catch (...) { + set_error(err_out, "unknown C++ exception"); + return CELEBORN_FFI_ERROR; + } +} + +} // namespace + +extern "C" { + +void celeborn_ffi_free_error(char* err) { + std::free(err); +} + +void celeborn_ffi_free_buffer(uint8_t* data) { + delete[] data; +} + +celeborn_ffi_handle* celeborn_ffi_create_client( + const char* app_id, + size_t app_id_len, + int32_t push_buffer_max_size, + const char* codec, + size_t codec_len, + char** err_out) { + try { + auto impl = std::make_unique<ClientImpl>(); + impl->app_id.assign(app_id, app_id_len); + impl->conf = std::make_shared<celeborn::conf::CelebornConf>(); + + if (push_buffer_max_size > 0) { + impl->conf->registerProperty( + celeborn::conf::CelebornConf::kClientPushBufferMaxSize, + std::to_string(push_buffer_max_size) + "b"); + } + if (codec_len > 0) { + impl->conf->registerProperty( + celeborn::conf::CelebornConf::kShuffleCompressionCodec, + std::string(codec, codec_len)); + } + + impl->endpoint = + std::make_shared<celeborn::client::ShuffleClientEndpoint>(impl->conf); + impl->client = celeborn::client::ShuffleClientImpl::create( + impl->app_id, impl->conf, *(impl->endpoint)); + return reinterpret_cast<celeborn_ffi_handle*>(impl.release()); + } catch (const std::exception& e) { + set_error(err_out, e.what()); + return nullptr; + } catch (...) { + set_error(err_out, "unknown C++ exception"); + return nullptr; + } +} + +celeborn_ffi_status celeborn_ffi_setup_lifecycle_manager( + celeborn_ffi_handle* handle, + const char* host, + size_t host_len, + int32_t port, + char** err_out) { + return guarded(err_out, [&] { + auto* impl = as_impl(handle); + impl->lifecycle_manager_host.assign(host, host_len); + impl->client->setupLifecycleManagerRef(impl->lifecycle_manager_host, port); + }); Review Comment: `celeborn_ffi_setup_lifecycle_manager` dereferences `handle` and `host` without validation. A null handle/host from a foreign-language caller will crash the process instead of returning `CELEBORN_FFI_ERROR` with a message. ########## cpp/celeborn/ffi/CelebornFfi.cc: ########## @@ -0,0 +1,281 @@ +#include "celeborn/ffi/CelebornFfi.h" + +#include <cstdlib> +#include <cstring> +#include <exception> +#include <memory> +#include <new> +#include <string> +#include <vector> + +#include "celeborn/client/ShuffleClient.h" +#include "celeborn/client/reader/CelebornInputStream.h" +#include "celeborn/conf/CelebornConf.h" + +namespace { + +struct ClientImpl { + std::shared_ptr<celeborn::conf::CelebornConf> conf; + std::shared_ptr<celeborn::client::ShuffleClientEndpoint> endpoint; + std::shared_ptr<celeborn::client::ShuffleClientImpl> client; + std::string app_id; + std::string lifecycle_manager_host; +}; + +// One open partition stream. The stream references the owning +// ShuffleClient, so callers must close every reader before shutting down +// the client. +struct PartitionReaderImpl { + std::unique_ptr<celeborn::client::CelebornInputStream> stream; +}; + +inline ClientImpl* as_impl(celeborn_ffi_handle* h) { + return reinterpret_cast<ClientImpl*>(h); +} + +inline PartitionReaderImpl* as_reader_impl(celeborn_ffi_partition_reader* r) { + return reinterpret_cast<PartitionReaderImpl*>(r); +} + +char* dup_message(const std::string& msg) { + char* out = static_cast<char*>(std::malloc(msg.size() + 1)); + if (!out) { + return nullptr; + } + std::memcpy(out, msg.data(), msg.size()); + out[msg.size()] = '\0'; + return out; +} + +void set_error(char** err_out, const std::string& msg) { + if (err_out) { + *err_out = dup_message("celeborn-ffi: " + msg); + } +} + +template <typename Fn> +celeborn_ffi_status guarded(char** err_out, Fn&& fn) { + try { + fn(); + return CELEBORN_FFI_OK; + } catch (const std::exception& e) { + set_error(err_out, e.what()); + return CELEBORN_FFI_ERROR; + } catch (...) { + set_error(err_out, "unknown C++ exception"); + return CELEBORN_FFI_ERROR; + } +} + +} // namespace + +extern "C" { + +void celeborn_ffi_free_error(char* err) { + std::free(err); +} + +void celeborn_ffi_free_buffer(uint8_t* data) { + delete[] data; +} + +celeborn_ffi_handle* celeborn_ffi_create_client( + const char* app_id, + size_t app_id_len, + int32_t push_buffer_max_size, + const char* codec, + size_t codec_len, + char** err_out) { + try { + auto impl = std::make_unique<ClientImpl>(); + impl->app_id.assign(app_id, app_id_len); + impl->conf = std::make_shared<celeborn::conf::CelebornConf>(); + + if (push_buffer_max_size > 0) { + impl->conf->registerProperty( + celeborn::conf::CelebornConf::kClientPushBufferMaxSize, + std::to_string(push_buffer_max_size) + "b"); + } + if (codec_len > 0) { + impl->conf->registerProperty( + celeborn::conf::CelebornConf::kShuffleCompressionCodec, + std::string(codec, codec_len)); + } + + impl->endpoint = + std::make_shared<celeborn::client::ShuffleClientEndpoint>(impl->conf); + impl->client = celeborn::client::ShuffleClientImpl::create( + impl->app_id, impl->conf, *(impl->endpoint)); + return reinterpret_cast<celeborn_ffi_handle*>(impl.release()); + } catch (const std::exception& e) { + set_error(err_out, e.what()); + return nullptr; + } catch (...) { + set_error(err_out, "unknown C++ exception"); + return nullptr; + } +} + +celeborn_ffi_status celeborn_ffi_setup_lifecycle_manager( + celeborn_ffi_handle* handle, + const char* host, + size_t host_len, + int32_t port, + char** err_out) { + return guarded(err_out, [&] { + auto* impl = as_impl(handle); + impl->lifecycle_manager_host.assign(host, host_len); + impl->client->setupLifecycleManagerRef(impl->lifecycle_manager_host, port); + }); +} + +celeborn_ffi_status celeborn_ffi_shutdown( + celeborn_ffi_handle* handle, + char** err_out) { + return guarded(err_out, [&] { as_impl(handle)->client->shutdown(); }); +} + +celeborn_ffi_status celeborn_ffi_push_data( + celeborn_ffi_handle* handle, + int32_t shuffle_id, + int32_t map_id, + int32_t attempt_id, + int32_t partition_id, + const uint8_t* data, + size_t data_len, + int32_t num_mappers, + int32_t num_partitions, + char** err_out) { + return guarded(err_out, [&] { + as_impl(handle)->client->pushData( + shuffle_id, + map_id, + attempt_id, + partition_id, + data, + 0, + static_cast<int>(data_len), + num_mappers, + num_partitions); + }); +} + +celeborn_ffi_status celeborn_ffi_mapper_end( + celeborn_ffi_handle* handle, + int32_t shuffle_id, + int32_t map_id, + int32_t attempt_id, + int32_t num_mappers, + char** err_out) { + return guarded(err_out, [&] { + as_impl(handle)->client->mapperEnd( + shuffle_id, map_id, attempt_id, num_mappers); + }); +} + +celeborn_ffi_status celeborn_ffi_update_reducer_file_group( + celeborn_ffi_handle* handle, + int32_t shuffle_id, + char** err_out) { + return guarded(err_out, [&] { + as_impl(handle)->client->updateReducerFileGroup(shuffle_id); + }); +} + +celeborn_ffi_status celeborn_ffi_read_partition_full( + celeborn_ffi_handle* handle, + int32_t shuffle_id, + int32_t partition_id, + int32_t attempt_number, + int32_t start_map_index, + int32_t end_map_index, + uint8_t** data_out, + size_t* len_out, + char** err_out) { + return guarded(err_out, [&] { + auto stream = as_impl(handle)->client->readPartition( + shuffle_id, + partition_id, + attempt_number, + start_map_index, + end_map_index); + + constexpr size_t kReadBufSize = 64 * 1024; + std::vector<uint8_t> accumulated; + accumulated.reserve(kReadBufSize); + std::vector<uint8_t> buf(kReadBufSize); + + while (true) { + int n = stream->read(buf.data(), 0, buf.size()); + if (n == -1) { + break; + } + if (n <= 0) { + throw std::runtime_error( + "CelebornInputStream::read returned unexpected non-positive " + + std::to_string(n)); + } + accumulated.insert(accumulated.end(), buf.data(), buf.data() + n); + } + + auto* out = new uint8_t[accumulated.size()]; + std::memcpy(out, accumulated.data(), accumulated.size()); + *data_out = out; + *len_out = accumulated.size(); + }); +} + +celeborn_ffi_status celeborn_ffi_open_partition_reader( + celeborn_ffi_handle* handle, + int32_t shuffle_id, + int32_t partition_id, + int32_t attempt_number, + int32_t start_map_index, + int32_t end_map_index, + celeborn_ffi_partition_reader** reader_out, + char** err_out) { + return guarded(err_out, [&] { + auto reader = std::make_unique<PartitionReaderImpl>(); + reader->stream = as_impl(handle)->client->readPartition( + shuffle_id, + partition_id, + attempt_number, + start_map_index, + end_map_index); + *reader_out = + reinterpret_cast<celeborn_ffi_partition_reader*>(reader.release()); + }); Review Comment: `celeborn_ffi_open_partition_reader` writes to `*reader_out` without checking that `reader_out` (or `handle`) is non-null. A null pointer here will segfault, which is especially risky since this C ABI is meant for non-Rust consumers too. ########## rust/celeborn-client-sys/build.rs: ########## @@ -0,0 +1,174 @@ +// 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. + +use std::path::{Path, PathBuf}; +use std::process::Command; + +/// Resolve the directory containing `libceleborn_client.{so,dylib}`. +/// +/// Lookup order (first match wins): +/// 1. `CELEBORN_CPP_PREFIX` env var → use `<prefix>/lib`. Intended for CI / +/// custom installs. +/// 2. In-repo prebuilt artifact at +/// `rust/resource/lib/<target-triple>/libceleborn_client.{so,dylib}`. +/// This mirrors how `alake_rust_pangu_lightsdk` ships its native blob: +/// drop the file in, no env var needed. +/// 3. Fall back to driving `cmake` against the in-repo `cpp/` source tree +/// and installing into `$OUT_DIR/celeborn-cpp-install/lib`. +fn resolve_lib_dir() -> PathBuf { + let target_os = std::env::var("CARGO_CFG_TARGET_OS").unwrap(); + let lib_filename = match target_os.as_str() { + "macos" => "libceleborn_client.dylib", + "linux" => "libceleborn_client.so", + other => panic!("unsupported target_os: {other} (only linux/macos supported)"), + }; + + // 1. Explicit override. + if let Ok(prefix) = std::env::var("CELEBORN_CPP_PREFIX") { + let lib_dir = PathBuf::from(&prefix).join("lib"); + if !lib_dir.join(lib_filename).exists() { + panic!( + "CELEBORN_CPP_PREFIX={prefix} but {} does not exist.", + lib_dir.join(lib_filename).display() + ); + } + eprintln!("cargo:warning=Using prebuilt Celeborn dylib from {prefix}"); + return lib_dir; + } + + // 2. In-repo prebuilt: rust/resource/lib/<target>/libceleborn_client.<ext> + let manifest_dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap()); + let target = std::env::var("TARGET").unwrap(); + let resource_lib_dir = manifest_dir + .join("../resource/lib") + .join(&target); + if resource_lib_dir.join(lib_filename).exists() { + eprintln!( + "cargo:warning=Using in-repo prebuilt Celeborn dylib at {}", + resource_lib_dir.display() + ); + return resource_lib_dir + .canonicalize() + .expect("failed to canonicalize resource lib dir"); + } + + // 3. Fall back to cmake from source. + let cpp_source_dir = manifest_dir + .join("../../cpp") + .canonicalize() + .unwrap_or_else(|_| { + panic!( + "No prebuilt {lib_filename} found at {}, CELEBORN_CPP_PREFIX is unset, \ + and the in-repo cpp/ directory is not reachable from {}.", + resource_lib_dir.display(), + manifest_dir.display(), + ) + }); + + eprintln!( + "cargo:warning=No prebuilt dylib at {}; building Celeborn C++ from source at {}", + resource_lib_dir.display(), + cpp_source_dir.display() + ); + + cmake_build_cpp(&cpp_source_dir).join("lib") +} + +fn cmake_build_cpp(source_dir: &Path) -> PathBuf { + let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap()); + let build_dir = out_dir.join("celeborn-cpp-build"); + let install_dir = out_dir.join("celeborn-cpp-install"); + + std::fs::create_dir_all(&build_dir).expect("failed to create cmake build directory"); + std::fs::create_dir_all(&install_dir).expect("failed to create cmake install directory"); + + let target_os = std::env::var("CARGO_CFG_TARGET_OS").unwrap(); + + let mut configure_cmd = Command::new("cmake"); + configure_cmd + .current_dir(&build_dir) + .arg(source_dir) + .arg(format!("-DCMAKE_INSTALL_PREFIX={}", install_dir.display())) + .arg("-DCMAKE_BUILD_TYPE=Release") + .arg("-DCELEBORN_BUILD_TESTS=OFF"); + + if target_os == "macos" { + let homebrew_prefix = std::env::var("HOMEBREW_PREFIX") + .unwrap_or_else(|_| "/opt/homebrew".to_string()); + configure_cmd.arg(format!("-DCMAKE_PREFIX_PATH={homebrew_prefix}")); + configure_cmd.env( + "OPENSSL_ROOT_DIR", + format!("{homebrew_prefix}/opt/openssl@3"), + ); + } + + let configure_status = configure_cmd + .status() + .expect("failed to execute `cmake` – is cmake installed?"); + if !configure_status.success() { + panic!("cmake configure step failed (exit code: {configure_status})"); + } + + let num_jobs = std::env::var("NUM_JOBS").unwrap_or_else(|_| num_cpus().to_string()); + let build_status = Command::new("cmake") + .current_dir(&build_dir) + .args(["--build", "."]) + .args(["--config", "Release"]) + .args(["--parallel", &num_jobs]) + .status() + .expect("failed to execute cmake --build"); + if !build_status.success() { + panic!("cmake build step failed (exit code: {build_status})"); + } + + let install_status = Command::new("cmake") + .current_dir(&build_dir) + .args(["--install", "."]) + .status() + .expect("failed to execute cmake --install"); + if !install_status.success() { + panic!("cmake install step failed (exit code: {install_status})"); + } + + install_dir +} + +fn num_cpus() -> usize { + std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(4) +} + +fn main() { + let lib_dir = resolve_lib_dir(); + let lib_dir_str = lib_dir.display().to_string(); + + // Single aggregated dylib: libceleborn_client.{so,dylib} bundles every + // internal static lib (including the celeborn_ffi C ABI shim) and pulls + // third-party deps via NEEDED entries that resolve at runtime. The Rust + // crate touches no C++ headers and links nothing else. + println!("cargo:rustc-link-search=native={lib_dir_str}"); + println!("cargo:rustc-link-lib=dylib=celeborn_client"); + + // Re-export lib_dir so the downstream `celeborn-client` crate can pick + // it up via DEP_CELEBORN_CLIENT_LIB_DIR and embed an rpath in its + // examples / tests / binaries. (`cargo:rustc-link-arg` does not + // propagate from a sys crate to dependent crates' artifacts.) + println!("cargo:lib_dir={lib_dir_str}"); + + println!("cargo:rerun-if-env-changed=CELEBORN_CPP_PREFIX"); + let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap(); + println!("cargo:rerun-if-changed={manifest_dir}/../resource/lib"); Review Comment: This build script may drive an in-repo CMake build of `cpp/`, but it doesn’t declare any `cargo:rerun-if-changed` dependency on the `cpp/` sources. That can lead to confusing incremental builds where edits under `cpp/` don’t trigger a rebuild of the dylib. ########## cpp/celeborn/ffi/CelebornFfi.cc: ########## @@ -0,0 +1,281 @@ +#include "celeborn/ffi/CelebornFfi.h" + +#include <cstdlib> +#include <cstring> +#include <exception> +#include <memory> +#include <new> +#include <string> +#include <vector> + +#include "celeborn/client/ShuffleClient.h" +#include "celeborn/client/reader/CelebornInputStream.h" +#include "celeborn/conf/CelebornConf.h" + +namespace { + +struct ClientImpl { + std::shared_ptr<celeborn::conf::CelebornConf> conf; + std::shared_ptr<celeborn::client::ShuffleClientEndpoint> endpoint; + std::shared_ptr<celeborn::client::ShuffleClientImpl> client; + std::string app_id; + std::string lifecycle_manager_host; +}; + +// One open partition stream. The stream references the owning +// ShuffleClient, so callers must close every reader before shutting down +// the client. +struct PartitionReaderImpl { + std::unique_ptr<celeborn::client::CelebornInputStream> stream; +}; + +inline ClientImpl* as_impl(celeborn_ffi_handle* h) { + return reinterpret_cast<ClientImpl*>(h); +} + +inline PartitionReaderImpl* as_reader_impl(celeborn_ffi_partition_reader* r) { + return reinterpret_cast<PartitionReaderImpl*>(r); +} + +char* dup_message(const std::string& msg) { + char* out = static_cast<char*>(std::malloc(msg.size() + 1)); + if (!out) { + return nullptr; + } + std::memcpy(out, msg.data(), msg.size()); + out[msg.size()] = '\0'; + return out; +} + +void set_error(char** err_out, const std::string& msg) { + if (err_out) { + *err_out = dup_message("celeborn-ffi: " + msg); + } +} + +template <typename Fn> +celeborn_ffi_status guarded(char** err_out, Fn&& fn) { + try { + fn(); + return CELEBORN_FFI_OK; + } catch (const std::exception& e) { + set_error(err_out, e.what()); + return CELEBORN_FFI_ERROR; + } catch (...) { + set_error(err_out, "unknown C++ exception"); + return CELEBORN_FFI_ERROR; + } +} + +} // namespace + +extern "C" { + +void celeborn_ffi_free_error(char* err) { + std::free(err); +} + +void celeborn_ffi_free_buffer(uint8_t* data) { + delete[] data; +} + +celeborn_ffi_handle* celeborn_ffi_create_client( + const char* app_id, + size_t app_id_len, + int32_t push_buffer_max_size, + const char* codec, + size_t codec_len, + char** err_out) { + try { + auto impl = std::make_unique<ClientImpl>(); + impl->app_id.assign(app_id, app_id_len); + impl->conf = std::make_shared<celeborn::conf::CelebornConf>(); + + if (push_buffer_max_size > 0) { + impl->conf->registerProperty( + celeborn::conf::CelebornConf::kClientPushBufferMaxSize, + std::to_string(push_buffer_max_size) + "b"); + } + if (codec_len > 0) { + impl->conf->registerProperty( + celeborn::conf::CelebornConf::kShuffleCompressionCodec, + std::string(codec, codec_len)); + } + + impl->endpoint = + std::make_shared<celeborn::client::ShuffleClientEndpoint>(impl->conf); + impl->client = celeborn::client::ShuffleClientImpl::create( + impl->app_id, impl->conf, *(impl->endpoint)); + return reinterpret_cast<celeborn_ffi_handle*>(impl.release()); + } catch (const std::exception& e) { + set_error(err_out, e.what()); + return nullptr; + } catch (...) { + set_error(err_out, "unknown C++ exception"); + return nullptr; + } +} + +celeborn_ffi_status celeborn_ffi_setup_lifecycle_manager( + celeborn_ffi_handle* handle, + const char* host, + size_t host_len, + int32_t port, + char** err_out) { + return guarded(err_out, [&] { + auto* impl = as_impl(handle); + impl->lifecycle_manager_host.assign(host, host_len); + impl->client->setupLifecycleManagerRef(impl->lifecycle_manager_host, port); + }); +} + +celeborn_ffi_status celeborn_ffi_shutdown( + celeborn_ffi_handle* handle, + char** err_out) { + return guarded(err_out, [&] { as_impl(handle)->client->shutdown(); }); +} + +celeborn_ffi_status celeborn_ffi_push_data( + celeborn_ffi_handle* handle, + int32_t shuffle_id, + int32_t map_id, + int32_t attempt_id, + int32_t partition_id, + const uint8_t* data, + size_t data_len, + int32_t num_mappers, + int32_t num_partitions, + char** err_out) { + return guarded(err_out, [&] { + as_impl(handle)->client->pushData( + shuffle_id, + map_id, + attempt_id, + partition_id, + data, + 0, + static_cast<int>(data_len), + num_mappers, + num_partitions); + }); +} + +celeborn_ffi_status celeborn_ffi_mapper_end( + celeborn_ffi_handle* handle, + int32_t shuffle_id, + int32_t map_id, + int32_t attempt_id, + int32_t num_mappers, + char** err_out) { + return guarded(err_out, [&] { + as_impl(handle)->client->mapperEnd( + shuffle_id, map_id, attempt_id, num_mappers); + }); +} + +celeborn_ffi_status celeborn_ffi_update_reducer_file_group( + celeborn_ffi_handle* handle, + int32_t shuffle_id, + char** err_out) { + return guarded(err_out, [&] { + as_impl(handle)->client->updateReducerFileGroup(shuffle_id); + }); +} + +celeborn_ffi_status celeborn_ffi_read_partition_full( + celeborn_ffi_handle* handle, + int32_t shuffle_id, + int32_t partition_id, + int32_t attempt_number, + int32_t start_map_index, + int32_t end_map_index, + uint8_t** data_out, + size_t* len_out, + char** err_out) { + return guarded(err_out, [&] { + auto stream = as_impl(handle)->client->readPartition( + shuffle_id, + partition_id, + attempt_number, + start_map_index, + end_map_index); Review Comment: Several FFI entrypoints blindly dereference `handle` and output pointers inside the guarded lambda. Passing a null `handle`/`data_out`/`len_out` should fail gracefully with `CELEBORN_FFI_ERROR` and an error string rather than crashing the process. ########## cpp/celeborn/ffi/CelebornFfi.cc: ########## @@ -0,0 +1,281 @@ +#include "celeborn/ffi/CelebornFfi.h" + +#include <cstdlib> +#include <cstring> +#include <exception> +#include <memory> +#include <new> +#include <string> +#include <vector> + +#include "celeborn/client/ShuffleClient.h" +#include "celeborn/client/reader/CelebornInputStream.h" +#include "celeborn/conf/CelebornConf.h" + +namespace { + +struct ClientImpl { + std::shared_ptr<celeborn::conf::CelebornConf> conf; + std::shared_ptr<celeborn::client::ShuffleClientEndpoint> endpoint; + std::shared_ptr<celeborn::client::ShuffleClientImpl> client; + std::string app_id; + std::string lifecycle_manager_host; +}; + +// One open partition stream. The stream references the owning +// ShuffleClient, so callers must close every reader before shutting down +// the client. +struct PartitionReaderImpl { + std::unique_ptr<celeborn::client::CelebornInputStream> stream; +}; + +inline ClientImpl* as_impl(celeborn_ffi_handle* h) { + return reinterpret_cast<ClientImpl*>(h); +} + +inline PartitionReaderImpl* as_reader_impl(celeborn_ffi_partition_reader* r) { + return reinterpret_cast<PartitionReaderImpl*>(r); +} + +char* dup_message(const std::string& msg) { + char* out = static_cast<char*>(std::malloc(msg.size() + 1)); + if (!out) { + return nullptr; + } + std::memcpy(out, msg.data(), msg.size()); + out[msg.size()] = '\0'; + return out; +} + +void set_error(char** err_out, const std::string& msg) { + if (err_out) { + *err_out = dup_message("celeborn-ffi: " + msg); + } +} + +template <typename Fn> +celeborn_ffi_status guarded(char** err_out, Fn&& fn) { + try { + fn(); + return CELEBORN_FFI_OK; + } catch (const std::exception& e) { + set_error(err_out, e.what()); + return CELEBORN_FFI_ERROR; + } catch (...) { + set_error(err_out, "unknown C++ exception"); + return CELEBORN_FFI_ERROR; + } +} + +} // namespace + +extern "C" { + +void celeborn_ffi_free_error(char* err) { + std::free(err); +} + +void celeborn_ffi_free_buffer(uint8_t* data) { + delete[] data; +} + +celeborn_ffi_handle* celeborn_ffi_create_client( + const char* app_id, + size_t app_id_len, + int32_t push_buffer_max_size, + const char* codec, + size_t codec_len, + char** err_out) { + try { + auto impl = std::make_unique<ClientImpl>(); + impl->app_id.assign(app_id, app_id_len); + impl->conf = std::make_shared<celeborn::conf::CelebornConf>(); + + if (push_buffer_max_size > 0) { + impl->conf->registerProperty( + celeborn::conf::CelebornConf::kClientPushBufferMaxSize, + std::to_string(push_buffer_max_size) + "b"); + } + if (codec_len > 0) { + impl->conf->registerProperty( + celeborn::conf::CelebornConf::kShuffleCompressionCodec, + std::string(codec, codec_len)); + } + + impl->endpoint = + std::make_shared<celeborn::client::ShuffleClientEndpoint>(impl->conf); + impl->client = celeborn::client::ShuffleClientImpl::create( + impl->app_id, impl->conf, *(impl->endpoint)); + return reinterpret_cast<celeborn_ffi_handle*>(impl.release()); + } catch (const std::exception& e) { + set_error(err_out, e.what()); + return nullptr; + } catch (...) { + set_error(err_out, "unknown C++ exception"); + return nullptr; + } +} + +celeborn_ffi_status celeborn_ffi_setup_lifecycle_manager( + celeborn_ffi_handle* handle, + const char* host, + size_t host_len, + int32_t port, + char** err_out) { + return guarded(err_out, [&] { + auto* impl = as_impl(handle); + impl->lifecycle_manager_host.assign(host, host_len); + impl->client->setupLifecycleManagerRef(impl->lifecycle_manager_host, port); + }); +} + +celeborn_ffi_status celeborn_ffi_shutdown( + celeborn_ffi_handle* handle, + char** err_out) { + return guarded(err_out, [&] { as_impl(handle)->client->shutdown(); }); +} + +celeborn_ffi_status celeborn_ffi_push_data( + celeborn_ffi_handle* handle, + int32_t shuffle_id, + int32_t map_id, + int32_t attempt_id, + int32_t partition_id, + const uint8_t* data, + size_t data_len, + int32_t num_mappers, + int32_t num_partitions, + char** err_out) { + return guarded(err_out, [&] { + as_impl(handle)->client->pushData( + shuffle_id, + map_id, + attempt_id, + partition_id, + data, + 0, + static_cast<int>(data_len), + num_mappers, + num_partitions); + }); +} + +celeborn_ffi_status celeborn_ffi_mapper_end( + celeborn_ffi_handle* handle, + int32_t shuffle_id, + int32_t map_id, + int32_t attempt_id, + int32_t num_mappers, + char** err_out) { + return guarded(err_out, [&] { + as_impl(handle)->client->mapperEnd( + shuffle_id, map_id, attempt_id, num_mappers); + }); +} + +celeborn_ffi_status celeborn_ffi_update_reducer_file_group( + celeborn_ffi_handle* handle, + int32_t shuffle_id, + char** err_out) { + return guarded(err_out, [&] { + as_impl(handle)->client->updateReducerFileGroup(shuffle_id); + }); +} + +celeborn_ffi_status celeborn_ffi_read_partition_full( + celeborn_ffi_handle* handle, + int32_t shuffle_id, + int32_t partition_id, + int32_t attempt_number, + int32_t start_map_index, + int32_t end_map_index, + uint8_t** data_out, + size_t* len_out, + char** err_out) { + return guarded(err_out, [&] { + auto stream = as_impl(handle)->client->readPartition( + shuffle_id, + partition_id, + attempt_number, + start_map_index, + end_map_index); + + constexpr size_t kReadBufSize = 64 * 1024; + std::vector<uint8_t> accumulated; + accumulated.reserve(kReadBufSize); + std::vector<uint8_t> buf(kReadBufSize); + + while (true) { + int n = stream->read(buf.data(), 0, buf.size()); + if (n == -1) { + break; + } + if (n <= 0) { + throw std::runtime_error( + "CelebornInputStream::read returned unexpected non-positive " + + std::to_string(n)); + } + accumulated.insert(accumulated.end(), buf.data(), buf.data() + n); + } + + auto* out = new uint8_t[accumulated.size()]; + std::memcpy(out, accumulated.data(), accumulated.size()); + *data_out = out; + *len_out = accumulated.size(); + }); +} + +celeborn_ffi_status celeborn_ffi_open_partition_reader( + celeborn_ffi_handle* handle, + int32_t shuffle_id, + int32_t partition_id, + int32_t attempt_number, + int32_t start_map_index, + int32_t end_map_index, + celeborn_ffi_partition_reader** reader_out, + char** err_out) { + return guarded(err_out, [&] { + auto reader = std::make_unique<PartitionReaderImpl>(); + reader->stream = as_impl(handle)->client->readPartition( + shuffle_id, + partition_id, + attempt_number, + start_map_index, + end_map_index); + *reader_out = + reinterpret_cast<celeborn_ffi_partition_reader*>(reader.release()); + }); +} + +celeborn_ffi_status celeborn_ffi_read_partition_chunk( + celeborn_ffi_partition_reader* reader, + uint8_t* buf, + size_t buf_len, + size_t* bytes_read, + char** err_out) { + return guarded(err_out, [&] { + if (buf_len == 0) { + *bytes_read = 0; + return; + } + int n = as_reader_impl(reader)->stream->read(buf, 0, buf_len); + if (n == -1) { + // EOF — surface as 0 bytes read to match std::io::Read semantics. + *bytes_read = 0; + return; Review Comment: `celeborn_ffi_read_partition_chunk` writes to `*bytes_read` even when `bytes_read` is null and dereferences `reader` without validation. This is a hard crash risk for any FFI consumer passing invalid pointers. ########## rust/celeborn-client-sys/build.rs: ########## @@ -0,0 +1,174 @@ +// 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. + +use std::path::{Path, PathBuf}; +use std::process::Command; + +/// Resolve the directory containing `libceleborn_client.{so,dylib}`. +/// +/// Lookup order (first match wins): +/// 1. `CELEBORN_CPP_PREFIX` env var → use `<prefix>/lib`. Intended for CI / +/// custom installs. +/// 2. In-repo prebuilt artifact at +/// `rust/resource/lib/<target-triple>/libceleborn_client.{so,dylib}`. +/// This mirrors how `alake_rust_pangu_lightsdk` ships its native blob: +/// drop the file in, no env var needed. +/// 3. Fall back to driving `cmake` against the in-repo `cpp/` source tree +/// and installing into `$OUT_DIR/celeborn-cpp-install/lib`. Review Comment: The doc comment references `alake_rust_pangu_lightsdk`, which appears to be an unrelated/internal project name and is confusing in an Apache Celeborn codebase. It would be better to describe the behavior generically without naming external proprietary projects. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
