ssn-installer/src/rateLimiter.c

68 lines
2.1 KiB
C

//To avoid reaching the write speed capacity of the hard disk, limit all writes
#include <stdbool.h>
#include <stdio.h>
#include <time.h>
#include <unistd.h>
#include "errorAndExit.h"
#include "rateLimiter.h"
#include "utils/min.h"
//How many bytes we can write to disk per second. Use 0 to specify no limit. Default is 30 MB/s.
unsigned long diskSpeed = 30UL * 1024UL * 1024UL;
/** If we are not writing to disk, how many seconds the capacity can keep increasing before it reaches the limit.
* Increasing this value can result in a spike of data once we are writing to disk again, but then allows the data
* to quickly be written without having to wait for capacity to be available. */
#define MAX_IDLE_TIME 3UL
//Capacity can be filled up to this maximum amount
#define CAPACITY_MAX diskSpeed * MAX_IDLE_TIME
//The current amount of free capacity
unsigned long capacity = 0UL;
time_t prevTime = 0;
/** Set disk speed to the given speed in bytes per second. */
void setDiskSpeed(unsigned long newSpeed) {
diskSpeed = newSpeed;
}
/** Increases the capacity based on time elapsed since the last time this function was called */
void increaseCapacity() {
//Get seconds since the last time this function was called, or default to 1 if this is the first time this function is called
unsigned long timeElapsed = 1;
if (prevTime == 0) {
time(&prevTime);
} else {
time_t prevTimeCache = prevTime;
time(&prevTime);
timeElapsed = prevTime - prevTimeCache;
}
capacity = min(CAPACITY_MAX, capacity + timeElapsed * diskSpeed);
}
/** Pauses execution until the given amount of capacity is available */
void consumeCapacity(unsigned long numBytes) {
//Exit early if there is no capacity limit
if (CAPACITY_MAX == 0UL) {
return;
}
if (numBytes > CAPACITY_MAX) {
fprintf(stderr, "Could not consume %lu bytes in rate limiter; maximum is %lu.\n", numBytes, CAPACITY_MAX);
errorAndExit();
}
increaseCapacity();
//While not enough capacity is available
while (capacity < numBytes) {
//Sleep for one second
sleep(1U);
increaseCapacity();
}
//Consume capacity
capacity -= numBytes;
}