mr_framework.cpp

#include "mapreduce.h"
#include "mr_comm_helpers.h"
#include "mr_data.pb.h"
#include "mr_data.grpc.pb.h"

#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <thread>
#include <memory>
#include <algorithm> // For std::find_if
#include <grpcpp/grpcpp.h>
#include <grpcpp/health_check_service_interface.h>
#include <grpc/grpc.h> // For GRPC_ARG_IPV6_V6ONLY
#include <hdfs.h>
#include <sys/stat.h> // For stat()
#include <arpa/inet.h> // For htonl, ntohl
#include <unicode/unistr.h> // For icu::UnicodeString
#include <unicode/normalizer2.h> // For icu::Normalizer2

// Thread function forward declarations 
void* mapper_thread_func(void* arg);
void* sender_thread_func(void* arg);

// Thread argument structure 
typedef struct {
    int id;
    struct map_reduce* mr;
} thread_arg_t;

// Thread-local variable 
thread_local int current_mapper_id = -1;

const char* HDFS_INPUT_PATH = "/mr_input.txt";

// For non-utf-8 laniguages, normalize strings to NFC form
std::string normalize_utf8(const std::string& input) {
    UErrorCode status = U_ZERO_ERROR;
    const icu::Normalizer2* normalizer = icu::Normalizer2::getNFCInstance(status);
    if (U_FAILURE(status)) return input;

    icu::UnicodeString u_in = icu::UnicodeString::fromUTF8(input);
    icu::UnicodeString u_out;
    normalizer->normalize(u_in, u_out, status);
    if (U_FAILURE(status)) return input;

    std::string out;
    u_out.toUTF8String(out);
    return out;
}

// HDFS Helper Functions 
// Reads a chunk of data from HDFS starting at 'offset' with size 'chunk_size'.
std::string read_hdfs_chunk(hdfsFS fs, const char* path, tOffset offset, tSize chunk_size) {
    if (!fs) return "";
    hdfsFile file = hdfsOpenFile(fs, path, O_RDONLY, 0, 0, 0);
    if (!file) {
        std::cerr << "Failed to open HDFS file: " << path << std::endl;
        return "";
    }
    std::vector<char> buffer(chunk_size);
    tSize bytes_read = hdfsPread(fs, file, offset, buffer.data(), chunk_size);
    std::string result;
    if (bytes_read > 0) {
        result.assign(buffer.data(), bytes_read);
    } else if (bytes_read < 0) {
        std::cerr << "Failed to read from HDFS file: " << path << std::endl;
    }
    hdfsCloseFile(fs, file);
    return result;
}

// Uploads a local file to HDFS at the specified path.
bool upload_to_hdfs(hdfsFS fs, const char* local_path, const char* hdfs_path) {
    if (hdfsExists(fs, hdfs_path) == 0) {
        if (hdfsDelete(fs, hdfs_path, 1) != 0) {
            std::cerr << "Failed to delete existing HDFS file: " << hdfs_path << std::endl;
            return false;
        }
    }
    FILE* local_file = fopen(local_path, "rb");
    if (!local_file) {
        std::cerr << "Failed to open local file for reading: " << local_path << std::endl;
        return false;
    }
    hdfsFile writeFile = hdfsOpenFile(fs, hdfs_path, O_WRONLY | O_CREAT, 0, 0, 0);
    if (!writeFile) {
        std::cerr << "Failed to open HDFS file for writing: " << hdfs_path << std::endl;
        fclose(local_file);
        return false;
    }
    const int buffer_size = 8192;
    char buffer[buffer_size];
    size_t bytes_read = 0;
    while ((bytes_read = fread(buffer, 1, buffer_size, local_file)) > 0) {
        if (hdfsWrite(fs, writeFile, buffer, bytes_read) != bytes_read) {
            std::cerr << "Failed to write a chunk to HDFS." << std::endl;
            fclose(local_file);
            hdfsCloseFile(fs, writeFile);
            return false;
        }
    }
    fclose(local_file);
    if (hdfsFlush(fs, writeFile) != 0) std::cerr << "Failed to flush HDFS file." << std::endl;
    hdfsCloseFile(fs, writeFile);
    return true;
}

// Trims leading and trailing whitespace from a string.
std::string trim(const std::string& str) {
    size_t first = str.find_first_not_of(" \t\n\r\f\v");
    if (std::string::npos == first) return "";
    size_t last = str.find_last_not_of(" \t\n\r\f\v");
    return str.substr(first, (last - first + 1));
}

