nixos/estuary: Implement 95% bandwidth limiter

nixos/boxes/colony/vms/estuary/bandwidth.nix

@@ -0,0 +1,93 @@
+{ lib, pkgs, config, assignments, allAssignments, ... }: {
+  config = {
+    systemd = {
+      services = {
+        # systemd-networkd doesn't support tc filtering
+        wan-filter-to-ifb =
+        let
+          waitOnline = [
+            "systemd-networkd-wait-online@wan.service"
+            "systemd-networkd-wait-online@ifb-wan.service"
+          ];
+        in
+        {
+          description = "Install tc filter to pass WAN traffic to IFB";
+          enable = true;
+          bindsTo = waitOnline;
+          after = waitOnline;
+          serviceConfig = {
+            Type = "oneshot";
+            RemainAfterExit = true;
+          };
+          script = ''
+            ${pkgs.iproute2}/bin/tc filter add dev wan parent ffff: u32 match u32 0 0 action mirred egress redirect dev ifb-wan
+          '';
+          wantedBy = [ "multi-user.target" ];
+        };
+
+        bandwidth-limiter =
+        let
+          deps = [ "wan-filter-to-ifb.service" ];
+        in
+        {
+          description = "WAN bandwidth limiter";
+          enable = true;
+          bindsTo = deps;
+          after = deps;
+          path = with pkgs; [ python310 iproute2 ];
+          environment = {
+            PYTHONUNBUFFERED = "1";
+          };
+          serviceConfig = {
+            ExecStart = [ "${./bandwidth.py} wan,ifb-wan 245 10000" ];
+            StateDirectory = "bandwidth-limiter";
+          };
+          wantedBy = [ "multi-user.target" ];
+        };
+      };
+
+      network = {
+        netdevs = {
+          "25-ifb-wan".netdevConfig = {
+            Name = "ifb-wan";
+            Kind = "ifb";
+          };
+        };
+
+        networks = {
+          "80-wan" = {
+            extraConfig = ''
+              [QDisc]
+              Parent=ingress
+              Handle=ffff
+
+              # Outbound traffic limiting
+              [TokenBucketFilter]
+              Parent=root
+              LatencySec=0.3
+              BurstBytes=512K
+              # *bits
+              Rate=245M
+            '';
+          };
+          "80-ifb-wan" = {
+            matchConfig.Name = "ifb-wan";
+            extraConfig = ''
+              # Inbound traffic limiting
+              [TokenBucketFilter]
+              Parent=root
+              LatencySec=0.3
+              BurstBytes=512K
+              # *bits
+              Rate=245M
+            '';
+          };
+        };
+      };
+    };
+
+    my = {
+      tmproot.persistence.config.directories = [ "/var/lib/bandwidth-limiter" ];
+    };
+  };
+}

