1 files changed, 44 insertions, 0 deletions
diff --git a/migration/migration-stats.c b/migration/migration-stats.c
index 2f2cea965c..0890220ba5 100644
--- a/migration/migration-stats.c
+++ b/migration/migration-stats.c
@@ -12,6 +12,50 @@
 
 #include "qemu/osdep.h"
 #include "qemu/stats64.h"
+#include "qemu-file.h"
 #include "migration-stats.h"
 
 MigrationAtomicStats mig_stats;
+
+bool migration_rate_exceeded(QEMUFile *f)
+{
+    if (qemu_file_get_error(f)) {
+        return true;
+    }
+
+    uint64_t rate_limit_used = stat64_get(&mig_stats.rate_limit_used);
+    uint64_t rate_limit_max = stat64_get(&mig_stats.rate_limit_max);
+
+    if (rate_limit_max == RATE_LIMIT_DISABLED) {
+        return false;
+    }
+    if (rate_limit_max > 0 && rate_limit_used > rate_limit_max) {
+        return true;
+    }
+    return false;
+}
+
+uint64_t migration_rate_get(void)
+{
+    return stat64_get(&mig_stats.rate_limit_max);
+}
+
+#define XFER_LIMIT_RATIO (1000 / BUFFER_DELAY)
+
+void migration_rate_set(uint64_t limit)
+{
+    /*
+     * 'limit' is per second.  But we check it each BUFER_DELAY miliseconds.
+     */
+    stat64_set(&mig_stats.rate_limit_max, limit / XFER_LIMIT_RATIO);
+}
+
+void migration_rate_reset(void)
+{
+    stat64_set(&mig_stats.rate_limit_used, 0);
+}
+
+void migration_rate_account(uint64_t len)
+{
+    stat64_add(&mig_stats.rate_limit_used, len);
+}