// Master's RPC Service (to receive "Job Complete" signal) 
class MasterServiceImpl final : public mrdata::MasterService::Service {
public:
    struct map_reduce* mr_master_glob = nullptr; 
    MasterServiceImpl(struct map_reduce* mr) : mr_master_glob(mr) {}
    // SignalJobComplete RPC Handler 
    grpc::Status SignalJobComplete(grpc::ServerContext* context, const mrdata::JobCompleteRequest* request,
                                   mrdata::TransferResponse* response) override {
        std::cout << "Master: Received SignalJobComplete from Reducer." << std::endl;
        pthread_mutex_lock(&mr_master_glob->job_lock);
        mr_master_glob->job_running = false; 
        pthread_cond_signal(&mr_master_glob->job_complete_cond);
        pthread_mutex_unlock(&mr_master_glob->job_lock);
        response->set_success(true);
        return grpc::Status::OK;
    }
};

// Master's background RPC server thread 
void* RunMasterServer(void* arg) {
    struct map_reduce* mr = (struct map_reduce*)arg;
    std::string server_address("0.0.0.0:50053");
    MasterServiceImpl service(mr);
    grpc::EnableDefaultHealthCheckService(true);
    grpc::ServerBuilder builder;
    builder.AddChannelArgument(std::string("grpc.ipv6_v6only"), 0); // IPv4/IPv6 fix
    builder.AddListeningPort(server_address, grpc::InsecureServerCredentials());
    builder.RegisterService(&service);
    std::unique_ptr<grpc::Server> server(builder.BuildAndStart());
    std::cout << "Master RPC Server listening on " << server_address << std::endl;
    
    // Wait for the job to be marked as complete (by mr_finish)
    while(mr->job_running) {
        sleep(1);
    }
    server->Shutdown();
    std::cout << "Master RPC Server shutting down." << std::endl;
    return NULL;
}

// C API for MapReduce Framework 
// These functions are linked into all executableqs using the framework.
struct map_reduce *mr_init(const char *application, int threads, const char *inpath, const char *outpath, const char *helper_args) {
    struct map_reduce *mr = (struct map_reduce *)calloc(1, sizeof(struct map_reduce));
    if (!mr) return NULL;

    mr->threads_per_unit = threads;
    mr->input_path = strdup(inpath);
    mr->output_path = strdup(outpath);
    mr->helper_args = strdup(helper_args);
    mr->app_name = strdup(application);
    mr->job_running = true;
    mr->map_fn = map;
    mr->reduce_fn = reduce;

    // Read Cluster Configuration 
    std::cout << "Reading cluster configuration from cluster.conf..." << std::endl;
    std::ifstream config_file("cluster.conf");
    if (!config_file.is_open()) {
        std::cerr << "FATAL: Could not open cluster.conf." << std::endl;
        free(mr);
        return NULL;
    }
    mr->mapper_buffer_capacity = 10240;
    int config_threads = threads;
    const std::string REDUCER_KEY = "reducer_ip=";
    const std::string MAPPER_KEY = "mapper_ip=";
    const std::string BUFFER_KEY = "mapper_buffer_size=";
    const std::string THREADS_KEY = "threads_per_unit=";
    std::string line;
    while (std::getline(config_file, line)) {
        line = trim(line);
        if (line.empty() || line[0] == '#') continue;
        if (line.rfind(REDUCER_KEY, 0) == 0) {
            mr->reducer_ip = trim(line.substr(REDUCER_KEY.length()));
        } else if (line.rfind(MAPPER_KEY, 0) == 0) {
            mr->mapper_ips.push_back(trim(line.substr(MAPPER_KEY.length())));
        } else if (line.rfind(BUFFER_KEY, 0) == 0) {
             mr->mapper_buffer_capacity = std::stoul(trim(line.substr(BUFFER_KEY.length())));
        } else if (line.rfind(THREADS_KEY, 0) == 0) {
             config_threads = std::stoi(trim(line.substr(THREADS_KEY.length())));
        }
    }
    config_file.close();
    mr->threads_per_unit = threads;
    if (mr->reducer_ip.empty() || mr->mapper_ips.empty()) {
        std::cerr << "FATAL: 'reducer_ip' or 'mapper_ip' not found in cluster.conf." << std::endl;
        mr_destroy(mr);
        return NULL;
    }
    mr->num_mapper_units = mr->mapper_ips.size();
    mr->total_mapper_threads = mr->num_mapper_units * mr->threads_per_unit;
    
