Restoring slurm paralel

#include <vector>

#include <cstdlib>   // For getenv

#include <pthread.h> // Include Pthreads header

#include <atomic>    // For atomic counters

using namespace std;

static pfind_runtime_options_t runtime;

static pfind_options_t* opt;

// Global rank/size for logging and work distribution

int sfind_rank = 0;

int sfind_size = 1;

[[noreturn]] void

pfind_abort(const string& str) {

    cerr << "ERROR: " << str << endl;

    cerr << "ERROR [Rank " << sfind_rank << "]: " << str << endl;

    exit(1);

}

// (pfind_print_help and pfind_parse_args are identical to your provided code)

#pragma region parsing

// ... Paste your pfind_print_help and pfind_parse_args functions here ...

static void

pfind_print_help(const pfind_options_t* res) {

    printf("pfind \nSynopsis:\n"

    auto res = new pfind_options_t();

    bool print_help = force_print_help;

    vector<char*> modified_argv(argv, argv + argc);

    for(int i = 1; i < argc; ++i) {

        if(strcmp(argv[i], "-newer") == 0 && i + 1 < argc) {

            res->timestamp_file = argv[i + 1];

        }

    }

    if(print_help) {

        if(sfind_rank == 0)

            pfind_print_help(res);

        exit(0);

    }

}

#pragma endregion

// Structure to pass data to and from each worker thread

// Structure to pass data to each worker thread

// Using atomics for results is simpler than passing a pointer and locking

struct ThreadData {

    int thread_id;

    const pfind_options_t* opt;

    const string* workdir;

    queue<int>* server_queue;

    pthread_mutex_t* queue_mutex;

    // Thread-local results to avoid race conditions

    unsigned long long local_found = 0;

    unsigned long long local_checked = 0;

    atomic<unsigned long long>* total_found;

    atomic<unsigned long long>* total_checked;

};

// The function each worker thread will execute

void*

worker_routine(void* arg) {

    ThreadData* data = static_cast<ThreadData*>(arg);

    unsigned long long local_found = 0;

    unsigned long long local_checked = 0;

    while(true) {

        int server_id = -1;

        // Safely get a server ID from the shared queue

        pthread_mutex_lock(data->queue_mutex);

        if(!data->server_queue->empty()) {

            server_id = data->server_queue->front();

        }

        pthread_mutex_unlock(data->queue_mutex);

        // If the queue was empty, our work is done

        if(server_id == -1) {

        if(server_id == -1)

            break;

        }

        struct dirent_extended* entries = nullptr;

        long n = gkfs_getsingleserverdir(data->workdir->c_str(), &entries,

                                         server_id);

        if(n <= 0) { // Handle error or empty directory from server

        if(n <= 0) {

            if(n < 0) {

                cerr << "Warning: Thread " << data->thread_id

                     << " received error from server " << server_id << endl;

                cerr << "Warning: Rank " << sfind_rank << " Thread "

                     << data->thread_id << " received error from server "

                     << server_id << endl;

            }

            if(entries)

                free(entries); // Ensure we free even on error if allocated

                free(entries);

            continue;

        }

        // Iterate through the results from this server

        char* ptr = reinterpret_cast<char*>(entries);

        int bytes_processed = 0;

        while(bytes_processed < n) {

            struct dirent_extended* temp =

                    reinterpret_cast<struct dirent_extended*>(ptr);

            if(temp->d_reclen == 0)

                break; // Safety break

                break;

            if(temp->d_type != 1) { // Process files, not directories

            if(temp->d_type != 1) {

                bool timeOK = data->opt->timestamp_file.empty() ||

                              ((uint64_t) temp->ctime >= runtime.ctime_min);

                bool sizeOK = (data->opt->size ==

                bool nameOK = data->opt->name_pattern.empty() ||

                              regex_search(temp->d_name, data->opt->name_regex);

                if(timeOK && sizeOK && nameOK) {

                    data->local_found++;

                }

                data->local_checked++;

                if(timeOK && sizeOK && nameOK)

                    local_found++;

                local_checked++;

            }

            bytes_processed += temp->d_reclen;

            ptr += temp->d_reclen;

        }

        free(entries); // Free the memory allocated by gkfs_getsingleserverdir

        free(entries);

    }

    // Atomically add local results to the global counters

    data->total_found->fetch_add(local_found);

    data->total_checked->fetch_add(local_checked);

    return nullptr;

}

int

process_parallel_pthreads(const pfind_options_t* opt) {

    unsigned long long global_found = 0;

    unsigned long long global_checked = 0;

    atomic<unsigned long long> global_found(0);

    atomic<unsigned long long> global_checked(0);

    runtime = {};

    if(!opt->timestamp_file.empty()) {

    if(workdir.empty())

        workdir = "/";

    // --- 1. Calculate this process's subset of servers ---

    int servers_per_proc = opt->num_servers / sfind_size;

    int remainder = opt->num_servers % sfind_size;

    int start_server =

            sfind_rank * servers_per_proc + min(sfind_rank, remainder);

    int end_server =

            start_server + servers_per_proc + (sfind_rank < remainder ? 1 : 0);

    int num_servers_for_this_rank = end_server - start_server;

    if(num_servers_for_this_rank <= 0) {

        cout << "[Rank " << sfind_rank << "] No servers to process. Exiting."

             << endl;

        return 0;

    }

    // --- 2. Populate the local work queue with this rank's servers ---

    queue<int> server_queue;

    for(int i = start_server; i < end_server; ++i) {

        server_queue.push(i);

    }

    // --- Pthreads Setup ---

    int num_threads = 20; // Default number of parallel requests

    const char* env_threads = getenv("GFIND_NUM_THREADS");

    const char* env_threads = getenv("SFIND_NUM_THREADS");

    if(env_threads) {

        try {

            num_threads = stoi(env_threads);

        } catch(...) { /* ignore, use default */

        } catch(...) { /* ignore */

        }

    }

    // Don't create more threads than there are servers to query

    num_threads = min(num_threads, opt->num_servers);

    num_threads = min(num_threads, num_servers_for_this_rank);

    cout << "Using " << num_threads << " threads to query " << opt->num_servers

         << " servers." << endl;

    cout << "[Rank " << sfind_rank << "] Processing servers " << start_server

         << "-" << end_server - 1 << " using " << num_threads << " threads."

         << endl;

    // 1. Create the shared work queue

    queue<int> server_queue;

    for(int i = 0; i < opt->num_servers; ++i) {

        server_queue.push(i);

    }

    // 2. Initialize the mutex

    pthread_mutex_t queue_mutex;

    pthread_mutex_init(&queue_mutex, nullptr);

    // 3. Create thread data structures

    vector<pthread_t> threads(num_threads);

    vector<ThreadData> thread_args(num_threads);

    // 4. Launch worker threads

    for(int i = 0; i < num_threads; ++i) {

        thread_args[i] = {.thread_id = i,

                          .opt = opt,

                          .workdir = &workdir,

                          .server_queue = &server_queue,

                          .queue_mutex = &queue_mutex};

        thread_args[i] = {i,

                          opt,

                          &workdir,

                          &server_queue,

                          &queue_mutex,

                          &global_found,

                          &global_checked};

        pthread_create(&threads[i], nullptr, worker_routine, &thread_args[i]);

    }

    // 5. Wait for all threads to complete

    for(int i = 0; i < num_threads; ++i) {

        pthread_join(threads[i], nullptr);

    }

    // 6. Aggregate results from all threads

    for(int i = 0; i < num_threads; ++i) {

        global_found += thread_args[i].local_found;

        global_checked += thread_args[i].local_checked;

    }

    // 7. Clean up

    pthread_mutex_destroy(&queue_mutex);

    // Results are already aggregated via atomics.

    // Each process prints its own sub-total. The final aggregation must be done

    // externally (e.g., with a script).

    cout << "MATCHED " << global_found << "/" << global_checked << endl;

    return 0;

}

int

main(int argc, char** argv) {

    // --- Get Rank/Size from SLURM Environment Variables ---

    const char* env_rank = getenv("SLURM_PROCID");

    const char* env_size = getenv("SLURM_NPROCS");

    if(env_rank && env_size) {

        try {

            sfind_rank = stoi(env_rank);

            sfind_size = stoi(env_size);

        } catch(const std::exception& e) {

            cerr << "Could not parse SLURM environment variables: " << e.what()

                 << ". Falling back to sequential mode." << endl;

            sfind_rank = 0;

            sfind_size = 1;

        }

    } else {

        if(getenv("SLURM_JOB_ID")) {

            cerr << "Warning: SLURM_JOB_ID is set, but SLURM_PROCID/SLURM_NPROCS are not. Are you running with `srun`?"

                 << endl;

        }

        cout << "SLURM variables not found. Running in sequential mode (rank 0 of 1)."

             << endl;

        sfind_rank = 0;

        sfind_size = 1;

    }

    // Each process parses its own arguments. This is simple and correct.

    opt = pfind_parse_args(argc, argv, false);

    if(gkfs_getsingleserverdir == nullptr) {