fillmem.cpp

#include <algorithm>
#include <chrono>
#include <cstring>
#include <iostream>
#include <memory>
#include <mutex>
#include <thread>
#include <vector>
#include <getopt.h>

int main(int argc, char** argv)
{
    int threads = 1;
    int delta = 1; // 1 MB
    int sleep_interval = 5;
    unsigned long long size_mem;

    struct option long_options[] =
    {
        {"threads",  required_argument, 0, 't'},
        {"sleep",  required_argument, 0, 's'},
        {"delta",  required_argument, 0, 'd'},
        {"help",  no_argument, 0, 'h'},
        {0, 0, 0, 0}
    };

    int i,j;

    while( (j = getopt_long (argc, argv, "ht:s:d:", long_options, &i)) != -1)
        switch(j)
        {
            case 'h':
            case '?':
                std::cout << "Usage: " << argv[0] << " [OPTIONS] MEGABYTES-TO-ALLOCATE\n"
                    "\n"
                    "    -t, --threads=num-threads       Number of threads to run (default: " << threads << ")\n"
                    "    -s, --sleep=seconds             Time in seconds to wait before exiting (default " << sleep_interval << " s)\n"
                    "    -d, --delta=megabytes           How many megabytes to allocate at every step (default: " << delta << " MB)\n"
                    "    -h, --help                      Display this help\n"
                    "\n"
                    "The final mandatory argument is the number of megabytes to allocate.\n"
                    "\n";
                return 0;
                break;
            case 't':
                threads = atoi(optarg);
                if(threads < 1)
                {
                    std::cerr << "Invalid number of threads!" << std::endl;
                    return 1;
                }
                break;
            case 's':
                sleep_interval = atoi(optarg);
                if(sleep_interval < 0)
                {
                    std::cerr << "Invalid sleep time!" << std::endl;
                    return 2;
                }
                break;
            case 'd':
                delta = atoi(optarg);
                if(delta <= 0)
                {
                    std::cerr << "Invalid delta in megabytes!" << std::endl;
                    return 3;
                }
                break;
        }

    if(optind >= argc)
    {
        std::cerr << "Missing size in megabytes to allocate." << std::endl;
        std::cerr << "Use --help view the correct arguments to pass. " << std::endl;
        return 4;
    }

    size_mem = atoi(argv[optind]);
    std::cout << "Allocating " << size_mem << " MB per thread (" << threads << " thread(s) used) ";
    std::cout << "in steps of " << delta << " MB." << std::endl;

    std::mutex print_mutex;
    auto alloc_mem = [&]() {
        auto this_id = std::this_thread::get_id();
        auto size = size_mem / delta;
        std::vector<std::unique_ptr<char []>> memlist(size);

        for(int c=0;c<size;c++)
        {
            try {
            //memlist[c] = std::make_unique<char []>(delta*1024*1024);
            memlist[c].reset(new char [delta*1024*1024]);
            std::memset(memlist[c].get(), '1', delta*1024*1024);
            } catch(std::bad_alloc &err) {
                print_mutex.lock();
                std::cerr << "Thread " << this_id << ": Failed to allocate " << delta << " MB: " << err.what() << std::endl;
                print_mutex.unlock();
                return;
            }

            print_mutex.lock();
            std::cout << "Thread " << this_id << ": Allocated " << (c+1)*delta << " MB" << std::endl;
            print_mutex.unlock();
        }

        print_mutex.lock();
        std::cout << "Thread " << this_id << " is sleeping " << sleep_interval << " s." << std::endl;
        print_mutex.unlock();
        std::this_thread::sleep_for(std::chrono::seconds(sleep_interval));
    };

    std::vector<std::thread> workers;

    for(i=0;i<threads;i++)
        workers.push_back(std::thread(alloc_mem));

    std::for_each(workers.begin(), workers.end(), [](std::thread &t)
    {
        t.join();
    });

    return 0;
}