nixos/atop: Add configuration for atop services, allow to enable netatop, gpuatop, allow setuid wrapper

3 files changed, 251 insertions, 8 deletions

diff --git a/nixos/modules/programs/atop.nix b/nixos/modules/programs/atop.nix
index 7ef8d687ca1..d1577b32adf 100644
--- a/nixos/modules/programs/atop.nix
+++ b/nixos/modules/programs/atop.nix
@@ -1,6 +1,6 @@
 # Global configuration for atop.
 
-{ config, lib, ... }:
+{ config, lib, pkgs, ... }:
 
 with lib;
 
@@ -12,11 +12,82 @@ in
 
   options = {
 
-    programs.atop = {
+    programs.atop = rec {
 
+      enable = mkEnableOption "Atop";
+
+      package = mkOption {
+        type = types.package;
+        default = pkgs.atop;
+        description = ''
+          Which package to use for Atop.
+        '';
+      };
+
+      netatop = {
+        enable = mkOption {
+          type = types.bool;
+          default = false;
+          description = ''
+            Whether to install and enable the netatop kernel module.
+          '';
+        };
+        package = mkOption {
+          type = types.package;
+          default = config.boot.kernelPackages.netatop;
+          description = ''
+            Which package to use for netatop.
+          '';
+        };
+      };
+
+      atopgpu.enable = mkOption {
+        type = types.bool;
+        default = false;
+        description = ''
+          Whether to install and enable the atopgpud daemon to get information about
+          NVIDIA gpus.
+        '';
+      };
+
+      setuidWrapper.enable = mkOption {
+        type = types.bool;
+        default = cfg.netatop.enable || cfg.atopgpu.enable;
+        description = ''
+          Whether to install a setuid wrapper for Atop. This is required to use some of
+          the features as non-root user (e.g.: ipc information, netatop, atopgpu).
+          Atop tries to drop the root privileges shortly after starting.
+        '';
+      };
+
+      atopsvc.enable = mkOption {
+        type = types.bool;
+        default = true;
+        description = ''
+          Whether to enable the atop service responsible for storing statistics for
+          long-term analysis.
+        '';
+      };
+      atopRotate.enable = mkOption {
+        type = types.bool;
+        default = true;
+        description = ''
+          Whether to enable the atop-rotate timer, which restarts the atop service
+          daily to make sure the data files are rotate.
+        '';
+      };
+      atopacct.enable = mkOption {
+        type = types.bool;
+        default = true;
+        description = ''
+          Whether to enable the atopacct service which manages process accounting.
+          This allows Atop to gather data about processes that disappeared in between
+          two refresh intervals.
+        '';
+      };
       settings = mkOption {
         type = types.attrs;
-        default = {};
+        default = { };
         example = {
           flags = "a1f";
           interval = 5;
@@ -25,12 +96,51 @@ in
           Parameters to be written to <filename>/etc/atoprc</filename>.
         '';
       };
-
     };
   };
 
-  config = mkIf (cfg.settings != {}) {
-    environment.etc.atoprc.text =
-      concatStrings (mapAttrsToList (n: v: "${n} ${toString v}\n") cfg.settings);
-  };
+  config = mkIf cfg.enable (
+    let
+      atop =
+        if cfg.atopgpu.enable then
+          (cfg.package.override { withAtopgpu = true; })
+        else
+          cfg.package;
+      packages = [ atop (lib.mkIf cfg.netatop.enable cfg.netatop.package) ];
+    in
+    {
+      environment.etc = mkIf (cfg.settings != { }) {
+        atoprc.text = concatStrings
+          (mapAttrsToList
+            (n: v: ''
+              ${n} ${toString v}
+            '')
+            cfg.settings);
+      };
+      environment.systemPackages = packages;
+      boot.extraModulePackages = [ (lib.mkIf cfg.netatop.enable cfg.netatop.package) ];
+      systemd =
+        let
+          mkSystemd = type: cond: name: {
+            ${name} = lib.mkIf cond {
+              restartTriggers = packages;
+              wantedBy = [ (if type == "services" then "multi-user.target" else if type == "timers" then "timers.target" else null) ];
+            };
+          };
+          mkService = mkSystemd "services";
+          mkTimer = mkSystemd "timers";
+        in
+        {
+          inherit packages;
+          services =
+            mkService cfg.atopsvc.enable "atop"
+            // mkService cfg.atopacct.enable "atopacct"
+            // mkService cfg.netatop.enable "netatop"
+            // mkService cfg.atopgpu.enable "atopgpu";
+          timers = mkTimer cfg.atopRotate.enable "atop-rotate";
+        };
+      security.wrappers =
+        lib.mkIf cfg.setuidWrapper.enable { atop = { source = "${atop}/bin/atop"; }; };
+    }
+  );
 }
diff --git a/nixos/tests/all-tests.nix b/nixos/tests/all-tests.nix
index 232d89052d4..e468c0da64c 100644
--- a/nixos/tests/all-tests.nix
+++ b/nixos/tests/all-tests.nix
@@ -28,6 +28,7 @@ in
   amazon-init-shell = handleTest ./amazon-init-shell.nix {};
   ammonite = handleTest ./ammonite.nix {};
   atd = handleTest ./atd.nix {};
+  atop = handleTest ./atop.nix {};
   avahi = handleTest ./avahi.nix {};
   avahi-with-resolved = handleTest ./avahi.nix { networkd = true; };
   awscli = handleTest ./awscli.nix { };
diff --git a/nixos/tests/atop.nix b/nixos/tests/atop.nix
new file mode 100644
index 00000000000..8cecf02d28f
--- /dev/null
+++ b/nixos/tests/atop.nix
@@ -0,0 +1,132 @@
+import ./make-test-python.nix ({ pkgs, ... }: {
+  name = "atop";
+
+  nodes = {
+    defaults = { ... }: {
+      programs.atop = {
+        enable = true;
+      };
+    };
+    minimal = { ... }: {
+      programs.atop = {
+        enable = true;
+        atopsvc.enable = false;
+        atopRotate.enable = false;
+        atopacct.enable = false;
+      };
+    };
+    minimal_with_setuid = { ... }: {
+      programs.atop = {
+        enable = true;
+        atopsvc.enable = false;
+        atopRotate.enable = false;
+        atopacct.enable = false;
+        setuidWrapper.enable = true;
+      };
+    };
+
+    atoprc_and_netatop = { ... }: {
+      programs.atop = {
+        enable = true;
+        netatop.enable = true;
+        settings = {
+          flags = "faf1";
+          interval = 2;
+        };
+      };
+    };
+
+    atopgpu = { lib, ... }: {
+      nixpkgs.config.allowUnfreePredicate = pkg: builtins.elem (lib.getName pkg) [
+        "cudatoolkit"
+      ];
+      programs.atop = {
+        enable = true;
+        atopgpu.enable = true;
+      };
+    };
+  };
+
+  testScript = ''
+    def a_version(m):
+      v = m.succeed("atop -V")
+      pkgver = "${pkgs.atop.version}"
+      assert v.startswith("Version: {}".format(pkgver)), "Version is {}, expected `{}`".format(v, pkgver)
+
+    def __exp_path(m, prg, expected):
+      p = m.succeed("type -p \"{}\" | head -c -1".format(prg))
+      assert p == expected, "{} is `{}`, expected `{}`".format(prg, p, expected)
+
+    def a_setuid(m, present=True):
+      if present:
+        __exp_path(m, "atop", "/run/wrappers/bin/atop")
+        stat = m.succeed("stat --printf '%a %u' /run/wrappers/bin/atop")
+        assert stat == "4511 0", "Wrapper stat is {}, expected `4511 0`".format(stat)
+      else:
+        __exp_path(m, "atop", "/run/current-system/sw/bin/atop")
+
+    def assert_no_netatop(m):
+      m.require_unit_state("netatop.service", "inactive")
+      m.fail("modprobe -n -v netatop")
+
+    def a_netatop(m, present=True):
+      m.require_unit_state("netatop.service", "active" if present else "inactive")
+      if present:
+        out = m.succeed("modprobe -n -v netatop")
+        assert out == "", "Module should be loaded, but modprobe would have done `{}`.".format(out)
+      else:
+        m.fail("modprobe -n -v netatop")
+
+    def a_atopgpu(m, present=True):
+      m.require_unit_state("atopgpu.service", "active" if present else "inactive")
+      if present:
+        __exp_path(m, "atopgpud", "/run/current-system/sw/bin/atopgpud")
+
+    # atop.service should log some data to /var/log/atop
+    def a_atopsvc(m, present=True):
+      m.require_unit_state("atop.service", "active" if present else "inactive")
+      if present:
+          files = int(m.succeed("ls -1 /var/log/atop | wc -l"))
+          assert files >= 1, "Expected at least 1 data file"
+        # def check_files(_):
+        #   files = int(m.succeed("ls -1 /var/log/atop | wc -l"))
+        #   return files >= 1
+        # retry(check_files)
+
+    def a_atoprotate(m, present=True):
+      m.require_unit_state("atop-rotate.timer", "active" if present else "inactive")
+
+    # atopacct.service should make kernel write to /run/pacct_source and make dir
+    # /run/pacct_shadow.d
+    def a_atopacct(m, present=True):
+      m.require_unit_state("atopacct.service", "active" if present else "inactive")
+      if present:
+        m.succeed("test -f /run/pacct_source")
+        files = int(m.succeed("ls -1 /run/pacct_shadow.d | wc -l"))
+        assert files >= 1, "Expected at least 1 pacct_shadow.d file"
+
+    def a_atoprc(m, contents):
+      if contents:
+        f = m.succeed("cat /etc/atoprc")
+        assert f == contents, "/etc/atoprc contents: `{}`, expected `{}`".format(f, contents)
+      else:
+        m.succeed("test ! -e /etc/atoprc")
+
+    def assert_all(m, setuid, atopsvc, atoprotate, atopacct, netatop, atopgpu, atoprc):
+      a_version(m)
+      a_setuid(m, setuid)
+      a_atopsvc(m, atopsvc)
+      a_atoprotate(m, atoprotate)
+      a_atopacct(m, atopacct)
+      a_netatop(m, netatop)
+      a_atopgpu(m, atopgpu)
+      a_atoprc(m, atoprc)
+
+    assert_all(defaults, False, True, True, True, False, False, False)
+    assert_all(minimal, False, False, False, False, False, False, False)
+    assert_all(minimal_with_setuid, True, False, False, False, False, False, False)
+    assert_all(atoprc_and_netatop, False, True, True, True, True, False,
+      "flags faf1\ninterval 2\n")
+    assert_all(atopgpu, False, True, True, True, False, True, False)
+  '';
+})