    std::cout << "Configured threads per unit: " << mr->threads_per_unit << std::endl;
    std::cout << "Reducer IP: " << mr->reducer_ip << std::endl;
    std::cout << "Mapper IPs: ";
    for(const auto& ip : mr->mapper_ips) { std::cout << ip << " "; }
    std::cout << std::endl;
    std::cout << "Total mapper threads: " << mr->total_mapper_threads << std::endl;
    std::cout << "Mapper buffer size: " << mr->mapper_buffer_capacity << " bytes." << std::endl;
    
    pthread_mutex_init(&mr->job_lock, NULL);
    pthread_cond_init(&mr->job_complete_cond, NULL);

    // HDFS Initialization
    mr->hdfs_connection = hdfsConnect("default", 0);
    if (!mr->hdfs_connection) {
        std::cerr << "Failed to connect to HDFS." << std::endl;
        mr_destroy(mr);
        return NULL;
    }
    struct stat file_stat;
    if (stat(mr->input_path, &file_stat) != 0) {
        std::cerr << "Input file does not exist: " << mr->input_path << std::endl;
        mr_destroy(mr);
        return NULL;
    }
    mr->input_file_size = file_stat.st_size;

    if (!upload_to_hdfs(mr->hdfs_connection, mr->input_path, HDFS_INPUT_PATH)) {
        std::cerr << "Failed to upload input file to HDFS." << std::endl;
        mr_destroy(mr);
        return NULL;
    }
    std::cout << "Successfully uploaded " << mr->input_path << " to " << HDFS_INPUT_PATH << std::endl;
    
    // Start the Master's background RPC server
    pthread_t master_server_tid;
    if (pthread_create(&master_server_tid, NULL, RunMasterServer, mr) != 0) {
        std::cerr << "FATAL: Failed to create Master RPC server thread!" << std::endl;
        mr_destroy(mr);
        return NULL;
    }
    pthread_detach(master_server_tid);
    return mr;
}

// Start the MapReduce job by instructing all Mapper workers to begin 
int mr_start(struct map_reduce *mr, const int barrierEnable) {
    mr->barrier_enabled = (barrierEnable != 0);
    int mapper_id_counter = 0; 

    // Send StartMappers RPC to each mapper worker
    for (const auto& worker_ip : mr->mapper_ips) {
        std::string worker_address = worker_ip + ":50052";
        auto channel = grpc::CreateChannel("ipv4:" + worker_address, grpc::InsecureChannelCredentials());
        auto stub = mrdata::WorkerService::NewStub(channel);

        // Prepare StartMappers Request 
        mrdata::StartMappersRequest request;
        request.set_app_name(mr->app_name);
        request.set_num_threads(mr->threads_per_unit);
        request.set_reducer_address(mr->reducer_ip);
        request.set_hdfs_input_path(HDFS_INPUT_PATH);
        request.set_helper_args(mr->helper_args);
        request.set_total_mapper_threads(mr->total_mapper_threads);
        request.set_total_file_size(mr->input_file_size);
        request.set_buffer_capacity(mr->mapper_buffer_capacity);
        request.set_mapper_id_start(mapper_id_counter);

        mapper_id_counter += mr->threads_per_unit;
        // Send RPC 
        mrdata::TransferResponse response;
        grpc::ClientContext context;
        context.set_deadline(std::chrono::system_clock::now() + std::chrono::seconds(30));

        std::cout << "Master: Sending StartMappers RPC to " << worker_address << std::endl;
        grpc::Status status = stub->StartMappers(&context, request, &response);

        if (!status.ok()) {
            std::cerr << "Master: Failed to start mappers on worker " << worker_ip << ": " << status.error_message() << std::endl;
            return -1;
        }
        std::cout << "Master: Successfully started mappers on " << worker_ip << std::endl;
    }
    return 0;
}

// Wait for the MapReduce job to complete 
int mr_finish(struct map_reduce *mr) {
    std::cout << "Master: mr_finish() called. Waiting for job completion signal from Reducer..." << std::endl;
    
    pthread_mutex_lock(&mr->job_lock);
    // Use a while loop to protect against spurious wakeups
    while (mr->job_running) {
        pthread_cond_wait(&mr->job_complete_cond, &mr->job_lock);
    }
    pthread_mutex_unlock(&mr->job_lock);
    
    // mr->job_running is now false (set by the MasterService handler)
    std::cout << "Master: Job complete. mr_finish() returning." << std::endl;
    return 0;
}

