ipc_debug/src/ipc_debug.cpp

#include <new>
#include <numeric>
#include <unistd.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cstdint>
#include <cassert>
#include <errno.h>
#include <iostream>
#include <vector>
#include <mpi.h>
#include <map>
#include <string>
#include <fcntl.h>

#include <ipc_debug.h>

bool debug_ipc_disabled = true;
bool debug_ipc_is_root = false;
int debug_ipc_file = -1;

int rank, cluster_size;
std::map<uint64_t, int> start_indices_map; // maps indices to MPI ranks
std::vector<int> start_indices_list; // maps ranks to start indices
std::vector<int> recv_counts;

std::vector<unsigned char> buffer;
std::vector<double> float_buffer;

constexpr size_t INITIAL_BUFFER_SIZE = 5 * 1024 * 1024; // 5 MiB

void print_backtrace();

void debug_ipc_init() {
    const char *root_env = std::getenv("IPC_DEBUG_ROOT");
    const char *file_env = std::getenv("IPC_DEBUG_FILE");
    debug_ipc_disabled = file_env == nullptr;

    int mpi_initialized;
    MPI_Initialized(&mpi_initialized);
    if (mpi_initialized) {
        MPI_Comm_rank(MPI_COMM_WORLD, &rank);
        MPI_Comm_size(MPI_COMM_WORLD, &cluster_size);
    }

    if (rank == 0) {
        
        if (file_env == nullptr) {
            return;
        }

        if (root_env != nullptr && 0 == strcmp(root_env, "1")) {
            debug_ipc_is_root = true;
            printf("[IPCDBG] I am host\n");
        } else {
            printf("[IPCDBG] I am client\n");
        }

        if (debug_ipc_is_root) {
            debug_ipc_file = open(file_env, O_RDONLY | O_CLOEXEC);
        } else {
            debug_ipc_file = open(file_env, O_WRONLY | O_CLOEXEC);
        }

        if (debug_ipc_file == -1) {
            printf("[IPCDBG] Error, could not open named pipe: %s", strerror(errno));
            exit(-1);
        }


        buffer.resize(INITIAL_BUFFER_SIZE);
    }

}

void endless_loop() {
    printf("Entering endless loop, attach debugger to PID %i \n", getpid());
    fflush(stdout);
    while (1) {
        sleep(1);
    }
}

template<typename T>
void debug_ipc_assert_equal(T value) {
    if (debug_ipc_file == -1 || debug_ipc_disabled) {
        return;
    }

    ssize_t expected_size = sizeof(T);
    if (debug_ipc_is_root) {
        T other_value;
        ssize_t num_read = read(debug_ipc_file, &other_value, expected_size);

        if (num_read != expected_size) {
            printf("[IPCDBG] Could not read enough bytes. Error: %s\n", strerror(errno));
            exit(-1);
        }

        if (other_value != value) {
            std::cout << "[IPCDBG] Assertion failed!"
                      << " Root has " << value << " but client has " << other_value;
            print_backtrace();
            std::cout << "Entering endless loop, attach debugger to PID " << getpid();
            fflush(stdout);
            while (1) {
                sleep(1);
            }
        } else {
#ifdef TRACE
            std::cout << "[IPCDBG] Assertion passed, value = " << value << std::endl;
#endif
        }
    } else {
        ssize_t written = write(debug_ipc_file, &value, expected_size);

        if (written != expected_size) {
            printf("[IPCDBG] Could not write enough bytes. Error: %s\n", strerror(errno));
            endless_loop();
        }
    }
}

void read_loop(int fd, unsigned char *data, size_t count) {
    size_t num_read = 0;

    do {
        ssize_t ret = read(fd, &data[num_read], count - num_read);

        if (ret == -1) {
            printf("[IPCDBG] Could not read enough bytes. Error: %s\n", strerror(errno));
            endless_loop();
        } else if (ret == 0) {
            printf("[IPCDBG] Could not read enough bytes, unexpected EOF.");
            endless_loop();
        }

        num_read += ret;
    } while (num_read < count);
}

void write_loop(int fd, unsigned char *data, size_t count) {
    size_t num_written = 0;

    do {
        ssize_t ret = write(fd, &data[num_written], count - num_written);

        if (ret == -1) {
            printf("[IPCDBG] Could not write enough bytes. Error: %s\n", strerror(errno));
            endless_loop();
        }
        num_written += ret;
    } while (num_written < count);
}


