chrome_throttle.cpp

#include "brt.h"

#include <windows.h>
#include <psapi.h>

struct WorkingSetParams {
    HANDLE process_handle;
    size_t min_working_set;
    size_t max_working_set;
    DWORD flags;
};

static void
GetWorkingSetParams(WorkingSetParams * params) {
    if (!GetProcessWorkingSetSizeEx(params->process_handle, 
                                    &params->min_working_set,
                                    &params->max_working_set,
                                    &params->flags)) 
    {
        // If the call failed, clear to zero.
        memset(params, 0, sizeof(*params));
    }
}

static void
SetWorkingSetParams(WorkingSetParams * params) {
    if (!params->process_handle ||
        (params->min_working_set == params->max_working_set && params->min_working_set == 0)) {
        // Invaild params, bail.
        return;
    }
    SetProcessWorkingSetSizeEx(params->process_handle, 
                               params->min_working_set, 
                               params->max_working_set, 
                               params->flags);
}

static bool
CheckProcessNameMatches(HANDLE process_handle, char * exe_name, u32 exe_name_len) {
    char process_name[2048] = {};
    bool success = false;

    DWORD process_name_len = array_count(process_name);
    QueryFullProcessImageNameA(process_handle, 0, process_name, &process_name_len);
    if (process_name_len >= exe_name_len) {
        char * process_name_cmp_start = process_name + (process_name_len - exe_name_len);
        success = true;
        for (u32 i = 0; i < exe_name_len; ++i) {
            if (exe_name[i] != process_name_cmp_start[i]) {
                success = false;
                break;
            }
        }
    }

    return success;
}

static WorkingSetParams stored_params[1024];
u32 stored_params_count;

static void
LimitExecutableMemUsage(char * exe_name, u32 exe_name_len, bool flush_only=false) {
    DWORD pids[1024] = {};
    DWORD pids_bytes_written = 0;
    if (EnumProcesses(pids, sizeof(pids), &pids_bytes_written)) {
        if (pids_bytes_written == sizeof(pids)) {
            fprintf(stderr, "Too many processes to enumerate, some may be missed.\n");
        }

        WorkingSetParams throttle_params = { 0, (size_t)-1, (size_t)-1, QUOTA_LIMITS_HARDWS_MIN_DISABLE | QUOTA_LIMITS_HARDWS_MAX_ENABLE };

        u32 access_mask = PROCESS_QUERY_INFORMATION | PROCESS_SET_QUOTA;

        u32 pid_count = pids_bytes_written >> 1;
        HANDLE process_handle;
        for (u32 i = 0; i < pid_count; ++i) {
            if (process_handle = OpenProcess(access_mask, false, pids[i])) {
                if (CheckProcessNameMatches(process_handle, exe_name, exe_name_len)) {
                    assert(stored_params_count < array_count(stored_params));
                    // Store the existing working set params.
                    WorkingSetParams * old_params = stored_params + stored_params_count++;
                    old_params->process_handle = process_handle;
                    GetWorkingSetParams(old_params);

                    // Set new params to throttle working set usage.
                    throttle_params.process_handle = process_handle;
                    if (flush_only) throttle_params.flags = old_params->flags;
                    SetWorkingSetParams(&throttle_params);

                    fprintf(stderr, "Set new working set params for pid %u\n", pids[i]);
                }
            }
        }
    }
}

static void
RestoreExecutableMemUsage() {
    for (u32 i = 0; i < stored_params_count; ++i) {
        SetWorkingSetParams(stored_params + i);
    }
}

int
main(int argc, char ** argv) {
    printf("This will flush the working set memory of all running chrome.exe processes.\n"
           "Proceed? (y/n): ");
    char yn = fgetc(stdin);
    if (yn == 'y' || yn == 'Y') {
        char chrome_name[] = "chrome.exe";
        LimitExecutableMemUsage(chrome_name, array_count(chrome_name) - 1, true);
    }
    // Consume newline
    fgetc(stdin);

    printf("Press any key to restore previous settings.\n");
    char cont = fgetc(stdin);
    if (cont != 'n') {
        RestoreExecutableMemUsage();
    }
    return 0;
}