// Clean up and free MapReduce structure 
void mr_destroy(struct map_reduce *mr) {
    if (!mr) return;

    std::cout << "Master: Telling all Mapper workers to shut down..." << std::endl;
    // Send shutdown RPCs to all mapper workers
    for (const auto& worker_ip : mr->mapper_ips) {
        std::string worker_address = worker_ip + ":50052";
        auto channel = grpc::CreateChannel("ipv4:" + worker_address, grpc::InsecureChannelCredentials());
        auto stub = mrdata::WorkerService::NewStub(channel);

        mrdata::ShutdownRequest request;
        request.set_reason("Job complete");
        mrdata::TransferResponse response;
        grpc::ClientContext context;
        context.set_deadline(std::chrono::system_clock::now() + std::chrono::seconds(5));
        stub->ShutdownWorker(&context, request, &response);
    }
    std::cout << "Master: Shutdown signals sent." << std::endl;

    // Cleanup
    pthread_mutex_destroy(&mr->job_lock);
    pthread_cond_destroy(&mr->job_complete_cond);
    if (mr->hdfs_connection) hdfsDisconnect(mr->hdfs_connection);
    free((void*)mr->input_path);
    free((void*)mr->output_path);
    free((void*)mr->helper_args);
    free((void*)mr->app_name);
    free(mr);
}

// Produce a key-value pair from Mapper to Reducer 
int mr_produce(struct map_reduce *mr, const struct kvpair *kv) {
    int current_mapper_id = ::current_mapper_id;
    if (current_mapper_id < 0 || current_mapper_id >= mr->total_mapper_threads) {
        std::cerr << "Invalid mapper ID in mr_produce: " << current_mapper_id << std::endl;
        return -1;
    }
    mapper_unit_t* unit = &mr->mapper_units[current_mapper_id];

    // Serialize the kvpair using Protocol Buffers
    mrdata::IntermediateKvPair kv_msg;
    kv_msg.set_key_data(kv->key, kv->keysz);
    kv_msg.set_value_data(kv->value, kv->valuesz); // Transports the 4-byte integer
    std::string serialized_kv;
    kv_msg.SerializeToString(&serialized_kv);

    // Prepare the message with length prefix
    uint32_t msg_size = static_cast<uint32_t>(serialized_kv.length());
    uint32_t net_msg_size = htonl(msg_size); // Convert to network byte order
    size_t required_size = sizeof(net_msg_size) + msg_size;

    pthread_mutex_lock(&unit->lock);
    
    // Wait until there is enough space in the buffer
    while (unit->capacity - unit->current_size < required_size) {
        unit->mapper_waiting_for_space = true;
        pthread_cond_signal(&unit->threshold_reached_cond);
        pthread_cond_wait(&unit->not_full_cond, &unit->lock);
        unit->mapper_waiting_for_space = false;
    }

    // Append the length-prefixed message to the buffer
    memcpy(unit->buffer_data + unit->current_size, &net_msg_size, sizeof(net_msg_size));
    unit->current_size += sizeof(net_msg_size);
    memcpy(unit->buffer_data + unit->current_size, serialized_kv.c_str(), msg_size);
    unit->current_size += msg_size;

    // If buffer exceeds threshold, signal the sender thread
    if (unit->current_size / (double)unit->capacity > MR_BUFFER_THRESHOLD) {
        pthread_cond_signal(&unit->threshold_reached_cond);
    }
    pthread_mutex_unlock(&unit->lock);
    return 1;
}

int mr_consume(struct map_reduce *mr, struct kvpair *kvset, unsigned long *num, const int barrierEnable) {
    // Stub for linking
    return 0;
}

// Output the final result from Reducer to HDFS 
int mr_output(struct map_reduce *mr, char *writeBuffer, unsigned long bufferLength) {
    const char* tmp_file = "/tmp/mr_final_output.txt";
    std::cout << "mr_output: Writing " << bufferLength << " bytes to " << tmp_file << std::endl;
    std::ofstream out(tmp_file, std::ios::binary);
    if (!out.is_open()) {
        std::cerr << "mr_output ERROR: Failed to open temp file " << tmp_file << std::endl;
        return -1;
    }
    out.write(writeBuffer, bufferLength);
    out.close();
    std::cout << "mr_output: Successfully wrote to temporary file." << std::endl;
    
    // Use libhdfs C API to upload the file
    if (!upload_to_hdfs(mr->hdfs_connection, tmp_file, mr->output_path)) {
        std::cerr << "mr_output failed! HDFS upload failed." << std::endl;
        return -1;
    }

    std::cout << "mr_output: Successfully uploaded output to HDFS at " << mr->output_path << std::endl;
    remove(tmp_file);
    return 0;
}