template<typename T>
void debug_ipc_assert_equal_vector(std::vector<T> value) {
    if (debug_ipc_file == -1 || debug_ipc_disabled) {
        return;
    }

    debug_ipc_assert_equal(value.size());
    const size_t array_byte_length = sizeof(T) * value.size();

    if (debug_ipc_is_root) {
        buffer.resize(array_byte_length);

        read_loop(debug_ipc_file, buffer.data(), array_byte_length);

        assert(reinterpret_cast<uint64_t>(buffer.data()) % 8 == 0); // Make sure the array is properly aligned
        T *local_array = value.data();
        T *other_array = reinterpret_cast<T *>(buffer.data());

        for (size_t i = 0; i < value.size(); ++i) {
            if (local_array[i] != other_array[i]) {
                std::cout << "[IPCDBG] Assertion failed in vector  at index " << i 
                    << ". Root has " << local_array[i] << " but client has " << other_array[i] << std::endl;
                print_backtrace();
                endless_loop();
            }
        }

    } else {
        write_loop(debug_ipc_file, reinterpret_cast<unsigned char *>(value.data()), array_byte_length);
    }
}

// Explicit template instantiation
template void debug_ipc_assert_equal<double>(double);
template void debug_ipc_assert_equal<float>(float);
template void debug_ipc_assert_equal<uint32_t>(uint32_t);
template void debug_ipc_assert_equal<int32_t>(int32_t);
template void debug_ipc_assert_equal<uint64_t>(uint64_t);
template void debug_ipc_assert_equal<int64_t>(int64_t);
template void debug_ipc_assert_equal<bool>(bool);
template void debug_ipc_assert_equal<std::string>(std::string);

void debug_ipc_assert_equal_double(double value) {
    debug_ipc_assert_equal<double>(value);
}
void debug_ipc_assert_equal_uint(uint32_t value) {
    debug_ipc_assert_equal<uint32_t>(value);
}
void debug_ipc_assert_equal_int(int32_t value) {
    debug_ipc_assert_equal<int32_t>(value);
}
void debug_ipc_assert_equal_int64(int64_t value) {
    debug_ipc_assert_equal<int64_t>(value);
}

void debug_ipc_assert_equal_array(void *value, size_t size) {
    static_assert(sizeof(unsigned char) == 1);
    if (debug_ipc_file == -1 || debug_ipc_disabled) {
        return;
    }
    debug_ipc_assert_equal(size); // Make sure arrays are the same size
    

    unsigned char *array = reinterpret_cast<unsigned char *>(value);
    if (debug_ipc_is_root) {
        buffer.resize(size);
        unsigned char *other_array = buffer.data();

        read_loop(debug_ipc_file, other_array, size);

        for (size_t i = 0; i < size; i++) {
            if (array[i] != other_array[i]) {
                printf("[IPCDBG] Assertion failed in byte %lu!\n", i);
                print_backtrace();
                endless_loop();
            }
        }
    } else {
        write_loop(debug_ipc_file, reinterpret_cast<unsigned char *>(value), size);
    }
}

void debug_ipc_assert_source_location(const char *source_file, const long int line_number) {
    std::string fname = source_file;

    debug_ipc_assert_equal_array(fname.data(), fname.size());
    debug_ipc_assert_equal(line_number);
}

#include <backward.hpp>

void print_backtrace()
{
    backward::StackTrace st;
    st.load_here(32);
    backward::Printer p;
    p.print(st);
}

void debug_ipc_mpi_set_data_distribution(int start_index) {
    if (rank == 0) {
        start_indices_map.clear();
        start_indices_list.resize(cluster_size + 1);
    }

    int my_start_index = start_index;
    MPI_Gather(&my_start_index, 1, MPI_INT,
            start_indices_list.data(), 1, MPI_INT,
            0, MPI_COMM_WORLD);

    if (rank == 0) {
        for (int i = 0; i < cluster_size; i++) {
            start_indices_map[start_indices_list[i]] = i;
        }
    }
}




void debug_ipc_assert_equal_mpi_double_array(double *array, size_t array_length) {
    if (debug_ipc_disabled) {
        return;
    }

    std::vector<int> recv_counts(cluster_size);
    std::vector<int> displacements(cluster_size);
    int local_array_length = static_cast<int>(array_length);

    // Gather local lengths
    MPI_Gather(&local_array_length, 1, MPI_INT,
            recv_counts.data(), 1, MPI_INT,
            0, MPI_COMM_WORLD);

    if (rank == 0) {
        // Setup displacements and receive counts
        int global_index = 0;
        // Iterate over start indices from small to large
        for (const auto& [start_index, m_rank] : start_indices_map) {
            displacements[m_rank] = global_index;
            global_index += recv_counts[m_rank];
        }
        float_buffer.resize(global_index); // Make sure buffer is appropriately sized.
    }

    MPI_Gatherv(array, array_length, MPI_DOUBLE,
            float_buffer.data(), // recvbuf
            recv_counts.data(),  // recv counts
            displacements.data(), // displacements
            MPI_DOUBLE, 0, MPI_COMM_WORLD);

    if (rank == 0) {
        debug_ipc_assert_equal_vector(float_buffer);
    }
}