From 8b7abe7c79ecf6028a472a6d1e18d1d42d7b45d1 Mon Sep 17 00:00:00 2001
From: C-3PO <c-3po@jedipedia.net>
Date: Sun, 30 Sep 2018 16:48:41 +0200
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Add=20rate=20limit=20for=20disk=20w?=
 =?UTF-8?q?rite=20speed?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/main.c        | 11 ++++++---
 src/rateLimiter.c | 58 +++++++++++++++++++++++++++++++++++++++++++++++
 src/rateLimiter.h |  4 ++++
 src/xdelta3.c     |  4 +++-
 4 files changed, 73 insertions(+), 4 deletions(-)
 create mode 100644 src/rateLimiter.c
 create mode 100644 src/rateLimiter.h

diff --git a/src/main.c b/src/main.c
index 3720f84..c0d835b 100644
--- a/src/main.c
+++ b/src/main.c
@@ -12,6 +12,7 @@
 #include "fileReader.h"
 #include "inflate.h"
 #include "parseArguments.h"
+#include "rateLimiter.h"
 #include "xdelta3.h"
 
 //The size of the buffers where the compressed and uncompressed data is stored
@@ -125,14 +126,18 @@ int main(int argc, char *argv[]) {
 
     //important: we must not modify uncompressedChunk since miniz may use it as dictionary and read from it during the next invocation of inflateInflate()
 
+    uint8_t* dataBuffer = uncompressedChunk + uncompressedPosition;
+    unsigned long numBytes = inflateResult.numBytesWrittenToOutput;
+
     //Optionally perform xdelta3
     if (state.prevFile) {
-      xdelta3AddInput(uncompressedChunk + uncompressedPosition, inflateResult.numBytesWrittenToOutput, remainingBytes == 0 && hasReachedEnd, targetFile);
+      xdelta3AddInput(dataBuffer, numBytes, remainingBytes == 0 && hasReachedEnd, targetFile);
     } else { //otherwise write to target file or stdout
+      consumeCapacity(numBytes);
       if (state.target) {
-        fwrite(uncompressedChunk + uncompressedPosition, 1, inflateResult.numBytesWrittenToOutput, targetFile);
+        fwrite(dataBuffer, 1, numBytes, targetFile);
       } else {
-        write(1, uncompressedChunk + uncompressedPosition, inflateResult.numBytesWrittenToOutput);
+        write(1, dataBuffer, numBytes);
       }
     }
 
diff --git a/src/rateLimiter.c b/src/rateLimiter.c
new file mode 100644
index 0000000..f4f29b6
--- /dev/null
+++ b/src/rateLimiter.c
@@ -0,0 +1,58 @@
+//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: 100 MB - TODO: need to use a better value for this
+#define DISK_SPEED 100UL * 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 5UL
+//Capacity can be filled up to this maximum amount
+#define CAPACITY_MAX DISK_SPEED * MAX_IDLE_TIME
+
+//The current amount of free capacity
+unsigned long capacity = 0UL;
+
+time_t prevTime = 0;
+
+/** 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 * DISK_SPEED);
+}
+
+/** Pauses execution until the given amount of capacity is available */
+void consumeCapacity(unsigned long numBytes) {
+  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;
+}
diff --git a/src/rateLimiter.h b/src/rateLimiter.h
new file mode 100644
index 0000000..890691d
--- /dev/null
+++ b/src/rateLimiter.h
@@ -0,0 +1,4 @@
+#pragma once
+
+/** Pauses execution until the given amount of capacity is available */
+void consumeCapacity(unsigned long numBytes);
diff --git a/src/xdelta3.c b/src/xdelta3.c
index 6a229ab..c6b7d60 100644
--- a/src/xdelta3.c
+++ b/src/xdelta3.c
@@ -1,4 +1,4 @@
-//To fix compile errors with xdelta3
+//Required to fix compile errors with xdelta3
 #define SIZEOF_SIZE_T 4
 #define static_assert(e,m) /* do nothing */
 #define XD3_ENCODER 0
@@ -15,6 +15,7 @@ typedef unsigned long long xoff_t;
 
 //Include our own files
 #include "errorAndExit.h"
+#include "rateLimiter.h"
 #include "xdelta3.h"
 
 //Various variables for Xdelta3
@@ -70,6 +71,7 @@ process:
     case XD3_OUTPUT: {
       //fprintf(stderr, "XD3_OUTPUT\n");
       totalOut += stream.avail_out;
+      consumeCapacity(stream.avail_out);
       if (targetFile) {
         fwrite(stream.next_out, 1, stream.avail_out, targetFile);
       } else {