// THREAD FUNCTIONS (Called by Worker Servers)
void* mapper_thread_func(void* arg) {
    // Get thread arguments
    thread_arg_t* t_arg = (thread_arg_t*)arg;
    struct map_reduce* mr = t_arg->mr;

    const int mapper_id = t_arg->id; // Unique global mapper ID
    current_mapper_id = mapper_id; // Set thread-local mapper ID

    std::cout << "Mapper " << mapper_id << ": Starting." << std::endl;

    // Compute nominal chunk bounds
    const long long total_file_size = mr->input_file_size;
    const long long base_chunk      = (mr->total_mapper_threads > 0)
                                      ? (total_file_size / mr->total_mapper_threads)
                                      : total_file_size;
    const long long start_offset    = mapper_id * base_chunk;
    const bool      is_last_mapper  = (mapper_id == mr->total_mapper_threads - 1);
    const long long nominal_end     = is_last_mapper ? total_file_size : (start_offset + base_chunk);

    // 128 KiB minimum or 1/8th of the base chunk, capped at 1 MiB
    const long long SAFETY_MARGIN = std::min(1024LL*1024LL, std::max(128LL*1024LL, base_chunk / 8));

    // Read up to nominal_end + SAFETY_MARGIN (capped at EOF), not simply base_chunk+margin
    const long long read_limit_end = std::min(nominal_end + SAFETY_MARGIN, total_file_size);
    const long long max_read       = (read_limit_end > start_offset) ? (read_limit_end - start_offset) : 0;

    // Read chunk from HDFS 
    std::string buf;
    if (max_read > 0) {
        buf = read_hdfs_chunk(mr->hdfs_connection, HDFS_INPUT_PATH, start_offset, (tSize)max_read);
    }
    if (buf.empty() && max_read > 0) {
        std::cerr << "Mapper " << mapper_id << ": Failed to read input chunk." << std::endl;
    }

    // Adjust START (drop leading partial line for non-first mapper) 
    size_t begin = 0;
    if (mapper_id > 0) {
        size_t first_nl = buf.find('\n'); // find first newline
        if (first_nl != std::string::npos) {
            begin = first_nl + 1; // start after the newline
        } else {
            begin = 0; // no newline found, process entire buffer
        }
    }

    // Adjust END (stop at the first newline at/after nominal_end) 
    size_t cutoff_in_buf = 0;
    if (nominal_end > start_offset) {
        long long diff = nominal_end - start_offset;
        cutoff_in_buf = (size_t)std::min<long long>(diff, (long long)buf.size());
    }

    // Find the end position
    size_t end_pos;
    // Not the last mapper: stop at first '\n' after cutoff_in_buf
    if (mapper_id < mr->total_mapper_threads - 1) {
        size_t first_nl_after = buf.find('\n', cutoff_in_buf);
        if (first_nl_after != std::string::npos) {
            end_pos = first_nl_after;    // stop BEFORE '\n'
        } else {
            end_pos = buf.size();        // include overhang to avoid dropping a long line
        }
    } else {
        end_pos = buf.size();            // last mapper: to EOF
    }

    // Build final input_chunk once 
    std::string input_chunk;
    // Check for valid range
    if (begin < end_pos && end_pos <= buf.size()) {
        input_chunk.assign(buf.data() + begin, end_pos - begin);
    } else {
        input_chunk.clear();
    }

    // Trim trailing '\r' if present
    if (!input_chunk.empty() && input_chunk.back() == '\r') {
        input_chunk.pop_back();
    }

    // Tokenize and emit 
    uint32_t lines_processed = 0;
    std::vector<kvpair> pairs;
    // Only tokenize if there's data to process
    if (!input_chunk.empty()) {
        uint32_t startLineNum = 1;
        // For "grep", compute actual starting line number
        if (strcmp(mr->app_name, "grep") == 0) {
            uint64_t newline_count = 0;
            const long long CHUNK = 65536;
            std::vector<char> buf(CHUNK);
            hdfsFile f = hdfsOpenFile(mr->hdfs_connection, HDFS_INPUT_PATH, O_RDONLY, 0, 0, 0);
            if (f) { // Count newlines up to start_offset
                long long pos = 0;
                while (pos < start_offset) { // Read in chunks
                    tSize to_read = (tSize)std::min<long long>(CHUNK, start_offset - pos);
                    tSize n = hdfsPread(mr->hdfs_connection, f, pos, buf.data(), to_read);
                    if (n <= 0) break;
                    for (int i = 0; i < n; ++i)
                        if (buf[i] == '\n') newline_count++;
                    pos += n;
                }
                hdfsCloseFile(mr->hdfs_connection, f);
            }
            if (start_offset == 0)
                startLineNum = static_cast<uint32_t>(newline_count + 1);
            else
                startLineNum = static_cast<uint32_t>(newline_count + 2);
        } else {
            // Default: line 1 for first mapper, line 2 for others
            startLineNum = (start_offset == 0) ? 1u : 2u;
        }

        // Normalize UTF-8 input (if applicable)
        std::string normalized_chunk = normalize_utf8(input_chunk); // keeps '\n' intact

        // Tokenize
        pairs = tokenization(normalized_chunk.c_str(),
                            startLineNum,
                            mr->helper_args,
                            &lines_processed,
                            mr->app_name);
    }

    std::cout << "Mapper " << mapper_id << ": Starting map loop for " << pairs.size() << " pairs." << std::endl;
    
    // Map each key-value pair
    for (auto& kv : pairs) {
        mr->map_fn(mr, &kv);
        free(kv.key);
        free(kv.value);
    }
    std::cout << "Mapper " << mapper_id << ": Finished map loop." << std::endl;

    // Signal mapper completion
    mapper_unit_t* unit = &mr->mapper_units[mapper_id]; // Get mapper unit
    pthread_mutex_lock(&unit->lock); // Lock
    unit->mapper_done = true; // Set done flag
    pthread_cond_signal(&unit->threshold_reached_cond); // Wake sender thread
    pthread_mutex_unlock(&unit->lock); // Unlock

    std::cout << "Mapper " << mapper_id << ": Exiting." << std::endl;
    return NULL;
}