nixos/boxes/colony/vms/estuary/bandwidth.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python
+import time
+import datetime
+import sys
+import os
+import subprocess
+import json
+import signal
+import shlex
+
+INTERVAL = 5 * 60
+HZ = 1000
+LATENCY = '300ms'
+
+def end_of_month(dt: datetime.datetime):
+    return datetime.datetime(dt.year, dt.month + 1 if dt.month != 12 else 1, 1) - datetime.timedelta(seconds=1)
+
+def start_of_month(dt: datetime.datetime):
+    return datetime.datetime(dt.year, dt.month, 1)
+
+def month_seconds(dt: datetime.datetime):
+    return (end_of_month(dt) - start_of_month(dt)).total_seconds()
+
+def month_fraction(dt: datetime.datetime):
+    return (dt - start_of_month(dt)).total_seconds() / month_seconds(dt)
+
+def main():
+    if len(sys.argv) != 4:
+        print(f'usage: {sys.argv[0]} <interfaces> <95% limit mbit> <hi limit mbit>')
+        sys.exit(1)
+
+    ifaces = sys.argv[1].split(',')
+    lo = int(sys.argv[2])
+    hi = int(sys.argv[3])
+
+    cutoff = int((lo / 8) * 1024 * 1024 * INTERVAL)
+
+    basedir = os.environ['STATE_DIRECTORY']
+
+    def sig_handler(signum, frame):
+        sys.exit(0)
+    signal.signal(signal.SIGTERM, sig_handler)
+
+    last_total = 0
+    while True:
+        now = datetime.datetime.now()
+
+        total = 0
+        for n in ifaces:
+            output = subprocess.check_output(['ip', '-j', '-s', 'link', 'show', 'dev', n], encoding='utf-8')
+            stats = json.loads(output)
+            total += stats[0]['stats64']['tx']['bytes']
+
+        if last_total == 0:
+            last_total = total
+
+        data_file = os.path.join(basedir, str(now.year), f'{now.month}.json')
+        os.makedirs(os.path.dirname(data_file), exist_ok=True)
+
+        data = {
+            'hi_fraction_used': 0.0,
+        }
+        if os.path.exists(data_file):
+            with open(data_file, 'r') as f:
+                data = json.load(f)
+
+        if total - last_total > cutoff:
+            print(f'used more than {lo}mbps over the last {INTERVAL}s')
+            data['hi_fraction_used'] += INTERVAL / month_seconds(now)
+
+        limit = hi
+        mf = month_fraction(now)
+        if data['hi_fraction_used'] >= mf:
+            print(f"warning: used too many 5% buckets so far {data['hi_fraction_used']} and we are {mf} into the month); applying bandwidth limit")
+            limit = lo
+
+        with open(data_file, 'w') as f:
+            json.dump(data, f)
+
+        qdisc_args = ['rate', f'{limit}mbit', 'burst', str(int(((limit*1000*1000)/HZ/8) * 4)), 'latency', LATENCY]
+        for n in ifaces:
+            subprocess.check_call(['tc', 'qdisc', 'change', 'dev', n, 'root', 'tbf'] + qdisc_args)
+
+        last_total = total
+        time.sleep(INTERVAL)
+
+if __name__ == '__main__':
+	main()

nixos/boxes/colony/vms/estuary/default.nix

@@ -39,7 +39,7 @@
         inherit (lib.my) networkdAssignment;
       in
       {
-        imports = [ "${modulesPath}/profiles/qemu-guest.nix" ./dns.nix ];
+        imports = [ "${modulesPath}/profiles/qemu-guest.nix" ./dns.nix ./bandwidth.nix ];
 
         config = mkMerge [
           {
@@ -92,29 +92,6 @@
                   '';
                   wantedBy = [ "multi-user.target" ];
                 };
-
-                # systemd-networkd doesn't support tc filtering
-                wan-filter-to-ifb =
-                let
-                  waitOnline = [
-                    "systemd-networkd-wait-online@wan.service"
-                    "systemd-networkd-wait-online@ifb-wan.service"
-                  ];
-                in
-                {
-                  description = "Install tc filter to pass WAN traffic to IFB";
-                  enable = true;
-                  bindsTo = waitOnline;
-                  after = waitOnline;
-                  serviceConfig = {
-                    Type = "oneshot";
-                    RemainAfterExit = true;
-                  };
-                  script = ''
-                    ${pkgs.iproute2}/bin/tc filter add dev wan parent ffff: u32 match u32 0 0 action mirred egress redirect dev ifb-wan
-                  '';
-                  wantedBy = [ "multi-user.target" ];
-                };
               };
             };
 
@@ -135,13 +112,6 @@
                 };
               };
 
-              netdevs = {
-                "25-ifb-wan".netdevConfig = {
-                  Name = "ifb-wan";
-                  Kind = "ifb";
-                };
-              };
-
               networks = {
                 "80-wan" = {
                   matchConfig.Name = "wan";
@@ -160,31 +130,6 @@
                     LinkLocalAddressing = "no";
                     IPv6AcceptRA = false;
                   };
-                  extraConfig = ''
-                    [QDisc]
-                    Parent=ingress
-                    Handle=ffff
-
-                    # Outbound traffic limiting
-                    [TokenBucketFilter]
-                    Parent=root
-                    LatencySec=0.3
-                    BurstBytes=512K
-                    # *bits
-                    Rate=245M
-                  '';
-                };
-                "80-ifb-wan" = {
-                  matchConfig.Name = "ifb-wan";
-                  extraConfig = ''
-                    # Inbound traffic limiting
-                    [TokenBucketFilter]
-                    Parent=root
-                    LatencySec=0.3
-                    BurstBytes=512K
-                    # *bits
-                    Rate=245M
-                  '';
                 };
 
                 "80-base" = mkMerge [