// Sender thread function to transmit intermediate data to Reducer 
void* sender_thread_func(void* arg) {
    thread_arg_t* t_arg = (thread_arg_t*)arg;
    struct map_reduce* mr = t_arg->mr;
    int id = t_arg->id;
    mapper_unit_t* unit = &mr->mapper_units[id];
    // Create Reducer RPC stub
    mrdata::ReducerService::Stub* stub = static_cast<mrdata::ReducerService::Stub*>(unit->reducer_rpc_client);
    std::cout << "Sender " << id << ": Starting." << std::endl;

    // Main sending loop 
    while (true) {
        pthread_mutex_lock(&unit->lock);
        // Wait until buffer exceeds threshold or mapper is done
        while (unit->current_size < (size_t)(unit->capacity * MR_BUFFER_THRESHOLD) &&
               !unit->mapper_done &&
               !unit->mapper_waiting_for_space)
        {
            pthread_cond_wait(&unit->threshold_reached_cond, &unit->lock);
        }
        // If mapper is done and buffer is empty, exit loop
        if (unit->mapper_done && unit->current_size == 0) {
            pthread_mutex_unlock(&unit->lock);
            break; 
        }

        std::string buffer_content;
        // Copy buffer content to local string and reset buffer
        if (unit->current_size > 0) {
             buffer_content.assign(unit->buffer_data, unit->current_size);
             unit->current_size = 0;
        } else {
             pthread_mutex_unlock(&unit->lock);
             continue; 
        }

        // Unlock *before* send, signal afterwards
        pthread_mutex_unlock(&unit->lock);

        if (!buffer_content.empty()) {
            std::cout << "Sender " << id << ": Sending data..." << std::endl;
            send_intermediate_data(stub, buffer_content, id);
            std::cout << "Sender " << id << ": Data sent." << std::endl;
        }

        // Re-lock to signal mapper
        pthread_mutex_lock(&unit->lock); // Lock
        pthread_cond_broadcast(&unit->not_full_cond); // Signal mapper waiting for space
        pthread_mutex_unlock(&unit->lock); // Unlock
    }

    std::cout << "Sender " << id << ": Signaling mapper finished to reducer." << std::endl;
    signal_mapper_finished(stub, id, 0);
    std::cout << "Sender " << id << ": Exiting." << std::endl;
    return NULL;
}